crysome/Crysome.Common/Crysome.Common.Network.Transport/QuicTransportSession.cs

using System;
using System.Buffers.Binary;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Net.Quic;
using System.Net.Security;
using System.Security.Cryptography.X509Certificates;
using System.Threading;
using System.Threading.Channels;
using System.Threading.Tasks;

namespace Crysome.Common.Network.Transport;

public sealed class QuicTransportSession : ITransportSession, IDisposable, IAsyncDisposable
{
	private sealed class FragBuf
	{
		public byte[][] Frags;

		public int RecvCount;

		public bool IsComplete
		{
			get
			{
				if (Frags != null)
				{
					return RecvCount >= Frags.Length;
				}
				return false;
			}
		}
	}

	public const string Alpn = "crysome";

	private readonly QuicConnection _conn;

	private readonly bool _isServer;

	private readonly QuicStream _control;

	private readonly QuicStream _screenOut;

	private readonly QuicStream _screenIn;

	private readonly QuicStream _bulkOut;

	private readonly QuicStream _bulkIn;

	private readonly CancellationTokenSource _cts = new CancellationTokenSource();

	private readonly object _controlWriteLock = new object();

	private readonly object _bulkWriteLock = new object();

	private readonly object _screenWriteLock = new object();

	private readonly Channel<byte[]> _screenChannel;

	private volatile bool _alive = true;

	private long _bytesSent;

	private long _bytesRecv;

	private double _rttMs = 100.0;

	private double _lossRate;

	private DateTime _lastCongestionEmit = DateTime.UtcNow;

	private readonly ConcurrentDictionary<uint, FragBuf> _screenFrags = new ConcurrentDictionary<uint, FragBuf>();

	private uint _screenLastSeq;

	private readonly object _screenSeqLock = new object();

	private bool _disposed;

	public IPEndPoint RemoteEndPoint => _conn?.RemoteEndPoint;

	public IPEndPoint LocalEndPoint => _conn?.LocalEndPoint;

	public bool IsAlive
	{
		get
		{
			if (_alive)
			{
				return !_disposed;
			}
			return false;
		}
	}

	public double RttMs => _rttMs;

	public double LossRate => _lossRate;

	public long BytesSent => _bytesSent;

	public long BytesRecv => _bytesRecv;

	public event Action<byte[]> ReliableReceived;

	public event Action<byte, byte[]> UnreliableFrameReceived;

	public event Action<double> CongestionChanged;

	public event Action Disconnected;

	private QuicTransportSession(QuicConnection conn, bool isServer, QuicStream control, QuicStream screenOut, QuicStream screenIn, QuicStream bulkOut, QuicStream bulkIn)
	{
		_conn = conn;
		_isServer = isServer;
		_control = control;
		_screenOut = screenOut;
		_screenIn = screenIn;
		_bulkOut = bulkOut;
		_bulkIn = bulkIn;
		_screenChannel = Channel.CreateBounded<byte[]>(new BoundedChannelOptions(1)
		{
			FullMode = BoundedChannelFullMode.DropOldest,
			SingleReader = true,
			SingleWriter = false
		});
	}

	public static async ValueTask<QuicTransportSession> ConnectClientAsync(EndPoint remote, CancellationToken ct)
	{
		if (!QuicConnection.IsSupported)
		{
			throw new InvalidOperationException("QUIC is not supported on this OS or runtime build.");
		}
		QuicConnection conn = await QuicConnection.ConnectAsync(new QuicClientConnectionOptions
		{
			RemoteEndPoint = remote,
			DefaultCloseErrorCode = 0L,
			DefaultStreamErrorCode = 0L,
			MaxInboundBidirectionalStreams = 256,
			MaxInboundUnidirectionalStreams = 256,
			ClientAuthenticationOptions = new SslClientAuthenticationOptions
			{
				ApplicationProtocols = new List<SslApplicationProtocol>
				{
					new SslApplicationProtocol("crysome")
				},
				RemoteCertificateValidationCallback = (object _, X509Certificate? _, X509Chain? _, SslPolicyErrors _) => true,
				TargetHost = "crysome"
			}
		}, ct).ConfigureAwait(continueOnCapturedContext: false);
		QuicTransportSession quicTransportSession = new QuicTransportSession(conn, isServer: false, await conn.OpenOutboundStreamAsync(QuicStreamType.Bidirectional, ct).ConfigureAwait(continueOnCapturedContext: false), await conn.OpenOutboundStreamAsync(QuicStreamType.Unidirectional, ct).ConfigureAwait(continueOnCapturedContext: false), null, null, await conn.AcceptInboundStreamAsync(ct).ConfigureAwait(continueOnCapturedContext: false));
		quicTransportSession.StartLoops();
		return quicTransportSession;
	}

