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 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>
    internal FakeShellSession(long bytesToProduce = 0, bool blockReads = false)
    {
        remaining = bytesToProduce;
        this.blockReads = blockReads;
    }

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

    /// <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);
        }

        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>Connections handed out, in order.</summary>
    internal List<FakeConnection> Connections { get; } = [];

    /// <inheritdoc />
    public Task<ISshConnection> ConnectAsync(
        SshConnectionRequest request,
        CancellationToken cancellationToken)
    {
        var connection = new FakeConnection(request, bytesPerShell, blockShellReads);
        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; }

    /// <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);

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

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

        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>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);
}