LobbyServer/LobbyClient/TcpClient.cs

using System;
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;

        TcpClient? tcpClient;
        NetworkStream? networkStream;
        CancellationTokenSource? cancellationTokenSource = new CancellationTokenSource();
        bool running = false;

        internal async Task Connect(string host, int port)
        {
            bool wasError = false;
            string error = string.Empty;

            try
            {
                cancellationTokenSource!.Token.ThrowIfCancellationRequested();

                running = true; 
                tcpClient = new TcpClient();
                using (cancellationTokenSource!.Token.Register(() => { tcpClient.Close(); }))
                {
                    await tcpClient.ConnectAsync(host, port);
                }
                networkStream = tcpClient.GetStream();

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

                int currentOffset = 0;
                int currentMessageRemainingLength = 0;
                int currentMessageLength = 0;
                bool validMessage = true;
                int currentReadOffset = 0;
                bool offsetSizeInt = false;
                Connected?.Invoke();

                while (running)
                {
                    int copyOffset = 0;
                    int receivedBytes = currentReadOffset;

                    if (currentReadOffset < 4)
                    {
                        receivedBytes += await networkStream.ReadAsync(buffer.Slice(currentReadOffset), cancellationTokenSource.Token) + currentReadOffset;
                    }

                    if (receivedBytes == 0 && running && !cancellationTokenSource.Token.IsCancellationRequested)
                    {
                        throw new Exception("Connection lost!");
                    }

                    if (receivedBytes > 3 || (currentMessageRemainingLength > 0 && receivedBytes > currentMessageRemainingLength))
                    {
                        currentReadOffset = 0;
                        if (currentMessageLength == 0)
                        {
                            currentMessageRemainingLength = BitConverter.ToInt32(buffer.Span);
                            currentMessageLength = currentMessageRemainingLength;
                            receivedBytes -= 4;
                            copyOffset += 4;
                            offsetSizeInt = true;
                        }
                        else
                            offsetSizeInt = false;

                        var receivedCount = Math.Min(receivedBytes, currentMessageRemainingLength);
                        receivedBytes -= receivedCount;
                        copyOffset += receivedCount;

                        if (validMessage && currentOffset + receivedCount > 0)
                        {
                            if (currentOffset + receivedCount < target.Length)
                                buffer.Slice(offsetSizeInt ? 4 : 0, receivedCount).CopyTo(target.Slice(currentOffset));
                            else
                                validMessage = false;
                        }

                        currentOffset += receivedCount;
                        currentMessageRemainingLength -= receivedCount;

                        if (currentMessageRemainingLength <= 0)
                        {
                            if (validMessage)
                                DataReceived?.Invoke(currentMessageLength, target);

                            if (receivedBytes > 0)
                            {
                                buffer.Slice(copyOffset, receivedBytes).CopyTo(buffer);
                                currentReadOffset += receivedBytes;
                            }

                            currentOffset = 0;
                            currentMessageLength = 0;
                            currentMessageRemainingLength = 0;
                            validMessage = true;
                        }
                    }
                    else if (receivedBytes > 0)
                    {
                        currentReadOffset += receivedBytes;
                    }
                }
            }
            catch(Exception e)
            {
                error = e.Message;
            }
            finally 
            {
                wasError = running;

                running = false;
                
                networkStream?.Dispose();
                tcpClient?.Dispose();

                tcpClient = null;
                networkStream = null;

                Disconnected?.Invoke(!wasError, error);
            }
        }

        internal async Task Send(byte[] buffer, int offset, int count)
        {
            try
            {
                if (running && networkStream != null)
                {
                    await networkStream.WriteAsync(BitConverter.GetBytes(count - offset), 0, 4, cancellationTokenSource!.Token);
                    await networkStream.WriteAsync(buffer, offset, count, cancellationTokenSource!.Token);
                }
            }
            catch { }
        }

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

        public void Dispose()
        {
            cancellationTokenSource?.Cancel();
            tcpClient?.Dispose();
            cancellationTokenSource?.Dispose();
        }
    }
}