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 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>
    internal FakeShellSession(long bytesToProduce = 0) => remaining = bytesToProduce;

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

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