DodoSSH/tests/DodoSSH.Api.Tests/RelayEndpointTests.cs

using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;
using System.Text;
using DodoSSH.Api.Features.Events;
using DodoSSH.Contracts;
using DodoSSH.Domain;
using DodoSSH.Infrastructure;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;

namespace DodoSSH.Api.Tests;

/// <summary>
/// The relay, over a real socket, through the real ticket, against a target the test owns.
/// </summary>
/// <remarks>
/// <para>
/// The network is the only thing replaced — see <see cref="FakeRelayNetwork"/> — and it is replaced
/// <em>under</em> the address policy, so every address these tests arrange is still judged by the real
/// rules. That is what lets the refusals here be meaningful: a loopback target is refused because the
/// policy refused it, not because the fake had nothing to route it to, and the tests say which by
/// asserting that nothing was dialled.
/// </para>
/// <para>
/// The refusals share one answer on purpose, and a test holds them to it. A distinct response for "that
/// name resolves internally, but no" is a scanner for the operator's DNS in the hands of anyone who can
/// edit a host.
/// </para>
/// </remarks>
[Collection(ApiCollection.Name)]
public sealed class RelayEndpointTests(ApiFixture fixture)
{
    private static readonly DateTimeOffset Now = new(2026, 9, 28, 12, 0, 0, TimeSpan.Zero);

    /// <summary>How long a test waits for a frame. Generous, because it bounds a broken build, not a measurement.</summary>
    private static readonly TimeSpan FrameTimeout = TimeSpan.FromSeconds(30);

    // ---- The ticket ----

    [Fact]
    public async Task ATicket_NamesTheSocketAndExpiresWithinThirtySeconds()
    {
        var (subject, hostId) = await SeedRelayHostAsync("10.40.0.1", 22);

        var ticket = await IssueAsync(subject, hostId);

        ticket.SubProtocol.ShouldBe(RelayProtocol.SubProtocol);
        ticket.WebSocketUrl.AbsolutePath.ShouldBe(RelayProtocol.ConnectPath);
        ticket.ExpiresAt.ShouldBeLessThanOrEqualTo(TimeProvider.System.GetUtcNow().AddSeconds(31));

        // It has to stand in a Sec-WebSocket-Protocol list, so it must be an HTTP token: no padding, no
        // slash, no plus — which is what base64url is for.
        ticket.Ticket.ShouldAllBe(c => char.IsAsciiLetterOrDigit(c) || c == '-' || c == '_');
    }

    [Fact]
    public async Task EveryTargetRefusal_IsTheSameAnswer()
    {
        var (subject, relayOff) = await SeedHostAsync(relay: null);
        var (_, others) = await SeedRelayHostAsync("10.40.0.2", 22);
        var (_, loopback) = await SeedRelayHostAsync("127.0.0.1", 22, subject);
        var (_, smtp) = await SeedRelayHostAsync("10.40.0.3", 25, subject);
        var (_, unresolvable) = await SeedRelayHostAsync("nowhere.invalid", 22, subject);

        var client = fixture.CreateClientFor(subject);
        var answers = new List<JsonProblem>();

        foreach (var hostId in new[] { relayOff, others, loopback, smtp, unresolvable, Guid.CreateVersion7() })
        {
            var response = await client.PostContractAsync(RelayProtocol.TicketsPath, new RelayTicketRequest(hostId, null));

            response.StatusCode.ShouldBe(HttpStatusCode.Forbidden);
            answers.Add((await response.Content.ReadProblemAsync()).ShouldNotBeNull());
        }

        answers.ShouldAllBe(problem => problem.Code == ProblemCodes.RelayTargetRejected);
        answers.Select(problem => problem.Detail).Distinct(StringComparer.Ordinal).Count().ShouldBe(1);
    }

