This commit is contained in:
@@ -0,0 +1,266 @@
|
||||
using System.Buffers.Binary;
|
||||
using System.Collections.Concurrent;
|
||||
using System.IO;
|
||||
using System.Net;
|
||||
using System.Net.Security;
|
||||
using System.Net.Sockets;
|
||||
using System.Security.Cryptography.X509Certificates;
|
||||
using System.Threading;
|
||||
using Cysharp.Threading.Tasks;
|
||||
|
||||
namespace ShrinkNetwork.ServerHost;
|
||||
|
||||
using ShrinkNetwork;
|
||||
|
||||
public sealed class TcpServerTransport : IShrinkNetworkAsyncTransport, IShrinkNetworkSessionControlTransport
|
||||
{
|
||||
private readonly ConcurrentDictionary<long, TcpClient> _clients = new();
|
||||
private readonly ConcurrentDictionary<long, SemaphoreSlim> _sendLocks = new();
|
||||
private readonly ConcurrentDictionary<long, Stream> _streams = new();
|
||||
private readonly TcpListener _listener;
|
||||
private readonly int _maxPacketSize;
|
||||
private readonly ShrinkTcpTlsOptions? _tlsOptions;
|
||||
private long _sessionIdGenerator;
|
||||
private CancellationTokenSource? _cts;
|
||||
private X509Certificate2? _serverCertificate;
|
||||
|
||||
public TcpServerTransport(IPAddress ipAddress, int port, int maxPacketSize = 64 * 1024,
|
||||
ShrinkTcpTlsOptions? tlsOptions = null)
|
||||
{
|
||||
if (maxPacketSize <= 0)
|
||||
throw new ArgumentOutOfRangeException(nameof(maxPacketSize));
|
||||
|
||||
_listener = new TcpListener(ipAddress, port);
|
||||
_maxPacketSize = maxPacketSize;
|
||||
_tlsOptions = tlsOptions?.Clone();
|
||||
}
|
||||
|
||||
public bool IsStarted { get; private set; }
|
||||
|
||||
public event Action<ShrinkNetworkTransportEvent>? OnEvent;
|
||||
|
||||
public void Start()
|
||||
{
|
||||
if (IsStarted)
|
||||
return;
|
||||
|
||||
IsStarted = true;
|
||||
_cts = new CancellationTokenSource();
|
||||
EnsureTlsCertificateLoaded();
|
||||
_listener.Start();
|
||||
_ = AcceptLoopAsync(_cts.Token);
|
||||
}
|
||||
|
||||
public void Stop()
|
||||
{
|
||||
if (!IsStarted)
|
||||
return;
|
||||
|
||||
IsStarted = false;
|
||||
_cts?.Cancel();
|
||||
_listener.Stop();
|
||||
|
||||
foreach (var pair in _clients)
|
||||
{
|
||||
pair.Value.Close();
|
||||
}
|
||||
|
||||
_clients.Clear();
|
||||
|
||||
foreach (var pair in _streams)
|
||||
pair.Value.Dispose();
|
||||
|
||||
_streams.Clear();
|
||||
|
||||
foreach (var pair in _sendLocks)
|
||||
pair.Value.Dispose();
|
||||
|
||||
_sendLocks.Clear();
|
||||
}
|
||||
|
||||
public void Send(long sessionId, byte[] packetData)
|
||||
{
|
||||
SendAsync(sessionId, packetData).Forget();
|
||||
}
|
||||
|
||||
public bool DisconnectSession(long sessionId, string? reason = null)
|
||||
{
|
||||
if (!_clients.TryGetValue(sessionId, out var client))
|
||||
return false;
|
||||
|
||||
try
|
||||
{
|
||||
if (_streams.TryRemove(sessionId, out var stream))
|
||||
stream.Dispose();
|
||||
client.Close();
|
||||
}
|
||||
catch
|
||||
{
|
||||
}
|
||||
|
||||
if (!string.IsNullOrWhiteSpace(reason))
|
||||
ShrinkNetworkLogger.Warn($"[ShrinkNetwork][TCP-Server] Disconnect session {sessionId}: {reason}");
|
||||
return true;
|
||||
}
|
||||
|
||||
public async UniTask SendAsync(long sessionId, byte[] packetData)
|
||||
{
|
||||
if (!_clients.TryGetValue(sessionId, out var client))
|
||||
throw new InvalidOperationException($"Session {sessionId} is not connected.");
|
||||
if (!_streams.TryGetValue(sessionId, out var stream))
|
||||
throw new InvalidOperationException($"Session {sessionId} stream is not initialized.");
|
||||
|
||||
if (!_sendLocks.TryGetValue(sessionId, out var sendLock))
|
||||
throw new InvalidOperationException($"Session {sessionId} send lock is not initialized.");
|
||||
if ((packetData?.Length ?? 0) > _maxPacketSize)
|
||||
throw new InvalidOperationException($"TCP packet is too large. Size={(packetData?.Length ?? 0)}, Limit={_maxPacketSize}.");
|
||||
|
||||
await SendInternalAsync(stream, sendLock, packetData ?? Array.Empty<byte>(), CancellationToken.None);
|
||||
}
|
||||
|
||||
private async Task AcceptLoopAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
while (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
TcpClient client;
|
||||
try
|
||||
{
|
||||
client = await _listener.AcceptTcpClientAsync(cancellationToken);
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
break;
|
||||
}
|
||||
catch
|
||||
{
|
||||
if (!IsStarted)
|
||||
break;
|
||||
|
||||
ShrinkNetworkLogger.Warn("[ShrinkNetwork][TCP-Server] Accept failed, retrying.");
|
||||
continue;
|
||||
}
|
||||
|
||||
var sessionId = Interlocked.Increment(ref _sessionIdGenerator);
|
||||
client.NoDelay = true;
|
||||
var stream = await CreateServerStreamAsync(client, cancellationToken);
|
||||
_clients[sessionId] = client;
|
||||
_streams[sessionId] = stream;
|
||||
_sendLocks[sessionId] = new SemaphoreSlim(1, 1);
|
||||
var remoteAddress = client.Client.RemoteEndPoint?.ToString() ?? "unknown";
|
||||
OnEvent?.Invoke(ShrinkNetworkTransportEvent.Connected(sessionId, remoteAddress));
|
||||
_ = ReceiveLoopAsync(sessionId, client, stream, cancellationToken);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<Stream> CreateServerStreamAsync(TcpClient client, CancellationToken cancellationToken)
|
||||
{
|
||||
Stream stream = client.GetStream();
|
||||
var tlsOptions = _tlsOptions;
|
||||
if (tlsOptions == null || !tlsOptions.Enabled)
|
||||
return stream;
|
||||
|
||||
if (_serverCertificate == null)
|
||||
throw new InvalidOperationException("TCP TLS is enabled, but the server certificate is not loaded.");
|
||||
|
||||
var sslStream = new SslStream(stream, false);
|
||||
await sslStream.AuthenticateAsServerAsync(_serverCertificate, false, tlsOptions.EnabledProtocols,
|
||||
tlsOptions.CheckCertificateRevocation);
|
||||
return sslStream;
|
||||
}
|
||||
|
||||
private void EnsureTlsCertificateLoaded()
|
||||
{
|
||||
var tlsOptions = _tlsOptions;
|
||||
if (tlsOptions == null || !tlsOptions.Enabled)
|
||||
{
|
||||
_serverCertificate = null;
|
||||
return;
|
||||
}
|
||||
|
||||
if (string.IsNullOrWhiteSpace(tlsOptions.ServerCertificatePath))
|
||||
throw new InvalidOperationException("TCP TLS is enabled, but ServerCertificatePath is empty.");
|
||||
|
||||
_serverCertificate = new X509Certificate2(tlsOptions.ServerCertificatePath, tlsOptions.ServerCertificatePassword);
|
||||
}
|
||||
|
||||
private async Task ReceiveLoopAsync(long sessionId, TcpClient client, Stream stream, CancellationToken cancellationToken)
|
||||
{
|
||||
Exception? disconnectException = null;
|
||||
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 (Exception ex)
|
||||
{
|
||||
disconnectException = ex;
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (_clients.TryRemove(sessionId, out var removed))
|
||||
{
|
||||
if (_streams.TryRemove(sessionId, out var ownedStream))
|
||||
ownedStream.Dispose();
|
||||
if (_sendLocks.TryRemove(sessionId, out var sendLock))
|
||||
sendLock.Dispose();
|
||||
|
||||
var remoteAddress = removed.Client.RemoteEndPoint?.ToString() ?? "unknown";
|
||||
if (disconnectException != null)
|
||||
{
|
||||
ShrinkNetworkLogger.Warn(
|
||||
$"[ShrinkNetwork][TCP-Server] Session {sessionId} {remoteAddress} disconnected: {disconnectException.GetType().Name}: {disconnectException.Message}");
|
||||
}
|
||||
|
||||
removed.Close();
|
||||
OnEvent?.Invoke(ShrinkNetworkTransportEvent.Disconnected(sessionId, remoteAddress));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static async Task 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 Task 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user