voice-cat/dotnet/src/VoiceCat.Server/VoiceServer.cs

using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;
using Google.Protobuf;
using VoiceCat.Crypto;
using VoiceCat.Server.Data;
using VoiceCat.Server.Transport;
using Voicecat.V1;

namespace VoiceCat.Server;

public sealed class VoiceServer : IAsyncDisposable
{
    private readonly Socket listener;
    private readonly MediaRelay media;
    private readonly ServerCredentials credentials;
    private readonly AccountStore accounts;
    private readonly IReadOnlyList<Voicecat.V1.Channel> channels;
    private readonly bool allowGuests;
    private readonly string name;
    private readonly VoiceServerOptions options;
    private readonly TimeProvider clock;
    private readonly CancellationTokenSource shutdown = new();
    private readonly object gate = new();
    private readonly Dictionary<ulong, Session> sessions = [];
    private readonly List<Task> connections = [];
    private ulong nextSession;
    private uint nextUser;
    private uint nextSsrc;
    private readonly Task accepting;
    private readonly Task reaping;
    private int disposed;

    public IPEndPoint EndPoint => (IPEndPoint)listener.LocalEndPoint!;
    public IPEndPoint MediaEndPoint => media.EndPoint;
    public event Action<Exception>? ConnectionFailed;

    public VoiceServer(string directory, IPEndPoint endpoint, bool allowGuests = true, string name = "VoiceCat Server")
        : this(directory, endpoint, new VoiceServerOptions { AllowGuests = allowGuests, Name = name }) { }

