DodoSSH/tests/DodoSSH.Client.Session.Tests/FakeSpelunk.cs

using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;
using System.Security.Cryptography;
using System.Text;

namespace DodoSSH.Client.Session.Tests;

/// <summary>
/// Spelunk's two tunnel endpoints, as far as the wire is concerned: HTTP/1.1 over a real loopback socket.
/// </summary>
/// <remarks>
/// <para>
/// Real sockets rather than a stubbed handler, because what is being tested is the client's own HTTP and
/// WebSocket stack talking to something: the header the key goes in, the subprotocols the upgrade offers,
/// a refused upgrade's status read back through <c>CollectHttpResponseDetails</c>, and a close code
/// arriving through <see cref="Client.Ssh.RelayBridge"/>. A handler stub would skip exactly those.
/// </para>
/// <para>
/// Hand-rolled rather than Kestrel, so this client suite takes no dependency on the ASP.NET Core framework
/// for one fixture. It speaks only what the client sends: one request per connection, answered with
/// <c>Connection: close</c>, and the one upgrade.
/// </para>
/// </remarks>
internal sealed class FakeSpelunk : IAsyncDisposable
{
    internal const string Ticket = "dGlja2V0LWZvci1hLXR1bm5lbA";

    private readonly TcpListener listener = new(IPAddress.Loopback, 0);
    private readonly CancellationTokenSource stopping = new();
    private readonly Task accepting;

    internal FakeSpelunk()
    {
        listener.Start();
        accepting = AcceptAsync();
    }

    /// <summary>The API base a client is given.</summary>
    internal string ApiBase => $"http://127.0.0.1:{((IPEndPoint)listener.LocalEndpoint).Port}";

    /// <summary>Every request that arrived, in order.</summary>
    internal List<Request> Requests { get; } = [];

    /// <summary>How the ticket endpoint answers. By default, with a ticket.</summary>
    internal Func<Request, Answer> OnTicket { get; set; } = request => Answer.Json(
        HttpStatusCode.OK,
        $$"""
        {"ticket":"{{Ticket}}","expiresAt":"2026-09-28T12:00:30Z","path":"{{request.Path.Replace("tunnel-tickets", "tunnel", StringComparison.Ordinal)}}","subprotocol":"spelunk.tunnel.v1"}
        """);

    /// <summary>How the upgrade is answered: null accepts it, anything else refuses it with that status.</summary>
    internal Func<Request, HttpStatusCode?> OnUpgrade { get; set; } = _ => null;

    /// <summary>What Spelunk does with an accepted tunnel — by default nothing, until the client closes.</summary>
    internal Func<WebSocket, Task> OnTunnel { get; set; } = DrainAsync;

    public async ValueTask DisposeAsync()
    {
        await stopping.CancelAsync();
        listener.Stop();

        try
        {
            await accepting;
        }
        catch (Exception exception) when (exception is OperationCanceledException or SocketException or ObjectDisposedException)
        {
            // Stopped.
        }

        stopping.Dispose();
    }

