Relay/RelayServer/Services/Chat/ChatSocketBehavior.cs

using System.Net.Http.Headers;
using System.Text.Json;
using RelayServer.Models;
using RelayServer.Services.Crypto;
using RelayServer.Services.Data;
using RelayServer.Services.Rtc;
using WebSocketSharp;
using WebSocketSharp.Server;
using ErrorEventArgs = WebSocketSharp.ErrorEventArgs;
using RelayShared.Services;

namespace RelayServer.Services.Chat;

public class ChatSocketBehavior : WebSocketBehavior
{
    public static ClientKeyService? ClientKeyService { get; set; }
    public static string? ServerPublicKey { get; set; }
    public static string? ServerPrivateKey { get; set; }
    public static string? ChannelDbKey { get; set; }
    public static ChannelCryptoService? ChannelCryptoService { get; set; }
    public static SurrealDb.Net.SurrealDbClient? Db { get; set; }

    protected override void OnMessage(MessageEventArgs e)
    {
        var msg = e.Data;

        try
        {
            using var doc = JsonDocument.Parse(msg);
            var root = doc.RootElement;

            if (root.TryGetProperty("Action", out var actionProp))
            {
                var action = (WsAction)actionProp.GetInt32();
                var control = JsonSerializer.Deserialize<WsControlMessage>(msg)!;
                DispatchControl(action, control);
                return;
            }

            if (root.TryGetProperty("Type", out var typeProp))
            {
                var type = (SignalType)typeProp.GetInt32();

                switch (type)
                {
                    case SignalType.EncryptedSignal:
                        HandleEncryptedRtcSignal(msg);
                        return;
                    case SignalType.ClientEncryptedChat:
                        HandleEncryptedChatMessage(msg);
                        return;
                }
            }
        }
        catch (Exception ex)
        {
            Console.WriteLine($"WebSocket message error: session={ID}, error={ex.Message}");
        }
    }

    private void DispatchControl(WsAction action, WsControlMessage control)
    {
        switch (action)
        {
            case WsAction.Authenticate:
                HandleAuthenticate(control);
                break;
            case WsAction.RegisterKey:
                HandleRegisterKey(control);
                break;
            case WsAction.GetServerKey:
                HandleGetServerKey();
                break;
            case WsAction.GetChannels:
                HandleGetChannels();
                break;
            case WsAction.GetHistory:
                HandleGetHistory(control);
                break;
            case WsAction.RtcJoin:
                HandleRtcJoinChannel(control);
                break;
            case WsAction.RtcLeave:
                HandleRtcLeaveChannel(control);
                break;
        }
    }

    protected override void OnClose(CloseEventArgs e)
    {
        ConnectedClientService.Unregister(ID);
        RtcChannelPresenceService.RemoveSession(ID);
        Console.WriteLine($"WebSocket closed: session={ID}, code={e.Code}, reason={e.Reason}");
        base.OnClose(e);
    }

    protected override void OnError(ErrorEventArgs e)
    {
        Console.WriteLine($"WebSocket error: session={ID}, message={e.Message}");
        base.OnError(e);
    }

    private async void HandleAuthenticate(WsControlMessage control)
    {
        if (string.IsNullOrWhiteSpace(control.Username) || string.IsNullOrWhiteSpace(control.Token))
        {
            Console.WriteLine("Invalid Authenticate payload.");
            return;
        }

        using var core = new HttpClient { BaseAddress = new Uri("http://192.168.1.85:1337") };
        core.DefaultRequestHeaders.Accept.Clear();
        core.DefaultRequestHeaders.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
        core.DefaultRequestHeaders.Add("User-Agent", "RelayServer");

        var response = await core.PostAsJsonAsync("/server/verify/user", new AuthUserVerify
        {
            Username = control.Username,
            Token = control.Token
        });

        Console.WriteLine($"Auth response for {control.Username}: {await response.Content.ReadAsStringAsync()}");

        var result = new WsEventMessage { Event = WsEvent.Authenticated, Detail = control.Username };
        Send(JsonSerializer.Serialize(result));
    }

