DodoSSH/tests/DodoSSH.Client.Ssh.Tests/RelayBridgeTests.cs

using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;

namespace DodoSSH.Client.Ssh.Tests;

/// <summary>
/// The relay bridge: a real SSH connection, through a real SOCKS5 handshake, over a real WebSocket.
/// </summary>
/// <remarks>
/// <para>
/// The relay server is stood in for by the other end of an in-process WebSocket pair that pipes to the
/// fixture's sshd — which is the whole of what the server's socket does once a ticket is spent, so the
/// bytes on the wire here are the bytes a relayed session carries. What is not here is the ticket, which
/// is the server's half and is tested there.
/// </para>
/// <para>
/// Three of these are the security properties the bridge has to keep, from
/// <c>docs/reaching-a-host-you-cannot-dial.md</c>: it binds loopback only, it accepts once, and a
/// connection that goes away takes the bridge — and so the relay session — with it. The failure each
/// guards against is silent: an open proxy, or a relay session left running behind a closed tab.
/// </para>
/// </remarks>
[Collection(SshCollection.Name)]
public sealed class RelayBridgeTests(SshServerFixture fixture)
{
    private static readonly TimeSpan Settles = TimeSpan.FromSeconds(20);

    private static CancellationToken Token => TestContext.Current.CancellationToken;

    [Fact]
    public async Task TheBridge_ListensOnLoopbackOnly()
    {
        var (client, _) = await PairAsync();
        await using var bridge = RelayBridge.Start(client);

        // Asserted on the address rather than trusted to the constructor: a bridge on every interface
        // would be a proxy into the relay for anyone on the network, and nothing would report it.
        bridge.Endpoint.Address.ShouldBe(IPAddress.Loopback);
        bridge.Proxy.Port.ShouldBe(bridge.Endpoint.Port);
        bridge.Proxy.Lifetime.ShouldBeSameAs(bridge);
    }

    [Fact]
    public async Task AShell_ThroughTheBridge_ReachesTheHost_AndGoesWithTheConnection()
    {
        var knownHosts = await TrustedStoreAsync();
        var (client, server) = await PairAsync();
        var relaying = RelayToSshdAsync(server);

        var bridge = RelayBridge.Start(client);

        // The host as the request names it, so the pin is the host's — the bridge is a SOCKS proxy, and
        // the CONNECT carries this name through it rather than 127.0.0.1.
        var connection = await new SshNetConnectionFactory(knownHosts)
            .ConnectAsync(Request(bridge.Proxy), progress: null, Token);

        connection.IsConnected.ShouldBeTrue();
        connection.HostKey.Host.ShouldBe(fixture.Host);

        // The connection owns the bridge. Closing the one is what closes the other, and the relay hears it
        // as a close rather than as a socket that stopped answering.
        await connection.DisposeAsync();

        await bridge.Completion.WaitAsync(Settles, Token);
        (await relaying.WaitAsync(Settles, Token)).ShouldBe(WebSocketMessageType.Close);
    }

    [Fact]
    public async Task TheBridge_AcceptsOnce()
    {
        var (client, _) = await PairAsync();
        await using var bridge = RelayBridge.Start(client);

        using var first = new TcpClient();
        await first.ConnectAsync(bridge.Endpoint, Token);

        // A greeting answered means the first connection has been accepted, which is when the listener
        // stops — so what the second attempt meets is a closed port, not a race with the backlog.
        var stream = first.GetStream();
        await stream.WriteAsync(new byte[] { 5, 1, 0 }, Token);
        var reply = new byte[2];
        await stream.ReadExactlyAsync(reply, Token);
        reply.ShouldBe(new byte[] { 5, 0 });

        using var second = new TcpClient();
        await Should.ThrowAsync<SocketException>(second.ConnectAsync(bridge.Endpoint, Token).AsTask());
    }

