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

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

namespace DodoSSH.Client.Terminal.Tests;

/// <summary>
/// What the workspace waits for outside the process, and how long a session lives once it has one.
/// </summary>
/// <remarks>
/// <para>
/// The renderer gate itself is not in question. <see cref="TerminalDataPlane.SendAsync"/> drops frames when
/// nothing is attached, so a session opened before the renderer arrives loses its <c>SessionOpened</c> frame
/// and streams output at a terminal that was never created — and that it opens when a renderer does attach
/// is covered by <see cref="TerminalDataPlaneTests.TheRenderer_Attaches"/>. What is worth a test here is the
/// half that used to be missing: waiting for a renderer that never arrives has to end.
/// </para>
/// <para>
/// The session count is here for a different reason. It is shown to a user on the unlock screen, as the
/// disclosure that locking the vault leaves shells running, and a wrong number there is not a cosmetic bug:
/// it either hides a live connection or invents one.
/// </para>
/// </remarks>
public sealed class TerminalWorkspaceTests
{
    /// <summary>How long <see cref="WaitUntilAsync"/> polls before calling it a failure.</summary>
    private static readonly TimeSpan PollTimeout = TimeSpan.FromSeconds(10);

    [Fact]
    public async Task WaitingForARendererThatNeverAttaches_GivesUp()
    {
        // The shipped failure this stands in for is a WebView2 that never initialises — no Evergreen
        // runtime, an install blocked by policy, an AppContainer that cannot reach loopback. From this
        // side they are identical and all look like the listener being up with nothing ever connecting
        // to it. Before the wait was bounded this test would have hung instead of failing.
        await using var workspace = CreateWorkspace(
            new FakeConnectionFactory(), rendererTimeout: TimeSpan.FromMilliseconds(250));

        workspace.Start();

        await Should.ThrowAsync<TimeoutException>(async () =>
            await workspace.WaitForRendererAsync(TestContext.Current.CancellationToken));
    }

    [Fact]
    public async Task WaitingForARenderer_ObeysItsCancellationToken()
    {
        // The timeout is the backstop; the caller's token is what makes a Connect the user gave up on
        // return at once rather than sitting out the rest of the wait. The timeout here is long enough
        // that only cancellation can end this.
        await using var workspace = CreateWorkspace(
            new FakeConnectionFactory(), rendererTimeout: TimeSpan.FromMinutes(5));

        workspace.Start();

        using var cancellation = new CancellationTokenSource();
        var wait = workspace.WaitForRendererAsync(cancellation.Token);

        await cancellation.CancelAsync();

        await Should.ThrowAsync<OperationCanceledException>(async () => await wait);
    }

    [Fact]
    public async Task AnOpenSessionIsReportedAsLive()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        workspace.LiveSessionCount.ShouldBe(0);

        await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        workspace.LiveSessionCount.ShouldBe(1);

