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();
        string secret = context.Request.Query["secret"].ToString();

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

        using var webSocket = await context.WebSockets.AcceptWebSocketAsync();
        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
            {
            }
        }

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

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