    /// <summary>Reads until the client closes, and answers the close.</summary>
    internal static async Task DrainAsync(WebSocket socket)
    {
        var buffer = new byte[4096];

        while ((await socket.ReceiveAsync(buffer, CancellationToken.None)).MessageType != WebSocketMessageType.Close)
        {
        }

        await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);
    }

    private async Task AcceptAsync()
    {
        while (!stopping.IsCancellationRequested)
        {
            var client = await listener.AcceptTcpClientAsync(stopping.Token);
            _ = ServeAsync(client);
        }
    }

    private async Task ServeAsync(TcpClient client)
    {
        using (client)
        {
            var stream = client.GetStream();
            var request = await ReadRequestAsync(stream);

            lock (Requests)
            {
                Requests.Add(request);
            }

            if (request.Headers.TryGetValue("Upgrade", out var upgrade)
                && upgrade.Equals("websocket", StringComparison.OrdinalIgnoreCase))
            {
                await UpgradeAsync(stream, request);
                return;
            }

            var answer = OnTicket(request);
            var body = Encoding.UTF8.GetBytes(answer.Body);

            await WriteAsync(
                stream,
                $"HTTP/1.1 {(int)answer.Status} {answer.Status}\r\nContent-Type: {answer.ContentType}\r\n"
                + $"Content-Length: {body.Length}\r\nConnection: close\r\n\r\n");
            await stream.WriteAsync(body);
        }
    }

    private async Task UpgradeAsync(NetworkStream stream, Request request)
    {
        if (OnUpgrade(request) is { } refusal)
        {
            await WriteAsync(
                stream,
                $"HTTP/1.1 {(int)refusal} {refusal}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
            return;
        }

        await WriteAsync(
            stream,
            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
            + $"Sec-WebSocket-Accept: {AcceptFor(request.Headers["Sec-WebSocket-Key"])}\r\n"
            + "Sec-WebSocket-Protocol: spelunk.tunnel.v1\r\n\r\n");

        using var socket = WebSocket.CreateFromStream(
            stream,
            new WebSocketCreationOptions
            {
                IsServer = true,
                SubProtocol = "spelunk.tunnel.v1",
                KeepAliveInterval = TimeSpan.Zero,
            });

        try
        {
            await OnTunnel(socket);
        }
        catch (Exception exception) when (exception is WebSocketException or IOException)
        {
            // The client went away; nothing here is under test.
        }
    }

    /// <summary>The handshake's accept value, as RFC 6455 section 4.2.2 defines it.</summary>
    private static string AcceptFor(string key)
    {
#pragma warning disable RS0030, CA5350 // RFC 6455 fixes SHA-1 for this value. It authenticates nothing; it proves the peer read the key.
        return Convert.ToBase64String(SHA1.HashData(Encoding.ASCII.GetBytes(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")));
#pragma warning restore RS0030, CA5350
    }

    private static async Task<Request> ReadRequestAsync(NetworkStream stream)
    {
        var head = new List<byte>();
        var one = new byte[1];

        while (!(head.Count >= 4 && head[^4] == '\r' && head[^3] == '\n' && head[^2] == '\r' && head[^1] == '\n'))
        {
            if (await stream.ReadAsync(one) == 0)
            {
                break;
            }

            head.Add(one[0]);
        }

        var lines = Encoding.ASCII.GetString([.. head]).Split("\r\n", StringSplitOptions.RemoveEmptyEntries);
        var requestLine = lines[0].Split(' ');
        var headers = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);

        foreach (var line in lines.Skip(1))
        {
            var colon = line.IndexOf(':', StringComparison.Ordinal);
            headers[line[..colon].Trim()] = line[(colon + 1)..].Trim();
        }

        var body = string.Empty;

        if (headers.TryGetValue("Content-Length", out var length) && int.Parse(length, System.Globalization.CultureInfo.InvariantCulture) is > 0 and var count)
        {
            var bytes = new byte[count];
            await stream.ReadExactlyAsync(bytes);
            body = Encoding.UTF8.GetString(bytes);
        }

        return new Request(requestLine[0], requestLine[1], headers, body);
    }

    private static Task WriteAsync(NetworkStream stream, string text) =>
        stream.WriteAsync(Encoding.ASCII.GetBytes(text)).AsTask();

    /// <summary>One request as it arrived.</summary>
    internal sealed record Request(string Method, string Path, Dictionary<string, string> Headers, string Body);

    /// <summary>What the ticket endpoint says back.</summary>
    internal sealed record Answer(HttpStatusCode Status, string ContentType, string Body)
    {
        internal static Answer Json(HttpStatusCode status, string body) => new(status, "application/json", body);

        internal static Answer Problem(HttpStatusCode status, string title, string? detail = null) => new(
            status,
            "application/problem+json",
            detail is null
                ? $$"""{"type":"about:blank","title":"{{title}}","status":{{(int)status}}}"""
                : $$"""{"type":"about:blank","title":"{{title}}","status":{{(int)status}},"detail":"{{detail}}"}""");
    }
}