DodoSSH/src/DodoSSH.Api/Features/Relay/RelaySessionRegistry.cs

using DodoSSH.Api.Setup;
using Microsoft.Extensions.Options;

namespace DodoSSH.Api.Features.Relay;

/// <summary>
/// Counts the relay sessions open on this node, and refuses the ones over ADR 0004's limits.
/// </summary>
/// <remarks>
/// <para>
/// Per node, and so per user <em>per node</em>: a user behind a round-robin balancer with three nodes
/// could hold three times the per-user cap. That is accepted rather than fixed, because fixing it means
/// shared state on the hot path of every connect, and the cap exists to stop one account exhausting a
/// node — which a per-node count does exactly.
/// </para>
/// <para>
/// A singleton, because a session outlives the request that opened it. It holds counts and nothing else;
/// the sockets belong to the handlers that opened them.
/// </para>
/// </remarks>
internal sealed class RelaySessionRegistry(IOptions<RelayOptions> options)
{
    private readonly Lock gate = new();
    private readonly Dictionary<Guid, int> perUser = [];
    private readonly HashSet<Lease> open = [];
    private int total;

    /// <summary>How many sessions are open on this node.</summary>
    internal int Count
    {
        get
        {
            lock (gate)
            {
                return total;
            }
        }
    }

    /// <summary>Whether a session for this user would be admitted now.</summary>
    /// <remarks>
    /// Advisory, for the ticket endpoint: it lets a caller already at the cap hear so as an HTTP status
    /// before it has spent a ticket, but only <see cref="TryAdmit"/> decides, since the answer here can
    /// be stale by the time the socket arrives — possibly at another node.
    /// </remarks>
    internal bool HasRoomFor(Guid userId)
    {
        var settings = options.Value;

        lock (gate)
        {
            return total < settings.MaxConcurrentSessionsTotal
                && perUser.GetValueOrDefault(userId) < settings.MaxConcurrentSessionsPerUser;
        }
    }

    /// <summary>Admits a session, or answers null when a limit is met.</summary>
    /// <returns>A lease that gives the place back when disposed.</returns>
    /// <summary>The sessions a user has open now, for re-checking their access.</summary>
    internal IReadOnlyList<Lease> SessionsOf(Guid userId)
    {
        lock (gate)
        {
            return [.. open.Where(lease => lease.UserId == userId)];
        }
    }

    internal Lease? TryAdmit(Guid userId, Guid vaultId)
    {
        var settings = options.Value;

        lock (gate)
        {
            var held = perUser.GetValueOrDefault(userId);

            if (total >= settings.MaxConcurrentSessionsTotal || held >= settings.MaxConcurrentSessionsPerUser)
            {
                return null;
            }

            perUser[userId] = held + 1;
            total++;

            var lease = new Lease(this, userId, vaultId);
            open.Add(lease);

            return lease;
        }
    }

    private void Release(Lease lease, Guid userId)
    {
        lock (gate)
        {
            open.Remove(lease);
            total--;

            // Removed at zero rather than left behind, so the dictionary is the size of the users with a
            // session open and not of every user who ever had one.
            if (perUser[userId] <= 1)
            {
                perUser.Remove(userId);
            }
            else
            {
                perUser[userId]--;
            }
        }
    }

    /// <summary>One admitted session's place. Idempotent to dispose.</summary>
    internal sealed class Lease(RelaySessionRegistry registry, Guid userId, Guid vaultId) : IDisposable
    {
        private readonly CancellationTokenSource revoked = new();

        private int released;

        /// <summary>Whose session this is.</summary>
        internal Guid UserId { get; } = userId;

        /// <summary>The vault the ticket's host is filed in, which is what the session's access hangs on.</summary>
        internal Guid VaultId { get; } = vaultId;

        /// <summary>Cancelled when the session's user loses access to <see cref="VaultId"/>.</summary>
        internal CancellationToken Revoked => revoked.Token;

        /// <summary>Ends the session, because its user can no longer reach the vault it was opened for.</summary>
        internal void Revoke()
        {
            try
            {
                revoked.Cancel();
            }
            catch (ObjectDisposedException)
            {
                // Already over.
            }
        }

        /// <inheritdoc />
        public void Dispose()
        {
            if (Interlocked.Exchange(ref released, 1) == 0)
            {
                registry.Release(this, UserId);
                revoked.Dispose();
            }
        }
    }
}