    [Theory]
    [InlineData("127.0.0.1")]
    [InlineData("::1")]
    [InlineData("::ffff:127.0.0.1")]
    [InlineData("169.254.169.254")]
    [InlineData("::ffff:169.254.169.254")]
    [InlineData("0.0.0.0")]
    public async Task AnAddressThePolicyRefuses_IsNeverDialled(string address)
    {
        var (subject, hostId) = await SeedRelayHostAsync(address, 22);
        var dialledBefore = fixture.RelayNetwork.Dialled.Count;

        var response = await fixture.CreateClientFor(subject)
            .PostContractAsync(RelayProtocol.TicketsPath, new RelayTicketRequest(hostId, null));

        response.StatusCode.ShouldBe(HttpStatusCode.Forbidden);
        fixture.RelayNetwork.Dialled.Count.ShouldBe(dialledBefore);
    }

    [Fact]
    public async Task ANameResolvingToABadAndAGoodAddress_IsDialledOnlyOnTheGoodOne()
    {
        // The rebinding shape: one answer the policy likes and one it must never reach. The ticket carries
        // only what passed, so the socket has nothing else it could try.
        using var target = LoopbackTarget.Echo();
        var good = new IPEndPoint(IPAddress.Parse("10.40.1.1"), 22);

        fixture.RelayNetwork.Resolve("mixed.test", IPAddress.Parse("169.254.169.254"), good.Address);
        fixture.RelayNetwork.Route(good, target.Endpoint);

        var (subject, hostId) = await SeedRelayHostAsync("mixed.test", 22);

        using var timeout = Timeout();
        using var socket = await ConnectAsync(await IssueAsync(subject, hostId), timeout.Token);

        await RoundTripAsync(socket, "hello", timeout.Token);

        fixture.RelayNetwork.Dialled.ShouldNotContain(endpoint => endpoint.Address.Equals(IPAddress.Parse("169.254.169.254")));
    }

    [Fact]
    public async Task APortForward_IsRefused()
    {
        // Its port lives in ciphertext, so honouring one would mean taking the port from the client.
        var (subject, hostId) = await SeedRelayHostAsync("10.40.0.4", 22);

        var response = await fixture.CreateClientFor(subject)
            .PostContractAsync(RelayProtocol.TicketsPath, new RelayTicketRequest(hostId, Guid.CreateVersion7()));

        response.StatusCode.ShouldBe(HttpStatusCode.Forbidden);
        (await response.Content.ReadProblemAsync()).ShouldNotBeNull().Code.ShouldBe(ProblemCodes.RelayTargetRejected);
    }

    // ---- The socket ----

    [Fact]
    public async Task BytesGoBothWays_AndTheSessionIsAudited()
    {
        using var target = LoopbackTarget.Echo();
        var (subject, hostId) = await SeedRoutedHostAsync(target, "10.40.2.1");
        var ticket = await IssueAsync(subject, hostId);

        using var timeout = Timeout();

        using (var socket = await ConnectAsync(ticket, timeout.Token))
        {
            socket.SubProtocol.ShouldBe(RelayProtocol.SubProtocol);

            await RoundTripAsync(socket, "SSH-2.0-test", timeout.Token);
            await RoundTripAsync(socket, "a second message", timeout.Token);

            await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, timeout.Token);
        }

        var session = await WaitForClosedAuditAsync(ticket.SessionId, timeout.Token);

        session.HostId.ShouldBe(hostId);
        session.TargetAddress.ShouldBe("10.40.2.1");
        session.TargetPort.ShouldBe(22);
        session.CloseReason.ShouldBe(RelayCloseReason.ClientClosed);

