DodoSSH/tests/DodoSSH.Client.Terminal.Tests/TerminalDataPlaneTests.cs

using System.Globalization;
using System.Net;
using System.Net.WebSockets;
using System.Text;
using DodoSSH.Client.Ssh;

namespace DodoSSH.Client.Terminal.Tests;

/// <summary>
/// The loopback transport end to end, with a real WebSocket client.
/// </summary>
/// <remarks>
/// Uses <see cref="ClientWebSocket"/> against the real listener rather than a stubbed transport,
/// because the hand-rolled HTTP upgrade is the part that could be subtly wrong — and a handshake that
/// no real client accepts would pass any test that skipped it.
/// </remarks>
public sealed class TerminalDataPlaneTests : IAsyncDisposable
{
    private const uint SessionId = 3;

    private static readonly byte[] PageTemplate = Encoding.UTF8.GetBytes(
        $"<html><body data-token=\"{TerminalDataPlane.TokenPlaceholder}\" "
        + $"data-socket=\"{TerminalDataPlane.SocketUrlPlaceholder}\"></body></html>");

    private readonly TerminalDataPlane plane = new(new InMemoryTerminalAssetProvider(
        new Dictionary<string, TerminalAsset>(StringComparer.Ordinal)
        {
            [TerminalDataPlane.PagePath] = new("text/html; charset=utf-8", PageTemplate),
            ["/xterm.js"] = new("text/javascript; charset=utf-8", "console.log(1)"u8.ToArray()),
        }));

    private readonly CancellationTokenSource lifetime = new();
    private Task? server;

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        await lifetime.CancelAsync();
        await plane.DisposeAsync();

        if (server is not null)
        {
            try
            {
                await server;
            }
            catch (OperationCanceledException)
            {
                // Expected: the accept loop is stopped by cancelling it.
            }
        }