        await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        workspace.LiveSessionCount.ShouldBe(2);
    }

    /// <remarks>
    /// The case a naive <c>sessions.Count</c> gets wrong. Nothing removes the entry when the remote
    /// closes the channel by itself — the session is still in the dictionary, holding a connection that
    /// is finished — so counting entries would report a shell that exited as still running. A user
    /// deciding whether it is safe to walk away is the person that lie is told to.
    /// </remarks>
    [Fact]
    public async Task ASessionWhoseRemoteHasExitedIsNotReportedAsLive()
    {
        // A shell with no output to give: its first read returns 0, which is a remote closing the
        // channel, so the pump finishes on its own with nobody asking it to.
        var connections = new FakeConnectionFactory(bytesPerShell: 0);

        await using var workspace = CreateWorkspace(connections);

        await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        await WaitUntilAsync(() => workspace.LiveSessionCount == 0);
    }

    [Fact]
    public async Task ClosingASessionEndsItAndDisposesItsConnection()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        await workspace.CloseSessionAsync(sessionId);

        workspace.LiveSessionCount.ShouldBe(0);
        connections.Connections.ShouldHaveSingleItem().IsDisposed.ShouldBeTrue();
    }

    /// <remarks>
    /// <para>
    /// Per-session liveness, which the tab strip's status dot is. The aggregate count answers "may I walk
    /// away"; this answers "is <em>this</em> tab still connected", and a tab that went on claiming a shell
    /// it no longer has would be the same lie one row down.
    /// </para>
    /// <para>
    /// The unknown-id case is asserted because it is what a stale tab asks. A session id this workspace
    /// never issued, or has already closed, is not live — it must not throw and must not be optimistic.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task LivenessIsAnsweredPerSession()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        var first = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);
        var second = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        workspace.IsSessionLive(first).ShouldBeTrue();
        workspace.IsSessionLive(second).ShouldBeTrue();
        workspace.IsSessionLive(9999).ShouldBeFalse("this workspace never issued that id");

        await workspace.CloseSessionAsync(first);

        workspace.IsSessionLive(first).ShouldBeFalse();
        workspace.IsSessionLive(second).ShouldBeTrue("closing one tab must not disturb another");
    }

    /// <remarks>
    /// The shell's connect path reads these back once, right after <see cref="TerminalWorkspace.OpenSessionAsync"/>
    /// returns, to fill in the status bar's cipher and host-key facts — see <c>VaultViewModel.ConnectAndAnnounceAsync</c>.
    /// Asserted the same way <see cref="LivenessIsAnsweredPerSession"/> asserts liveness: per session, and
    /// null rather than thrown for an id this workspace never issued.
    /// </remarks>
    [Fact]
    public async Task SessionFactsAreReadPerSession()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        var facts = workspace.GetSessionFacts(sessionId).ShouldNotBeNull();
        var connection = connections.Connections.ShouldHaveSingleItem();

        facts.Cipher.ShouldBe(connection.Cipher);
        facts.HostKeyAlgorithm.ShouldBe(connection.HostKey.Algorithm);

        workspace.GetSessionFacts(9999).ShouldBeNull("this workspace never issued that id");
    }

    /// <remarks>
    /// <para>
    /// Inserting a snippet has to be able to say whether it arrived, and the transport cannot: it drops
    /// frames for a session nothing is listening to, so a send at a dead tab succeeds exactly as loudly as a
    /// send at a live one. That is why <c>PasteAsync</c> answers rather than returning void — and why the
    /// answer is a <see cref="bool"/> and not an exception, since a tab whose remote hung up an hour ago is
    /// still on screen and still clickable.
    /// </para>
    /// <para>
    /// What the renderer does with the frame — bracketed paste, and the Enter deliberately outside it —
    /// cannot be reached from here at all. It is JavaScript inside a WebView, and it is in
    /// <c>docs/manual-checks.md</c> for that reason.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task PastingIntoADeadSessionSaysSoRatherThanDroppingIt()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        (await workspace.PasteAsync(
            sessionId, "uptime", execute: false, TestContext.Current.CancellationToken))
            .ShouldBeTrue("the session is live");

        (await workspace.PasteAsync(
            9999, "uptime", execute: false, TestContext.Current.CancellationToken))
            .ShouldBeFalse("this workspace never issued that id");

        await workspace.CloseSessionAsync(sessionId);

        (await workspace.PasteAsync(
            sessionId, "uptime", execute: false, TestContext.Current.CancellationToken))
            .ShouldBeFalse("the tab it names is gone");
    }

    /// <remarks>
    /// The event the tab strip and the phone's keep-alive listen to, so a dot can go out — and a foreground
    /// notification can come down — the moment a shell exits rather than at the next thing that happens to
    /// repaint. The count captured inside the handler is the sharper half of this test: the announcement
    /// used to fire from inside the run's own finally block, where the run task is not yet complete, so
    /// <c>LiveSessionCount</c> read from the handler still said 1 — and the phone's notification went on
    /// claiming a shell that was gone, with nothing left to fire and correct it.
    /// </remarks>
    [Fact]
    public async Task ASessionEndingOnItsOwnIsAnnounced_AfterTheCountStoppedIncludingIt()
    {
        // A shell with no output to give: its first read returns 0, which is a remote closing the channel,
        // so the pump finishes with nobody asking it to.
        var connections = new FakeConnectionFactory(bytesPerShell: 0);

        await using var workspace = CreateWorkspace(connections);

        var ended = new List<uint>();
        var liveAtAnnouncement = -1;
        workspace.SessionEnded += (_, e) =>
        {
            lock (ended)
            {
                liveAtAnnouncement = workspace.LiveSessionCount;
                ended.Add(e.SessionId);
            }
        };

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        await WaitUntilAsync(() =>
        {
            lock (ended)
            {
                return ended.Contains(sessionId);
            }
        });

        lock (ended)
        {
            liveAtAnnouncement.ShouldBe(0, "the announcement must wait for the run to actually complete");
        }
    }

    /// <remarks>
    /// The reversal of a recorded decision, and the event's own remark carries why: a close used to be
    /// announced to nobody, on the theory that the caller already knew — but the phone's keep-alive is not
    /// the caller, and a close it never heard about left the foreground notification claiming a shell that
    /// was gone. Announced once, after the drain, so the count a handler reads is already honest.
    /// </remarks>
    [Fact]
    public async Task ClosingASessionIsAnnounced_OnceItHasDrained()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        var announcements = 0;
        var liveAtAnnouncement = -1;
        workspace.SessionEnded += (_, _) =>
        {
            liveAtAnnouncement = workspace.LiveSessionCount;
            Interlocked.Increment(ref announcements);
        };

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        await workspace.CloseSessionAsync(sessionId);

        Volatile.Read(ref announcements)
            .ShouldBe(1, "a close is news to the keep-alive even though it is an echo to the closer");
        liveAtAnnouncement.ShouldBe(0, "announced after the drain, so the count already excludes it");
    }

    /// <remarks>
    /// Disposal is the other path that closes sessions, because it is process shutdown. Asserted so
    /// that the SSH connections are known to be released rather than assumed to be — and that these closes,
    /// unlike a deliberate one, are announced to nobody: shutdown is dismantling every subscriber along
    /// with the sessions, and news nobody is left to hear is not news.
    /// </remarks>
    [Fact]
    public async Task DisposingTheWorkspaceClosesEverySession()
    {
        var connections = new FakeConnectionFactory();

        var workspace = CreateWorkspace(connections);

        var announcements = 0;
        workspace.SessionEnded += (_, _) => Interlocked.Increment(ref announcements);

        await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);
        await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        await workspace.DisposeAsync();

        workspace.LiveSessionCount.ShouldBe(0);
        connections.Connections.Count.ShouldBe(2);
        connections.Connections.ShouldAllBe(connection => connection.IsDisposed);
        Volatile.Read(ref announcements).ShouldBe(0, "shutdown closes are not announced");
    }

    // ---- Reattach ----

    /// <remarks>
    /// <para>
    /// The scenario the whole fix exists for: a page that lost its socket — killed WebView renderer, or
    /// simply a reload — reattaches, and the session that was already running has to come back rather than
    /// sit there forever with its output going nowhere.
    /// </para>
    /// <para>
    /// The session's shell blocks on every read rather than producing output, which is what an idle prompt
    /// looks like and — for this test — is what keeps its <c>Run</c> live without a background read loop
    /// competing with this test over the credit window's exact value.
    /// </para>
    /// <para>
    /// Neither assertion below polls, deliberately. The "before" one does not need to: reserving credit is
    /// a synchronous call, so it is true the instant it returns. The "after" one does not need to either,
    /// for a subtler reason — <see cref="TerminalWorkspace.ReplayAfterAttachAsync"/> calls
    /// <c>Credits.Reset()</c> and only then awaits sending the replay frame for that same session, with no
    /// suspension between the two, so by the time this test has received that frame the reset has
    /// necessarily already happened. A poll here would only have hidden a real ordering bug behind a
    /// generous timeout instead of catching it.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task ANewRenderer_ReplaysTheLiveSessionAndResetsItsCredits()
    {
        var connections = new FakeConnectionFactory(blockShellReads: true);

        await using var workspace = CreateWorkspace(connections);
        workspace.Start();

        using var first = await ConnectRendererAsync(workspace);

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        // The session's own opening frame, sent as soon as the pump starts running. Not a replay, and not
        // what this test is about — read and discarded so it cannot be confused for one below.
        await ReceiveFrameAsync(first);

        var credits = workspace.CreditsFor(sessionId).ShouldNotBeNull();
        credits.TryReserve(4096);
        credits.Outstanding.ShouldBeGreaterThanOrEqualTo(
            4096, "the pump's own read loop may have reserved a buffer's worth on top of this");

        using var second = await ConnectRendererAsync(workspace);

        var replay = await ReceiveFrameAsync(second);
        replay.Opcode.ShouldBe((byte)TerminalServerOpcode.SessionOpened);
        replay.SessionId.ShouldBe(sessionId);
        replay.Payload.ShouldBe(new byte[] { 1 }, "a replay is flagged so the page can tell it apart from a fresh open");

        credits.Outstanding.ShouldBe(0, "the replay frame above cannot have been sent before the reset that precedes it");
    }

    /// <remarks>
    /// The other half of a reattach: the workspace has replayed what it owns, and this is the seam the
    /// shell uses to replay what it owns instead — the font size and the selected tab, neither of which a
    /// terminal session knows anything about. <c>MainWindowViewModel</c>'s subscription is what actually
    /// does that; this only asserts that the workspace hands it the chance to.
    /// </remarks>
    [Fact]
    public async Task ANewRenderer_RaisesRendererReattached()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);
        workspace.Start();

        using var first = await ConnectRendererAsync(workspace);

        var reattachedCount = 0;
        workspace.RendererReattached += (_, _) => Interlocked.Increment(ref reattachedCount);

        using var second = await ConnectRendererAsync(workspace);

        await WaitUntilAsync(() => Volatile.Read(ref reattachedCount) > 0);
    }

    // ---- Helpers ----

    /// <remarks>
    /// <para>
    /// Input that did not come from the keyboard. The Android head's accessory key row is what needs this —
    /// a software keyboard has no Ctrl, Esc, Tab or arrows — and what it sends has to arrive at the remote
    /// byte for byte, because an escape sequence that loses a byte is not a degraded arrow key, it is a
    /// stray character in somebody's shell.
    /// </para>
    /// <para>
    /// Ordinary typing does not come this way and is not what is being tested: that goes from the renderer
    /// down the socket, which the pump's own tests cover.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task SendingInput_ReachesTheRemoteUnchanged()
    {
        var connections = new FakeConnectionFactory();

        await using var workspace = CreateWorkspace(connections);

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        // The up-arrow, as a PTY expects it. Three bytes, and all three matter.
        byte[] upArrow = [0x1B, (byte)'[', (byte)'A'];

        await workspace.SendInputAsync(sessionId, upArrow, TestContext.Current.CancellationToken);

        await WaitUntilAsync(() =>
            connections.Connections.SingleOrDefault()?.Shell?.Written.SequenceEqual(upArrow) == true);
    }

    /// <remarks>
    /// A tab can close while a key is still in flight, which on a phone is one mis-tap rather than a rare
    /// race — the close cross sits inside the tab and the accessory row is directly under it. Throwing
    /// would turn that into a crash on a keystroke that no longer matters.
    /// </remarks>
    [Fact]
    public async Task SendingInputToASessionThatIsGone_IsIgnored()
    {
        await using var workspace = CreateWorkspace(new FakeConnectionFactory());

        var sessionId = await workspace.OpenSessionAsync(
            Request(), TerminalSize.Default, TestContext.Current.CancellationToken);

        await workspace.CloseSessionAsync(sessionId);

        await Should.NotThrowAsync(async () =>
            await workspace.SendInputAsync(sessionId, "x"u8.ToArray(), TestContext.Current.CancellationToken));
    }

    /// <param name="connections">
    /// How connections are made. The renderer-gate tests never reach it — they stop at the gate — but they
    /// take a fake anyway, because reaching a real host from here would make this a network test.
    /// </param>
    /// <param name="rendererTimeout">The gate's bound, or null for the shipped default.</param>
    private static TerminalWorkspace CreateWorkspace(
        ISshConnectionFactory connections,
        TimeSpan? rendererTimeout = null) =>
        new(
            StubAssets(),
            connections,
            TimeProvider.System,
            rendererTimeout is { } timeout
                ? new TerminalWorkspaceOptions { RendererTimeout = timeout }
                : null);

    private static InMemoryTerminalAssetProvider StubAssets() =>
        new(new Dictionary<string, TerminalAsset>(StringComparer.Ordinal)
        {
            // The placeholders, not a token and URL already filled in — the reattach tests below have to
            // connect a real renderer, and doing that by reading them back out of the served page is what
            // proves the workspace serves a page a real renderer could actually attach with, rather than
            // one that merely looks servable.
            [TerminalDataPlane.PagePath] = new(
                "text/html; charset=utf-8",
                Encoding.UTF8.GetBytes(
                    $"<html><body data-token=\"{TerminalDataPlane.TokenPlaceholder}\" "
                    + $"data-socket=\"{TerminalDataPlane.SocketUrlPlaceholder}\"></body></html>")),
        });

    private static SshConnectionRequest Request() =>
        new("host.invalid", 22, "dodo", new SshPasswordCredential("irrelevant"));

    /// <remarks>
    /// Attaches the way the real page does: by fetching the served page, reading the token and socket URL
    /// back out of it, and presenting them on the upgrade — rather than reaching into the workspace for a
    /// token it does not expose. A shortcut here would prove only that a socket can be opened, not that the
    /// workspace serves a page a renderer could actually attach with.
    /// </remarks>
    private static async Task<ClientWebSocket> ConnectRendererAsync(TerminalWorkspace workspace)
    {
        using var http = new HttpClient();
        var page = await http.GetStringAsync(workspace.PageUrl, TestContext.Current.CancellationToken);

        var token = ExtractAttribute(page, "data-token");
        var socketUrl = ExtractAttribute(page, "data-socket");

        var client = new ClientWebSocket();
        client.Options.AddSubProtocol(TerminalDataPlane.SubProtocol);
        client.Options.AddSubProtocol($"token.{token}");
        client.Options.SetRequestHeader(
            "Origin",
            string.Create(CultureInfo.InvariantCulture, $"http://127.0.0.1:{workspace.PageUrl.Port}"));

        try
        {
            await client.ConnectAsync(new Uri(socketUrl), TestContext.Current.CancellationToken);
        }
        catch
        {
            client.Dispose();
            throw;
        }

        return client;
    }

    private static string ExtractAttribute(string html, string name)
    {
        var marker = $"{name}=\"";
        var start = html.IndexOf(marker, StringComparison.Ordinal) + marker.Length;
        var end = html.IndexOf('"', start);

        return html[start..end];
    }

    private static async Task<(byte Opcode, uint SessionId, byte[] Payload)> ReceiveFrameAsync(
        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>
    /// Polled rather than awaited on a task, because the point is what an observer of the property
    /// sees: the pump ends on a thread of its own, and the count has to catch up without anyone
    /// telling it to.
    /// </remarks>
    private static async Task WaitUntilAsync(Func<bool> condition)
    {
        var deadline = TimeProvider.System.GetUtcNow() + PollTimeout;

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

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

        throw new TimeoutException($"The condition was still false after {PollTimeout}.");
    }
}