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

using DodoSSH.Client.Ssh;

namespace DodoSSH.Client.Terminal.Tests;

/// <summary>A shell session that produces output on demand, for exercising the pump.</summary>
internal sealed class FakeShellSession : ISshShellSession
{
    private readonly List<byte> written = [];
    private readonly Lock gate = new();

    private readonly bool blockReads;
    private readonly TimeSpan? pacePerRead;

    private long remaining;
    private byte pattern;

    /// <param name="bytesToProduce">
    /// How many bytes to emit before reporting end of stream. <see cref="long.MaxValue"/> for an
    /// endless producer, which is what a runaway remote process looks like — those sessions are ended
    /// by disposing the pump rather than by running out of data.
    /// </param>
    /// <param name="blockReads">
    /// True for a shell that is open and live but has nothing to say — an idle prompt, rather than either
    /// end of the "produces bytes" and "hit end of stream" spectrum <paramref name="bytesToProduce"/>
    /// covers. <see cref="ReadAsync"/> then blocks until cancelled, which is what a real idle SSH channel's
    /// read does. Exists for tests that need a session whose <c>Run</c> stays live without a background
    /// read loop racing the test for control of the pump's credit window — see the reattach tests in
    /// <c>TerminalWorkspaceTests</c>.
    /// </param>
    /// <param name="pacePerRead">
    /// How long each read takes, or null for one that returns as fast as the CPU can fill a buffer.
    /// </param>
    /// <remarks>
    /// <para>
    /// The default producer emits at memory bandwidth, which is faster than any network and faster than
    /// the pump can frame and send — useful for the throughput harness, and useless for asking whether a
    /// terminal that <em>can</em> keep up loses anything, because nothing can keep up with it. A pace makes
    /// the fake a link instead of a memcpy: 8 KiB every millisecond is 8 MB/s, which is a fast remote
    /// rather than an impossible one.
    /// </para>
    /// </remarks>
    internal FakeShellSession(
        long bytesToProduce = 0,
        bool blockReads = false,
        TimeSpan? pacePerRead = null)
    {
        remaining = bytesToProduce;
        this.blockReads = blockReads;
        this.pacePerRead = pacePerRead;
    }

    /// <inheritdoc />
    public bool IsOpen { get; private set; } = true;

    /// <inheritdoc />
    /// <remarks>
    /// Settable, because it is the whole of what tells a shell that exited apart from one whose connection
    /// broke — a test sets it for the second case and leaves it null for the first, which is exactly what the
    /// real session's latch does off SSH.NET's error event.
    /// </remarks>
    public string? TransportFailure { get; set; }

    /// <summary>Reads issued against this session, to detect a pump that kept reading.</summary>
    public int ReadCount { get; private set; }

    /// <summary>Bytes the pump wrote toward the remote.</summary>
    public byte[] Written
    {
        get
        {
            lock (gate)
            {
                return [.. written];
            }
        }
    }

    /// <summary>The last size the pump forwarded, or null if it forwarded none.</summary>
    public TerminalSize? LastResize { get; private set; }

    /// <summary>How many resizes were forwarded, so a dropped one is observable.</summary>
    public int ResizeCount { get; private set; }

    /// <inheritdoc />
    /// <remarks>
    /// Returns 0 once the configured budget is spent, which is what a remote closing the channel looks
    /// like. Not synchronous: a fake that never yields would let the pump's read loop monopolise the
    /// thread and hide any ordering problem between reading and flushing.
    /// </remarks>
    public async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken)
    {
        ReadCount++;

        if (blockReads)
        {
            // Never completes on its own. The only way out is the same way a real blocked read ends: the
            // token being cancelled, which is what disposing the pump does.
            await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false);
        }

        if (pacePerRead is { } pace)
        {
            await Task.Delay(pace, cancellationToken).ConfigureAwait(false);
        }

        await Task.Yield();
        cancellationToken.ThrowIfCancellationRequested();

        if (remaining <= 0)
        {
            return 0;
        }

        var count = (int)Math.Min(buffer.Length, remaining);
        buffer.Span[..count].Fill(unchecked(pattern++));
        remaining -= count;

        return count;
    }

    /// <inheritdoc />
    public ValueTask WriteAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
    {
        lock (gate)
        {
            written.AddRange(data.ToArray());
        }

        return ValueTask.CompletedTask;
    }

    /// <inheritdoc />
    public void Resize(TerminalSize size)
    {
        // Mirrors the real session: an unusable size is dropped rather than forwarded.
        if (!size.IsUsable)
        {
            return;
        }

        ResizeCount++;
        LastResize = size;
    }

    /// <inheritdoc />
    public ValueTask DisposeAsync()
    {
        IsOpen = false;
        remaining = 0;

        return ValueTask.CompletedTask;
    }
}