    [Fact]
    public async Task ARelayThatClosesWithAReason_FailsTheConnect_AndKeepsTheReason()
    {
        // What the server does when it cannot reach the host: accept the socket, then close it with 4502.
        var (client, server) = await PairAsync();
        await server.CloseOutputAsync((WebSocketCloseStatus)4502, "unreachable", Token);

        var bridge = RelayBridge.Start(client);

        var failure = await Should.ThrowAsync<Exception>(async () =>
            await new SshNetConnectionFactory(new InMemoryKnownHostStore())
                .ConnectAsync(Request(bridge.Proxy), progress: null, Token));

        // A connection error, not a question about a host key nobody presented.
        failure.ShouldNotBeOfType<SshHostKeyUnknownException>();

        bridge.CloseStatus.ShouldBe((WebSocketCloseStatus)4502);

        // And the failed connect disposed the bridge it was handed, rather than leaving it listening.
        await bridge.Completion.WaitAsync(Settles, Token);
    }

    [Fact]
    public async Task NobodyArriving_ClosesTheRelay()
    {
        var (client, server) = await PairAsync();
        await using var bridge = RelayBridge.Start(client, TimeSpan.FromMilliseconds(200));

        await bridge.Completion.WaitAsync(Settles, Token);

        var received = await server.ReceiveAsync(new byte[16], Token);
        received.MessageType.ShouldBe(WebSocketMessageType.Close);
    }

    /// <summary>Two ends of one WebSocket, over a loopback TCP connection this test owns.</summary>
    private static async Task<(WebSocket Client, WebSocket Server)> PairAsync()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();

        var dialling = new TcpClient();
        var connecting = dialling.ConnectAsync((IPEndPoint)listener.LocalEndpoint, Token);
        var accepted = await listener.AcceptTcpClientAsync(Token);
        await connecting;

        var client = WebSocket.CreateFromStream(
            dialling.GetStream(),
            new WebSocketCreationOptions { IsServer = false, KeepAliveInterval = TimeSpan.Zero });

        var server = WebSocket.CreateFromStream(
            accepted.GetStream(),
            new WebSocketCreationOptions { IsServer = true, KeepAliveInterval = TimeSpan.Zero });

        return (client, server);
    }

    /// <summary>
    /// The relay's half: bytes from the socket to sshd and back, until one side ends.
    /// </summary>
    /// <returns>The type of the last message the relay received — a close, when the client said goodbye.</returns>
    private Task<WebSocketMessageType> RelayToSshdAsync(WebSocket server) => Task.Run(async () =>
    {
        using var target = new TcpClient();
        await target.ConnectAsync(fixture.Host, fixture.Port, Token);
        var stream = target.GetStream();

        var down = Task.Run(async () =>
        {
            var buffer = new byte[32 * 1024];
            int read;

            try
            {
                while ((read = await stream.ReadAsync(buffer, Token)) > 0)
                {
                    await server.SendAsync(buffer.AsMemory(0, read), WebSocketMessageType.Binary, true, Token);
                }
            }
            catch (Exception exception) when (exception is IOException or WebSocketException or ObjectDisposedException)
            {
                // The test's connection is over.
            }
        });

        var inbound = new byte[32 * 1024];

        while (true)
        {
            var received = await server.ReceiveAsync(inbound.AsMemory(), Token);

            if (received.MessageType == WebSocketMessageType.Close)
            {
                await server.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, Token);
                target.Close();
                await down;

                return received.MessageType;
            }

            await stream.WriteAsync(inbound.AsMemory(0, received.Count), Token);
        }
    });

    private SshConnectionRequest Request(SshLoopbackProxy proxy) =>
        new(
            fixture.Host,
            fixture.Port,
            SshServerFixture.Username,
            new SshPasswordCredential(SshServerFixture.Password),
            ConnectTimeout: null,
            proxy);

    /// <summary>A store that already trusts the container's host key, learned directly.</summary>
    private async Task<InMemoryKnownHostStore> TrustedStoreAsync()
    {
        var knownHosts = new InMemoryKnownHostStore();

        var unknown = await Should.ThrowAsync<SshHostKeyUnknownException>(async () =>
            await new SshNetConnectionFactory(knownHosts).ConnectAsync(
                new SshConnectionRequest(
                    fixture.Host,
                    fixture.Port,
                    SshServerFixture.Username,
                    new SshPasswordCredential(SshServerFixture.Password)),
                progress: null,
                Token));

        await knownHosts.TrustAsync(unknown.Presentation, Token);

        return knownHosts;
    }
}