LobbyServer/LobbyClient/TcpClient.cs

using System;
using System.IO;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Lobbies
{
    internal class TcpLobbbyClient : IDisposable
    {
        internal delegate void DataReceivedEventArgs(int bytes, Memory<byte> data);
        internal event DataReceivedEventArgs? DataReceived;

        internal delegate void DisconnectedEventArgs(bool clean, string error);
        internal event DisconnectedEventArgs? Disconnected;

        internal delegate void ConnectedEventArgs();
        internal event ConnectedEventArgs? Connected;

        private const int HeaderSize = 4;
        private const int MaxMessageSize = 4096;
        private static readonly TimeSpan PingInterval = TimeSpan.FromSeconds(10);

        private TcpClient? tcpClient;
        private NetworkStream? networkStream;
        private CancellationTokenSource? cancellationTokenSource = new CancellationTokenSource();
        private bool running = false;
        private readonly SemaphoreSlim sendLock = new SemaphoreSlim(1, 1);

        internal async Task Connect(string host, int port)
        {
            bool cleanDisconnect = true;
            string error = string.Empty;
            Task? pingTask = null;

            try
            {
                if (cancellationTokenSource == null || cancellationTokenSource.IsCancellationRequested)
                    cancellationTokenSource = new CancellationTokenSource();

                var token = cancellationTokenSource.Token;

                token.ThrowIfCancellationRequested();

                running = true;
                tcpClient = new TcpClient();

                using (token.Register(() =>
                {
                    try { tcpClient.Close(); } catch { }
                }))
                {
                    await tcpClient.ConnectAsync(host, port);
                }

                networkStream = tcpClient.GetStream();

                Connected?.Invoke();

                pingTask = Task.Run(() => PingLoop(token), token);

                Memory<byte> buffer = new byte[MaxMessageSize];
                Memory<byte> target = new byte[MaxMessageSize];

                int bufferedBytes = 0;
                int currentMessageLength = -1;
                int currentMessageOffset = 0;

                while (running && !token.IsCancellationRequested)
                {
                    if (bufferedBytes == buffer.Length)
                        throw new InvalidDataException("Receive buffer overflow.");

                    int bytesRead = await networkStream.ReadAsync(buffer.Slice(bufferedBytes), token);
                    if (bytesRead == 0)
                        break;

                    bufferedBytes += bytesRead;

                    while (true)
                    {
                        if (currentMessageLength < 0)
                        {
                            if (bufferedBytes < HeaderSize)
                                break;

                            currentMessageLength = BitConverter.ToInt32(buffer.Span.Slice(0, HeaderSize));

                            if (currentMessageLength < 0)
                                throw new InvalidDataException("Negative message length received.");

                            if (currentMessageLength > MaxMessageSize)
                                throw new InvalidDataException($"Message too large: {currentMessageLength} > {MaxMessageSize}.");

                            if (bufferedBytes > HeaderSize)
                                buffer.Slice(HeaderSize, bufferedBytes - HeaderSize).CopyTo(buffer);

                            bufferedBytes -= HeaderSize;
                            currentMessageOffset = 0;

                            if (currentMessageLength == 0)
                            {
                                currentMessageLength = -1;
                                continue;
                            }
                        }

                        int remainingMessageBytes = currentMessageLength - currentMessageOffset;
                        if (remainingMessageBytes <= 0)
                        {
                            DataReceived?.Invoke(currentMessageLength, target.Slice(0, currentMessageLength));
                            currentMessageLength = -1;
                            currentMessageOffset = 0;
                            continue;
                        }

                        if (bufferedBytes == 0)
                            break;

                        int chunkSize = Math.Min(bufferedBytes, remainingMessageBytes);

                        buffer.Slice(0, chunkSize).CopyTo(target.Slice(currentMessageOffset));
                        currentMessageOffset += chunkSize;

                        if (bufferedBytes > chunkSize)
                            buffer.Slice(chunkSize, bufferedBytes - chunkSize).CopyTo(buffer);

                        bufferedBytes -= chunkSize;

                        if (currentMessageOffset == currentMessageLength)
                        {
                            DataReceived?.Invoke(currentMessageLength, target.Slice(0, currentMessageLength));
                            currentMessageLength = -1;
                            currentMessageOffset = 0;
                        }
                    }
                }
            }
            catch (OperationCanceledException)
            {
                cleanDisconnect = true;
            }
            catch (Exception e)
            {
                cleanDisconnect = false;
                error = e.Message;
            }
            finally
            {
                running = false;

                try { cancellationTokenSource?.Cancel(); } catch { }

                if (pingTask != null)
                {
                    try { await pingTask; } catch { }
                }

                try { networkStream?.Dispose(); } catch { }
                try { tcpClient?.Close(); } catch { }
                try { tcpClient?.Dispose(); } catch { }

                networkStream = null;
                tcpClient = null;

                Disconnected?.Invoke(cleanDisconnect, error);
            }
        }

        private async Task PingLoop(CancellationToken token)
        {
            while (running && !token.IsCancellationRequested)
            {
                await Task.Delay(PingInterval, token);

                if (!running || token.IsCancellationRequested)
                    break;

                await SendPing(token);
            }
        }

        private async Task SendPing(CancellationToken token)
        {
            if (!running || networkStream == null)
                return;

            await sendLock.WaitAsync(token);
            try
            {
                await networkStream.WriteAsync(BitConverter.GetBytes(0), 0, 4, token);
            }
            finally
            {
                sendLock.Release();
            }
        }

        internal async Task Send(byte[] buffer, int offset, int count)
        {
            try
            {
                if (!running || networkStream == null || cancellationTokenSource == null)
                    return;

                await sendLock.WaitAsync(cancellationTokenSource.Token);
                try
                {
                    await networkStream.WriteAsync(BitConverter.GetBytes(count), 0, 4, cancellationTokenSource.Token);
                    await networkStream.WriteAsync(buffer, offset, count, cancellationTokenSource.Token);
                }
                finally
                {
                    sendLock.Release();
                }
            }
            catch { }
        }

        internal void Stop()
        {
            running = false;
            cancellationTokenSource?.Cancel();
        }

        public void Dispose()
        {
            running = false;

            try { cancellationTokenSource?.Cancel(); } catch { }
            try { networkStream?.Dispose(); } catch { }
            try { tcpClient?.Close(); } catch { }
            try { tcpClient?.Dispose(); } catch { }
            try { cancellationTokenSource?.Dispose(); } catch { }

            networkStream = null;
            tcpClient = null;
            cancellationTokenSource = null;
        }
    }
}