voice-cat/src/VoiceCat.Core/VoiceCatClient.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Threading.Channels;
using VoiceCat.Crypto;
using VoiceCat.Transport;
using VoiceCat.Protocol;
using VoiceCat.Audio;
using Voicecat.V1;
using Channel = Voicecat.V1.Channel;

namespace VoiceCat.Core;

public enum ClientConnectionState { Disconnected, Connecting, VerifyingIdentity, Authenticating, Connected }
public sealed record ServerIdentityChallenge(string Host, ushort Port, string CertificateFingerprint, TofuStatus Status);

public sealed partial class VoiceCatClient : IAsyncDisposable
{
    private readonly string clientName;
    private readonly string clientVersion;
    private readonly TofuStore pins;
    private readonly SemaphoreSlim lifecycle = new(1);
    private readonly CancellationTokenSource disposed = new();
    private readonly object stateGate = new();
    private readonly ConcurrentDictionary<ulong, TaskCompletionSource<Envelope>> pending = new();
    private readonly System.Threading.Channels.Channel<Envelope> events = System.Threading.Channels.Channel.CreateBounded<Envelope>(128);
    private readonly Dictionary<uint, Channel> channels = [];
    private readonly Dictionary<uint, User> users = [];
    private readonly Dictionary<uint, StreamInfo> localStreams = [];
    public AudioEngine Audio { get; }
    public IReadOnlyList<StreamInfo> LocalStreams { get { lock (stateGate) return localStreams.Values.Select(s => s.Clone()).ToArray(); } }
    private TlsControlConnection? control;
    private MediaSessionCrypto? mediaCrypto;
    private ClientMediaTransport? media;
    private Task keepalive = Task.CompletedTask;
    public event EncodedVoiceHandler? VoiceReceived;
    private CancellationTokenSource? connectionLifetime;
    private Task reader = Task.CompletedTask;
    private long nextRequest;
    private AuthResult? authentication;
    private ServerHello? hello;
    private ClientConnectionState state;

    public event Action<ClientConnectionState>? ConnectionStateChanged;
    public ClientConnectionState State { get { lock (stateGate) return state; } }
    public Exception? ConnectionFailure { get; private set; }
    public Task Completion => reader;
    public AuthResult? Authentication { get { lock (stateGate) return authentication?.Clone(); } }
    public ServerHello? ServerHello { get { lock (stateGate) return hello?.Clone(); } }
    public IReadOnlyList<Channel> Channels { get { lock (stateGate) return channels.Values.Select(c => c.Clone()).ToArray(); } }
    public IReadOnlyList<User> Users { get { lock (stateGate) return users.Values.Select(u => u.Clone()).ToArray(); } }
    public bool TryReadEvent(out Envelope? envelope) => events.Reader.TryRead(out envelope);
    public IAsyncEnumerable<Envelope> ReadEventsAsync(CancellationToken cancellationToken = default) => events.Reader.ReadAllAsync(cancellationToken);

