ShrinkNetwork/Runtime/Transport/Tcp/ShrinkTcpClientTransport.cs

#nullable enable
using System;
using System.Buffers.Binary;
using System.IO;
using System.Net.Sockets;
using System.Net.Security;
using System.Security.Cryptography.X509Certificates;
using System.Threading;
using Cysharp.Threading.Tasks;

namespace ShrinkNetwork
{
    public sealed class ShrinkTcpClientTransport : IShrinkNetworkAsyncTransport
    {
        private readonly string _host;
        private readonly int _port;
        private readonly long _sessionId;
        private readonly int _maxPacketSize;
        private readonly ShrinkTcpTlsOptions? _tlsOptions;
        private readonly SemaphoreSlim _sendLock = new(1, 1);

        private TcpClient? _client;
        private Stream? _stream;
        private CancellationTokenSource? _cts;
        private int _started;

        public ShrinkTcpClientTransport(string host, int port, long sessionId = 1, int maxPacketSize = 64 * 1024,
            ShrinkTcpTlsOptions? tlsOptions = null)
        {
            if (string.IsNullOrWhiteSpace(host))
                throw new ArgumentException("Host cannot be empty.", nameof(host));
            if (port <= 0 || port > 65535)
                throw new ArgumentOutOfRangeException(nameof(port));
            if (maxPacketSize <= 0)
                throw new ArgumentOutOfRangeException(nameof(maxPacketSize));

            _host = host.Trim();
            _port = port;
            _sessionId = sessionId;
            _maxPacketSize = maxPacketSize;
            _tlsOptions = tlsOptions?.Clone();
        }

        public bool IsStarted => Volatile.Read(ref _started) == 1;

        public event Action<ShrinkNetworkTransportEvent>? OnEvent;

        public void Start()
        {
            if (Interlocked.Exchange(ref _started, 1) == 1)
                return;

            _cts = new CancellationTokenSource();
            ConnectAsync(_cts.Token).Forget();
        }

        public void Stop()
        {
            if (Interlocked.Exchange(ref _started, 0) == 0)
                return;

            _cts?.Cancel();

            try
            {
                _stream?.Dispose();
            }
            catch
            {
            }

            try
            {
                _client?.Close();
            }
            catch
            {
            }

            _client = null;
            _stream = null;
        }

        public void Send(long sessionId, byte[] packetData)
        {
            SendAsync(sessionId, packetData).Forget();
        }

        public UniTask SendAsync(long sessionId, byte[] packetData)
        {
            if (!IsStarted)
                throw new InvalidOperationException("Transport is not started.");
            if (sessionId != _sessionId)
                throw new InvalidOperationException($"Unsupported session id {sessionId}. This transport only supports {_sessionId}.");
            if (_client == null || !_client.Connected || _stream == null)
                throw new InvalidOperationException("TCP client is not connected.");
            if ((packetData?.Length ?? 0) > _maxPacketSize)
                throw new InvalidOperationException($"TCP packet is too large. Size={(packetData?.Length ?? 0)}, Limit={_maxPacketSize}.");

            return SendInternalAsync(_stream, _sendLock, packetData ?? Array.Empty<byte>(),
                _cts?.Token ?? CancellationToken.None);
        }

        private async UniTaskVoid ConnectAsync(CancellationToken cancellationToken)
        {
            try
            {
                _client = new TcpClient();
                _client.NoDelay = true;
                await _client.ConnectAsync(_host, _port);
                var stream = await CreateConnectedStreamAsync(_client, cancellationToken);
                _stream = stream;

                if (!IsStarted)
                {
                    stream.Dispose();
                    _client.Close();
                    return;
                }

                var remote = _client.Client.RemoteEndPoint?.ToString() ?? $"{_host}:{_port}";
                OnEvent?.Invoke(ShrinkNetworkTransportEvent.Connected(_sessionId, remote));
                ReceiveLoopAsync(_client, stream, cancellationToken).Forget();
            }
            catch (Exception ex)
            {
                ShrinkNetworkLogger.Exception(ex);
                ShrinkNetworkLogger.Error($"[ShrinkNetwork] TCP connect failed: {_host}:{_port} {ex.Message}");
                Stop();
            }
        }

