h-school/src/HSchool.Server/Net/GameSocketHandler.cs

using System.Buffers;
using System.Net.WebSockets;
using HSchool.Protocol;
using HSchool.Server.Game;
using HSchool.Simulation;

namespace HSchool.Server.Net;

/// <summary>
/// Drives one WebSocket connection: version handshake, then the receive loop that turns frames
/// into commands. Everything it reads from the wire is untrusted, so frames are validated before
/// anything reaches a school worker.
/// </summary>
internal sealed class GameSocketHandler(
    ClientRegistry clients,
    GameCommandQueue commands,
    GameLoopService loop,
    GameMetrics metrics,
    ILogger<GameSocketHandler> logger)
{
    private static readonly TimeSpan HandshakeTimeout = TimeSpan.FromSeconds(5);

    public async Task HandleAsync(WebSocket socket, string userName, CancellationToken cancellationToken)
    {
        var client = clients.Add(socket);
        var buffer = ArrayPool<byte>.Shared.Rent(ProtocolConstants.MaxMessageSize);
        using var connectionCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);

        Task? sendLoop = null;
        metrics.ClientConnected();

        try
        {
            if (!clients.TryClaimUserName(client, userName))
            {
                await CloseAsync(
                    socket,
                    WebSocketCloseStatus.PolicyViolation,
                    "Name is already online.",
                    cancellationToken).ConfigureAwait(false);
                return;
            }

            using var handshakeCts = CancellationTokenSource.CreateLinkedTokenSource(connectionCts.Token);
            handshakeCts.CancelAfter(HandshakeTimeout);

            var helloLength = await ReceiveFrameAsync(socket, buffer, handshakeCts.Token).ConfigureAwait(false);
            if (helloLength <= 0)
            {
                return;
            }

            var hello = ProtocolCodec.ReadHello(buffer.AsSpan(0, helloLength));
            if (hello.ProtocolVersion != ProtocolConstants.Version)
            {
                logger.LogWarning(
                    "Rejecting client {PlayerId}: protocol v{ClientVersion}, server speaks v{ServerVersion}.",
                    client.PlayerId,
                    hello.ProtocolVersion,
                    ProtocolConstants.Version);

                await CloseAsync(
                    socket,
                    WebSocketCloseStatus.ProtocolError,
                    $"Protocol v{ProtocolConstants.Version} required.",
                    cancellationToken).ConfigureAwait(false);
                return;
            }

            client.Locale = hello.Locale;

            await SendWelcomeAsync(socket, connectionCts.Token).ConfigureAwait(false);
            client.MarkReady();

            // From here on every outbound frame goes through the outbox, so there is
            // exactly one writer on the socket.
            sendLoop = client.RunSendLoopAsync(connectionCts.Token);

            await ReceiveLoopAsync(client, buffer, connectionCts.Token).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // Client went away or the host is shutting down.
        }
        catch (ProtocolException ex)
        {
            logger.LogWarning(ex, "Malformed frame from client {PlayerId}.", client.PlayerId);
            await CloseAsync(socket, WebSocketCloseStatus.InvalidPayloadData, "Malformed frame.", CancellationToken.None)
                .ConfigureAwait(false);
        }
        catch (WebSocketException ex)
        {
            logger.LogDebug(ex, "Connection {PlayerId} dropped.", client.PlayerId);
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
            clients.Remove(client.PlayerId);
            client.CompleteOutbox();
            metrics.ClientDisconnected();

            // Read after the removal above: an open that lands later finds no client and is
            // dropped, so this is the last chance to see the school this connection was watching.
            if (client.OpenSchoolId is { } watchedSchoolId)
            {
                commands.Enqueue(new GameCommand.CloseSchool(client.PlayerId, watchedSchoolId));
            }

            // Cancel first, then wait. The other order hung here whenever the send loop sat inside
            // SendAsync on a half-open connection: nothing was left to cancel it.
            await connectionCts.CancelAsync().ConfigureAwait(false);

            if (sendLoop is not null)
            {
                try
                {
                    await sendLoop.ConfigureAwait(false);
                }
                catch (Exception ex) when (ex is OperationCanceledException or WebSocketException)
                {
                    // Expected while tearing the connection down.
                }
            }
        }
    }

    private async Task ReceiveLoopAsync(GameClient client, byte[] buffer, CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            var length = await ReceiveFrameAsync(client.Socket, buffer, cancellationToken).ConfigureAwait(false);
            if (length <= 0)
            {
                return;
            }

            var frame = buffer.AsSpan(0, length);
            switch (ProtocolCodec.PeekMessageType(frame))
            {
                case MessageType.ClientPing:
                    SendPong(client, ProtocolCodec.ReadPing(frame).ClientTimeMs);
                    break;

                case MessageType.ClientOpenSchool:
                    var open = ProtocolCodec.ReadOpenSchool(frame);
                    commands.Enqueue(new GameCommand.OpenSchool(client.PlayerId, open.SchoolId));
                    break;

                case MessageType.ClientCloseSchool:
                    if (client.OpenSchoolId is { } openSchoolId)
                    {
                        commands.Enqueue(new GameCommand.CloseSchool(client.PlayerId, openSchoolId));
                    }

                    break;

                case MessageType.ClientSetRunning:
                    var setRunning = ProtocolCodec.ReadSetRunning(frame);
                    if (TryNormalizedUser(client, out var runningUser))
                    {
                        commands.Enqueue(new GameCommand.SetRunning(client.PlayerId, setRunning.Running, runningUser));
                    }

                    break;

                case MessageType.ClientSetSpeed:
                    var setSpeed = ProtocolCodec.ReadSetSpeed(frame);
                    if (TryNormalizedUser(client, out var speedUser))
                    {
                        commands.Enqueue(new GameCommand.SetSpeed(client.PlayerId, setSpeed.SpeedIndex, speedUser));
                    }

                    break;

                case MessageType.ClientSkipEmpty:
                    if (TryNormalizedUser(client, out var skipUser))
                    {
                        commands.Enqueue(new GameCommand.SkipEmpty(client.PlayerId, skipUser));
                    }

                    break;

                case MessageType.ClientDismissNotice:
                    var dismiss = ProtocolCodec.ReadDismissNotice(frame);
                    commands.Enqueue(new GameCommand.DismissNotice(client.PlayerId, dismiss.Id));
                    break;

                default:
                    logger.LogDebug(
                        "Ignoring unexpected frame 0x{MessageType:X2} from client {PlayerId}.",
                        frame[0],
                        client.PlayerId);
                    break;
            }
        }
    }

    /// <summary>Reads one whole message. Returns 0 on close, -1 on an oversized or non-binary frame.</summary>
    private async Task<int> ReceiveFrameAsync(WebSocket socket, byte[] buffer, CancellationToken cancellationToken)
    {
        var offset = 0;

        while (true)
        {
            var result = await socket
                .ReceiveAsync(new ArraySegment<byte>(buffer, offset, buffer.Length - offset), cancellationToken)
                .ConfigureAwait(false);

            if (result.MessageType == WebSocketMessageType.Close)
            {
                return 0;
            }

            if (result.MessageType != WebSocketMessageType.Binary)
            {
                logger.LogDebug("Dropping non-binary frame.");
                return -1;
            }

            offset += result.Count;

            if (result.EndOfMessage)
            {
                return offset;
            }

            if (offset >= buffer.Length)
            {
                logger.LogWarning("Frame exceeds {Limit} bytes; closing.", ProtocolConstants.MaxMessageSize);
                await CloseAsync(socket, WebSocketCloseStatus.MessageTooBig, "Frame too large.", cancellationToken)
                    .ConfigureAwait(false);
                return -1;
            }
        }
    }

    private async Task SendWelcomeAsync(WebSocket socket, CancellationToken cancellationToken)
    {
        var options = loop.Options;
        var frame = new byte[ProtocolCodec.MaxFrameSize];
        var length = ProtocolCodec.WriteWelcome(
            frame,
            new ServerWelcomeMessage(ProtocolConstants.Version, (byte)options.TickRate, (byte)options.MaxSchools));

        await socket
            .SendAsync(frame.AsMemory(0, length), WebSocketMessageType.Binary, endOfMessage: true, cancellationToken)
            .ConfigureAwait(false);
    }

    private void SendPong(GameClient client, long clientTimeMs)
    {
        var frame = new byte[ProtocolCodec.MaxFrameSize];
        var length = ProtocolCodec.WritePong(frame, new ServerPongMessage(clientTimeMs, loop.CurrentTick));
        client.TrySend(frame.AsMemory(0, length));
    }

    private static bool TryNormalizedUser(GameClient client, out string normalized) =>
        SchoolNames.TryNormalize(client.UserName, out normalized);

    private static async Task CloseAsync(
        WebSocket socket,
        WebSocketCloseStatus status,
        string description,
        CancellationToken cancellationToken)
    {
        if (socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
        {
            try
            {
                await socket.CloseAsync(status, description, cancellationToken).ConfigureAwait(false);
            }
            catch (WebSocketException)
            {
                // The peer may already be gone; nothing left to do.
            }
        }
    }
}