#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? 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(), _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 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; } } } }