using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;
using System.Text;
namespace OpenAC.MosswartMassacre.Tests.Backend;
///
/// 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.
///
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 Frames { get; } = new();
public ConcurrentQueue 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(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 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(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;
}
}