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

using System.Net;
using System.Net.Sockets;
using System.Text;
using DodoSSH.Client.Terminal;

namespace DodoSSH.Client.Ssh.Tests;

/// <summary>
/// The pump against a real shell, with the transport replaced by a recorder.
/// </summary>
/// <remarks>
/// Sits between the pump's unit tests, which use a fake shell, and the full end-to-end test, which
/// adds the loopback socket. Its value is diagnostic: when output does not reach a renderer, this
/// says whether the pump and SSH.NET are producing anything at all.
/// </remarks>
[Collection(SshCollection.Name)]
public sealed class PumpOverRealSshTests(SshServerFixture fixture)
{
    /// <summary>Trusts the container's host key, then connects. First contact is refused by design.</summary>
    private async Task<ISshConnection> ConnectTrustedAsync()
    {
        var knownHosts = new InMemoryKnownHostStore();
        var factory = new SshNetConnectionFactory(knownHosts);

        var request = new SshConnectionRequest(
            fixture.Host,
            fixture.Port,
            SshServerFixture.Username,
            new SshPasswordCredential(SshServerFixture.Password));

        var unknown = await Should.ThrowAsync<SshHostKeyUnknownException>(async () =>
            await factory.ConnectAsync(request, progress: null, TestContext.Current.CancellationToken));

        await knownHosts.TrustAsync(unknown.Presentation, TestContext.Current.CancellationToken);

        return await factory.ConnectAsync(request, progress: null, TestContext.Current.CancellationToken);
    }

    /// <summary>
    /// <c>exit</c> against a real sshd is reported as the remote closing the session, not as a lost
    /// connection.
    /// </summary>
    /// <remarks>
    /// <para>
    /// ◆ THIS IS THE TEST THE RACE WOULD FAIL. The obvious way to tell an ordinary exit from a dropped
    /// connection is to wait for the read to end and then ask whether the transport is still up — and against
    /// a real server it is wrong most of the time, because this was the only channel on the connection and
    /// sshd hangs up within a millisecond of closing it. By the time anything here could ask, the transport
    /// is gone. A fake shell cannot show that: it has no server behind it deciding to disconnect.
    /// </para>
    /// <para>
    /// What makes the answer settled instead is the failure being latched when it happens rather than
    /// inferred afterwards — see <see cref="ISshShellSession.TransportFailure"/>. Nothing latches one here,
    /// because nothing failed.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task AShellThatExits_IsReportedAsTheRemoteClosingItRatherThanAsALostConnection()
    {
        await using var connection = await ConnectTrustedAsync();

        var shell = await connection.OpenShellAsync(
            new TerminalSize(100, 30, 1000, 750), TestContext.Current.CancellationToken);

        var transport = new CountingTransport();

        await using var pump = new TerminalSessionPump(
            1,
            shell,
            transport,
            TimeProvider.System,
            new TerminalPumpOptions { FlushInterval = TimeSpan.FromMilliseconds(20) });

        transport.Pump = pump;

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

        await pump.WriteInputAsync("exit\n"u8.ToArray(), TestContext.Current.CancellationToken);

        // No disposal and no cancellation: the run has to end because the shell did, which is the whole
        // point. A timeout here would mean the session hangs on `exit`.
        await run.WaitAsync(TimeSpan.FromSeconds(30), TestContext.Current.CancellationToken);

        pump.Outcome.End.ShouldBe(
            SshSessionEnd.RemoteClosed,
            $"exit is not a dropped connection; got <{pump.Outcome.Description}>");
    }