    private void HandleRegisterKey(WsControlMessage control)
    {
        if (string.IsNullOrWhiteSpace(control.Username) || string.IsNullOrWhiteSpace(control.PublicKey))
        {
            Console.WriteLine("Invalid RegisterKey payload.");
            return;
        }

        if (ClientKeyService is null)
        {
            Console.WriteLine("ClientKeyService is not initialized.");
            return;
        }

        RegisterOrUpdateClientKeySync(control.Username, control.PublicKey);
        ConnectedClientService.Register(ID, control.Username);

        var response = new WsEventMessage { Event = WsEvent.KeyRegistered, Detail = control.Username };
        Send(JsonSerializer.Serialize(response));
    }

    private void HandleGetServerKey()
    {
        if (string.IsNullOrWhiteSpace(ServerPublicKey))
        {
            Console.WriteLine("Server public key is not initialized.");
            return;
        }

        var payload = new ServerPublicKeyMessage
        {
            Type = SignalType.ServerPublicKey,
            PublicKey = ServerPublicKey
        };

        Send(JsonSerializer.Serialize(payload));
    }

    private void HandleGetChannels()
    {
        if (Db is null)
        {
            Console.WriteLine("Db is not initialized.");
            return;
        }

        var channels = GetChannelsSync()
            .OrderBy(c => c.CreatedAt)
            .Select(c => new ChannelItem
            {
                ChannelId = GetRecordId(c.Id),
                Name = c.Name,
                CreatedAt = c.CreatedAt
            })
            .ToList();

        var payload = new SocketChannelList
        {
            Type = SignalType.ChannelList,
            Channels = channels
        };

        Send(JsonSerializer.Serialize(payload));
    }