        private async UniTask<Stream> CreateConnectedStreamAsync(TcpClient client, CancellationToken cancellationToken)
        {
            Stream stream = client.GetStream();
            var tlsOptions = _tlsOptions;
            if (tlsOptions == null || !tlsOptions.Enabled)
                return stream;

            var targetHost = string.IsNullOrWhiteSpace(tlsOptions.TargetHost) ? _host : tlsOptions.TargetHost.Trim();
            var sslStream = new SslStream(stream, false, (_, _, _, errors) =>
            {
                if (tlsOptions.AllowInvalidServerCertificate)
                    return true;

                return errors == SslPolicyErrors.None;
            });
            await sslStream.AuthenticateAsClientAsync(targetHost, null, tlsOptions.EnabledProtocols,
                tlsOptions.CheckCertificateRevocation);
            return sslStream;
        }

        private async UniTaskVoid ReceiveLoopAsync(TcpClient client, Stream stream, CancellationToken cancellationToken)
        {
            try
            {
                var lengthBuffer = new byte[4];

                while (!cancellationToken.IsCancellationRequested)
                {
                    await ReadExactlyAsync(stream, lengthBuffer, cancellationToken);
                    var length = BinaryPrimitives.ReadInt32LittleEndian(lengthBuffer);
                    if (length <= 0)
                        throw new InvalidDataException($"Invalid packet length: {length}");
                    if (length > _maxPacketSize)
                        throw new InvalidDataException($"TCP packet length exceeded limit. Length={length}, Limit={_maxPacketSize}");

                    var payload = new byte[length];
                    await ReadExactlyAsync(stream, payload, cancellationToken);
                    OnEvent?.Invoke(ShrinkNetworkTransportEvent.Packet(_sessionId, payload));
                }
            }
            catch (OperationCanceledException)
            {
            }
            catch (IOException ex) when (IsExpectedDisconnect(ex))
            {
                ShrinkNetworkLogger.Warn($"[ShrinkNetwork] TCP receive loop ended: {ex.Message}");
            }
            catch (SocketException ex) when (IsExpectedDisconnect(ex))
            {
                ShrinkNetworkLogger.Warn($"[ShrinkNetwork] TCP receive loop ended: {ex.Message}");
            }
            catch (Exception ex)
            {
                ShrinkNetworkLogger.Exception(ex);
                ShrinkNetworkLogger.Warn($"[ShrinkNetwork] TCP receive loop ended: {ex.Message}");
            }
            finally
            {
                var remote = client.Client.RemoteEndPoint?.ToString() ?? $"{_host}:{_port}";
                OnEvent?.Invoke(ShrinkNetworkTransportEvent.Disconnected(_sessionId, remote));
                Stop();
            }
        }

        private static bool IsExpectedDisconnect(IOException ex)
        {
            if (string.Equals(ex.Message, "Remote closed.", StringComparison.Ordinal))
                return true;

            return ex.InnerException is SocketException socketException && IsExpectedDisconnect(socketException);
        }

        private static bool IsExpectedDisconnect(SocketException ex)
        {
            return ex.SocketErrorCode == SocketError.ConnectionReset ||
                   ex.SocketErrorCode == SocketError.ConnectionAborted ||
                   ex.SocketErrorCode == SocketError.Shutdown;
        }

        private static async UniTask SendInternalAsync(Stream stream, SemaphoreSlim sendLock, byte[] packetData,
            CancellationToken cancellationToken)
        {
            await sendLock.WaitAsync(cancellationToken);
            try
            {
                var header = new byte[4];
                BinaryPrimitives.WriteInt32LittleEndian(header, packetData.Length);
                await stream.WriteAsync(header, cancellationToken);
                if (packetData.Length > 0)
                    await stream.WriteAsync(packetData, cancellationToken);
                await stream.FlushAsync(cancellationToken);
            }
            finally
            {
                sendLock.Release();
            }
        }

        private static async UniTask ReadExactlyAsync(Stream stream, byte[] buffer, CancellationToken cancellationToken)
        {
            var offset = 0;
            while (offset < buffer.Length)
            {
                var read = await stream.ReadAsync(buffer.AsMemory(offset, buffer.Length - offset), cancellationToken);
                if (read <= 0)
                    throw new IOException("Remote closed.");
                offset += read;
            }
        }
    }
}