    /// <summary>
    /// A transport that goes away under a live shell is reported as a lost connection, and the read does not
    /// hang.
    /// </summary>
    /// <remarks>
    /// <para>
    /// ◆ THE REBOOT CASE, produced without rebooting anything. The connection is dialled through a plain TCP
    /// relay this test owns and then the relay is torn down, which takes the socket out from under a live
    /// shell exactly as a machine going down does. The container is shared by every test in this collection,
    /// so killing <c>sshd</c> itself was never an option; killing something this test owns is.
    /// </para>
    /// <para>
    /// <b>What is not reproduced here is the waiting.</b> A relay that is closed sends a FIN, so the socket
    /// reports itself at once; a machine that has been switched off sends nothing at all, and finding that
    /// out requires writing to it — which is what the keep-alive in <see cref="SshNetConnectionFactory"/>
    /// exists to do, and which no test can hurry along without a clock this code does not own. So this covers
    /// the half that can be checked, and it is the half everything above the socket depends on: a transport
    /// that has failed is <em>reported as one</em> rather than being called an ordinary exit or leaving the
    /// read blocked for ever.
    /// </para>
    /// <para>
    /// A relay rather than the SOCKS forward <c>LoopbackProxyTests</c> uses, which was the first attempt:
    /// disconnecting the carrier of a <c>ForwardedPortDynamic</c> does not close the loopback socket it
    /// accepted, so the inner connection was left holding a socket with nothing behind it and the read did
    /// not return within thirty seconds. That is a property of SSH.NET's forwarder rather than of anything
    /// this repository ships, and it is not what a dying machine looks like.
    /// </para>
    /// </remarks>
    [Fact]
    public async Task AShellWhoseTransportDies_IsReportedAsALostConnection()
    {
        using var relay = new TcpRelay(fixture.Host, fixture.Port);

        await using var connection = await ConnectThroughAsync(relay);

        var shell = await connection.OpenShellAsync(
            new TerminalSize(100, 30, 1000, 750), TestContext.Current.CancellationToken);

        var transport = new CountingTransport();

        await using var pump = new TerminalSessionPump(
            1,
            shell,
            transport,
            TimeProvider.System,
            new TerminalPumpOptions { FlushInterval = TimeSpan.FromMilliseconds(20) });

        transport.Pump = pump;

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

        // The shell is up and talking before anything is broken, so what follows is a live session losing its
        // transport rather than one that never started.
        var deadline = TimeProvider.System.GetUtcNow() + TimeSpan.FromSeconds(20);

        while (TimeProvider.System.GetUtcNow() < deadline && pump.BytesRead == 0)
        {
            await Task.Delay(50, TestContext.Current.CancellationToken);
        }

        pump.BytesRead.ShouldBeGreaterThan(0, "the shell should have produced a prompt before the relay goes");

        // And the ground goes out from under it.
        relay.Kill();

        // No disposal and no cancellation: the run has to end because the transport did.
        await run.WaitAsync(TimeSpan.FromSeconds(30), TestContext.Current.CancellationToken);

        pump.Outcome.End.ShouldBe(
            SshSessionEnd.ConnectionLost,
            $"a transport that died is not a shell that exited; got <{pump.Outcome.Description}>");

        pump.Outcome.Description.ShouldNotBeNullOrWhiteSpace(
            "the strip under the terminal prints this sentence, so an empty one is a blank status line");
    }

    /// <summary>Trusts the relay's identity, then connects through it.</summary>
    /// <remarks>
    /// Pinned under the relay's own address, which is the honest thing here: what is under examination is a
    /// transport dying, and this store has one entry in it either way.
    /// </remarks>
    private static async Task<ISshConnection> ConnectThroughAsync(TcpRelay relay)
    {
        var knownHosts = new InMemoryKnownHostStore();
        var factory = new SshNetConnectionFactory(knownHosts);

        var request = new SshConnectionRequest(
            "127.0.0.1",
            relay.Port,
            SshServerFixture.Username,
            new SshPasswordCredential(SshServerFixture.Password));

        var unknown = await Should.ThrowAsync<SshHostKeyUnknownException>(async () =>
            await factory.ConnectAsync(request, progress: null, TestContext.Current.CancellationToken));

        await knownHosts.TrustAsync(unknown.Presentation, TestContext.Current.CancellationToken);

        return await factory.ConnectAsync(request, progress: null, TestContext.Current.CancellationToken);
    }

    /// <summary>A loopback TCP forwarder that can be cut, standing in for a machine that goes away.</summary>
    /// <remarks>
    /// Deliberately dumb: it accepts one connection, opens one to the far end and copies bytes both ways
    /// until something stops. SSH is opaque to it, which is the point — what is being simulated is the wire
    /// disappearing, and a relay that understood the protocol could not simulate that without also deciding
    /// how to say goodbye.
    /// </remarks>
    private sealed class TcpRelay : IDisposable
    {
        private readonly TcpListener listener;
        private readonly CancellationTokenSource lifetime = new();
        private readonly List<TcpClient> sockets = [];
        private readonly Lock gate = new();

        internal TcpRelay(string host, int port)
        {
            listener = new TcpListener(IPAddress.Loopback, 0);
            listener.Start();

            Port = ((IPEndPoint)listener.LocalEndpoint).Port;

            _ = AcceptAsync(host, port);
        }

        /// <summary>The loopback port to dial instead of the far end's.</summary>
        internal int Port { get; }

