DodoSSH/tests/DodoSSH.Client.App.Tests/ShellFlowTests.Relay.cs

using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;
using System.Text;
using DodoSSH.Client.Api;
using DodoSSH.Client.Session;
using DodoSSH.Client.Shell.ViewModels;
using DodoSSH.Client.Ssh;
using DodoSSH.Contracts;

namespace DodoSSH.Client.App.Tests;

/// <summary>
/// A host set to connect through the server relay, from the Connect button to the request that reaches SSH.
/// </summary>
/// <remarks>
/// <para>
/// The relay here is the real <see cref="ServerRelay"/> — ticket, subprotocols and bridge — over an
/// in-process WebSocket pair, with only the API call and the far end of the socket stood in for. What these
/// hold is the promise the checkbox makes and used not to keep: ticking it changes how the host is reached,
/// and a relay that cannot be used is a stated failure rather than a quiet direct dial.
/// </para>
/// <para>
/// The bridge's own behaviour against a real sshd is <c>RelayBridgeTests</c>, and the server's is
/// <c>RelayEndpointTests</c>.
/// </para>
/// </remarks>
public sealed partial class ShellFlowTests
{
    [Fact]
    public async Task ARelayedHost_IsDialledThroughTheBridge_AndStillAsItself()
    {
        var tickets = new FakeRelayApi();
        server.Relay = new ServerRelay(tickets, ConnectPairedAsync);

        var vault = await ReadyToConnectRelayedAsync();

        await using var renderer = await FakeRenderer.AttachAsync(workspace, Token);
        await vault.ConnectCommand.ExecuteAsync(null);

        tickets.HostIds.ShouldHaveSingleItem().ShouldBe(vault.SelectedHost.ShouldNotBeNull().EntityId);

        // Through a loopback proxy — and still named as the host, because the name is what the host key is
        // pinned under. The bridge carries the name in its CONNECT; it does not replace it.
        var request = ssh.Requests.ShouldHaveSingleItem();
        request.Proxy.ShouldNotBeNull().Lifetime.ShouldBeOfType<RelayBridge>();
        request.Host.ShouldBe("db.internal");
        request.Port.ShouldBe(22);

        // The ticket travels as the second subprotocol and never as anything else.
        offeredProtocols.ShouldBe([RelayProtocol.SubProtocol, RelayProtocol.TicketProtocolPrefix + FakeRelayApi.Ticket]);

        shell.IsTerminalShowing.ShouldBeTrue();

        await request.Proxy.Lifetime.DisposeAsync();
    }

    [Fact]
    public async Task ARelayedHost_OnAServerWithoutARelay_IsRefusedWithoutDialling()
    {
        // The fake server's default: a server that does not advertise the feature.
        var vault = await ReadyToConnectRelayedAsync();

        await using var renderer = await FakeRenderer.AttachAsync(workspace, Token);
        await vault.ConnectCommand.ExecuteAsync(null);

        // No fallback to the direct address. The user ticked the relay because that path does not work, or
        // goes somewhere they do not mean.
        ssh.Requests.ShouldBeEmpty();
        vault.Status.ShouldNotBeNull().ShouldContain("does not offer one");
    }

    [Theory]
    [InlineData(ProblemCodes.RelayTargetRejected, "refused to relay")]
    [InlineData(ProblemCodes.RelayLimitReached, "as many relay sessions")]
    [InlineData(ProblemCodes.RelayUnavailable, "does not offer one")]
    public async Task ARefusedTicket_IsExplained_AndNothingIsDialled(string code, string explanation)
    {
        server.Relay = new ServerRelay(
            new FakeRelayApi { Refusal = new DodoSshApiException(HttpStatusCode.Forbidden, code, "no") },
            ConnectPairedAsync);

        var vault = await ReadyToConnectRelayedAsync();

        await using var renderer = await FakeRenderer.AttachAsync(workspace, Token);
        await vault.ConnectCommand.ExecuteAsync(null);

        ssh.Requests.ShouldBeEmpty();
        vault.Status.ShouldNotBeNull().ShouldContain(explanation);
    }

    [Fact]
    public async Task ARelayThatCouldNotReachTheHost_SaysSo_InsteadOfTheTransportError()
    {
        // What the server does when it cannot dial the host: accept the socket, then close it with 4502.
        server.Relay = new ServerRelay(new FakeRelayApi(), async (url, protocols, cancellationToken) =>
        {
            var socket = await ConnectPairedAsync(url, protocols, cancellationToken);
            await relayEnd!.CloseOutputAsync(
                (WebSocketCloseStatus)RelayProtocol.TargetUnreachableCloseCode,
                "unreachable",
                cancellationToken);

            return socket;
        });

        // SSH.NET arriving at the bridge and then losing its stream, which is how a relay that hung up looks
        // from inside the connect.
        ssh.BeforeConnect = request => ThroughSocksToEndOfStreamAsync(request.Proxy.ShouldNotBeNull().Port);
        ssh.Failure = new InvalidOperationException("An established connection was aborted by the server.");

        var vault = await ReadyToConnectRelayedAsync();

        await using var renderer = await FakeRenderer.AttachAsync(workspace, Token);
        await vault.ConnectCommand.ExecuteAsync(null);

        vault.Status.ShouldNotBeNull().ShouldContain("could not reach this host");
        vault.Status.ShouldNotContain("aborted");
    }