    public VoiceCatClient(string clientName = "VoiceCat .NET", string clientVersion = "0.1.0", string? tofuStorePath = null)
    {
        this.clientName = clientName;
        this.clientVersion = clientVersion;
        pins = new(tofuStorePath ?? Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), "VoiceCat", "tofu.txt"));
        Audio = new(TrySendEncodedVoice);
        VoiceReceived += Audio.Receive;
    }

    public async Task ConnectAsync(string host, ushort port, Func<ServerIdentityChallenge, CancellationToken, ValueTask<bool>>? confirmIdentity = null, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(host);
        ArgumentOutOfRangeException.ThrowIfZero(port);
        await lifecycle.WaitAsync(cancellationToken).ConfigureAwait(false);
        Socket? socket = null;
        bool started = false;
        try
        {
            ObjectDisposedException.ThrowIf(disposed.IsCancellationRequested, this);
            if (control is not null) throw new InvalidOperationException("Disconnect before reconnecting.");
            started = true;
            ConnectionFailure = null;
            connectionLifetime = CancellationTokenSource.CreateLinkedTokenSource(disposed.Token);
            using var connecting = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, connectionLifetime.Token);
            CancellationToken token = connecting.Token;
            SetState(ClientConnectionState.Connecting);
            socket = new(SocketType.Stream, ProtocolType.Tcp) { NoDelay = true };
            await socket.ConnectAsync(host, port, token).ConfigureAwait(false);
            string? fingerprint = null;
            control = new(socket, TlsSession.CreateClient(value => { fingerprint = value; return true; }), connectionLifetime.Token);
            socket = null; // Transport owns it from here.
            mediaCrypto = await control.TakeMediaCryptoAsync(token).ConfigureAwait(false);
            string certificatePin = fingerprint ?? throw new IOException("TLS did not report a certificate fingerprint.");
            TofuStatus pinStatus = pins.Check(host, port, certificatePin);
            if (pinStatus != TofuStatus.Matched)
            {
                SetState(ClientConnectionState.VerifyingIdentity);
                if (confirmIdentity is null || !await confirmIdentity(new(host, port, certificatePin, pinStatus), token).ConfigureAwait(false))
                    throw new System.Security.Authentication.AuthenticationException("Server identity was rejected.");
                pins.Pin(host, port, certificatePin);
            }
            reader = ReadAsync(control, connectionLifetime.Token);
            Envelope response = await RequestAsync(new() { ClientHello = new() { ProtoVersion = 2, ClientName = clientName, ClientVersion = clientVersion } }, token).ConfigureAwait(false);
            if (response.ServerHello?.ProtoVersion != 2) throw new IOException("Unsupported server protocol.");
            lock (stateGate) hello = response.ServerHello.Clone();
            keepalive = KeepaliveAsync(connectionLifetime.Token);
            SetState(ClientConnectionState.Authenticating);
        }
        catch
        {
            socket?.Dispose();
            if (started) await CloseAsync().ConfigureAwait(false);
            throw;
        }
        finally { lifecycle.Release(); }
    }

    public Task<AuthResult> AuthenticateGuestAsync(string nickname, CancellationToken cancellationToken = default) =>
        AuthenticateAsync(new() { Guest = new() { Nickname = nickname } }, cancellationToken);
    public Task<AuthResult> AuthenticateUserAsync(string username, string password, CancellationToken cancellationToken = default) =>
        AuthenticateAsync(new() { Password = new() { Username = username, Password = password } }, cancellationToken);

    private async Task<AuthResult> AuthenticateAsync(AuthRequest request, CancellationToken cancellationToken)
    {
        if (State != ClientConnectionState.Authenticating) throw new InvalidOperationException("Authentication requires a connected TLS session.");
        Envelope response = await RequestAsync(new() { AuthRequest = request }, cancellationToken).ConfigureAwait(false);
        AuthResult result = response.AuthResult ?? throw new IOException("Unexpected authentication response.");
        if (result.Ok)
        {
            var endpoint = (IPEndPoint)control!.RemoteEndPoint;
            IPAddress address = endpoint.Address.IsIPv4MappedToIPv6 ? endpoint.Address.MapToIPv4() : endpoint.Address;
            media = new(new(address, checked((int)ServerHello!.UdpPort)), result.UdpToken.Span, mediaCrypto!, connectionLifetime!.Token);
            media.Received += (header, packet) => VoiceReceived?.Invoke(header, packet);
            SetState(ClientConnectionState.Connected);
        }
        return result;
    }

    public async Task<VoiceSubscriptionResult> SubscribeVoiceAsync(bool subscribe = true, CancellationToken cancellationToken = default)
    {
        if (subscribe && media is not null) await media.Bound.WaitAsync(TimeSpan.FromSeconds(5), cancellationToken).ConfigureAwait(false);
        return (await RequestAsync(subscribe ? new() { SubscribeVoice = new() } : new() { UnsubscribeVoice = new() }, cancellationToken).ConfigureAwait(false)).VoiceSubscriptionResult;
    }

    public bool TrySendEncodedVoice(uint ssrc, uint timestamp, ReadOnlySpan<byte> payload, VoiceFrameFlags flags = VoiceFrameFlags.None) =>
        media?.TrySend(new(MediaFrameType.Voice, flags, 0, ssrc, 0, timestamp), payload) == true;

    public async Task<StreamInfo> StartStreamAsync(StreamKind kind, string label = "", int captureChannels = 1, CancellationToken cancellationToken = default)
    {
        if (State != ClientConnectionState.Connected) throw new InvalidOperationException("Client is disconnected.");
        var response = (await RequestAsync(new() { StreamAnnounce = new() { Kind = kind, Label = label } }, cancellationToken).ConfigureAwait(false)).StreamAnnounceResult;
        if (!response.Ok) throw new InvalidOperationException(response.Error);
        var info = new StreamInfo { StreamId = response.StreamId, Ssrc = response.Ssrc, Kind = kind, Audio = response.EffectiveAudio.Clone(), Label = label };
        try
        {
            lock (stateGate) { if (State != ClientConnectionState.Connected) throw new InvalidOperationException("Client disconnected during stream negotiation."); Audio.AddLocalStream(info, captureChannels); localStreams[info.StreamId] = info; }
            return info.Clone();
        }
        catch { if (State == ClientConnectionState.Connected) Send(new() { StreamStop = new() { StreamId = info.StreamId } }); throw; }
    }

    public void StopStream(uint streamId)
    {
        lock (stateGate) { localStreams.Remove(streamId); Audio.RemoveLocalStream(streamId); }
        Send(new() { StreamStop = new() { StreamId = streamId } });
    }

    private async Task KeepaliveAsync(CancellationToken cancellationToken)
    {
        try
        {
            using var timer = new PeriodicTimer(TimeSpan.FromSeconds(10));
            while (await timer.WaitForNextTickAsync(cancellationToken).ConfigureAwait(false)) Send(new() { Ping = new() { Nonce = checked((ulong)Environment.TickCount64) } });
        }
        catch (Exception exception) when (exception is OperationCanceledException or IOException or InvalidOperationException) { }
    }

    public async Task<Envelope> RequestAsync(Envelope request, CancellationToken cancellationToken = default)
    {
        TlsControlConnection connection = control ?? throw new InvalidOperationException("Client is disconnected.");
        var completion = new TaskCompletionSource<Envelope>(TaskCreationOptions.RunContinuationsAsynchronously);
        ulong id = checked((ulong)Interlocked.Increment(ref nextRequest));
        Envelope outbound = request.Clone(); outbound.RequestId = id;
        if (!pending.TryAdd(id, completion)) throw new InvalidOperationException("Request ids exhausted.");
        try
        {
            if (!connection.TrySend(outbound)) throw new IOException("Control queue is full or closed.");
            return await completion.Task.WaitAsync(TimeSpan.FromSeconds(15), cancellationToken).ConfigureAwait(false);
        }
        finally { pending.TryRemove(id, out _); }
    }

    public void Send(Envelope message)
    {
        TlsControlConnection connection = control ?? throw new InvalidOperationException("Client is disconnected.");
        if (!connection.TrySend(message.Clone())) throw new IOException("Control queue is full or closed.");
    }

    private async Task ReadAsync(TlsControlConnection connection, CancellationToken cancellationToken)
    {
        Exception? failure = null;
        try
        {
            await foreach (Envelope message in connection.ReadAsync(cancellationToken).ConfigureAwait(false))
            {
                Apply(message);
                if (message.RequestId != 0 && pending.TryRemove(message.RequestId, out var completion)) completion.TrySetResult(message.Clone());
                if (!events.Writer.TryWrite(message.Clone())) throw new IOException("Client event queue exhausted; consume events regularly.");
                if (message.Disconnect is not null) { connection.CompleteWrites(); break; }
            }
        }
        catch (Exception exception) { failure = exception; ConnectionFailure = exception; }
        finally
        {
            connectionLifetime?.Cancel();
            foreach (var operation in pending.Values) operation.TrySetException(failure ?? new IOException("Connection closed."));
            SetState(ClientConnectionState.Disconnected);
        }
    }

    private void Apply(Envelope message)
    {
        lock (stateGate)
        {
            if (message.AuthResult?.Ok == true) authentication = message.AuthResult.Clone();
            if (message.ServerState is not null)
            {
                channels.Clear(); users.Clear();
                foreach (var channel in message.ServerState.Channels) channels[channel.Id] = channel.Clone();
                foreach (var user in message.ServerState.Users) users[user.Id] = user.Clone();
            }
            if (message.ChannelEvent is not null)
            {
                if (message.ChannelEvent.Kind == ChannelEvent.Types.Kind.Deleted) channels.Remove(message.ChannelEvent.DeletedId);
                else if (message.ChannelEvent.Channel is not null) channels[message.ChannelEvent.Channel.Id] = message.ChannelEvent.Channel.Clone();
            }
            if (message.UserEvent is not null)
            {
                if (message.UserEvent.Kind == UserEvent.Types.Kind.Left) users.Remove(message.UserEvent.LeftId);
                else if (message.UserEvent.User is not null) users[message.UserEvent.User.Id] = message.UserEvent.User.Clone();
            }
            if (authentication is not null && (message.ServerState is not null || message.UserEvent is not null))
            {
                User self = users.GetValueOrDefault(authentication.Self.Id, authentication.Self);
                Audio.SetRemoteStreams(users.Values.ToArray(), self.Id, self.ChannelId);
                foreach (var id in localStreams.Keys.Where(id => !self.Streams.Any(s => s.StreamId == id)).ToArray())
                { Audio.RemoveLocalStream(id); localStreams.Remove(id); }
            }
        }
    }

    private void SetState(ClientConnectionState value)
    {
        lock (stateGate) state = value;
        ConnectionStateChanged?.Invoke(value);
    }

    public async Task DisconnectAsync()
    {
        connectionLifetime?.Cancel();
        await lifecycle.WaitAsync().ConfigureAwait(false);
        try { await CloseAsync().ConfigureAwait(false); }
        finally { lifecycle.Release(); }
    }

    private async Task CloseAsync()
    {
        connectionLifetime?.Cancel();
        try { await reader.ConfigureAwait(false); }
        finally
        {
            try { await keepalive.ConfigureAwait(false); }
            catch (Exception exception) when (exception is IOException or OperationCanceledException or SocketException or ObjectDisposedException) { }
            try { if (media is not null) await media.DisposeAsync().ConfigureAwait(false); }
            catch (Exception exception) when (exception is IOException or OperationCanceledException or SocketException or ObjectDisposedException) { }
            try { if (control is not null) await control.DisposeAsync().ConfigureAwait(false); }
            catch (Exception exception) when (exception is IOException or OperationCanceledException or SocketException or ObjectDisposedException) { }
            control = null;
                media = null;
                mediaCrypto?.Dispose(); mediaCrypto = null;
                connectionLifetime?.Dispose(); connectionLifetime = null;
                lock (stateGate) { authentication = null; hello = null; channels.Clear(); users.Clear(); }
                lock (stateGate)
                {
                    foreach (var id in localStreams.Keys) Audio.RemoveLocalStream(id);
                    localStreams.Clear(); Audio.SetRemoteStreams([], 0, 0);
                }
            SetState(ClientConnectionState.Disconnected);
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (disposed.IsCancellationRequested) return;
        disposed.Cancel();
        await DisconnectAsync().ConfigureAwait(false);
        events.Writer.TryComplete();
        Audio.Dispose();
    }
}