using System.Buffers;
using System.Net;
using System.Net.WebSockets;
using DodoSSH.Api.Setup;
using DodoSSH.Contracts;
using DodoSSH.Domain;
using DodoSSH.Infrastructure;
using FastEndpoints;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Options;
using Npgsql;
namespace DodoSSH.Api.Features.Relay;
///
/// The relay socket: a ticket in, raw bytes to one host and back.
///
///
///
/// Anonymous at the HTTP layer, because the ticket is the authorization and it is the whole of it. A
/// ticket carries no API authority — it names one host, one set of addresses and one session — so this
/// endpoint can accept it from a client that cannot set an Authorization header on an upgrade,
/// which is the constraint ADR 0004 was written around. It is verified, checked for freshness and
/// redeemed before the upgrade is accepted, so a refused ticket costs a problem document and no socket.
///
///
/// Nothing here is scoped to the request. A session lasts as long as a shell does, which is hours,
/// and anything scoped held across it is a pooled database connection held across it. The database is
/// touched three times — redeem, open the audit row, close it — each from a scope of its own that is
/// disposed before the next byte moves. That is ADR 0004's consequence about DbContext lifetime.
///
///
/// Backpressure is the await. Each direction reads one buffer, writes it, and does not read again until
/// the write has completed, so a slow reader on either side stops this side reading, which fills the
/// kernel's receive window, which throttles the sender at the far end. Nothing is queued in between, and
/// a session costs its two buffers whatever the peers do.
///
///
internal sealed class RelayConnectEndpoint(
IServiceScopeFactory scopes,
IRelayNetwork network,
RelayTicketCodec codec,
RelaySessionRegistry sessions,
IOptions options,
TimeProvider clock,
IHostApplicationLifetime lifetime,
ILogger logger)
: EndpointWithoutRequest
{
///
/// One direction's buffer.
///
///
/// Room for an SSH packet at the size OpenSSH sends them, so a bulk transfer moves in one read and
/// one write per packet rather than several.
///
private const int BufferBytes = 32 * 1024;
/// How long a close handshake, or a peer's last write, may take before the socket is dropped.
private static readonly TimeSpan CloseTimeout = TimeSpan.FromSeconds(5);
///
public override void Configure()
{
Get(RelayProtocol.ConnectPath);
// The ticket is the authorization. See the remarks on this class, and ADR 0004.
AllowAnonymous();
Description(b => b
.WithName("RelayConnect")
.WithSummary("Pipes raw bytes between a WebSocket and the host a relay ticket names.")
.WithTags("Relay"));
}
///
public override async Task HandleAsync(CancellationToken ct)
{
var ticket = await AdmitTicketAsync(ct).ConfigureAwait(false);
if (ticket is null)
{
return;
}
var lease = sessions.TryAdmit(ticket.UserId, ticket.VaultId);
try
{
using var socket = await HttpContext.WebSockets
.AcceptWebSocketAsync(new WebSocketAcceptContext { SubProtocol = RelayProtocol.SubProtocol })
.ConfigureAwait(false);
if (lease is null)
{
RelayLog.LimitReached(logger, ticket.UserId, sessions.Count);
await CloseAsync(
socket,
new Ending(
RelayCloseReason.Unspecified,
RelayProtocol.TooManySessionsCloseCode,
"Too many relay sessions are open."))
.ConfigureAwait(false);
return;
}
if (lifetime.ApplicationStopping.IsCancellationRequested)
{
await CloseAsync(socket, ShuttingDown).ConfigureAwait(false);
return;
}
await RelayAsync(socket, ticket, lease.Revoked, ct).ConfigureAwait(false);
}
finally
{
// Inside a try that starts before the upgrade, for the reason VaultEventsEndpoint gives: an
// accept that throws would otherwise hold a place against this user's cap forever.
lease?.Dispose();
}
}
private static Ending ShuttingDown { get; } = new(
RelayCloseReason.ServerShutdown,
(int)WebSocketCloseStatus.EndpointUnavailable,
"The server is shutting down.");
///
/// Answers everything that is not a redeemable ticket, as ordinary HTTP, before any socket exists.
///
/// The redeemed ticket, or null when a response has been sent.
///
/// A ticket that is forged, expired, spent, or simply absent all answer 401
/// relay-ticket-invalid, with one remedy — ask for another — and no hint of which it was.
///
private async Task AdmitTicketAsync(CancellationToken cancellationToken)
{
if (!options.Value.Enabled)
{
await Send.ResultAsync(Problems.Coded(
StatusCodes.Status404NotFound,
ProblemCodes.RelayUnavailable,
"This server does not relay. GET /api/v1/meta lists the features it does offer."))
.ConfigureAwait(false);
return null;
}
var offered = HttpContext.WebSockets.WebSocketRequestedProtocols;
if (!HttpContext.WebSockets.IsWebSocketRequest
|| !offered.Contains(RelayProtocol.SubProtocol, StringComparer.Ordinal))
{
await Send.ResultAsync(Problems.Coded(
StatusCodes.Status400BadRequest,
ProblemCodes.MalformedRequest,
"This endpoint is a WebSocket. Send an upgrade request offering the "
+ $"'{RelayProtocol.SubProtocol}' subprotocol and the ticket as a second one, "
+ $"'{RelayProtocol.TicketProtocolPrefix}'."))
.ConfigureAwait(false);
return null;
}
var sealedTicket = offered
.FirstOrDefault(protocol => protocol.StartsWith(RelayProtocol.TicketProtocolPrefix, StringComparison.Ordinal))
?[RelayProtocol.TicketProtocolPrefix.Length..];
var ticket = sealedTicket is { Length: > 0 } ? codec.TryOpen(sealedTicket) : null;
if (ticket is null || !await RedeemAsync(ticket, cancellationToken).ConfigureAwait(false))
{
RelayLog.TicketRefused(logger, ticket?.SessionId);
await Send.ResultAsync(Problems.Coded(
StatusCodes.Status401Unauthorized,
ProblemCodes.RelayTicketInvalid,
"This relay ticket is not valid: it is expired, already used, or was not issued here. "
+ $"Ask for another from POST {RelayProtocol.TicketsPath}."))
.ConfigureAwait(false);
return null;
}
return ticket;
}
///
/// Spends the ticket, or answers false when it has been spent already.
///
///
///
/// A unique insert, which is the only form of "once" that holds across nodes: two redemptions racing
/// on two nodes both attempt the insert and the database lets exactly one of them win.
///
///
/// Rows for tickets that can no longer verify are swept on the way. They guard nothing, and sweeping
/// here — on a path that runs once per session rather than once per request — keeps the table the
/// size of the last thirty seconds without a background job to forget about.
///
///
private async Task RedeemAsync(RelayTicket ticket, CancellationToken cancellationToken)
{
var scope = scopes.CreateAsyncScope();
await using var _ = scope.ConfigureAwait(false);
var database = scope.ServiceProvider.GetRequiredService();
var now = clock.GetUtcNow();
await database.RelayTicketUses
.Where(use => use.ExpiresAtUtc < now)
.ExecuteDeleteAsync(cancellationToken)
.ConfigureAwait(false);
database.RelayTicketUses.Add(new RelayTicketUse
{
TicketId = ticket.TicketId,
ExpiresAtUtc = ticket.ExpiresAt,
UsedAtUtc = now,
});
try
{
await database.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
return true;
}
catch (DbUpdateException exception)
when (string.Equals(
(exception.InnerException as PostgresException)?.SqlState,
PostgresErrorCodes.UniqueViolation,
StringComparison.Ordinal))
{
return false;
}
}
/// Dials the target, pipes until something ends it, and writes down what happened.
private async Task RelayAsync(
WebSocket socket,
RelayTicket ticket,
CancellationToken revoked,
CancellationToken requestAborted)
{
var started = clock.GetUtcNow();
await OpenAuditAsync(ticket, started).ConfigureAwait(false);
var dialled = await DialAsync(ticket, requestAborted).ConfigureAwait(false);
if (dialled is not var (target, address))
{
RelayLog.TargetUnreachable(logger, ticket.SessionId, ticket.HostId);
var unreachable = new Ending(
RelayCloseReason.TargetUnreachable,
RelayProtocol.TargetUnreachableCloseCode,
"The relay could not reach the host.");
await CloseAsync(socket, unreachable).ConfigureAwait(false);
await CloseAuditAsync(ticket.SessionId, null, new Counters(clock, started), unreachable.Reason).ConfigureAwait(false);
return;
}
var counters = new Counters(clock, started);
Ending ending;
await using (target.ConfigureAwait(false))
{
RelayLog.SessionOpened(logger, ticket.SessionId, ticket.UserId, ticket.HostId, sessions.Count);
ending = await PumpAsync(socket, target, counters, started, revoked, requestAborted).ConfigureAwait(false);
}
RelayLog.SessionClosed(
logger,
ticket.SessionId,
ending.Reason,
counters.ClientToTarget,
counters.TargetToClient);
await CloseAuditAsync(ticket.SessionId, address, counters, ending.Reason).ConfigureAwait(false);
}
///
/// Connects to the first of the ticket's addresses that answers, within the connect timeout overall.
///
///
/// Overall rather than per address, so a name with sixteen dead addresses costs five seconds rather
/// than eighty. In the order resolved, which is the order the operating system's resolver prefers.
///
private async Task<(Stream Target, IPAddress Address)?> DialAsync(
RelayTicket ticket,
CancellationToken cancellationToken)
{
using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
timeout.CancelAfter(options.Value.ConnectTimeout);
foreach (var address in ticket.TargetAddresses)
{
try
{
var target = await network
.ConnectAsync(address, ticket.TargetPort, timeout.Token)
.ConfigureAwait(false);
return (target, address);
}
catch (OperationCanceledException)
{
return null;
}
catch (Exception exception) when (exception is IOException or System.Net.Sockets.SocketException)
{
// The next address, if there is one.
}
}
return null;
}
///
/// Runs both directions and the clock until one of them ends the session, then ends it politely.
///
///
///
/// No operation on the WebSocket is ever cancelled. Cancelling a pending receive or send on a
/// .NET WebSocket aborts it, and an aborted socket cannot send the close frame that is the only way
/// a client learns why its session ended — which, for an idle timeout or a drain, is the
/// difference between "reconnect" and "something is broken". So the token stops the target side and
/// the clock; the sender is let finish its last write; the close goes out; and the receiver ends on
/// the client's answering close. Each wait is bounded, and the socket is aborted only when a peer
/// has stopped answering altogether.
///
///
private async Task PumpAsync(
WebSocket socket,
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, revoked, pump.Token);
var first = await Task.WhenAny(upstream, downstream, minding).ConfigureAwait(false);
var ending = await first.ConfigureAwait(false) ?? new Ending(
RelayCloseReason.ClientClosed, (int)WebSocketCloseStatus.NormalClosure, string.Empty);
await pump.CancelAsync().ConfigureAwait(false);
// The only sender must be done before the close frame is written: a WebSocket permits one send
// at a time and faults permanently on a second.
if (!await SettlesAsync(downstream).ConfigureAwait(false))
{
socket.Abort();
}
await CloseAsync(socket, ending).ConfigureAwait(false);
if (!await SettlesAsync(upstream).ConfigureAwait(false))
{
socket.Abort();
await SettlesAsync(upstream).ConfigureAwait(false);
}
await SettlesAsync(minding).ConfigureAwait(false);
return ending;
}
/// Client to target: every binary message is the next bytes of the stream.
private static async Task ClientToTargetAsync(
WebSocket socket,
Stream target,
Counters counters,
CancellationToken cancellationToken)
{
var buffer = ArrayPool.Shared.Rent(BufferBytes);
try
{
while (true)
{
ValueWebSocketReceiveResult received;
try
{
// Not cancellable, deliberately. See PumpAsync.
received = await socket.ReceiveAsync(buffer.AsMemory(), CancellationToken.None)
.ConfigureAwait(false);
}
catch (Exception exception) when (exception is WebSocketException or OperationCanceledException)
{
return ClientGone;
}
if (received.MessageType == WebSocketMessageType.Close)
{
return ClientGone;
}
// A byte pipe has no text. A client sending some has the wrong idea about this endpoint,
// and forwarding it into an SSH stream would corrupt that stream at best.
if (received.MessageType != WebSocketMessageType.Binary)
{
return new Ending(
RelayCloseReason.Error,
(int)WebSocketCloseStatus.InvalidMessageType,
"The relay carries binary messages only.");
}
try
{
await target.WriteAsync(buffer.AsMemory(0, received.Count), cancellationToken)
.ConfigureAwait(false);
}
catch (OperationCanceledException)
{
return null;
}
catch (IOException)
{
return TargetGone;
}
counters.AddClientToTarget(received.Count);
}
}
finally
{
ArrayPool.Shared.Return(buffer);
}
}
/// Target to client: every read becomes one binary message.
private static async Task TargetToClientAsync(
Stream target,
WebSocket socket,
Counters counters,
CancellationToken cancellationToken)
{
var buffer = ArrayPool.Shared.Rent(BufferBytes);
try
{
while (true)
{
int read;
try
{
read = await target.ReadAsync(buffer.AsMemory(0, BufferBytes), cancellationToken)
.ConfigureAwait(false);
}
catch (OperationCanceledException)
{
return null;
}
catch (IOException)
{
return TargetGone;
}
if (read == 0)
{
return TargetGone;
}
try
{
// Not cancellable, deliberately. See PumpAsync.
await socket
.SendAsync(
buffer.AsMemory(0, read),
WebSocketMessageType.Binary,
endOfMessage: true,
CancellationToken.None)
.ConfigureAwait(false);
}
catch (Exception exception) when (exception is WebSocketException or ObjectDisposedException)
{
return ClientGone;
}
counters.AddTargetToClient(read);
}
}
finally
{
ArrayPool.Shared.Return(buffer);
}
}
///
/// Ends the session at the idle timeout, the lifetime cap, or the end of a shutdown's drain.
///
///
///
/// A shutdown does not end a session; it starts a countdown. ADR 0004 asks that a
/// docker compose up -d not guillotine live shells, so a session in progress when the server is
/// told to stop is given the drain timeout to finish — most will, because the user was already
/// closing it or the deployment waits — and is closed with 1001 only when that runs out, which tells
/// the client to reconnect rather than to report a fault. The host's own shutdown timeout is set
/// longer than the drain for this to be possible at all.
///
///
/// Idle means no bytes in either direction. WebSocket pings are not bytes of the stream and do not
/// count; SSH's own keepalives do, which is how a client that wants an idle shell kept open says so.
///
///
private async Task MindAsync(
Counters counters,
DateTimeOffset started,
CancellationToken revoked,
CancellationToken cancellationToken)
{
var settings = options.Value;
var hardDeadline = started + settings.MaxSessionDuration;
DateTimeOffset? drainDeadline = null;
while (!cancellationToken.IsCancellationRequested)
{
var now = clock.GetUtcNow();
if (drainDeadline is null && lifetime.ApplicationStopping.IsCancellationRequested)
{
drainDeadline = now + settings.DrainTimeout;
}
var idleDeadline = counters.LastActivity + settings.IdleTimeout;
if (Due(now, revoked.IsCancellationRequested, hardDeadline, drainDeadline, idleDeadline) is { } ending)
{
return ending;
}
var next = new[] { hardDeadline, idleDeadline, drainDeadline ?? DateTimeOffset.MaxValue }.Min();
// 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, revoked, lifetime.ApplicationStopping)
: CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, revoked);
try
{
await Task.Delay(next - now, clock, wake.Token).ConfigureAwait(false);
}
catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
{
// The shutdown began, or access was withdrawn. Round again to act on whichever it was.
}
catch (OperationCanceledException)
{
return null;
}
}
return null;
}
/// Whichever reason to end the session has come due, or null while none has.
///
/// 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.
///
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);
private static Ending TargetGone { get; } =
new(RelayCloseReason.TargetClosed, (int)WebSocketCloseStatus.NormalClosure, "The host closed the connection.");
///
/// Opens the audit row, before the target is dialled, so a session that never reached its host is
/// still one the operator can see.
///
private async Task OpenAuditAsync(RelayTicket ticket, DateTimeOffset started)
{
var scope = scopes.CreateAsyncScope();
await using var _ = scope.ConfigureAwait(false);
var database = scope.ServiceProvider.GetRequiredService();
database.RelaySessions.Add(new RelaySession
{
Id = ticket.SessionId,
UserId = ticket.UserId,
HostId = ticket.HostId,
VaultId = ticket.VaultId,
TargetHost = ticket.TargetHost,
TargetPort = ticket.TargetPort,
ClientAddress = HttpContext.Connection.RemoteIpAddress?.ToString(),
StartedAtUtc = started,
});
// Not the request's token: the row has to be written even for a client that has already gone,
// or a session that was cut off would be the one with no record.
await database.SaveChangesAsync(CancellationToken.None).ConfigureAwait(false);
}
///
/// Closes the audit row. A failure is logged and swallowed: the bytes have already moved, and
/// throwing here would only turn a finished session into an unhandled exception.
///
private async Task CloseAuditAsync(
Guid sessionId,
IPAddress? address,
Counters counters,
RelayCloseReason reason)
{
try
{
var scope = scopes.CreateAsyncScope();
await using var _ = scope.ConfigureAwait(false);
var database = scope.ServiceProvider.GetRequiredService();
var ended = clock.GetUtcNow();
var dialled = address?.ToString();
await database.RelaySessions
.Where(session => session.Id == sessionId)
.ExecuteUpdateAsync(
set => set
.SetProperty(session => session.EndedAtUtc, ended)
.SetProperty(session => session.TargetAddress, dialled)
.SetProperty(session => session.BytesClientToTarget, counters.ClientToTarget)
.SetProperty(session => session.BytesTargetToClient, counters.TargetToClient)
.SetProperty(session => session.CloseReason, reason),
CancellationToken.None)
.ConfigureAwait(false);
}
catch (Exception exception) when (exception is DbUpdateException or NpgsqlException or InvalidOperationException)
{
RelayLog.AuditCloseFailed(logger, sessionId, exception);
}
}
private static async Task CloseAsync(WebSocket socket, Ending ending)
{
if (socket.State is not (WebSocketState.Open or WebSocketState.CloseReceived))
{
return;
}
using var timeout = new CancellationTokenSource(CloseTimeout);
try
{
await socket
.CloseOutputAsync((WebSocketCloseStatus)ending.Code, ending.Description, timeout.Token)
.ConfigureAwait(false);
}
catch (Exception exception)
when (exception is OperationCanceledException or WebSocketException or ObjectDisposedException)
{
// The peer is already gone. There is nothing to tell it.
}
}
/// Waits for a loop to finish, for at most the close timeout.
/// Whether it finished.
private static async Task SettlesAsync(Task loop)
{
try
{
await loop.WaitAsync(CloseTimeout).ConfigureAwait(false);
return true;
}
catch (TimeoutException)
{
return false;
}
catch (Exception exception)
when (exception is OperationCanceledException or WebSocketException or IOException or ObjectDisposedException)
{
return true;
}
}
/// How a session ended, and what the client is told.
private sealed record Ending(RelayCloseReason Reason, int Code, string Description);
///
/// Byte counts and the time of the last byte, written by the two directions and read by the clock.
///
private sealed class Counters(TimeProvider clock, DateTimeOffset started)
{
private long clientToTarget;
private long targetToClient;
private long lastActivityTicks = started.UtcTicks;
internal long ClientToTarget => Interlocked.Read(ref clientToTarget);
internal long TargetToClient => Interlocked.Read(ref targetToClient);
internal DateTimeOffset LastActivity => new(Interlocked.Read(ref lastActivityTicks), TimeSpan.Zero);
internal void AddClientToTarget(int count)
{
Interlocked.Add(ref clientToTarget, count);
Interlocked.Exchange(ref lastActivityTicks, clock.GetUtcNow().UtcTicks);
}
internal void AddTargetToClient(int count)
{
Interlocked.Add(ref targetToClient, count);
Interlocked.Exchange(ref lastActivityTicks, clock.GetUtcNow().UtcTicks);
}
}
}