natsdotnet/src/NATS.Server/Auth/Account.cs

using System.Collections.Concurrent;
using NATS.Server.Imports;
using NATS.Server.Subscriptions;

namespace NATS.Server.Auth;

public sealed class Account : IDisposable
{
    public const string GlobalAccountName = "$G";

    public string Name { get; }
    public SubList SubList { get; } = new();
    public Permissions? DefaultPermissions { get; set; }
    public int MaxConnections { get; set; } // 0 = unlimited
    public int MaxSubscriptions { get; set; } // 0 = unlimited
    public ExportMap Exports { get; } = new();
    public ImportMap Imports { get; } = new();
    public int MaxJetStreamStreams { get; set; } // 0 = unlimited
    public string? JetStreamTier { get; set; }

    // JWT fields
    public string? Nkey { get; set; }
    public string? Issuer { get; set; }
    public Dictionary<string, object>? SigningKeys { get; set; }
    private readonly ConcurrentDictionary<string, long> _revokedUsers = new(StringComparer.Ordinal);

    public void RevokeUser(string userNkey, long issuedAt) => _revokedUsers[userNkey] = issuedAt;

    public bool IsUserRevoked(string userNkey, long issuedAt)
    {
        if (_revokedUsers.TryGetValue(userNkey, out var revokedAt))
            return issuedAt <= revokedAt;
        // Check "*" wildcard for all-user revocation
        if (_revokedUsers.TryGetValue("*", out revokedAt))
            return issuedAt <= revokedAt;
        return false;
    }

    private readonly ConcurrentDictionary<ulong, byte> _clients = new();
    private int _subscriptionCount;
    private int _jetStreamStreamCount;

    public Account(string name)
    {
        Name = name;
    }

    public int ClientCount => _clients.Count;
    public int SubscriptionCount => Volatile.Read(ref _subscriptionCount);
    public int JetStreamStreamCount => Volatile.Read(ref _jetStreamStreamCount);

    /// <summary>Returns false if max connections exceeded.</summary>
    public bool AddClient(ulong clientId)
    {
        if (MaxConnections > 0 && _clients.Count >= MaxConnections)
            return false;
        _clients[clientId] = 0;
        return true;
    }

    public void RemoveClient(ulong clientId) => _clients.TryRemove(clientId, out _);

    public bool IncrementSubscriptions()
    {
        if (MaxSubscriptions > 0 && Volatile.Read(ref _subscriptionCount) >= MaxSubscriptions)
            return false;
        Interlocked.Increment(ref _subscriptionCount);
        return true;
    }

    public void DecrementSubscriptions()
    {
        Interlocked.Decrement(ref _subscriptionCount);
    }

    public bool TryReserveStream()
    {
        if (MaxJetStreamStreams > 0 && Volatile.Read(ref _jetStreamStreamCount) >= MaxJetStreamStreams)
            return false;

        Interlocked.Increment(ref _jetStreamStreamCount);
        return true;
    }

    public void ReleaseStream()
    {
        if (Volatile.Read(ref _jetStreamStreamCount) == 0)
            return;

        Interlocked.Decrement(ref _jetStreamStreamCount);
    }

    // Per-account message/byte stats
    private long _inMsgs;
    private long _outMsgs;
    private long _inBytes;
    private long _outBytes;

    public long InMsgs => Interlocked.Read(ref _inMsgs);
    public long OutMsgs => Interlocked.Read(ref _outMsgs);
    public long InBytes => Interlocked.Read(ref _inBytes);
    public long OutBytes => Interlocked.Read(ref _outBytes);

    public void IncrementInbound(long msgs, long bytes)
    {
        Interlocked.Add(ref _inMsgs, msgs);
        Interlocked.Add(ref _inBytes, bytes);
    }

    public void IncrementOutbound(long msgs, long bytes)
    {
        Interlocked.Add(ref _outMsgs, msgs);
        Interlocked.Add(ref _outBytes, bytes);
    }

    // Internal (ACCOUNT) client for import/export message routing
    private InternalClient? _internalClient;

    public InternalClient GetOrCreateInternalClient(ulong clientId)
    {
        if (_internalClient != null) return _internalClient;
        _internalClient = new InternalClient(clientId, ClientKind.Account, this);
        return _internalClient;
    }

    public void AddServiceExport(string subject, ServiceResponseType responseType, IEnumerable<Account>? approved)
    {
        var auth = new ExportAuth
        {
            ApprovedAccounts = approved != null ? new HashSet<string>(approved.Select(a => a.Name)) : null,
        };
        Exports.Services[subject] = new ServiceExport
        {
            Auth = auth,
            Account = this,
            ResponseType = responseType,
        };
    }

    public void AddStreamExport(string subject, IEnumerable<Account>? approved)
    {
        var auth = new ExportAuth
        {
            ApprovedAccounts = approved != null ? new HashSet<string>(approved.Select(a => a.Name)) : null,
        };
        Exports.Streams[subject] = new StreamExport { Auth = auth };
    }

    public ServiceImport AddServiceImport(Account destination, string from, string to)
    {
        if (!destination.Exports.Services.TryGetValue(to, out var export))
            throw new InvalidOperationException($"No service export found for '{to}' on account '{destination.Name}'");

        if (!export.Auth.IsAuthorized(this))
            throw new UnauthorizedAccessException($"Account '{Name}' not authorized to import '{to}' from '{destination.Name}'");

        var si = new ServiceImport
        {
            DestinationAccount = destination,
            From = from,
            To = to,
            Export = export,
            ResponseType = export.ResponseType,
        };

        Imports.AddServiceImport(si);
        return si;
    }

    public void AddStreamImport(Account source, string from, string to)
    {
        if (!source.Exports.Streams.TryGetValue(from, out var export))
            throw new InvalidOperationException($"No stream export found for '{from}' on account '{source.Name}'");

        if (!export.Auth.IsAuthorized(this))
            throw new UnauthorizedAccessException($"Account '{Name}' not authorized to import '{from}' from '{source.Name}'");

        var si = new StreamImport
        {
            SourceAccount = source,
            From = from,
            To = to,
        };

        Imports.Streams.Add(si);
    }

    public void Dispose() => SubList.Dispose();
}