	public static async ValueTask<QuicTransportSession> AcceptServerAsync(QuicConnection conn, CancellationToken ct)
	{
		QuicTransportSession quicTransportSession = new QuicTransportSession(conn, isServer: true, await conn.AcceptInboundStreamAsync(ct).ConfigureAwait(continueOnCapturedContext: false), null, await conn.AcceptInboundStreamAsync(ct).ConfigureAwait(continueOnCapturedContext: false), await conn.OpenOutboundStreamAsync(QuicStreamType.Unidirectional, ct).ConfigureAwait(continueOnCapturedContext: false), null);
		quicTransportSession.StartLoops();
		return quicTransportSession;
	}

	public static async ValueTask<QuicListener> StartListenerAsync(int port, X509Certificate2 certificate, CancellationToken ct)
	{
		if (!QuicListener.IsSupported)
		{
			throw new InvalidOperationException("QUIC listener is not supported on this OS or runtime build.");
		}
		return await QuicListener.ListenAsync(new QuicListenerOptions
		{
			ListenEndPoint = new IPEndPoint(IPAddress.Any, port),
			ListenBacklog = 512,
			ApplicationProtocols = new List<SslApplicationProtocol>
			{
				new SslApplicationProtocol("crysome")
			},
			ConnectionOptionsCallback = (QuicConnection connection, SslClientHelloInfo clientHello, CancellationToken token) => ValueTask.FromResult(new QuicServerConnectionOptions
			{
				DefaultStreamErrorCode = 0L,
				DefaultCloseErrorCode = 0L,
				MaxInboundBidirectionalStreams = 256,
				MaxInboundUnidirectionalStreams = 256,
				ServerAuthenticationOptions = new SslServerAuthenticationOptions
				{
					ApplicationProtocols = new List<SslApplicationProtocol>
					{
						new SslApplicationProtocol("crysome")
					},
					ServerCertificate = certificate
				}
			})
		}, ct).ConfigureAwait(continueOnCapturedContext: false);
	}

	private void StartLoops()
	{
		CancellationToken ct = _cts.Token;
		Task.Run(() => ControlReadLoop(ct), ct);
		Task.Run(() => CongestionEmitLoop(ct), ct);
		if (_isServer && _screenIn != null)
		{
			Task.Run(() => ScreenReadLoop(ct), ct);
		}
		if (!_isServer && _bulkIn != null)
		{
			Task.Run(() => BulkReadLoop(ct), ct);
		}
		if (!_isServer && _screenOut != null)
		{
			Task.Run(() => ScreenWriterLoop(ct), ct);
		}
	}

	private async Task CongestionEmitLoop(CancellationToken ct)
	{
		while (!ct.IsCancellationRequested && _alive)
		{
			try
			{
				await Task.Delay(2000, ct).ConfigureAwait(continueOnCapturedContext: false);
				EmitCongestion();
			}
			catch
			{
				break;
			}
		}
	}

	private void EmitCongestion()
	{
		DateTime utcNow = DateTime.UtcNow;
		if ((utcNow - _lastCongestionEmit).TotalMilliseconds < 500.0)
		{
			return;
		}
		_lastCongestionEmit = utcNow;
		double num = Math.Min(1.0, _rttMs / 2000.0);
		double lossRate = _lossRate;
		double obj = Math.Min(1.0, num * 0.5 + lossRate * 0.5);
		try
		{
			this.CongestionChanged?.Invoke(obj);
		}
		catch
		{
		}
	}