        var sent = Encoding.UTF8.GetByteCount("SSH-2.0-test") + Encoding.UTF8.GetByteCount("a second message");
        session.BytesClientToTarget.ShouldBe(sent);
        session.BytesTargetToClient.ShouldBe(sent);
    }

    [Fact]
    public async Task ATicket_IsSpentByItsFirstUse()
    {
        using var target = LoopbackTarget.Echo();
        var (subject, hostId) = await SeedRoutedHostAsync(target, "10.40.3.1");
        var ticket = await IssueAsync(subject, hostId);

        using var timeout = Timeout();
        using var first = await ConnectAsync(ticket, timeout.Token);

        var replay = await Should.ThrowAsync<InvalidOperationException>(ConnectAsync(ticket, timeout.Token));

        replay.Message.ShouldContain("401");
    }

    [Theory]
    [InlineData(null)]
    [InlineData("")]
    [InlineData("not-a-ticket")]
    public async Task AnUpgradeWithoutAGenuineTicket_IsRefused(string? ticket)
    {
        var client = fixture.Server.CreateWebSocketClient();
        client.SubProtocols.Add(RelayProtocol.SubProtocol);

        if (ticket is not null)
        {
            client.SubProtocols.Add(RelayProtocol.TicketProtocolPrefix + ticket);
        }

        var refused = await Should.ThrowAsync<InvalidOperationException>(client.ConnectAsync(
            new Uri(fixture.Server.BaseAddress, RelayProtocol.ConnectPath.TrimStart('/')),
            TestContext.Current.CancellationToken));

        refused.Message.ShouldContain("401");
    }

    [Fact]
    public async Task APlainGet_SaysItIsAWebSocket()
    {
        var response = await fixture.CreateClient().GetAsync(
            new Uri(RelayProtocol.ConnectPath, UriKind.Relative),
            TestContext.Current.CancellationToken);

        response.StatusCode.ShouldBe(HttpStatusCode.BadRequest);
        (await response.Content.ReadProblemAsync()).ShouldNotBeNull().Code.ShouldBe(ProblemCodes.MalformedRequest);
    }

    [Fact]
    public async Task ATargetThatDoesNotAnswer_ClosesWithItsOwnCode()
    {
        // Routed nowhere, so the fake refuses the connect the way a closed port does. The user has to be
        // told the *server* could not reach the machine, which is not what a direct failure means.
        var (subject, hostId) = await SeedRelayHostAsync("10.40.4.1", 22);
        var ticket = await IssueAsync(subject, hostId);

        using var timeout = Timeout();
        using var socket = await ConnectAsync(ticket, timeout.Token);

        await ReadUntilCloseAsync(socket, timeout.Token);

        socket.CloseStatus.ShouldBe((WebSocketCloseStatus)RelayProtocol.TargetUnreachableCloseCode);
        (await WaitForClosedAuditAsync(ticket.SessionId, timeout.Token)).CloseReason
            .ShouldBe(RelayCloseReason.TargetUnreachable);
    }

    [Fact]
    public async Task ATargetThatHangsUp_EndsTheSession()
    {
        using var target = LoopbackTarget.SayThenHangUp("bye");
        var (subject, hostId) = await SeedRoutedHostAsync(target, "10.40.5.1");
        var ticket = await IssueAsync(subject, hostId);

        using var timeout = Timeout();
        using var socket = await ConnectAsync(ticket, timeout.Token);

        (await ReadTextAsync(socket, timeout.Token)).ShouldBe("bye");
        await ReadUntilCloseAsync(socket, timeout.Token);

        socket.CloseStatus.ShouldBe(WebSocketCloseStatus.NormalClosure);
        (await WaitForClosedAuditAsync(ticket.SessionId, timeout.Token)).CloseReason
            .ShouldBe(RelayCloseReason.TargetClosed);
    }

    [Fact]
    public async Task AnIdleSession_IsClosedWithTheIdleCode()
    {
        // The close code has to arrive, not merely the close. Cancelling a pending receive aborts a .NET
        // WebSocket, and an aborted one sends nothing — so a relay that stopped its reader by cancelling
        // it would pass a test that only checked the socket ended.
        using var target = LoopbackTarget.Echo();
        var (subject, hostId) = await SeedRoutedHostAsync(target, "10.40.6.1");
        var ticket = await IssueAsync(subject, hostId);

        using var timeout = Timeout();
        using var socket = await ConnectAsync(ticket, timeout.Token);

        await ReadUntilCloseAsync(socket, timeout.Token);

        socket.CloseStatus.ShouldBe((WebSocketCloseStatus)RelayProtocol.IdleTimeoutCloseCode);
        (await WaitForClosedAuditAsync(ticket.SessionId, timeout.Token)).CloseReason
            .ShouldBe(RelayCloseReason.IdleTimeout);
    }

    /// <remarks>
    /// <para>
    /// A session is authorised once, when its ticket is issued, and used to run on for as long as it carried
    /// bytes — removing somebody from a team left their shells into its machines open for up to the session's
    /// maximum lifetime. The announcement that tells their event socket to re-read its vaults now has the
    /// relay ask again too.
    /// </para>
    /// <para>
    /// Access is withdrawn here by deleting the vault rather than by a team removal, because what is under
    /// test is the relay hearing the announcement and asking the access model again — any way of making the
    /// answer change will do, and this one needs no team. The removal path announces through the same call;
    /// see <c>TeamService.RemoveMemberAsync</c>.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task ASessionWhoseUserLosesAccess_IsClosedWithItsOwnCode()
    {
        using var target = LoopbackTarget.Echo();
        var (subject, hostId) = await SeedRoutedHostAsync(target, "10.40.8.1");
        var ticket = await IssueAsync(subject, hostId);

        using var timeout = Timeout();
        using var socket = await ConnectAsync(ticket, timeout.Token);

        await RoundTripAsync(socket, "still here", timeout.Token);

        Guid userId;

        await using (var scope = fixture.CreateScope())
        {
            var database = scope.ServiceProvider.GetRequiredService<DodoDbContext>();

            userId = (await database.Users.SingleAsync(u => u.Subject == subject, timeout.Token)).Id;

            var vaultId = (await database.Hosts.SingleAsync(h => h.Id == hostId, timeout.Token)).VaultId;
            var vault = await database.Vaults.SingleAsync(v => v.Id == vaultId, timeout.Token);

            vault.DeletedAtUtc = Now;
            await database.SaveChangesAsync(timeout.Token);

            scope.ServiceProvider.GetRequiredService<IVaultEventPublisher>().VaultAccessChanged(userId);
        }

        await ReadUntilCloseAsync(socket, timeout.Token);

        socket.CloseStatus.ShouldBe((WebSocketCloseStatus)RelayProtocol.AccessRevokedCloseCode);
        (await WaitForClosedAuditAsync(ticket.SessionId, timeout.Token)).CloseReason
            .ShouldBe(RelayCloseReason.AccessRevoked);
    }

    [Fact]
    public async Task AUserAtTheCap_IsRefusedATicket_AndASocketOverItIsClosed()
    {
        using var target = LoopbackTarget.Echo();
        var (subject, hostId) = await SeedRoutedHostAsync(target, "10.40.7.1");

        // Three tickets while nothing is open, so the third reaches the socket and is refused there —
        // the authoritative check, which the ticket endpoint's is only an early warning of.
        var tickets = new[]
        {
            await IssueAsync(subject, hostId),
            await IssueAsync(subject, hostId),
            await IssueAsync(subject, hostId),
        };

        using var timeout = Timeout();
        using var first = await ConnectAsync(tickets[0], timeout.Token);
        using var second = await ConnectAsync(tickets[1], timeout.Token);

        await RoundTripAsync(first, "one", timeout.Token);
        await RoundTripAsync(second, "two", timeout.Token);

        var early = await fixture.CreateClientFor(subject)
            .PostContractAsync(RelayProtocol.TicketsPath, new RelayTicketRequest(hostId, null));

        early.StatusCode.ShouldBe(HttpStatusCode.TooManyRequests);
        (await early.Content.ReadProblemAsync()).ShouldNotBeNull().Code.ShouldBe(ProblemCodes.RelayLimitReached);

        using var third = await ConnectAsync(tickets[2], timeout.Token);
        await ReadUntilCloseAsync(third, timeout.Token);

        third.CloseStatus.ShouldBe((WebSocketCloseStatus)RelayProtocol.TooManySessionsCloseCode);
    }

    // ---- Helpers ----

    private static CancellationTokenSource Timeout()
    {
        var source = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
        source.CancelAfter(FrameTimeout);
        return source;
    }

    private static string NewSubject() => $"relay-{Guid.CreateVersion7():N}";

    private async Task<RelayTicketResponse> IssueAsync(string subject, Guid hostId)
    {
        var response = await fixture.CreateClientFor(subject)
            .PostContractAsync(RelayProtocol.TicketsPath, new RelayTicketRequest(hostId, null));

        response.StatusCode.ShouldBe(HttpStatusCode.OK);

        return (await response.Content.ReadContractAsync<RelayTicketResponse>()).ShouldNotBeNull();
    }

    private Task<WebSocket> ConnectAsync(RelayTicketResponse ticket, CancellationToken cancellationToken)
    {
        var client = fixture.Server.CreateWebSocketClient();
        client.SubProtocols.Add(RelayProtocol.SubProtocol);
        client.SubProtocols.Add(RelayProtocol.TicketProtocolPrefix + ticket.Ticket);

        return client.ConnectAsync(
            new Uri(fixture.Server.BaseAddress, RelayProtocol.ConnectPath.TrimStart('/')),
            cancellationToken);
    }

    private static async Task RoundTripAsync(WebSocket socket, string text, CancellationToken cancellationToken)
    {
        await socket.SendAsync(Encoding.UTF8.GetBytes(text), WebSocketMessageType.Binary, true, cancellationToken);

        (await ReadTextAsync(socket, cancellationToken, Encoding.UTF8.GetByteCount(text))).ShouldBe(text);
    }

    /// <summary>
    /// Reads until at least <paramref name="expected"/> bytes have arrived, because message boundaries
    /// carry no meaning on the relay and an echo may come back split.
    /// </summary>
    private static async Task<string> ReadTextAsync(
        WebSocket socket,
        CancellationToken cancellationToken,
        int expected = 1)
    {
        var buffer = new byte[64 * 1024];
        var received = 0;

        while (received < expected)
        {
            var result = await socket.ReceiveAsync(buffer.AsMemory(received), cancellationToken);

            result.MessageType.ShouldBe(WebSocketMessageType.Binary);
            received += result.Count;
        }

        return Encoding.UTF8.GetString(buffer, 0, received);
    }

    private static async Task ReadUntilCloseAsync(WebSocket socket, CancellationToken cancellationToken)
    {
        var buffer = new byte[64 * 1024];

        while (true)
        {
            var result = await socket.ReceiveAsync(buffer.AsMemory(), cancellationToken);

            if (result.MessageType == WebSocketMessageType.Close)
            {
                return;
            }
        }
    }

    /// <summary>
    /// Waits for the audit row's second write, which lands after the close frame has gone out.
    /// </summary>
    private async Task<RelaySession> WaitForClosedAuditAsync(Guid sessionId, CancellationToken cancellationToken)
    {
        while (true)
        {
            await using (var scope = fixture.CreateScope())
            {
                var database = scope.ServiceProvider.GetRequiredService<DodoDbContext>();
                var session = await database.RelaySessions
                    .AsNoTracking()
                    .SingleOrDefaultAsync(s => s.Id == sessionId, cancellationToken);

                if (session?.EndedAtUtc is not null)
                {
                    return session;
                }
            }

            await Task.Delay(TimeSpan.FromMilliseconds(50), cancellationToken);
        }
    }

    private async Task<(string Subject, Guid HostId)> SeedRoutedHostAsync(LoopbackTarget target, string address)
    {
        fixture.RelayNetwork.Route(new IPEndPoint(IPAddress.Parse(address), 22), target.Endpoint);

        return await SeedRelayHostAsync(address, 22);
    }

    private Task<(string Subject, Guid HostId)> SeedRelayHostAsync(
        string hostname,
        int port,
        string? subject = null) =>
        SeedHostAsync((hostname, port), subject);

    /// <summary>A user, their personal vault, and one host in it — relayed when an address is given.</summary>
    private async Task<(string Subject, Guid HostId)> SeedHostAsync(
        (string Hostname, int Port)? relay,
        string? subject = null)
    {
        await using var scope = fixture.CreateScope();
        var database = scope.ServiceProvider.GetRequiredService<DodoDbContext>();

        var user = await FindOrAddUserAsync(database, subject);

        var vault = await database.Vaults.FirstOrDefaultAsync(
                v => v.OwnerUserId == user.Id,
                TestContext.Current.CancellationToken)
            ?? database.Vaults.Add(new Vault
            {
                Id = Guid.CreateVersion7(),
                Name = "Personal",
                OwnerKind = VaultOwnerKind.Personal,
                OwnerUserId = user.Id,
                KeyGeneration = 1,
                CreatedAtUtc = Now,
                UpdatedAtUtc = Now,
            }).Entity;

        var host = new SshHost
        {
            Id = Guid.CreateVersion7(),
            VaultId = vault.Id,
            Payload = [1, 2, 3],
            KeyGeneration = 1,
            PayloadAadVersion = 1,
            RelayEnabled = relay is not null,
            Hostname = relay?.Hostname,
            Port = relay?.Port,
            Version = 1,
            CreatedAtUtc = Now,
            UpdatedAtUtc = Now,
            CreatedByUserId = user.Id,
            UpdatedByUserId = user.Id,
        };

        database.Hosts.Add(host);
        await database.SaveChangesAsync(TestContext.Current.CancellationToken);

        return (user.Subject, host.Id);
    }

    private async Task<UserAccount> FindOrAddUserAsync(DodoDbContext database, string? subject)
    {
        if (subject is not null
            && await database.Users.SingleOrDefaultAsync(
                u => u.Subject == subject,
                TestContext.Current.CancellationToken) is { } existing)
        {
            return existing;
        }

        var user = new UserAccount
        {
            Id = Guid.CreateVersion7(),
            Issuer = fixture.IdentityProvider.Authority,
            Subject = NewSubject(),
            Status = UserStatus.Active,
            CreatedAtUtc = Now,
            UpdatedAtUtc = Now,
        };

        database.Users.Add(user);
        database.UserKeys.Add(Seed.CurrentKey(user.Id, Now));

        return user;
    }

    /// <summary>A TCP server on loopback, standing in for the machine the relay reaches.</summary>
    private sealed class LoopbackTarget : IDisposable
    {
        private readonly TcpListener listener = new(IPAddress.Loopback, 0);
        private readonly CancellationTokenSource stopping = new();

        private LoopbackTarget(Func<NetworkStream, CancellationToken, Task> serve)
        {
            listener.Start();
            _ = AcceptAsync(serve);
        }

        internal IPEndPoint Endpoint => (IPEndPoint)listener.LocalEndpoint;

        /// <summary>Writes back everything it reads, on every connection.</summary>
        internal static LoopbackTarget Echo() => new(static async (stream, token) =>
        {
            var buffer = new byte[16 * 1024];
            int read;

            while ((read = await stream.ReadAsync(buffer, token)) > 0)
            {
                await stream.WriteAsync(buffer.AsMemory(0, read), token);
            }
        });

        /// <summary>Writes one line and closes, as a server does that refuses the client.</summary>
        internal static LoopbackTarget SayThenHangUp(string text) => new(async (stream, token) =>
            await stream.WriteAsync(Encoding.UTF8.GetBytes(text), token));

        public void Dispose()
        {
            stopping.Cancel();
            listener.Stop();
            stopping.Dispose();
        }

        private async Task AcceptAsync(Func<NetworkStream, CancellationToken, Task> serve)
        {
            try
            {
                while (!stopping.IsCancellationRequested)
                {
                    var client = await listener.AcceptTcpClientAsync(stopping.Token);

                    _ = Task.Run(
                        async () =>
                        {
                            using (client)
                            {
                                try
                                {
                                    await serve(client.GetStream(), stopping.Token);
                                }
                                catch (Exception exception) when (exception is IOException or OperationCanceledException or ObjectDisposedException)
                                {
                                    // The relay hung up, or the test is over.
                                }
                            }
                        },
                        CancellationToken.None);
                }
            }
            catch (Exception exception) when (exception is SocketException or OperationCanceledException or ObjectDisposedException)
            {
                // Stopped.
            }
        }
    }
}