voice-cat/dotnet/src/VoiceCat.Server/Transport/TlsControlConnection.cs

using System.Buffers;
using System.Buffers.Binary;
using System.Net.Sockets;
using System.Threading.Channels;
using Google.Protobuf;
using VoiceCat.Crypto;
using VoiceCat.Protocol;
using Voicecat.V1;

namespace VoiceCat.Server.Transport;

internal sealed class TlsControlConnection : IAsyncDisposable
{
    internal const int MaximumPayloadLength = 65536;
    private readonly Socket socket;
    private readonly TlsSession tls;
    private readonly CancellationTokenSource lifetime;
    private readonly Channel<byte[]> outgoing = System.Threading.Channels.Channel.CreateBounded<byte[]>(64);
    private readonly Channel<Envelope> incoming = System.Threading.Channels.Channel.CreateBounded<Envelope>(32);
    private readonly byte[] prefix = new byte[4];
    private int prefixBytes;
    private byte[]? payload;
    private int payloadBytes;

    public Task Completion { get; }
    public CancellationToken CancellationToken => lifetime.Token;

    internal TlsControlConnection(Socket socket, TlsSession tls, CancellationToken cancellationToken)
    {
        this.socket = socket;
        this.tls = tls;
        lifetime = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        lifetime.CancelAfter(TimeSpan.FromSeconds(15));
        Completion = RunAsync();
    }

    public IAsyncEnumerable<Envelope> ReadAsync(CancellationToken cancellationToken) => incoming.Reader.ReadAllAsync(cancellationToken);

    public bool TrySend(Envelope envelope)
    {
        if (envelope.CalculateSize() > MaximumPayloadLength) throw new InvalidDataException("Server control payload exceeds 64 KiB.");
        var framed = new ArrayBufferWriter<byte>();
        ControlFraming.WriteEnvelope(framed, envelope);
        if (outgoing.Writer.TryWrite(framed.WrittenSpan.ToArray())) return true;
        lifetime.Cancel();
        return false;
    }

    public void CompleteWrites() => outgoing.Writer.TryComplete();

    private async Task RunAsync()
    {
        byte[] ciphertext = new byte[16384];
        byte[] plaintext = new byte[16384];
        byte[] sendBuffer = new byte[16384];
        CancellationToken cancellationToken = lifetime.Token;
        Task<int>? receive = null;
        Task<bool>? ready = null;
        Exception? error = null;
        try
        {
            await FlushAsync(sendBuffer, cancellationToken).ConfigureAwait(false);
            receive = socket.ReceiveAsync(ciphertext, SocketFlags.None, cancellationToken).AsTask();
            while (true)
            {
                if (tls.IsReady)
                {
                    while (outgoing.Reader.TryRead(out byte[]? frame)) tls.WritePlaintext(frame);
                    await FlushAsync(sendBuffer, cancellationToken).ConfigureAwait(false);
                    ready ??= outgoing.Reader.WaitToReadAsync(cancellationToken).AsTask();
                }
                Task winner = ready is null ? receive : await Task.WhenAny(receive, ready).ConfigureAwait(false);
                if (winner == receive)
                {
                    int count = await receive.ConfigureAwait(false);
                    if (count == 0)
                    {
                        tls.CompleteInput();
                        if (prefixBytes != 0 || payload is not null) throw new InvalidDataException("Truncated control frame.");
                        break;
                    }
                    tls.ReceiveCiphertext(ciphertext.AsSpan(0, count));
                    if (tls.IsReady) lifetime.CancelAfter(TimeSpan.FromSeconds(60));
                    while ((count = tls.ReadPlaintext(plaintext)) > 0) Parse(plaintext.AsSpan(0, count));
                    await FlushAsync(sendBuffer, cancellationToken).ConfigureAwait(false);
                    receive = socket.ReceiveAsync(ciphertext, SocketFlags.None, cancellationToken).AsTask();
                }
                else
                {
                    bool hasOutgoing = await ready!.ConfigureAwait(false);
                    ready = null;
                    if (!hasOutgoing)
                    {
                        tls.Close();
                        await FlushAsync(sendBuffer, cancellationToken).ConfigureAwait(false);
                        break;
                    }
                }
            }
        }
        catch (Exception exception) when (exception is IOException or SocketException or OperationCanceledException or ObjectDisposedException)
        {
            if (!cancellationToken.IsCancellationRequested) error = exception;
        }
        finally
        {
            lifetime.Cancel();
            socket.Dispose();
            if (receive is not null)
            {
                try { await receive.ConfigureAwait(false); }
                catch (Exception exception) when (exception is SocketException or OperationCanceledException or ObjectDisposedException) { }
            }
            tls.Dispose();
            incoming.Writer.TryComplete(error);
            outgoing.Writer.TryComplete(error);
        }
    }

    private async Task FlushAsync(byte[] buffer, CancellationToken cancellationToken)
    {
        int count;
        while ((count = tls.DrainCiphertext(buffer)) > 0)
        {
            int sent = 0;
            while (sent < count)
            {
                int written = await socket.SendAsync(buffer.AsMemory(sent, count - sent), SocketFlags.None, cancellationToken).ConfigureAwait(false);
                if (written == 0) throw new IOException("Socket closed during TLS send.");
                sent += written;
            }
        }
    }

    private void Parse(ReadOnlySpan<byte> input)
    {
        while (!input.IsEmpty)
        {
            if (payload is null)
            {
                int count = Math.Min(4 - prefixBytes, input.Length);
                input[..count].CopyTo(prefix.AsSpan(prefixBytes));
                prefixBytes += count;
                input = input[count..];
                if (prefixBytes != 4) continue;
                uint length = BinaryPrimitives.ReadUInt32BigEndian(prefix);
                if (length > MaximumPayloadLength) throw new InvalidDataException("Server control payload exceeds 64 KiB.");
                payload = new byte[length];
                prefixBytes = 0;
            }
            int consumed = Math.Min(payload.Length - payloadBytes, input.Length);
            input[..consumed].CopyTo(payload.AsSpan(payloadBytes));
            payloadBytes += consumed;
            input = input[consumed..];
            if (payloadBytes != payload.Length) continue;
            Envelope envelope = Envelope.Parser.ParseFrom(payload);
            payload = null;
            payloadBytes = 0;
            if (!incoming.Writer.TryWrite(envelope)) throw new IOException("Control consumer exceeded its bounded queue.");
        }
    }

    public async ValueTask DisposeAsync()
    {
        lifetime.Cancel();
        await Completion.ConfigureAwait(false);
        lifetime.Dispose();
    }
}