Relay: end a user's sessions into a vault the moment they lose access to it

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

@@ -22,6 +22,7 @@ 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>
@@ -55,7 +56,16 @@ internal sealed class RelaySessionRegistry(IOptions<RelayOptions> options)
 
     /// <summary>Admits a session, or answers null when a limit is met.</summary>
     /// <returns>A lease that gives the place back when disposed.</returns>
-    internal Lease? TryAdmit(Guid userId)
+    /// <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;
 
@@ -71,14 +81,18 @@ internal sealed class RelaySessionRegistry(IOptions<RelayOptions> options)
             perUser[userId] = held + 1;
             total++;
 
-            return new Lease(this, userId);
+            var lease = new Lease(this, userId, vaultId);
+            open.Add(lease);
+
+            return lease;
         }
     }
 
-    private void Release(Guid userId)
+    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
@@ -95,16 +109,41 @@ internal sealed class RelaySessionRegistry(IOptions<RelayOptions> options)
     }
 
     /// <summary>One admitted session's place. Idempotent to dispose.</summary>
-    internal sealed class Lease(RelaySessionRegistry registry, Guid userId) : IDisposable
+    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(userId);
+                registry.Release(this, UserId);
+                revoked.Dispose();
             }
         }
     }