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; }

    /// <summary>Per-account JetStream resource limits (storage, consumers, ack pending).</summary>
    public AccountLimits JetStreamLimits { get; set; } = AccountLimits.Unlimited;

    // 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;
    private int _consumerCount;
    private long _storageUsed;

    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);
    public int ConsumerCount => Volatile.Read(ref _consumerCount);
    public long StorageUsed => Interlocked.Read(ref _storageUsed);

    /// <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);
    }

    /// <summary>
    /// Reserves a stream slot, checking both <see cref="MaxJetStreamStreams"/> (legacy)
    /// and <see cref="JetStreamLimits"/>.<see cref="AccountLimits.MaxStreams"/>.
    /// </summary>
    public bool TryReserveStream()
    {
        var effectiveMax = JetStreamLimits.MaxStreams > 0
            ? JetStreamLimits.MaxStreams
            : MaxJetStreamStreams;

        if (effectiveMax > 0 && Volatile.Read(ref _jetStreamStreamCount) >= effectiveMax)
            return false;

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

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

        Interlocked.Decrement(ref _jetStreamStreamCount);
    }

    /// <summary>Reserves a consumer slot. Returns false if <see cref="AccountLimits.MaxConsumers"/> is exceeded.</summary>
    public bool TryReserveConsumer()
    {
        var max = JetStreamLimits.MaxConsumers;
        if (max > 0 && Volatile.Read(ref _consumerCount) >= max)
            return false;

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

    public void ReleaseConsumer()
    {
        if (Volatile.Read(ref _consumerCount) == 0)
            return;

        Interlocked.Decrement(ref _consumerCount);
    }

    /// <summary>
    /// Adjusts the tracked storage usage by <paramref name="deltaBytes"/>.
    /// Returns false if the positive delta would exceed <see cref="AccountLimits.MaxStorage"/>.
    /// A negative delta always succeeds.
    /// </summary>
    public bool TrackStorageDelta(long deltaBytes)
    {
        var maxStorage = JetStreamLimits.MaxStorage;

        if (deltaBytes > 0 && maxStorage > 0)
        {
            var current = Interlocked.Read(ref _storageUsed);
            if (current + deltaBytes > maxStorage)
                return false;
        }

        Interlocked.Add(ref _storageUsed, deltaBytes);
        return true;
    }

    // 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 };
    }

    /// <summary>
    /// Adds a service import with cycle detection.
    /// Go reference: accounts.go addServiceImport with checkForImportCycle.
    /// </summary>
    /// <exception cref="InvalidOperationException">Thrown if no export found or import would create a cycle.</exception>
    /// <exception cref="UnauthorizedAccessException">Thrown if this account is not authorized.</exception>
    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}'");

        // Cycle detection: check if adding this import from destination would
        // create a path back to this account.
        if (AccountImportExport.DetectCycle(destination, this))
            throw new InvalidOperationException("Import would create a cycle");

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

        Imports.AddServiceImport(si);
        return si;
    }

    /// <summary>Removes a service import by its 'from' subject.</summary>
    /// <returns>True if the import was found and removed.</returns>
    public bool RemoveServiceImport(string from)
    {
        return Imports.Services.Remove(from);
    }

    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);
    }

    /// <summary>Removes a stream import by its 'from' subject.</summary>
    /// <returns>True if the import was found and removed.</returns>
    public bool RemoveStreamImport(string from)
    {
        var idx = Imports.Streams.FindIndex(s => string.Equals(s.From, from, StringComparison.Ordinal));
        if (idx < 0)
            return false;
        Imports.Streams.RemoveAt(idx);
        return true;
    }

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