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

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Threading.Channels;
using VoiceCat.Protocol;

namespace VoiceCat.Server.Transport;

internal sealed class MediaPeer(byte[] token, MediaSessionCrypto crypto)
{
    public byte[] Token { get; } = token;
    public MediaSessionCrypto Crypto { get; } = crypto;
    // Only the UDP loop reads or changes the endpoint and binding state.
    public SocketAddress? Endpoint { get; set; }
    public void Dispose() { Crypto.Dispose(); CryptographicOperations.ZeroMemory(Token); }
}

internal sealed record MediaRoute(MediaPeer Peer, uint ChannelId, bool Subscribed, bool Muted, bool Deafened, uint[] Sources);

internal sealed class MediaRelay : IAsyncDisposable
{
    private readonly Socket socket;
    private readonly CancellationTokenSource shutdown = new();
    private readonly ConcurrentQueue<MediaPeer> retired = new();
    private readonly Channel<byte> changed = Channel.CreateBounded<byte>(1);
    private MediaRoute[] routes = [];
    private readonly byte[] input = new byte[65535];
    private readonly MediaFanout fanout = new();
    private readonly Task receiving;
    public IPEndPoint EndPoint { get; }
    public event Action<Exception>? Failed;

    public MediaRelay(IPEndPoint endpoint)
    {
        socket = new(endpoint.AddressFamily, SocketType.Dgram, ProtocolType.Udp);
        try { socket.Bind(endpoint); EndPoint = (IPEndPoint)socket.LocalEndPoint!; }
        catch { socket.Dispose(); shutdown.Dispose(); throw; }
        receiving = ReceiveAsync();
    }

    // Publications are serialized by the server's session gate. Crypto ownership transfers here.
    public void Publish(MediaRoute[] next)
    {
        MediaRoute[] previous = Volatile.Read(ref routes);
        Volatile.Write(ref routes, next);
        foreach (MediaRoute route in previous)
            if (!next.Any(candidate => ReferenceEquals(candidate.Peer, route.Peer))) retired.Enqueue(route.Peer);
        changed.Writer.TryWrite(0);
    }

    private void DrainRetired()
    {
        while (retired.TryDequeue(out MediaPeer? peer)) peer.Dispose();
    }

    private async Task ReceiveAsync()
    {
        var sender = new SocketAddress(socket.AddressFamily);
        Task<int>? receive = null;
        Task<bool>? update = null;
        try
        {
            while (true)
            {
                receive ??= socket.ReceiveFromAsync(input, SocketFlags.None, sender, shutdown.Token).AsTask();
                update ??= changed.Reader.WaitToReadAsync(shutdown.Token).AsTask();
                await Task.WhenAny(receive, update).ConfigureAwait(false);
                if (update.IsCompleted)
                {
                    await update.ConfigureAwait(false);
                    while (changed.Reader.TryRead(out _)) { }
                    update = null;
                    DrainRetired();
                }
                if (!receive.IsCompleted) continue;
                int length;
                try { length = await receive.ConfigureAwait(false); }
                catch (SocketException exception) when (exception.SocketErrorCode is SocketError.MessageSize or SocketError.ConnectionReset) { continue; }
                finally { receive = null; }
                DrainRetired();
                MediaRoute[] current = Volatile.Read(ref routes);
                if (!VoiceFrameHeader.TryRead(input.AsSpan(0, length), out var header)) continue;
                MediaRoute? source = null;
                foreach (MediaRoute route in current)
                    if (route.Peer.Endpoint?.Equals(sender) == true) { source = route; break; }

                if (header.Type == MediaFrameType.UdpBinding)
                {
                    if (length != VoiceFrameHeader.Size + 16 || source is not null) continue;
                    foreach (MediaRoute route in current)
                    {
                        if (route.Peer.Endpoint is not null || !CryptographicOperations.FixedTimeEquals(route.Peer.Token, input.AsSpan(VoiceFrameHeader.Size, 16))) continue;
                        var bound = new SocketAddress(sender.Family, sender.Size);
                        for (int index = 0; index < sender.Size; index++) bound[index] = sender[index];
                        route.Peer.Endpoint = bound;
                        break;
                    }
                    continue;
                }
                if (source is null) continue;
                if (header.Type == MediaFrameType.Keepalive)
                {
                    if (length == VoiceFrameHeader.Size) await SendAsync(input.AsMemory(0, length), sender).ConfigureAwait(false);
                    continue;
                }
                if (!fanout.TryStart(input.AsSpan(0, length), source, current)) continue;
                while (fanout.TryNext(out ReadOnlyMemory<byte> packet, out SocketAddress? endpoint))
                    await SendAsync(packet, endpoint!).ConfigureAwait(false);
            }
        }
        catch (Exception exception) when (shutdown.IsCancellationRequested && exception is OperationCanceledException or SocketException or ObjectDisposedException) { }
        catch (Exception exception) { Failed?.Invoke(exception); throw; }
        finally
        {
            shutdown.Cancel();
            socket.Dispose();
            fanout.Dispose();
            if (receive is not null)
            {
                try { await receive.ConfigureAwait(false); }
                catch (Exception exception) when (exception is OperationCanceledException or SocketException or ObjectDisposedException) { }
            }
            if (update is not null)
            {
                try { await update.ConfigureAwait(false); }
                catch (OperationCanceledException) { }
            }
        }
    }

    private async ValueTask SendAsync(ReadOnlyMemory<byte> packet, SocketAddress endpoint)
    {
        try { await socket.SendToAsync(packet, SocketFlags.None, endpoint, shutdown.Token).ConfigureAwait(false); }
        catch (SocketException exception) when (exception.SocketErrorCode is SocketError.ConnectionReset or SocketError.HostUnreachable or SocketError.NetworkUnreachable) { }
    }

    public async ValueTask DisposeAsync()
    {
        shutdown.Cancel();
        socket.Dispose();
        try { await receiving.ConfigureAwait(false); }
        finally
        {
            DrainRetired();
            foreach (MediaRoute route in Volatile.Read(ref routes)) route.Peer.Dispose();
            shutdown.Dispose();
        }
    }
}