	public void SendReliable(ReadOnlySpan<byte> payload)
	{
		if (_disposed || !_alive || _control == null)
		{
			return;
		}
		lock (_controlWriteLock)
		{
			try
			{
				WriteFrameSync(_control, payload);
				Interlocked.Add(ref _bytesSent, 4 + payload.Length);
			}
			catch
			{
				TriggerDisconnect();
			}
		}
	}

	public void SendReliableBulk(ReadOnlySpan<byte> payload)
	{
		if (!_isServer || _bulkOut == null || _disposed || !_alive)
		{
			return;
		}
		lock (_bulkWriteLock)
		{
			try
			{
				WriteFrameSync(_bulkOut, payload);
				Interlocked.Add(ref _bytesSent, 4 + payload.Length);
			}
			catch
			{
				TriggerDisconnect();
			}
		}
	}

	public void SendPacketUnreliable(ReadOnlySpan<byte> payload)
	{
		SendReliable(payload);
	}

	public void SendUnreliableScreen(byte packetTypeId, ReadOnlySpan<byte> payload)
	{
		if (!_isServer && _screenOut != null)
		{
			uint frameSeq;
			lock (_screenSeqLock)
			{
				frameSeq = ++_screenLastSeq;
			}
			byte[] array = payload.ToArray();
			int num = Math.Max(1, (array.Length + 60000 - 1) / 60000);
			for (ushort num2 = 0; num2 < num; num2++)
			{
				int num3 = num2 * 60000;
				int payloadLength = Math.Min(60000, array.Length - num3);
				byte[] item = UdpTransport.BuildFragment(0u, packetTypeId, frameSeq, num2, (ushort)num, array, num3, payloadLength);
				_screenChannel.Writer.TryWrite(item);
			}
		}
	}

	private async Task ScreenWriterLoop(CancellationToken ct)
	{
		_ = 1;
		try
		{
			await foreach (byte[] item in _screenChannel.Reader.ReadAllAsync(ct).ConfigureAwait(continueOnCapturedContext: false))
			{
				if (_disposed || !_alive)
				{
					break;
				}
				lock (_screenWriteLock)
				{
					try
					{
						WriteFrameSync(_screenOut, item);
						Interlocked.Add(ref _bytesSent, 4 + item.Length);
					}
					catch
					{
					}
				}
			}
		}
		catch
		{
		}
	}

	private void WriteFrameSync(QuicStream stream, ReadOnlySpan<byte> payload)
	{
		byte[] array = new byte[4];
		BinaryPrimitives.WriteInt32LittleEndian(array, payload.Length);
		stream.Write(array);
		stream.Write(payload);
		stream.Flush();
	}

	private async Task ControlReadLoop(CancellationToken ct)
	{
		try
		{
			while (!ct.IsCancellationRequested && _alive)
			{
				byte[] array = await ReadFrameAsync(_control, ct).ConfigureAwait(continueOnCapturedContext: false);
				Interlocked.Add(ref _bytesRecv, 4 + array.Length);
				try
				{
					this.ReliableReceived?.Invoke(array);
				}
				catch
				{
				}
			}
		}
		catch
		{
			TriggerDisconnect();
		}
	}

	private async Task ScreenReadLoop(CancellationToken ct)
	{
		try
		{
			while (!ct.IsCancellationRequested && _alive)
			{
				byte[] array = await ReadFrameAsync(_screenIn, ct).ConfigureAwait(continueOnCapturedContext: false);
				Interlocked.Add(ref _bytesRecv, 4 + array.Length);
				ProcessScreenDatagram(array);
			}
		}
		catch
		{
			TriggerDisconnect();
		}
	}