/// <summary>
/// Hands out <see cref="FakeShellSession"/>s, so a workspace can be driven with no network.
/// </summary>
/// <remarks>
/// Exists for the session-lifetime tests. Everything else in this suite works on a pump directly; the
/// workspace is the layer that decides when a session is over, and that decision is what needs a
/// connection whose shell can be made to end on cue.
/// </remarks>
internal sealed class FakeConnectionFactory(long bytesPerShell = long.MaxValue, bool blockShellReads = false)
    : ISshConnectionFactory
{
    /// <summary>Armed on every shell handed out, so a workspace-level test can produce a dropped connection.</summary>
    internal string? TransportFailure { get; set; }

    /// <summary>Connections handed out, in order.</summary>
    internal List<FakeConnection> Connections { get; } = [];

    /// <inheritdoc />
    public Task<ISshConnection> ConnectAsync(
        SshConnectionRequest request,
        IProgress<SshConnectionPhase>? progress,
        CancellationToken cancellationToken)
    {
        // The real factory's own order, so that a test watching this fake is watching the same sequence a
        // real handshake produces. It cannot report CheckingHostKey — there is no key exchange here to
        // produce a key — and inventing one would make this the only place that phase came from.
        progress?.Report(SshConnectionPhase.Reaching);
        progress?.Report(SshConnectionPhase.Authenticating);

        var connection = new FakeConnection(request, bytesPerShell, blockShellReads)
        {
            TransportFailure = TransportFailure,
        };

        Connections.Add(connection);

        return Task.FromResult<ISshConnection>(connection);
    }
}

/// <summary>A connection that opens fake shells and records its own disposal.</summary>
internal sealed class FakeConnection(SshConnectionRequest request, long bytesPerShell, bool blockShellReads = false)
    : ISshConnection
{
    /// <inheritdoc />
    public bool IsConnected { get; private set; } = true;

    /// <inheritdoc />
    public HostKeyPresentation HostKey { get; } =
        new(request.Host, request.Port, "ssh-ed25519", "SHA256:fake");

    /// <inheritdoc />
    public string Cipher { get; } = "aes256-gcm@openssh.com";

    /// <summary>The shell this connection opened, if it opened one.</summary>
    internal FakeShellSession? Shell { get; private set; }

    /// <inheritdoc cref="ISshShellSession.TransportFailure" />
    /// <remarks>Handed to whatever shell this connection opens, so the arming happens before the pump runs.</remarks>
    internal string? TransportFailure { get; init; }

    /// <summary>Whether the connection was disposed, which is what closing a session must do.</summary>
    internal bool IsDisposed { get; private set; }

    /// <inheritdoc />
    public Task<ISshShellSession> OpenShellAsync(TerminalSize size, CancellationToken cancellationToken)
    {
        Shell = new FakeShellSession(bytesPerShell, blockShellReads) { TransportFailure = TransportFailure };

        return Task.FromResult<ISshShellSession>(Shell);
    }

    /// <summary>The forwards started on this connection, in order.</summary>
    internal List<SshPortForwardSpec> Forwards { get; } = [];

    /// <inheritdoc />
    public Task<ISshPortForward> StartForwardAsync(SshPortForwardSpec spec, CancellationToken cancellationToken)
    {
        Forwards.Add(spec);

        return Task.FromResult<ISshPortForward>(new FakePortForward(spec));
    }

    /// <inheritdoc />
    public ValueTask DisposeAsync()
    {
        IsDisposed = true;
        IsConnected = false;

        return ValueTask.CompletedTask;
    }
}

/// <summary>
/// Counts output without keeping it, and acknowledges it as a renderer that is keeping up would.
/// </summary>
/// <remarks>
/// <see cref="RecordingTransport"/>'s cheap sibling, and the throughput harness needs it to be the one
/// measuring rather than the one measured: recording copies every frame into a list that grows to the size
/// of the whole run, which at the rates that harness drives makes the <em>test double</em> the slowest thing
/// in the pipeline. A harness whose consumer cannot keep up is measuring itself.
/// </remarks>
internal sealed class CountingTransport : ITerminalTransport
{
    private long outputBytes;

