openac-mosswartmassacre/tests/OpenAC.MosswartMassacre.Tests/Backend/TestSocketServer.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;
using System.Text;

namespace OpenAC.MosswartMassacre.Tests.Backend;

/// <summary>
/// A throwaway WebSocket server on a free local port, standing in for the
/// backend: it records the secret header and every frame, and can push frames
/// and drop the connection.
/// </summary>
internal sealed class TestSocketServer : IDisposable
{
    private readonly HttpListener _listener = new();
    private readonly CancellationTokenSource _cts = new();
    private readonly Task _accept;
    private WebSocket? _current;
    private HttpListenerContext? _currentContext;

    public TestSocketServer()
    {
        int port = FreePort();
        Uri = new Uri($"ws://localhost:{port}/websocket/");
        _listener.Prefixes.Add($"http://localhost:{port}/websocket/");
        _listener.Start();
        _accept = Task.Run(AcceptLoopAsync);
    }

    public Uri Uri { get; }

    public BlockingCollection<string> Frames { get; } = new();

    public ConcurrentQueue<string?> Secrets { get; } = new();

    public int Connections => Secrets.Count;

    public async Task PushAsync(string text)
    {
        WebSocket socket = await WaitForSocketAsync();
        byte[] bytes = Encoding.UTF8.GetBytes(text);
        // Split into small fragments so the client has to reassemble.
        const int chunk = 1000;
        for (int offset = 0; offset < bytes.Length; offset += chunk)
        {
            int count = Math.Min(chunk, bytes.Length - offset);
            bool last = offset + count >= bytes.Length;
            await socket.SendAsync(new ArraySegment<byte>(bytes, offset, count), WebSocketMessageType.Text, last, CancellationToken.None);
        }
    }

    public async Task DropAsync()
    {
        WebSocket socket = await WaitForSocketAsync();
        socket.Abort();
        // Aborting the socket alone leaves the TCP connection open on some
        // platforms' listeners, so the client never sees the drop; ending the
        // response closes the connection everywhere, as a real server's drop does.
        _currentContext?.Response.Abort();
        _current = null;
        _currentContext = null;
    }

    public string NextFrame(int timeoutMs = 5000)
    {
        if (!Frames.TryTake(out string? frame, timeoutMs))
            throw new TimeoutException("No frame arrived.");
        return frame;
    }

    public void Dispose()
    {
        _cts.Cancel();
        _current?.Abort();
        _listener.Stop();
        _listener.Close();
        try
        {
            _accept.Wait(1000);
        }
        catch (AggregateException)
        {
        }
    }

    private async Task<WebSocket> WaitForSocketAsync()
    {
        for (int i = 0; i < 200; i++)
        {
            if (_current is { State: WebSocketState.Open } socket)
                return socket;
            await Task.Delay(25);
        }
        throw new TimeoutException("No client connected.");
    }

    private async Task AcceptLoopAsync()
    {
        while (!_cts.IsCancellationRequested)
        {
            HttpListenerContext context;
            try
            {
                context = await _listener.GetContextAsync();
            }
            catch
            {
                return;
            }
            if (!context.Request.IsWebSocketRequest)
            {
                context.Response.StatusCode = 400;
                context.Response.Close();
                continue;
            }
            Secrets.Enqueue(context.Request.Headers["X-Plugin-Secret"]);
            HttpListenerWebSocketContext ws = await context.AcceptWebSocketAsync(null);
            _current = ws.WebSocket;
            _currentContext = context;
            _ = Task.Run(() => ReadLoopAsync(ws.WebSocket));
        }
    }

    private async Task ReadLoopAsync(WebSocket socket)
    {
        var buffer = new byte[65536];
        using var message = new MemoryStream();
        try
        {
            while (socket.State == WebSocketState.Open)
            {
                WebSocketReceiveResult result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), _cts.Token);
                if (result.MessageType == WebSocketMessageType.Close)
                    return;
                message.Write(buffer, 0, result.Count);
                if (!result.EndOfMessage)
                    continue;
                Frames.Add(Encoding.UTF8.GetString(message.ToArray()));
                message.SetLength(0);
            }
        }
        catch
        {
        }
    }

    private static int FreePort()
    {
        var probe = new TcpListener(IPAddress.Loopback, 0);
        probe.Start();
        int port = ((IPEndPoint)probe.LocalEndpoint).Port;
        probe.Stop();
        return port;
    }
}