    private void HandleGetHistory(WsControlMessage control)
    {
        if (string.IsNullOrWhiteSpace(control.Username) || string.IsNullOrWhiteSpace(control.ChannelId))
        {
            Console.WriteLine("Invalid GetHistory payload.");
            return;
        }

        if (!EnsureCoreReady() || ChannelCryptoService is null || string.IsNullOrWhiteSpace(ChannelDbKey))
        {
            Console.WriteLine("History dependencies are not initialized.");
            return;
        }

        var targetClient = GetClientPublicKeyByUsernameSync(control.Username);

        if (targetClient is null)
        {
            Console.WriteLine($"No public key found for history request user {control.Username}");
            return;
        }

        var channelMessages = GetChannelMessagesSync()
            .Where(m => m.ChannelId == control.ChannelId)
            .OrderBy(m => m.CreatedAt)
            .ToList();

        Console.WriteLine($"Sending {channelMessages.Count} history messages to {control.Username}");

        foreach (var dbMessage in channelMessages)
        {
            string plainText;

            try
            {
                plainText = ChannelCryptoService.Decrypt(
                    dbMessage.CipherText,
                    dbMessage.Nonce,
                    dbMessage.Tag,
                    ChannelDbKey
                );
            }
            catch (Exception ex)
            {
                Console.WriteLine($"Failed to decrypt DB history row {dbMessage.Id}: {ex.Message}");
                continue;
            }

            var encrypted = E2EeHelper.EncryptForRecipient(plainText, targetClient.PublicKey);

            var outbound = new SocketEncryptedMessage
            {
                Type = SignalType.EncryptedChat,
                SenderUsername = ExtractUsernameFromUserId(dbMessage.SenderUserId),
                RecipientUsername = control.Username,
                ChannelId = control.ChannelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Send(JsonSerializer.Serialize(outbound));
        }
    }

    private void HandleRtcJoinChannel(WsControlMessage control)
    {
        if (string.IsNullOrWhiteSpace(control.Username) || string.IsNullOrWhiteSpace(control.ChannelId))
        {
            Console.WriteLine("Invalid RtcJoin payload.");
            return;
        }

        RtcChannelPresenceService.SetUser(ID, control.Username);
        RtcChannelPresenceService.JoinChannel(ID, control.ChannelId);

        Console.WriteLine($"RTC presence joined: session={ID}, user={control.Username}, channel={control.ChannelId}");
    }

    private void HandleRtcLeaveChannel(WsControlMessage control)
    {
        if (string.IsNullOrWhiteSpace(control.Username) || string.IsNullOrWhiteSpace(control.ChannelId))
        {
            Console.WriteLine("Invalid RtcLeave payload.");
            return;
        }

        if (RtcChannelPresenceService.IsInChannel(ID, control.ChannelId))
            RtcChannelPresenceService.LeaveChannel(ID);

        Console.WriteLine($"RTC presence left: session={ID}, user={control.Username}, channel={control.ChannelId}");
    }

    private void HandleEncryptedChatMessage(string msg)
    {
        SocketEncryptedMessage? clientPayload;

        try
        {
            clientPayload = JsonSerializer.Deserialize<SocketEncryptedMessage>(msg);
        }
        catch
        {
            Console.WriteLine("Failed to parse encrypted client payload.");
            return;
        }

        if (clientPayload is null || clientPayload.Type != SignalType.ClientEncryptedChat)
            return;

        if (!EnsureCoreReady() || !EnsureCryptoReady())
            return;

        string plainText;

        try
        {
            plainText = E2EeHelper.DecryptForRecipient(
                new EncryptedPayload
                {
                    CipherText = clientPayload.CipherText,
                    Nonce = clientPayload.Nonce,
                    Tag = clientPayload.Tag,
                    EncryptedKey = clientPayload.EncryptedKey
                },
                ServerPrivateKey
            );
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to decrypt client payload: {ex.Message}");
            return;
        }

        Console.WriteLine($"Server decrypted message from {clientPayload.SenderUsername}: {plainText}");

        try
        {
            var dbEncrypted = ChannelCryptoService!.Encrypt(plainText, ChannelDbKey);

            var savedMessage = CreateChannelMessageSync(new ChannelMessages
            {
                ChannelId = clientPayload.ChannelId,
                SenderUserId = $"users:{clientPayload.SenderUsername.ToLower()}",
                CipherText = dbEncrypted.cipherText,
                Nonce = dbEncrypted.nonce,
                Tag = dbEncrypted.tag,
                CreatedAt = DateTime.UtcNow
            });

            Console.WriteLine($"Live message saved to DB: {JsonSerializer.Serialize(savedMessage)}");
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to save live message to DB: {ex.Message}");
            return;
        }

        var members = GetServerMembersSync();

        foreach (var member in members)
        {
            var username = ExtractUsernameFromUserId(member.UserId);
            var sessionIds = ConnectedClientService.GetSessionsForUser(username);

            if (!sessionIds.Any())
                continue;

            // Preserve the exact casing the client registered with
            var properUsername = sessionIds
                .Select(ConnectedClientService.GetUsernameForSession)
                .FirstOrDefault(u => u is not null) ?? username;

            var clientKey = GetClientPublicKeyByUsernameSync(properUsername);

            if (clientKey is null)
                continue;

            var encrypted = E2EeHelper.EncryptForRecipient(plainText, clientKey.PublicKey);

            Console.WriteLine($"Routing message from {clientPayload.SenderUsername} to {properUsername}");

            var outbound = new SocketEncryptedMessage
            {
                Type = SignalType.EncryptedChat,
                SenderUsername = clientPayload.SenderUsername,
                RecipientUsername = properUsername,
                ChannelId = clientPayload.ChannelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            var json = JsonSerializer.Serialize(outbound);

            foreach (var sessionId in sessionIds)
                Sessions.SendTo(json, sessionId);
        }
    }

    private void HandleEncryptedRtcSignal(string msg)
    {
        Console.WriteLine("RTC SIGNAL HIT");
        SocketRtcSignalMessage? clientPayload;

        try
        {
            clientPayload = JsonSerializer.Deserialize<SocketRtcSignalMessage>(msg);
        }
        catch
        {
            Console.WriteLine("Failed to parse encrypted RTC signal payload.");
            return;
        }

        if (clientPayload is null || clientPayload.Type != SignalType.EncryptedSignal)
            return;

        if (string.IsNullOrWhiteSpace(clientPayload.ChannelId))
        {
            Console.WriteLine("Encrypted RTC signal missing channel id.");
            return;
        }

        string plainText;

        try
        {
            plainText = E2EeHelper.DecryptForRecipient(
                new EncryptedPayload
                {
                    CipherText = clientPayload.CipherText,
                    Nonce = clientPayload.Nonce,
                    Tag = clientPayload.Tag,
                    EncryptedKey = clientPayload.EncryptedKey
                },
                ServerPrivateKey
            );
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to decrypt RTC signal: {ex.Message}");
            return;
        }

        var sessionIds = RtcChannelPresenceService.GetSessionsInChannel(clientPayload.ChannelId);

        foreach (var sessionId in sessionIds)
        {
            if (sessionId == ID)
                continue;

            var username = RtcChannelPresenceService.GetUsernameForSession(sessionId);
            if (string.IsNullOrWhiteSpace(username))
                continue;

            var clientKey = GetClientPublicKeyByUsernameSync(username);
            if (clientKey is null)
                continue;

            var encrypted = E2EeHelper.EncryptForRecipient(plainText, clientKey.PublicKey);

            var outbound = new SocketRtcSignalMessage
            {
                Type = SignalType.EncryptedSignal,
                SenderUsername = clientPayload.SenderUsername,
                ChannelId = clientPayload.ChannelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Sessions.SendTo(JsonSerializer.Serialize(outbound), sessionId);
        }

        Console.WriteLine($"Forwarded encrypted RTC signal from {clientPayload.SenderUsername} to channel {clientPayload.ChannelId}");
    }

    private static string ExtractUsernameFromUserId(string senderUserId)
    {
        if (string.IsNullOrWhiteSpace(senderUserId))
            return "Unknown";

        var parts = senderUserId.Split(':', 2);
        return parts.Length == 2 ? parts[1] : senderUserId;
    }

    private static string GetRecordId(object? id)
    {
        if (id is null)
            return string.Empty;

        var json = JsonSerializer.Serialize(id);

        using var doc = JsonDocument.Parse(json);
        var root = doc.RootElement;

        var recordId = root.GetProperty("Id").GetString() ?? string.Empty;
        var table = root.GetProperty("Table").GetString() ?? string.Empty;

        return $"{table}:{recordId}";
    }

    private void RegisterOrUpdateClientKeySync(string username, string publicKey)
    {
        Task.Run(async () => await ClientKeyService!.RegisterOrUpdateKeyAsync(username, publicKey))
            .GetAwaiter()
            .GetResult();
    }

    private List<Channels> GetChannelsSync()
    {
        return Task.Run(async () => await Db!.Select<Channels>("channels"))
            .GetAwaiter()
            .GetResult()
            .ToList();
    }

    private ClientPublicKeys? GetClientPublicKeyByUsernameSync(string username)
    {
        return Task.Run(async () => await ClientKeyService!.GetByUsernameAsync(username))
            .GetAwaiter()
            .GetResult();
    }

    private List<ChannelMessages> GetChannelMessagesSync()
    {
        return Task.Run(async () => await Db!.Select<ChannelMessages>("channel_messages"))
            .GetAwaiter()
            .GetResult()
            .ToList();
    }

    private ChannelMessages CreateChannelMessageSync(ChannelMessages message)
    {
        return Task.Run(async () => await Db!.Create("channel_messages", message))
            .GetAwaiter()
            .GetResult();
    }

    private List<ServerMembers> GetServerMembersSync()
    {
        return Task.Run(async () => await Db!.Select<ServerMembers>("server_members"))
            .GetAwaiter()
            .GetResult()
            .ToList();
    }

    private bool EnsureCoreReady()
    {
        if (ClientKeyService is null || Db is null)
        {
            Console.WriteLine("Core services not initialized.");
            return false;
        }

        return true;
    }

    private bool EnsureCryptoReady()
    {
        if (string.IsNullOrWhiteSpace(ServerPrivateKey) || string.IsNullOrWhiteSpace(ChannelDbKey))
        {
            Console.WriteLine("Crypto keys not initialized.");
            return false;
        }

        if (ChannelCryptoService is null)
        {
            Console.WriteLine("ChannelCryptoService is not initialized.");
            return false;
        }

        return true;
    }
}