mrgameeng/src/MrGameEng.Net/WebSocketServer.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using MrGameEng.Core;

namespace MrGameEng.Net;

/// <summary>
/// A dependency-free WebSocket server (RFC 6455 over <see cref="TcpListener"/>) for
/// dedicated servers. Accepting and reading happen on background tasks; the simulation
/// drains new connections with <see cref="TryAcceptConnection"/> and reads messages by
/// polling each connection — nothing here touches the ECS world from another thread.
/// Binary messages only; pings are answered automatically.
/// </summary>
public sealed class WebSocketServer : IDisposable
{
    /// <summary>The port the server listens on.</summary>
    public int Port { get; }

    /// <summary>Snapshot of currently open connections.</summary>
    public IReadOnlyList<INetConnection> Connections
    {
        get
        {
            lock (_connections)
            {
                return _connections.Where(c => c.IsOpen).Cast<INetConnection>().ToArray();
            }
        }
    }

    private readonly TcpListener _listener;
    private readonly List<ServerConnection> _connections = [];
    private readonly ConcurrentQueue<ServerConnection> _accepted = new();
    private readonly CancellationTokenSource _shutdown = new();
    private int _nextConnectionId;
    private bool _started;

    /// <summary>Creates a server for <paramref name="port"/> on all interfaces. Call <see cref="Start"/> to listen.</summary>
    public WebSocketServer(int port)
    {
        Port = port;
        _listener = new TcpListener(IPAddress.Any, port);
    }

    /// <summary>Starts listening and accepting connections in the background.</summary>
    public void Start()
    {
        if (_started)
        {
            return;
        }

        _started = true;
        _listener.Start();
        Task.Run(AcceptLoop);
        Log.Info($"WebSocketServer listening on port {Port}");
    }

    /// <summary>Dequeues a connection that completed its handshake since the last call.</summary>
    public bool TryAcceptConnection(out INetConnection connection)
    {
        if (_accepted.TryDequeue(out var accepted))
        {
            connection = accepted;
            return true;
        }

        connection = null!;
        return false;
    }

    /// <summary>Stops listening and closes every connection.</summary>
    public void Dispose()
    {
        _shutdown.Cancel();
        _listener.Stop();
        lock (_connections)
        {
            foreach (var connection in _connections)
            {
                connection.Close();
            }

            _connections.Clear();
        }
    }

    private async Task AcceptLoop()
    {
        while (!_shutdown.IsCancellationRequested)
        {
            TcpClient client;
            try
            {
                client = await _listener.AcceptTcpClientAsync(_shutdown.Token);
            }
            catch (Exception) when (_shutdown.IsCancellationRequested)
            {
                return;
            }
            catch (Exception exception)
            {
                Log.Warning($"WebSocketServer accept failed: {exception.Message}");
                continue;
            }

            _ = Task.Run(() => Handshake(client));
        }
    }

    private void Handshake(TcpClient client)
    {
        try
        {
            client.NoDelay = true;
            var stream = client.GetStream();
            if (!WebSocketProtocol.TryReadHandshakeKey(stream, out var key))
            {
                client.Dispose();
                return;
            }

            WebSocketProtocol.WriteHandshakeResponse(stream, key);
            var connection = new ServerConnection(
                Interlocked.Increment(ref _nextConnectionId),
                client
            );
            lock (_connections)
            {
                _connections.RemoveAll(c => !c.IsOpen);
                _connections.Add(connection);
            }

            _accepted.Enqueue(connection);
            connection.StartReceiveLoop();
            Log.Info($"WebSocketServer: connection #{connection.Id} accepted");
        }
        catch (Exception exception)
        {
            Log.Warning($"WebSocketServer handshake failed: {exception.Message}");
            client.Dispose();
        }
    }

    private sealed class ServerConnection : INetConnection
    {
        public int Id { get; }
        public bool IsOpen => !_closed;

        private readonly TcpClient _client;
        private readonly NetworkStream _stream;
        private readonly ConcurrentQueue<byte[]> _inbox = new();
        private readonly object _sendLock = new();
        private volatile bool _closed;

        internal ServerConnection(int id, TcpClient client)
        {
            Id = id;
            _client = client;
            _stream = client.GetStream();
        }

        internal void StartReceiveLoop() => Task.Run(ReceiveLoop);

        public void Send(ReadOnlySpan<byte> message)
        {
            if (_closed)
            {
                return;
            }

            var frame = WebSocketProtocol.EncodeFrame(message, WebSocketOpcode.Binary);
            try
            {
                lock (_sendLock)
                {
                    _stream.Write(frame, 0, frame.Length);
                }
            }
            catch (Exception)
            {
                Close();
            }
        }

        public bool TryReceive(out byte[] message) => _inbox.TryDequeue(out message!);

        public void Close()
        {
            if (_closed)
            {
                return;
            }

            _closed = true;
            try
            {
                _client.Dispose();
            }
            catch (Exception)
            {
                // соединение уже мертво — закрытие не должно бросать
            }
        }

        private void ReceiveLoop()
        {
            var pending = new List<byte>();
            var pendingOpcode = WebSocketOpcode.Binary;
            try
            {
                while (!_closed)
                {
                    if (
                        !WebSocketProtocol.TryReadFrame(
                            _stream,
                            out var opcode,
                            out var fin,
                            out var payload
                        )
                    )
                    {
                        break;
                    }

                    switch (opcode)
                    {
                        case WebSocketOpcode.Ping:
                            lock (_sendLock)
                            {
                                var pong = WebSocketProtocol.EncodeFrame(
                                    payload,
                                    WebSocketOpcode.Pong
                                );
                                _stream.Write(pong, 0, pong.Length);
                            }

                            continue;
                        case WebSocketOpcode.Pong:
                            continue;
                        case WebSocketOpcode.Close:
                            lock (_sendLock)
                            {
                                var close = WebSocketProtocol.EncodeFrame(
                                    [],
                                    WebSocketOpcode.Close
                                );
                                _stream.Write(close, 0, close.Length);
                            }

                            return;
                    }

                    if (opcode != WebSocketOpcode.Continuation)
                    {
                        pendingOpcode = opcode;
                        pending.Clear();
                    }

                    pending.AddRange(payload);
                    if (fin && pendingOpcode == WebSocketOpcode.Binary)
                    {
                        _inbox.Enqueue(pending.ToArray());
                        pending.Clear();
                    }
                }
            }
            catch (Exception)
            {
                // обрыв соединения — штатный путь завершения цикла
            }
            finally
            {
                Close();
            }
        }
    }
}