LarpixServer/LarpixVoice/Requests.cs

using System.Buffers;
using System.Collections.Concurrent;
using System.Net.WebSockets;

using static LarpixVoice.Program;

namespace LarpixVoice;

public class Requests
{
    public static async Task Websocket(HttpContext context)
    {
        if (!context.WebSockets.IsWebSocketRequest)
        {
            context.Response.StatusCode = 400;
            return;
        }

        string roomId = context.Request.Query["room"].ToString();
        string userIdStr = context.Request.Query["userId"].ToString();

        if (string.IsNullOrEmpty(roomId) || !ulong.TryParse(userIdStr, out ulong clientId))
        {
            context.Response.StatusCode = 400;
            return;
        }

        using var webSocket = await context.WebSockets.AcceptWebSocketAsync();

        // --- Ticket handshake: first frame must be "<publicId>.<iv>.<ct>".
        // Unauthenticated connections never touch the room roster; the blob is
        // validated by the main server, which consumes it atomically.
        // Tiny buffer: unauthenticated clients cannot pin large allocations. ---
        var authBuffer = ArrayPool<byte>.Shared.Rent(2048);
        string ticket;
        using (var authCts = new CancellationTokenSource(TimeSpan.FromSeconds(5)))
        {
            WebSocketReceiveResult firstFrame;
            try
            {
                firstFrame = await webSocket.ReceiveAsync(new ArraySegment<byte>(authBuffer), authCts.Token);
            }
            catch (OperationCanceledException)
            {
                Console.WriteLine($"Voice ws: ticket timeout, room={roomId}");
                await SafeClose(webSocket, "Ticket timeout");
                return;
            }
            catch (WebSocketException)
            {
                return;
            }

            if (firstFrame.MessageType != WebSocketMessageType.Text || !firstFrame.EndOfMessage || firstFrame.Count == 0)
            {
                await SafeClose(webSocket, "Invalid ticket");
                return;
            }
            ticket = System.Text.Encoding.UTF8.GetString(authBuffer, 0, firstFrame.Count);
        }
        ArrayPool<byte>.Shared.Return(authBuffer);

        string[] ticketParts = ticket.Split('.');
        bool frameShapeOk = ticketParts.Length == 3
            && ticketParts[0].Length > 0 && ticketParts[0].Length <= 64
            && ticketParts[1].Length == 16
            && ticketParts[2].Length > 0 && ticketParts[2].Length <= 128;

        Console.WriteLine($"Voice ws: validating {ticketParts[0][..Math.Min(8, ticketParts[0].Length)]} room={roomId}");
        if (!frameShapeOk ||
            !await Utils.AuthorizeConnection(ticketParts[0], ticketParts[1], ticketParts[2], roomId, clientId))
        {
            Console.WriteLine($"Voice ws: invalid ticket, room={roomId}");
            await SafeClose(webSocket, "Invalid ticket");
            return;
        }

        Session currentSession = null;
        var roomClients = Rooms.GetOrAdd(roomId, _ => new ConcurrentDictionary<ulong, Session>());

        if (roomClients.TryGetValue(clientId, out var oldSession) && oldSession.Socket.State == WebSocketState.Open)
        {
            try
            {
                await oldSession.Socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Kicked", CancellationToken.None);
            }
            catch
            {
            }
        }

        currentSession = new Session(webSocket);
        roomClients[clientId] = currentSession;

        _ = Task.Run(async () =>
        {
            await foreach (var packet in currentSession.SendQueue.Reader.ReadAllAsync())
            {
                try
                {
                    if (currentSession.Socket.State == WebSocketState.Open)
                    {
                        await currentSession.Socket.SendAsync(new ArraySegment<byte>(packet.Buffer, 0, packet.Length),
                            WebSocketMessageType.Binary, true, CancellationToken.None);

                        Interlocked.Add(ref TotalBytesSent, (ulong)packet.Length);
                    }
                }
                catch
                {
                }
                finally
                {
                    ArrayPool<byte>.Shared.Return(packet.Buffer);
                }
            }
        });

        // Big buffer only for authenticated media traffic.
        var receiveBuffer = ArrayPool<byte>.Shared.Rent(1024 * 1024);
        try
        {
            while (webSocket.State == WebSocketState.Open)
            {
                int totalBytes = 0;
                WebSocketReceiveResult result;

                do
                {
                    result = await webSocket.ReceiveAsync(
                        new ArraySegment<byte>(receiveBuffer, totalBytes, receiveBuffer.Length - totalBytes),
                        CancellationToken.None);

                    totalBytes += result.Count;
                    Interlocked.Add(ref TotalBytesReceived, (ulong)result.Count);

                } while (!result.EndOfMessage);

                if (result.MessageType == WebSocketMessageType.Close) break;
                if (result.MessageType == WebSocketMessageType.Binary)
                {
                    int payloadLength = totalBytes + 8;

                    foreach (var (id, client) in roomClients)
                    {
                        if (id != clientId && client.Socket.State == WebSocketState.Open)
                        {
                            byte[] payload = ArrayPool<byte>.Shared.Rent(payloadLength);
                            BitConverter.TryWriteBytes(payload, clientId);
                            Array.Copy(receiveBuffer, 0, payload, 8, totalBytes);

                            if (!client.SendQueue.Writer.TryWrite((payload, payloadLength)))
                                ArrayPool<byte>.Shared.Return(payload);
                        }
                    }
                }
            }
        }
        catch (WebSocketException)
        {
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(receiveBuffer);
            if (roomClients.TryGetValue(clientId, out var session) && session == currentSession)
            {
                roomClients.TryRemove(clientId, out _);
                session.SendQueue.Writer.Complete();
            }

            if (roomClients.IsEmpty) Rooms.TryRemove(roomId, out _);
        }
    }

    private static async Task SafeClose(WebSocket socket, string reason)
    {
        try
        {
            if (socket.State == WebSocketState.Open)
                await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, reason, CancellationToken.None);
        }
        catch
        {
        }
    }
}