        lifetime.Dispose();
    }

    // ---- Serving the page ----

    [Fact]
    public async Task ThePage_IsServedWithTheTokenAndSocketUrlSubstituted()
    {
        // Substituted at serve time rather than written into the file, so the token never touches disk
        // and never appears in a URL that a log or a browser history could keep.
        Start();

        using var client = new HttpClient();
        var body = await client.GetStringAsync(plane.PageUrl, TestContext.Current.CancellationToken);

        body.ShouldContain(plane.Token);
        body.ShouldContain(
            string.Create(CultureInfo.InvariantCulture, $"ws://127.0.0.1:{plane.Port}/socket"));

        body.ShouldNotContain(TerminalDataPlane.TokenPlaceholder);
        body.ShouldNotContain(TerminalDataPlane.SocketUrlPlaceholder);
    }

    [Fact]
    public async Task ThePage_IsNotCached()
    {
        // It carries a connection token.
        Start();

        using var client = new HttpClient();
        using var response = await client.GetAsync(plane.PageUrl, TestContext.Current.CancellationToken);

        response.Headers.CacheControl!.NoStore.ShouldBeTrue();
    }

    [Fact]
    public async Task AnUnknownPath_Is404()
    {
        Start();

        using var client = new HttpClient();
        using var response = await client.GetAsync(
            new Uri($"http://127.0.0.1:{plane.Port}/nope"),
            TestContext.Current.CancellationToken);

        response.StatusCode.ShouldBe(HttpStatusCode.NotFound);
    }

    // ---- Attaching ----

    [Fact]
    public async Task TheRenderer_Attaches()
    {
        Start();

        using var socket = await ConnectAsync();

        socket.State.ShouldBe(WebSocketState.Open);
        socket.SubProtocol.ShouldBe(TerminalDataPlane.SubProtocol);

        await plane.RendererAttached;
    }

    [Fact]
    public async Task AConnectionWithoutTheToken_IsRejected()
    {
        Start();

        await Should.ThrowAsync<WebSocketException>(async () =>
            await ConnectAsync(token: null));
    }

    [Fact]
    public async Task AConnectionWithTheWrongToken_IsRejected()
    {
        Start();

        await Should.ThrowAsync<WebSocketException>(async () =>
            await ConnectAsync(token: "not-the-token"));
    }

    [Fact]
    public async Task AConnectionFromAnotherOrigin_IsRejected()
    {
        // The threat this actually addresses: a page open in the user's browser can attempt WebSocket
        // connections to loopback ports, and would otherwise reach a terminal.
        Start();

        await Should.ThrowAsync<WebSocketException>(async () =>
            await ConnectAsync(origin: "https://evil.example"));
    }

    /// <remarks>
    /// The truth this replaced: a second valid attach used to be a 409, on the theory that one renderer
    /// lives for the whole process. Android's WebView does not honour that theory — its renderer process is
    /// routinely killed and the page reloads with a fresh socket — so a second valid attach is now a
    /// takeover. This asserts both halves: the newcomer gets the connection, and the displaced socket
    /// actually goes rather than lingering as a phantom nothing is reading from.
    /// </remarks>
    [Fact]
    public async Task ASecondRenderer_TakesOver_AndTheFirstSocketIsDropped()
    {
        Start();

        using var first = await ConnectAsync();
        first.State.ShouldBe(WebSocketState.Open);

        using var second = await ConnectAsync();
        second.State.ShouldBe(WebSocketState.Open);

        // The first socket was aborted rather than closed gracefully — Abort skips the close handshake
        // entirely, so there is no Close frame for this side to see coming. What a receive on it sees
        // instead is the connection simply gone, which the client surfaces as an exception rather than as
        // a state that quietly flips on its own; nothing here reads from the socket otherwise, so the
        // state alone would not move.
        var firstBuffer = new byte[16];
        await Should.ThrowAsync<Exception>(async () =>
            await first.ReceiveAsync(firstBuffer.AsMemory(), TestContext.Current.CancellationToken));

        await using var session = new FakeShellSession(bytesToProduce: 64);
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        var run = pump.RunAsync(TestContext.Current.CancellationToken);

        var opened = await ReceiveAsync(second);
        opened.Opcode.ShouldBe((byte)TerminalServerOpcode.SessionOpened);

        var output = await ReceiveAsync(second);
        output.Opcode.ShouldBe((byte)TerminalServerOpcode.Output);
        output.Payload.Length.ShouldBe(64);

        await SendAsync(
            second,
            (byte)TerminalClientOpcode.Acknowledge,
            TerminalFrame.CreateAcknowledgementPayload((uint)output.Payload.Length));

        await run;
    }

    /// <remarks>
    /// <para>
    /// The other half of the takeover: a renderer process that dies without a close handshake — which is
    /// what a killed Android WebView actually does, no FIN, nothing — must not fault the send path. A
    /// faulted send would propagate into <see cref="TerminalSessionPump"/>'s flush loop and freeze a live
    /// session; see <see cref="TerminalDataPlane.SendAsync"/>'s remark for why. Disposing the client socket
    /// abruptly, with no close handshake sent, is the closest this harness gets to that: the server-side
    /// socket is left believing itself open until it actually tries to write to it.
    /// </para>
    /// <para>
    /// Driven straight through <see cref="TerminalDataPlane.SendAsync"/> rather than through a pump, because
    /// a pump adds nothing here — the point is entirely about the transport's own contract, and a session
    /// layered on top would only leave it unclear whether a passing test proved the transport never threw or
    /// merely that the frames never happened to need a live socket.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task SendAsync_DoesNotThrow_WhenTheAttachedRendererDiedWithoutClosing_AndAFreshAttachStillReceives()
    {
        Start();

        var first = await ConnectAsync();
        first.State.ShouldBe(WebSocketState.Open);
        first.Dispose();

        // Whether this particular send lands on the OS's send buffer before the peer's absence is noticed,
        // or fails immediately, is not the point — either way it must not throw.
        await Should.NotThrowAsync(async () =>
            await plane.SendAsync(
                TerminalFrame.Create((byte)TerminalServerOpcode.Output, SessionId, "before"u8.ToArray()),
                TestContext.Current.CancellationToken));

        using var second = await ConnectAsync();

        await Should.NotThrowAsync(async () =>
            await plane.SendAsync(
                TerminalFrame.Create((byte)TerminalServerOpcode.Output, SessionId, "after"u8.ToArray()),
                TestContext.Current.CancellationToken));

        var output = await ReceiveAsync(second);
        output.Opcode.ShouldBe((byte)TerminalServerOpcode.Output);
        Encoding.UTF8.GetString(output.Payload).ShouldBe("after");
    }

    // ---- Frames ----

    [Fact]
    public async Task Output_ReachesTheRenderer()
    {
        Start();

        using var socket = await ConnectAsync();

        await using var session = new FakeShellSession(bytesToProduce: 128);
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        var run = pump.RunAsync(TestContext.Current.CancellationToken);

        // SessionOpened first, then the output.
        var opened = await ReceiveAsync(socket);
        opened.Opcode.ShouldBe((byte)TerminalServerOpcode.SessionOpened);
        opened.SessionId.ShouldBe(SessionId);

        var output = await ReceiveAsync(socket);
        output.Opcode.ShouldBe((byte)TerminalServerOpcode.Output);
        output.Payload.Length.ShouldBe(128);

        // Acknowledge it, exactly as the page does from term.write's callback.
        await SendAsync(
            socket,
            (byte)TerminalClientOpcode.Acknowledge,
            TerminalFrame.CreateAcknowledgementPayload((uint)output.Payload.Length));

        await run;
    }

    [Fact]
    public async Task Input_ReachesTheRemote()
    {
        Start();

        using var socket = await ConnectAsync();

        await using var session = new FakeShellSession();
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        await SendAsync(socket, (byte)TerminalClientOpcode.Input, "whoami\r"u8.ToArray());

        await WaitUntilAsync(() => session.Written.Length > 0);

        Encoding.UTF8.GetString(session.Written).ShouldBe("whoami\r");
    }

    [Fact]
    public async Task AResize_ReachesTheRemote()
    {
        Start();

        using var socket = await ConnectAsync();

        await using var session = new FakeShellSession();
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        await SendAsync(
            socket,
            (byte)TerminalClientOpcode.Resize,
            TerminalFrame.CreateResizePayload(new TerminalSize(132, 43, 1320, 1075)));

        await WaitUntilAsync(() => session.ResizeCount > 0);

        session.LastResize.ShouldBe(new TerminalSize(132, 43, 1320, 1075));
    }

    [Fact]
    public async Task AnAcknowledgement_ReturnsCredit()
    {
        Start();

        using var socket = await ConnectAsync();

        await using var session = new FakeShellSession();
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        pump.Credits.TryReserve(1024);
        pump.Credits.Outstanding.ShouldBe(1024);

        await SendAsync(
            socket,
            (byte)TerminalClientOpcode.Acknowledge,
            TerminalFrame.CreateAcknowledgementPayload(1024));

        await WaitUntilAsync(() => pump.Credits.Outstanding == 0);
    }

    /// <remarks>
    /// The point of this one is the <em>absence</em> of a registered session. Every other client frame is
    /// about a terminal and is dropped once its session has gone; a text-size chord is about the person
    /// reading, and the page sends whichever id it had to hand — often one whose shell has just ended,
    /// which is exactly when somebody may be squinting at the message it left behind. Dispatching this
    /// opcode before the session lookup is what makes that work.
    /// </remarks>
    [Fact]
    public async Task AFontSizeStep_IsHeardWithNoSessionRegistered()
    {
        Start();

        using var socket = await ConnectAsync();

        var steps = new List<int>();
        plane.FontSizeStepRequested += (_, e) => steps.Add(e.Step);

        await SendAsync(
            socket,
            (byte)TerminalClientOpcode.FontSizeStep,
            TerminalFrame.CreateFontSizeStepPayload(-1),
            sessionId: 0);

        await WaitUntilAsync(() => steps.Count > 0);

        steps[0].ShouldBe(-1);
    }

    /// <remarks>
    /// Dispatched before the session lookup for the reason the text-size chord is, and one of its own: a
    /// clipboard is not a property of any terminal. The absence of a registered session is the point of the
    /// test — copying the last thing a remote said before it hung up is an ordinary thing to want, and a
    /// dead tab must not be the one place the gesture silently does nothing.
    /// </remarks>
    [Fact]
    public async Task ACopy_IsHeardWithNoSessionRegistered()
    {
        Start();

        using var socket = await ConnectAsync();

        var copied = new List<string>();
        plane.ClipboardCopyRequested += (_, e) => copied.Add(e.Text);

        await SendAsync(
            socket,
            (byte)TerminalClientOpcode.ClipboardCopy,
            Encoding.UTF8.GetBytes("error: nothing was written"));

        await WaitUntilAsync(() => copied.Count > 0);

        copied[0].ShouldBe("error: nothing was written");
    }

    /// <remarks>
    /// An empty copy is dropped rather than forwarded. The clipboard is shared with everything else on the
    /// machine, so "copy nothing" would be a gesture that quietly destroys whatever was on it — and the page
    /// does not send one, which makes this a guard against a frame from something that is not the page.
    /// </remarks>
    [Fact]
    public async Task AnEmptyCopy_IsDropped()
    {
        Start();

        using var socket = await ConnectAsync();

        var copies = 0;
        plane.ClipboardCopyRequested += (_, _) => Interlocked.Increment(ref copies);

        await SendAsync(socket, (byte)TerminalClientOpcode.ClipboardCopy, []);

        // Followed by one that must be heard, so the assertion below is about ordering rather than about
        // having waited long enough: the receive loop is sequential, so a copy arriving after this proves
        // the empty one was seen and passed over.
        var copied = new List<string>();
        plane.ClipboardCopyRequested += (_, e) => copied.Add(e.Text);

        await SendAsync(socket, (byte)TerminalClientOpcode.ClipboardCopy, "uptime"u8.ToArray());
        await WaitUntilAsync(() => copied.Count > 0);

        copies.ShouldBe(1, "only the second frame carried anything");
    }

    /// <remarks>
    /// The session id is the whole payload of a paste request — it says which pane the pointer was over, and
    /// therefore where the answer has to land. Heard with nothing registered for the reason a copy is.
    /// </remarks>
    [Fact]
    public async Task APasteRequest_CarriesTheSessionItWasMadeIn()
    {
        Start();

        using var socket = await ConnectAsync();

        var asked = new List<uint>();
        plane.ClipboardPasteRequested += (_, e) => asked.Add(e.SessionId);

        await SendAsync(socket, (byte)TerminalClientOpcode.ClipboardPaste, [], sessionId: 4242);

        await WaitUntilAsync(() => asked.Count > 0);

        asked[0].ShouldBe(4242u);
    }

    [Fact]
    public async Task AFrameForAnUnregisteredSession_IsIgnored()
    {
        // Ordinary during teardown: the page can still have frames in flight for a session that just
        // ended. Dropping them must not disturb anything else.
        Start();

        using var socket = await ConnectAsync();

        await using var session = new FakeShellSession();
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        // A session id that is never registered. Using the registered one and relying on ordering
        // would be a race: frames are dispatched on the receive loop, so the orphan can land after
        // registration and legitimately be delivered.
        await SendAsync(socket, (byte)TerminalClientOpcode.Input, "orphan"u8.ToArray(), sessionId: 999);

        await SendAsync(socket, (byte)TerminalClientOpcode.Input, "real"u8.ToArray());
        await WaitUntilAsync(() => session.Written.Length > 0);

        Encoding.UTF8.GetString(session.Written).ShouldBe("real");
    }

    [Fact]
    public async Task AMalformedFrame_IsIgnored()
    {
        Start();

        using var socket = await ConnectAsync();

        // Shorter than the header.
        await socket.SendAsync(
            new byte[] { 1, 2 },
            WebSocketMessageType.Binary,
            endOfMessage: true,
            TestContext.Current.CancellationToken);

        await using var session = new FakeShellSession();
        await using var pump = CreatePump(session);
        plane.Register(SessionId, pump);

        await SendAsync(socket, (byte)TerminalClientOpcode.Input, "still here"u8.ToArray());
        await WaitUntilAsync(() => session.Written.Length > 0);

        Encoding.UTF8.GetString(session.Written).ShouldBe("still here");
    }

    // ---- Helpers ----

    private void Start() => server = plane.RunAsync(lifetime.Token);

    private TerminalSessionPump CreatePump(FakeShellSession session) =>
        new(
            SessionId,
            session,
            plane,
            TimeProvider.System,
            new TerminalPumpOptions { FlushInterval = TimeSpan.FromMilliseconds(10) });

    /// <summary>
    /// Opens a renderer's socket and returns once the plane is actually holding it.
    /// </summary>
    /// <remarks>
    /// <para>
    /// <b>Connected and attached are two different moments, and the gap between them is where this used to
    /// flake.</b> <c>ClientWebSocket.ConnectAsync</c> completes on the 101, which
    /// <see cref="TerminalDataPlane.UpgradeAsync"/> writes before it has a <see cref="WebSocket"/> to
    /// attach — it builds one from the stream and swaps it in a few instructions later, on the accept
    /// thread. A frame sent in between is dropped, by design rather than by accident: the transport has
    /// nowhere to put a frame for a renderer that is not there, and queueing it is the unbounded growth the
    /// credit window exists to prevent.
    /// </para>
    /// <para>
    /// So a test that connected and immediately expected a frame was racing that window on every run. It
    /// lost one on CI — <c>Output_ReachesTheRenderer</c> read the output frame first and asked why it was
    /// not the session's opening one, the opening one having been dropped a moment earlier — which is a
    /// scheduling accident on a loaded machine and says nothing whatever about the transport.
    /// </para>
    /// <para>
    /// Production does not race it and needs no change: everything that opens a session waits on
    /// <c>TerminalWorkspace.WaitForRendererAsync</c> first, and that resolves from the same few lines this
    /// event is raised from.
    /// </para>
    /// <para>
    /// Subscribed before the connection rather than after it, because the event is raised on the accept
    /// thread and can be over before <c>ConnectAsync</c> has returned here. Bounded, so that a socket that
    /// never attaches fails this helper rather than hanging the suite in a later receive.
    /// </para>
    /// </remarks>
    private async Task<ClientWebSocket> ConnectAsync(
        string? token = "",
        string? origin = null)
    {
        var socket = new ClientWebSocket();

        socket.Options.AddSubProtocol(TerminalDataPlane.SubProtocol);

        if (token is not null)
        {
            socket.Options.AddSubProtocol($"token.{(token.Length == 0 ? plane.Token : token)}");
        }

        socket.Options.SetRequestHeader(
            "Origin",
            origin ?? string.Create(CultureInfo.InvariantCulture, $"http://127.0.0.1:{plane.Port}"));

        var attached = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);

        void OnAttached(object? sender, EventArgs e) => attached.TrySetResult();

        plane.SocketAttached += OnAttached;

        try
        {
            await socket.ConnectAsync(
                new Uri($"ws://127.0.0.1:{plane.Port}{TerminalDataPlane.SocketPath}"),
                TestContext.Current.CancellationToken);

            await attached.Task.WaitAsync(
                TimeSpan.FromSeconds(5),
                TestContext.Current.CancellationToken);
        }
        catch
        {
            socket.Dispose();
            throw;
        }
        finally
        {
            plane.SocketAttached -= OnAttached;
        }

        return socket;
    }

    private static async Task SendAsync(
        ClientWebSocket socket,
        byte opcode,
        byte[] payload,
        uint sessionId = SessionId) =>
        await socket.SendAsync(
            TerminalFrame.Create(opcode, sessionId, payload),
            WebSocketMessageType.Binary,
            endOfMessage: true,
            TestContext.Current.CancellationToken);

    private static async Task<(byte Opcode, uint SessionId, byte[] Payload)> ReceiveAsync(
        ClientWebSocket socket)
    {
        var buffer = new byte[64 * 1024];

        var result = await socket.ReceiveAsync(
            buffer.AsMemory(),
            TestContext.Current.CancellationToken);

        TerminalFrame.TryRead(buffer.AsSpan(0, result.Count), out var opcode, out var sessionId, out var payload)
            .ShouldBeTrue();

        return (opcode, sessionId, payload.ToArray());
    }

    /// <remarks>
    /// Polls rather than awaiting a signal: inbound frames are dispatched on the transport's receive
    /// loop, which has no completion for the caller to await. The bound keeps a genuine failure a
    /// failure rather than a hang.
    /// </remarks>
    private static async Task WaitUntilAsync(Func<bool> condition)
    {
        var deadline = TimeProvider.System.GetUtcNow() + TimeSpan.FromSeconds(5);

        while (TimeProvider.System.GetUtcNow() < deadline)
        {
            if (condition())
            {
                return;
            }

            await Task.Delay(10, TestContext.Current.CancellationToken);
        }

        throw new TimeoutException("The expected state was not reached within 5 seconds.");
    }
}