	private void ProcessScreenDatagram(byte[] dg)
	{
		if (!UdpTransport.ParseHeader(dg, dg.Length, out var _, out var packetTypeId, out var frameSeq, out var fragIdx, out var fragCount, out var payloadOffset, out var payloadLength))
		{
			return;
		}
		if (fragCount <= 1)
		{
			byte[] array = new byte[payloadLength];
			if (payloadLength > 0)
			{
				Buffer.BlockCopy(dg, payloadOffset, array, 0, payloadLength);
			}
			try
			{
				this.UnreliableFrameReceived?.Invoke(packetTypeId, array);
				return;
			}
			catch
			{
				return;
			}
		}
		byte[] array2 = new byte[payloadLength];
		if (payloadLength > 0)
		{
			Buffer.BlockCopy(dg, payloadOffset, array2, 0, payloadLength);
		}
		FragBuf orAdd = _screenFrags.GetOrAdd(frameSeq, (uint _) => new FragBuf
		{
			Frags = new byte[fragCount][]
		});
		if (fragIdx >= orAdd.Frags.Length || orAdd.Frags[fragIdx] != null)
		{
			return;
		}
		orAdd.Frags[fragIdx] = array2;
		Interlocked.Increment(ref orAdd.RecvCount);
		if (!orAdd.IsComplete)
		{
			return;
		}
		_screenFrags.TryRemove(frameSeq, out var _);
		int num = 0;
		byte[][] frags = orAdd.Frags;
		foreach (byte[] array3 in frags)
		{
			if (array3 != null)
			{
				num += array3.Length;
			}
		}
		byte[] array4 = new byte[num];
		int num3 = 0;
		frags = orAdd.Frags;
		foreach (byte[] array5 in frags)
		{
			if (array5 != null)
			{
				Buffer.BlockCopy(array5, 0, array4, num3, array5.Length);
				num3 += array5.Length;
			}
		}
		try
		{
			this.UnreliableFrameReceived?.Invoke(packetTypeId, array4);
		}
		catch
		{
		}
	}

	private async Task BulkReadLoop(CancellationToken ct)
	{
		try
		{
			while (!ct.IsCancellationRequested && _alive)
			{
				byte[] array = await ReadFrameAsync(_bulkIn, ct).ConfigureAwait(continueOnCapturedContext: false);
				Interlocked.Add(ref _bytesRecv, 4 + array.Length);
				try
				{
					this.ReliableReceived?.Invoke(array);
				}
				catch
				{
				}
			}
		}
		catch
		{
			TriggerDisconnect();
		}
	}

	private static async Task<byte[]> ReadFrameAsync(QuicStream stream, CancellationToken ct)
	{
		byte[] lenBuf = new byte[4];
		await ReadExactAsync(stream, lenBuf, ct).ConfigureAwait(continueOnCapturedContext: false);
		int num = BinaryPrimitives.ReadInt32LittleEndian(lenBuf);
		if (num < 0 || num > 104857600)
		{
			throw new IOException("Invalid frame length");
		}
		byte[] buf = new byte[num];
		await ReadExactAsync(stream, buf, ct).ConfigureAwait(continueOnCapturedContext: false);
		return buf;
	}

	private static async Task ReadExactAsync(QuicStream stream, Memory<byte> buffer, CancellationToken ct)
	{
		int num;
		for (int o = 0; o < buffer.Length; o += num)
		{
			num = await stream.ReadAsync(buffer.Slice(o), ct).ConfigureAwait(continueOnCapturedContext: false);
			if (num == 0)
			{
				throw new IOException("Stream ended");
			}
		}
	}

	public void SendFin()
	{
		try
		{
			_cts.Cancel();
		}
		catch
		{
		}
		_alive = false;
	}

	private void TriggerDisconnect()
	{
		if (_disposed)
		{
			return;
		}
		_alive = false;
		try
		{
			_cts.Cancel();
		}
		catch
		{
		}
		try
		{
			this.Disconnected?.Invoke();
		}
		catch
		{
		}
	}

	public void Dispose()
	{
		if (_disposed)
		{
			return;
		}
		_disposed = true;
		_alive = false;
		try
		{
			_cts.Cancel();
		}
		catch
		{
		}
		try
		{
			_screenChannel.Writer.TryComplete();
		}
		catch
		{
		}
		try
		{
			_control?.Dispose();
		}
		catch
		{
		}
		try
		{
			_screenOut?.Dispose();
		}
		catch
		{
		}
		try
		{
			_screenIn?.Dispose();
		}
		catch
		{
		}
		try
		{
			_bulkOut?.Dispose();
		}
		catch
		{
		}
		try
		{
			_bulkIn?.Dispose();
		}
		catch
		{
		}
		try
		{
			_conn?.DisposeAsync().AsTask().Wait(2000);
		}
		catch
		{
		}
	}

	public async ValueTask DisposeAsync()
	{
		Dispose();
		await Task.CompletedTask;
	}
}