    /// <summary>Set to acknowledge every output frame immediately, as a keeping-up renderer would.</summary>
    internal TerminalSessionPump? AutoAcknowledge { get; set; }

    /// <summary>Payload bytes of every output frame sent so far.</summary>
    internal long OutputBytes => Interlocked.Read(ref outputBytes);

    /// <inheritdoc />
    public ValueTask SendAsync(ReadOnlyMemory<byte> frame, CancellationToken cancellationToken)
    {
        if (TerminalFrame.TryRead(frame.Span, out var opcode, out _, out var payload)
            && opcode == (byte)TerminalServerOpcode.Output)
        {
            _ = Interlocked.Add(ref outputBytes, payload.Length);

            AutoAcknowledge?.Acknowledge((uint)payload.Length);
        }

        return ValueTask.CompletedTask;
    }
}

/// <summary>Records frames, and can acknowledge them to keep credit flowing.</summary>
internal sealed class RecordingTransport : ITerminalTransport
{
    private readonly List<byte[]> frames = [];
    private readonly Lock gate = new();

    /// <summary>Set to acknowledge every output frame immediately, as a keeping-up renderer would.</summary>
    internal TerminalSessionPump? AutoAcknowledge { get; set; }

    /// <summary>Frames sent so far.</summary>
    internal IReadOnlyList<byte[]> Frames
    {
        get
        {
            lock (gate)
            {
                return [.. frames];
            }
        }
    }

    /// <inheritdoc />
    public ValueTask SendAsync(ReadOnlyMemory<byte> frame, CancellationToken cancellationToken)
    {
        var copy = frame.ToArray();

        lock (gate)
        {
            frames.Add(copy);
        }

        if (AutoAcknowledge is { } pump
            && TerminalFrame.TryRead(copy, out var opcode, out _, out var payload)
            && opcode == (byte)TerminalServerOpcode.Output)
        {
            pump.Acknowledge((uint)payload.Length);
        }

        return ValueTask.CompletedTask;
    }

    /// <summary>Concatenated payloads of every output frame.</summary>
    internal byte[] OutputBytes()
    {
        var output = new List<byte>();

        foreach (var frame in Frames)
        {
            if (TerminalFrame.TryRead(frame, out var opcode, out _, out var payload)
                && opcode == (byte)TerminalServerOpcode.Output)
            {
                output.AddRange(payload);
            }
        }

        return [.. output];
    }

    /// <summary>Concatenated payloads of every frame of one opcode.</summary>
    internal byte[] PayloadsOf(TerminalServerOpcode opcode)
    {
        var payloads = new List<byte>();

        foreach (var frame in Frames)
        {
            if (TerminalFrame.TryRead(frame, out var actual, out _, out var payload)
                && actual == (byte)opcode)
            {
                payloads.AddRange(payload);
            }
        }

        return [.. payloads];
    }

    /// <summary>Counts frames of one opcode.</summary>
    internal int CountOf(TerminalServerOpcode opcode) =>
        Frames.Count(frame =>
            TerminalFrame.TryRead(frame, out var actual, out _, out _)
            && actual == (byte)opcode);
}

/// <summary>An <see cref="IProgress{T}"/> that runs its callback on the thread that reported.</summary>
/// <remarks>
/// <c>System.Progress&lt;T&gt;</c> would post to a captured synchronisation context, or to the thread pool
/// when there is none — which is what a test has here — so a list it appended to would be asserted on before
/// it had been written. This is the same reason the shell does not use it either; see
/// <c>VaultViewModel.ReporterFor</c>.
/// </remarks>
internal sealed class DelegateProgress<T>(Action<T> report) : IProgress<T>
{
    public void Report(T value) => report(value);
}

/// <summary>A forward that listens on nothing and says it is listening until disposed.</summary>
internal sealed class FakePortForward(SshPortForwardSpec spec) : ISshPortForward
{
    /// <inheritdoc />
    public event EventHandler<SshPortForwardErrorEventArgs>? ErrorOccurred
    {
        add { }
        remove { }
    }

    /// <inheritdoc />
    public SshPortForwardSpec Spec { get; } = spec;

    /// <inheritdoc />
    public int BoundPort => Spec.BindPort;

    /// <inheritdoc />
    public bool IsListening { get; private set; } = true;

    /// <inheritdoc />
    public string? LastError => null;

    /// <inheritdoc />
    public ValueTask DisposeAsync()
    {
        IsListening = false;
        return ValueTask.CompletedTask;
    }
}