    public VoiceServer(string directory, IPEndPoint endpoint, VoiceServerOptions options, TimeProvider? timeProvider = null)
    {
        ArgumentNullException.ThrowIfNull(options);
        options.Validate();
        this.options = options;
        clock = timeProvider ?? TimeProvider.System;
        allowGuests = options.AllowGuests;
        name = options.Name;
        credentials = ServerCredentials.LoadOrCreate(directory, name);
        try
        {
            accounts = new AccountStore(Path.Combine(directory, "voicecat.db"));
            channels = accounts.LoadChannels();
            listener = new Socket(endpoint.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
            listener.Bind(endpoint);
            listener.Listen(options.MaximumConnections);
            media = new((IPEndPoint)listener.LocalEndPoint!);
            media.Failed += exception => ConnectionFailed?.Invoke(exception);
        }
        catch
        {
            listener?.Dispose();
            accounts?.Dispose();
            credentials.Dispose();
            shutdown.Dispose();
            throw;
        }
        accepting = AcceptAsync();
        reaping = ReapAsync();
    }

    private async Task AcceptAsync()
    {
        try
        {
            while (!shutdown.IsCancellationRequested)
            {
                Socket socket = await listener.AcceptAsync(shutdown.Token).ConfigureAwait(false);
                lock (gate)
                {
                    if (sessions.Count >= options.MaximumConnections) { socket.Dispose(); continue; }
                    socket.NoDelay = true;
                    string address = ((IPEndPoint)socket.RemoteEndPoint!).Address.ToString();
                    var connection = new TlsControlConnection(socket, credentials.CreateTlsSession(), shutdown.Token, options.HandshakeTimeout);
                    var session = new Session(++nextSession, connection, address, new(clock));
                    sessions.Add(session.Id, session);
                    connections.RemoveAll(task => task.IsCompleted);
                    connections.Add(HandleAsync(session));
                }
            }
        }
        catch (Exception exception) when (shutdown.IsCancellationRequested && exception is OperationCanceledException or SocketException or ObjectDisposedException) { }
    }

    private async Task HandleAsync(Session session)
    {
        try
        {
            await foreach (Envelope envelope in session.Connection.ReadAsync(shutdown.Token).ConfigureAwait(false))
            {
                session.Activity.Touch();
                if (envelope.Ping is not null)
                {
                    session.Connection.TrySend(new() { RequestId = envelope.RequestId, Pong = new() { Nonce = envelope.Ping.Nonce } });
                    continue;
                }
                if (envelope.Disconnect is not null) { session.Connection.CompleteWrites(); break; }
                if (!session.HelloReceived)
                {
                    if (envelope.ClientHello?.ProtoVersion != 2 || accounts.IsBanned("ip", session.Address))
                    {
                        Reject(session, "Unsupported protocol version or banned address.");
                        break;
                    }
                    session.Media = new(RandomNumberGenerator.GetBytes(16), await session.Connection.TakeMediaCryptoAsync(shutdown.Token).ConfigureAwait(false), session.Activity);
                    var hello = new ServerHello { ProtoVersion = 2, ServerName = name, ServerVersion = "0.1.0-dotnet", UdpPort = checked((uint)media.EndPoint.Port), ServerIdentityFingerprint = ByteString.CopyFrom(SHA256.HashData(credentials.Identity.PublicKey)) };
                    if (allowGuests) hello.AuthMethods.Add("guest");
                    hello.AuthMethods.Add("password");
                    session.Connection.TrySend(new() { RequestId = envelope.RequestId, ServerHello = hello });
                    session.HelloReceived = true;
                    continue;
                }
                if (session.User is null)
                {
                    if (envelope.AuthRequest is null) { Reject(session, "Authentication required."); break; }
                    await AuthenticateAsync(session, envelope.RequestId, envelope.AuthRequest).ConfigureAwait(false);
                    continue;
                }
                switch (envelope.BodyCase)
                {
                    case Envelope.BodyOneofCase.TextMessage: RelayText(session, envelope.TextMessage); break;
                    case Envelope.BodyOneofCase.Subscribe: SendSnapshot(session); break;
                    case Envelope.BodyOneofCase.JoinChannel: Join(session, envelope.RequestId, envelope.JoinChannel.ChannelId); break;
                    case Envelope.BodyOneofCase.SubscribeVoice: SubscribeVoice(session, envelope.RequestId, true); break;
                    case Envelope.BodyOneofCase.UnsubscribeVoice: SubscribeVoice(session, envelope.RequestId, false); break;
                    case Envelope.BodyOneofCase.StreamAnnounce: AnnounceStream(session, envelope.RequestId, envelope.StreamAnnounce); break;
                    case Envelope.BodyOneofCase.StreamStop: StopStream(session, envelope.StreamStop.StreamId); break;
                    case Envelope.BodyOneofCase.StreamState: UpdateStream(session, envelope.StreamState); break;
                    case Envelope.BodyOneofCase.UdpBinding:
                        if (!envelope.UdpBinding.Ack && CryptographicOperations.FixedTimeEquals(envelope.UdpBinding.UdpToken.Span, session.Media!.Token))
                            session.Connection.TrySend(new() { RequestId = envelope.RequestId, UdpBinding = new() { Ack = true } });
                        break;
                    default:
                        session.Connection.TrySend(new() { RequestId = envelope.RequestId, GenericResult = new() { Code = 1, Message = "Operation is not implemented by this server checkpoint." } });
                        break;
                }
            }
            await session.Connection.Completion.ConfigureAwait(false);
        }
        catch (Exception exception) when (exception is IOException or SocketException or OperationCanceledException or ObjectDisposedException)
        {
            if (!shutdown.IsCancellationRequested && exception is not OperationCanceledException) ConnectionFailed?.Invoke(exception);
        }
        finally
        {
            lock (gate)
            {
                sessions.Remove(session.Id);
                if (session.User is null) session.Media?.Dispose();
                else PublishMedia();
                if (session.User is not null) Broadcast(new() { UserEvent = new() { Kind = UserEvent.Types.Kind.Left, LeftId = session.User.Id } });
            }
            await session.Connection.DisposeAsync().ConfigureAwait(false);
        }
    }

    private static void Reject(Session session, string reason)
    {
        session.Connection.TrySend(new() { Disconnect = new() { Code = 1, Reason = reason } });
        session.Connection.CompleteWrites();
    }

    private async Task ReapAsync()
    {
        if (options.IdleTimeout == TimeSpan.Zero) return;
        using var timer = new PeriodicTimer(options.ReaperInterval, clock);
        try
        {
            while (await timer.WaitForNextTickAsync(shutdown.Token).ConfigureAwait(false))
            {
                lock (gate)
                {
                    foreach (Session session in sessions.Values)
                    {
                        if (session.Closing || !session.Activity.IsExpired(options.IdleTimeout)) continue;
                        session.Closing = true;
                        Reject(session, "Receive idle timeout.");
                    }
                }
            }
        }
        catch (OperationCanceledException) when (shutdown.IsCancellationRequested) { }
    }

    private async Task AuthenticateAsync(Session session, ulong requestId, AuthRequest request)
    {
        User? user = null;
        bool admin = false;
        if (request.Guest is not null && allowGuests && request.Guest.Nickname.Length <= 128)
            user = new() { Nickname = request.Guest.Nickname.Length == 0 ? "Guest" : request.Guest.Nickname, IsGuest = true, ChannelId = 1 };
        else if (request.Password is not null && request.Password.Username.Length <= 128 && request.Password.Password.Length <= 1024 && !accounts.IsBanned("username", request.Password.Username))
        {
            Account? account = await accounts.AuthenticateAsync(request.Password.Username, request.Password.Password, session.Connection.CancellationToken).ConfigureAwait(false);
            if (account is not null) { user = new() { Nickname = account.Username, ChannelId = 1 }; admin = account.IsAdmin; }
        }
        shutdown.Token.ThrowIfCancellationRequested();
        session.Connection.CancellationToken.ThrowIfCancellationRequested();
        lock (gate)
        {
            var lobby = channels.FirstOrDefault(channel => channel.Id == 1);
            if (user is null || lobby is null || lobby.PasswordProtected || lobby.MaxUsers != 0 && sessions.Values.Count(peer => peer.User?.ChannelId == 1) >= lobby.MaxUsers)
            {
                session.Connection.TrySend(new() { RequestId = requestId, AuthResult = new() { Error = "Invalid credentials or lobby unavailable." } });
                return;
            }
            user.Id = checked(++nextUser);
            session.User = user;
            session.Connection.TrySend(new() { RequestId = requestId, AuthResult = new()
            {
                Ok = true, SessionId = session.Id, Self = user.Clone(), UdpToken = ByteString.CopyFrom(session.Media!.Token),
                Permissions = new() { IsAdmin = admin, CanAdminAccounts = admin, CanBan = admin, CanKick = admin, CanMoveUsers = admin, CanCreateTempChannel = admin }
            } });
            PublishMedia();
            Broadcast(new() { UserEvent = new() { Kind = UserEvent.Types.Kind.Joined, User = user.Clone() } }, session.Id);
            SendSnapshot(session);
        }
    }

    private void SendSnapshot(Session session)
    {
        lock (gate)
        {
            var snapshot = new ServerStateSnapshot();
            snapshot.Channels.Add(channels.Select(channel => channel.Clone()));
            snapshot.Users.Add(sessions.Values.Where(peer => peer.User is not null).Select(peer => peer.User!.Clone()));
            session.Connection.TrySend(new() { ServerState = snapshot });
        }
    }

    private void Join(Session session, ulong requestId, uint channelId)
    {
        lock (gate)
        {
            var channel = channels.FirstOrDefault(candidate => candidate.Id == channelId);
            if (channel is null || channel.PasswordProtected || channel.MaxUsers != 0 && sessions.Values.Count(peer => peer.Id != session.Id && peer.User?.ChannelId == channelId) >= channel.MaxUsers)
            {
                session.Connection.TrySend(new() { RequestId = requestId, JoinChannelResult = new() { Error = "Channel unavailable." } });
                return;
            }
            if (session.User!.ChannelId != channelId) session.User.Streams.Clear();
            session.User.ChannelId = channelId;
            PublishMedia();
            var result = new JoinChannelResult { Ok = true, ChannelId = channelId, Audio = channel.Audio.Clone() };
            result.Members.Add(sessions.Values.Where(peer => peer.User?.ChannelId == channelId).Select(peer => peer.User!.Clone()));
            session.Connection.TrySend(new() { RequestId = requestId, JoinChannelResult = result });
            Broadcast(new() { UserEvent = new() { Kind = UserEvent.Types.Kind.Updated, User = session.User.Clone() } });
        }
    }

    private void RelayText(Session sender, TextMessage message)
    {
        lock (gate)
        {
            bool permitted = Encoding.UTF8.GetByteCount(message.Body) <= 4096 && message.ClientMsgId.Length <= 128 &&
                (message.Scope == TextScope.TextServer || message.Scope == TextScope.TextChannel && message.TargetId == sender.User!.ChannelId ||
                 message.Scope == TextScope.TextPrivate && sessions.Values.Any(peer => peer.User?.Id == message.TargetId));
            if (permitted)
            {
                var relay = message.Clone();
                relay.SenderId = sender.User!.Id;
                relay.SentAtUnixMs = checked((ulong)DateTimeOffset.UtcNow.ToUnixTimeMilliseconds());
                var envelope = new Envelope { TextMessage = relay };
                foreach (Session recipient in sessions.Values.Where(peer => peer.User is not null))
                    if (message.Scope == TextScope.TextServer || message.Scope == TextScope.TextChannel && recipient.User!.ChannelId == message.TargetId ||
                        message.Scope == TextScope.TextPrivate && (recipient.User!.Id == message.TargetId || recipient.Id == sender.Id))
                        recipient.Connection.TrySend(envelope);
            }
            sender.Connection.TrySend(new() { TextMessageAck = new() { ClientMsgId = message.ClientMsgId, Ok = permitted } });
        }
    }

    private void PublishMedia()
    {
        media.Publish(sessions.Values.Where(peer => peer.User is not null).Select(peer => new MediaRoute(
            peer.Media!, peer.User!.ChannelId, peer.User.VoiceSubscribed, peer.User.ServerMuted, peer.User.SelfDeafened || peer.User.ServerDeafened,
            peer.User.Streams.Select(stream => stream.Ssrc).ToArray())).ToArray());
    }

    private void BroadcastUser(Session session) => Broadcast(new() { UserEvent = new() { Kind = UserEvent.Types.Kind.Updated, User = session.User!.Clone() } });

    private void SubscribeVoice(Session session, ulong requestId, bool subscribed)
    {
        lock (gate)
        {
            session.User!.VoiceSubscribed = subscribed;
            if (!subscribed) session.User.Streams.Clear();
            PublishMedia();
            session.Connection.TrySend(new() { RequestId = requestId, VoiceSubscriptionResult = new() { Ok = true, Subscribed = subscribed } });
            BroadcastUser(session);
        }
    }

    private void AnnounceStream(Session session, ulong requestId, StreamAnnounce request)
    {
        lock (gate)
        {
            if (!session.User!.VoiceSubscribed || !Enum.IsDefined(request.Kind) || request.Label.Length > 128 || session.User.Streams.Count >= 16 ||
                nextSsrc == uint.MaxValue || session.NextStream == uint.MaxValue || request.RequestedAudio?.BitrateBps is > 0 and < 500)
            {
                session.Connection.TrySend(new() { RequestId = requestId, StreamAnnounceResult = new() { Error = "Voice subscription required, invalid stream, or stream limit reached." } });
                return;
            }
            AudioConfig audio = channels.First(channel => channel.Id == session.User.ChannelId).Audio.Clone();
            if (request.RequestedAudio?.BitrateBps > 0) audio.BitrateBps = Math.Min(audio.BitrateBps, request.RequestedAudio.BitrateBps);
            var stream = new StreamInfo { StreamId = ++session.NextStream, Ssrc = ++nextSsrc, Kind = request.Kind, Label = request.Label, Audio = audio };
            session.User.Streams.Add(stream);
            PublishMedia();
            session.Connection.TrySend(new() { RequestId = requestId, StreamAnnounceResult = new() { Ok = true, StreamId = stream.StreamId, Ssrc = stream.Ssrc, EffectiveAudio = audio.Clone() } });
            BroadcastUser(session);
        }
    }

    private void StopStream(Session session, uint streamId)
    {
        lock (gate)
        {
            StreamInfo? stream = session.User!.Streams.FirstOrDefault(candidate => candidate.StreamId == streamId);
            if (stream is null) return;
            session.User.Streams.Remove(stream);
            PublishMedia();
            BroadcastUser(session);
        }
    }

    private void UpdateStream(Session session, StreamStateUpdate update)
    {
        lock (gate)
        {
            StreamInfo? stream = session.User!.Streams.FirstOrDefault(candidate => candidate.StreamId == update.StreamId);
            if (stream is null) return;
            Broadcast(new() { StreamState = new() { UserId = session.User.Id, StreamId = stream.StreamId, Muted = update.Muted, Talking = update.Talking } });
        }
    }

    private void Broadcast(Envelope envelope, ulong excluded = 0)
    {
        foreach (Session recipient in sessions.Values.Where(peer => peer.Id != excluded && peer.User is not null)) recipient.Connection.TrySend(envelope);
    }

    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref disposed, 1) != 0) return;
        shutdown.Cancel();
        listener.Dispose();
        try
        {
            await Task.WhenAll(accepting, reaping).ConfigureAwait(false);
        }
        finally
        {
            try
            {
                Task[] pending;
                lock (gate) pending = connections.ToArray();
                await Task.WhenAll(pending).ConfigureAwait(false);
            }
            finally
            {
                try { await media.DisposeAsync().ConfigureAwait(false); }
                finally { accounts.Dispose(); credentials.Dispose(); shutdown.Dispose(); }
            }
        }
    }

    private sealed class Session(ulong id, TlsControlConnection connection, string address, SessionActivity activity)
    {
        public ulong Id { get; } = id;
        public TlsControlConnection Connection { get; } = connection;
        public string Address { get; } = address;
        public SessionActivity Activity { get; } = activity;
        public bool Closing { get; set; }
        public bool HelloReceived { get; set; }
        public User? User { get; set; }
        public MediaPeer? Media { get; set; }
        public uint NextStream;
    }
}