        /// <summary>
        /// Drops every socket, in both directions, without warning either end first.
        /// </summary>
        /// <remarks>
        /// Disposing the sockets is what closes them, and closing is what puts a FIN on the wire. Nothing
        /// here speaks to the SSH session running over it, which is the whole idea: this is the network
        /// going away, not a disconnect being negotiated.
        /// </remarks>
        internal void Kill()
        {
            lifetime.Cancel();
            listener.Stop();

            lock (gate)
            {
                foreach (var socket in sockets)
                {
                    socket.Dispose();
                }

                sockets.Clear();
            }
        }

        public void Dispose()
        {
            Kill();
            lifetime.Dispose();
        }

        private async Task AcceptAsync(string host, int port)
        {
            try
            {
                while (!lifetime.Token.IsCancellationRequested)
                {
                    var inbound = await listener.AcceptTcpClientAsync(lifetime.Token).ConfigureAwait(false);
                    var outbound = new TcpClient();

                    await outbound.ConnectAsync(host, port, lifetime.Token).ConfigureAwait(false);

                    lock (gate)
                    {
                        sockets.Add(inbound);
                        sockets.Add(outbound);
                    }

                    _ = CopyAsync(inbound, outbound);
                    _ = CopyAsync(outbound, inbound);
                }
            }
            catch (Exception exception) when (exception is not OutOfMemoryException)
            {
                // Killed, or the listener stopped. There is no caller on this thread to hand anything to.
            }
        }

        private async Task CopyAsync(TcpClient from, TcpClient to)
        {
            try
            {
                await from.GetStream().CopyToAsync(to.GetStream(), lifetime.Token).ConfigureAwait(false);
            }
            catch (Exception exception) when (exception is not OutOfMemoryException)
            {
                // Either end going is the ordinary way this ends, and is what the test is arranging.
            }
        }
    }

    [Fact]
    public async Task ThePump_ForwardsRealShellOutput()
    {
        await using var connection = await ConnectTrustedAsync();

        var shell = await connection.OpenShellAsync(
            new TerminalSize(100, 30, 1000, 750), TestContext.Current.CancellationToken);

        var transport = new CountingTransport();

        await using var pump = new TerminalSessionPump(
            1,
            shell,
            transport,
            TimeProvider.System,
            new TerminalPumpOptions { FlushInterval = TimeSpan.FromMilliseconds(20) });

        transport.Pump = pump;

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

        // Marker split so the PTY's echo of the command line does not satisfy the match.
        await pump.WriteInputAsync(
            Encoding.UTF8.GetBytes("echo \"DODO\"\"SSH-OK\"\n"),
            TestContext.Current.CancellationToken);

        var deadline = TimeProvider.System.GetUtcNow() + TimeSpan.FromSeconds(20);

        while (TimeProvider.System.GetUtcNow() < deadline
            && !transport.Text.Contains("DODOSSH-OK", StringComparison.Ordinal))
        {
            await Task.Delay(50, TestContext.Current.CancellationToken);
        }

        transport.Text.Contains("DODOSSH-OK", StringComparison.Ordinal).ShouldBeTrue(
            $"BytesRead={pump.BytesRead}, FramesSent={pump.FramesSent}, "
            + $"OutputFrames={transport.OutputFrames}, Text=<{transport.Text}>");

        await pump.DisposeAsync();

        try
        {
            await run;
        }
        catch (OperationCanceledException)
        {
            // Expected: disposing the pump cancels its run.
        }
    }

    /// <summary>Accumulates output and acknowledges it, as a keeping-up renderer would.</summary>
    private sealed class CountingTransport : ITerminalTransport
    {
        private readonly StringBuilder text = new();
        private readonly Lock gate = new();

        internal TerminalSessionPump? Pump { get; set; }

        internal int OutputFrames { get; private set; }

        internal string Text
        {
            get
            {
                lock (gate)
                {
                    return text.ToString();
                }
            }
        }

        public ValueTask SendAsync(ReadOnlyMemory<byte> frame, CancellationToken cancellationToken)
        {
            if (!TerminalFrame.TryRead(frame.Span, out var opcode, out _, out var payload)
                || opcode != (byte)TerminalServerOpcode.Output)
            {
                return ValueTask.CompletedTask;
            }

            lock (gate)
            {
                OutputFrames++;
                text.Append(Encoding.UTF8.GetString(payload));
            }

            Pump?.Acknowledge((uint)payload.Length);

            return ValueTask.CompletedTask;
        }
    }
}