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

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

@@ -83,7 +83,7 @@ internal sealed class RelayConnectEndpoint(
             return;
         }
 
-        var lease = sessions.TryAdmit(ticket.UserId);
+        var lease = sessions.TryAdmit(ticket.UserId, ticket.VaultId);
 
         try
         {
@@ -113,7 +113,7 @@ internal sealed class RelayConnectEndpoint(
                 return;
             }
 
-            await RelayAsync(socket, ticket, ct).ConfigureAwait(false);
+            await RelayAsync(socket, ticket, lease.Revoked, ct).ConfigureAwait(false);
         }
         finally
         {
@@ -239,7 +239,11 @@ internal sealed class RelayConnectEndpoint(
     }
 
     /// <summary>Dials the target, pipes until something ends it, and writes down what happened.</summary>
-    private async Task RelayAsync(WebSocket socket, RelayTicket ticket, CancellationToken requestAborted)
+    private async Task RelayAsync(
+        WebSocket socket,
+        RelayTicket ticket,
+        CancellationToken revoked,
+        CancellationToken requestAborted)
     {
         var started = clock.GetUtcNow();
 
@@ -269,7 +273,7 @@ internal sealed class RelayConnectEndpoint(
         {
             RelayLog.SessionOpened(logger, ticket.SessionId, ticket.UserId, ticket.HostId, sessions.Count);
 
-            ending = await PumpAsync(socket, target, counters, started, requestAborted).ConfigureAwait(false);
+            ending = await PumpAsync(socket, target, counters, started, revoked, requestAborted).ConfigureAwait(false);
         }
 
         RelayLog.SessionClosed(
@@ -338,13 +342,14 @@ internal sealed class RelayConnectEndpoint(
         Stream target,
         Counters counters,
         DateTimeOffset started,
+        CancellationToken revoked,
         CancellationToken requestAborted)
     {
         using var pump = CancellationTokenSource.CreateLinkedTokenSource(requestAborted);
 
         var upstream = ClientToTargetAsync(socket, target, counters, pump.Token);
         var downstream = TargetToClientAsync(target, socket, counters, pump.Token);
-        var minding = MindAsync(counters, started, pump.Token);
+        var minding = MindAsync(counters, started, revoked, pump.Token);
 
         var first = await Task.WhenAny(upstream, downstream, minding).ConfigureAwait(false);
         var ending = await first.ConfigureAwait(false) ?? new Ending(
@@ -512,7 +517,11 @@ internal sealed class RelayConnectEndpoint(
     /// count; SSH's own keepalives do, which is how a client that wants an idle shell kept open says so.
     /// </para>
     /// </remarks>
-    private async Task<Ending?> MindAsync(Counters counters, DateTimeOffset started, CancellationToken cancellationToken)
+    private async Task<Ending?> MindAsync(
+        Counters counters,
+        DateTimeOffset started,
+        CancellationToken revoked,
+        CancellationToken cancellationToken)
     {
         var settings = options.Value;
         var hardDeadline = started + settings.MaxSessionDuration;
@@ -527,27 +536,11 @@ internal sealed class RelayConnectEndpoint(
                 drainDeadline = now + settings.DrainTimeout;
             }
 
-            if (now >= hardDeadline)
-            {
-                return new Ending(
-                    RelayCloseReason.DurationLimit,
-                    RelayProtocol.DurationLimitCloseCode,
-                    "This relay session reached its maximum lifetime.");
-            }
-
-            if (drainDeadline is { } drain && now >= drain)
-            {
-                return ShuttingDown;
-            }
-
             var idleDeadline = counters.LastActivity + settings.IdleTimeout;
 
-            if (now >= idleDeadline)
+            if (Due(now, revoked.IsCancellationRequested, hardDeadline, drainDeadline, idleDeadline) is { } ending)
             {
-                return new Ending(
-                    RelayCloseReason.IdleTimeout,
-                    RelayProtocol.IdleTimeoutCloseCode,
-                    "This relay session carried nothing for the idle timeout.");
+                return ending;
             }
 
             var next = new[] { hardDeadline, idleDeadline, drainDeadline ?? DateTimeOffset.MaxValue }.Min();
@@ -555,8 +548,8 @@ internal sealed class RelayConnectEndpoint(
             // Woken early by the start of a shutdown, so the drain is measured from when it began rather
             // than from whenever the idle clock next came round.
             using var wake = drainDeadline is null
-                ? CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, lifetime.ApplicationStopping)
-                : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+                ? CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, revoked, lifetime.ApplicationStopping)
+                : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, revoked);
 
             try
             {
@@ -564,7 +557,7 @@ internal sealed class RelayConnectEndpoint(
             }
             catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
             {
-                // The shutdown began. Round again to start the drain.
+                // The shutdown began, or access was withdrawn. Round again to act on whichever it was.
             }
             catch (OperationCanceledException)
             {
@@ -575,6 +568,50 @@ internal sealed class RelayConnectEndpoint(
         return null;
     }
 
+    /// <summary>Whichever reason to end the session has come due, or null while none has.</summary>
+    /// <remarks>
+    /// Revocation first, because it is the one reason here that is about the user rather than the session:
+    /// they were removed from the team, or the vault was taken from them, after this session was admitted.
+    /// It used to run on for as long as it carried bytes, up to the maximum lifetime.
+    /// </remarks>
+    private static Ending? Due(
+        DateTimeOffset now,
+        bool revoked,
+        DateTimeOffset hardDeadline,
+        DateTimeOffset? drainDeadline,
+        DateTimeOffset idleDeadline)
+    {
+        if (revoked)
+        {
+            return AccessRevoked;
+        }
+
+        if (now >= hardDeadline)
+        {
+            return new Ending(
+                RelayCloseReason.DurationLimit,
+                RelayProtocol.DurationLimitCloseCode,
+                "This relay session reached its maximum lifetime.");
+        }
+
+        if (drainDeadline is { } drain && now >= drain)
+        {
+            return ShuttingDown;
+        }
+
+        return now >= idleDeadline
+            ? new Ending(
+                RelayCloseReason.IdleTimeout,
+                RelayProtocol.IdleTimeoutCloseCode,
+                "This relay session carried nothing for the idle timeout.")
+            : null;
+    }
+
+    private static Ending AccessRevoked { get; } = new(
+        RelayCloseReason.AccessRevoked,
+        RelayProtocol.AccessRevokedCloseCode,
+        "You no longer have access to this host.");
+
     private static Ending ClientGone { get; } =
         new(RelayCloseReason.ClientClosed, (int)WebSocketCloseStatus.NormalClosure, string.Empty);