    [Fact]
    public async Task ARelayedHostsFileTransfer_GoesThroughTheRelayToo()
    {
        // Its own connection, so it needs its own relay — or the transfers screen would fail for exactly the
        // hosts the relay exists for.
        server.Relay = new ServerRelay(new FakeRelayApi(), ConnectPairedAsync);

        var vault = await ReadyToConnectRelayedAsync();

        shell.Transfers.Attach(vault, knownHosts);
        shell.Transfers.SelectedHost = shell.Transfers.Hosts[0];

        await shell.Transfers.ConnectCommand.ExecuteAsync(null);

        var request = ssh.SftpRequests.ShouldHaveSingleItem();
        request.Proxy.ShouldNotBeNull().Lifetime.ShouldBeOfType<RelayBridge>();

        await request.Proxy.Lifetime.DisposeAsync();
    }

    /// <summary>The protocols the last relay socket was opened offering.</summary>
    private IReadOnlyList<string>? offeredProtocols;

    /// <summary>The far end of the last relay socket, which a test may close as the server would.</summary>
    private WebSocket? relayEnd;

    private async Task<VaultViewModel> ReadyToConnectRelayedAsync()
    {
        await UnlockedAsync();

        var vault = shell.Vault!;

        vault.NewHostCommand.Execute(null);
        vault.EditorLabel = "behind-the-server";
        vault.EditorHostname = "db.internal";
        vault.EditorUsername = "deploy";

        // A relayed host needs a port of its own; the server stores it and dials it.
        vault.EditorPort = 22;
        vault.EditorRelayEnabled = true;

        await vault.SaveHostCommand.ExecuteAsync(null);

        vault.SelectedHost = vault.Hosts.Single(row => row.Host.RelayEnabled);

        return vault;
    }

    /// <summary>Opens both ends of a WebSocket over loopback, and keeps the far one.</summary>
    private async Task<WebSocket> ConnectPairedAsync(
        Uri url,
        IReadOnlyList<string> protocols,
        CancellationToken cancellationToken)
    {
        url.ShouldBe(FakeRelayApi.SocketUrl);
        offeredProtocols = protocols;

        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();

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

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

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

    /// <summary>
    /// Speaks SOCKS5 to the bridge as SSH.NET would, then reads until the bridge closes the stream — which
    /// it does once the relay has closed, so the relay's reason is recorded by the time this returns.
    /// </summary>
    private static async Task ThroughSocksToEndOfStreamAsync(int port)
    {
        using var client = new TcpClient();
        await client.ConnectAsync(IPAddress.Loopback, port, Token);

        var stream = client.GetStream();
        var name = Encoding.ASCII.GetBytes("db.internal");

        await stream.WriteAsync(new byte[] { 5, 1, 0 }, Token);
        await stream.ReadExactlyAsync(new byte[2], Token);

        await stream.WriteAsync((byte[])[5, 1, 0, 3, (byte)name.Length, .. name, 0, 22], Token);
        await stream.ReadExactlyAsync(new byte[10], Token);

        var rest = new byte[256];
        while (await stream.ReadAsync(rest, Token) > 0)
        {
        }
    }

    /// <summary>The ticket half of the API, answering or refusing as a test arranges.</summary>
    private sealed class FakeRelayApi : IRelayApi
    {
        internal const string Ticket = "dGlja2V0";

        internal static readonly Uri SocketUrl = new("ws://relay.test" + RelayProtocol.ConnectPath);

        internal List<Guid> HostIds { get; } = [];

        internal DodoSshApiException? Refusal { get; init; }

        public Task<RelayTicketResponse> IssueRelayTicketAsync(Guid hostId, CancellationToken cancellationToken)
        {
            HostIds.Add(hostId);

            return Refusal is { } refusal
                ? Task.FromException<RelayTicketResponse>(refusal)
                : Task.FromResult(new RelayTicketResponse(
                    Ticket,
                    SocketUrl,
                    RelayProtocol.SubProtocol,
                    Guid.CreateVersion7(),
                    TimeProvider.System.GetUtcNow().AddSeconds(30)));
        }
    }
}