LobbyServer/LobbyClient/TcpClient.cs

using LobbyServerDto;
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 BufferRental bufferRental = new BufferRental(MaxMessageSize);
        private async Task ReadExact(NetworkStream stream, Memory<byte> buffer, int length, CancellationToken token)
        {
            int offset = 0;

            while (offset < length)
            {
                int read = await stream.ReadAsync(buffer.Slice(offset), token);
                if (read == 0)
                    throw new EndOfStreamException();

                offset += read;
            }
        }

        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];

                while (running && !token.IsCancellationRequested)
                {
                    await ReadExact(networkStream, buffer, HeaderSize, token);                    

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

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

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

                    if (length == 0)
                        continue;

                    await ReadExact(networkStream, buffer, length, token);

                    DataReceived?.Invoke(length, buffer.Slice(0, length));
                }                         
            }
            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;

            try
            {
                await networkStream.WriteAsync(BitConverter.GetBytes(0), 0, 4, token);
            }
            catch
            {
            }
        }

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

                if (count < 0 || count > MaxMessageSize)
                    throw new ArgumentOutOfRangeException(nameof(count));

                if (offset < 0 || offset + count > buffer.Length)
                    throw new ArgumentOutOfRangeException(nameof(offset));

                BitConverter.GetBytes(count).CopyTo(frame, 0);
                Buffer.BlockCopy(buffer, offset, frame, HeaderSize, count);

                await networkStream.WriteAsync(frame, 0, HeaderSize + count, cancellationTokenSource.Token);
            }
            catch
            {
            }
            finally
            {
                bufferRental.Return(frame);
            }
        }      

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