Relay/RelayClient/Services/RelaySocketClient.cs

using System.Text.Json;
using RelayClient.Crypto;
using RelayShared.Services;
using WebSocketSharp;

namespace RelayClient.Services;

public sealed class RelaySocketClient
{
    private readonly string _username;
    private readonly WebSocket _socket;

    public string? ServerPublicKey { get; private set; }

    public event Action<string>? RawMessageReceived;
    public event Action<SocketChannelList>? ChannelListReceived;
    public event Action<SocketEncryptedMessage>? EncryptedChatReceived;
    public event Action<SocketRtcSignalMessage>? EncryptedRtcSignalReceived;
    public event Action<string>? ServerPublicKeyReceived;
    public event Action<string>? Log;

    public RelaySocketClient(string username, string url = "ws://127.0.0.1:5001/")
    {
        _username = username;
        _socket = new WebSocket(url);
        _socket.OnMessage += OnMessage;
    }

    public void Connect()
    {
        _socket.Connect();

        var publicKey = KeyStorage.LoadPublicKey(_username);

        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.Authenticate,
            Username = _username,
            Token = MainPage._userToken
        });

        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.RegisterKey,
            Username = _username,
            PublicKey = publicKey
        });

        SendControlMessage(new WsControlMessage { Action = WsAction.GetServerKey });
        SendControlMessage(new WsControlMessage { Action = WsAction.GetChannels });
    }

    public void SendControlMessage(WsControlMessage message)
    {
        SendRaw(JsonSerializer.Serialize(message));
    }

    public void SendGetHistory(string channelId)
    {
        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.GetHistory,
            Username = _username,
            ChannelId = channelId
        });
    }

    public void SendRtcJoinChannel(string channelId)
    {
        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.RtcJoin,
            Username = _username,
            ChannelId = channelId
        });
    }

    public void SendRtcLeaveChannel(string channelId)
    {
        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.RtcLeave,
            Username = _username,
            ChannelId = channelId
        });
    }

    public void SendRaw(string message)
    {
        if (_socket.ReadyState == WebSocketState.Open)
            _socket.Send(message);
    }

    public void SendJson<T>(T payload)
    {
        SendRaw(JsonSerializer.Serialize(payload));
    }

    public void Disconnect()
    {
        _socket.OnMessage -= OnMessage;

        if (_socket.ReadyState == WebSocketState.Open)
            _socket.Close();
    }

    private void OnMessage(object? sender, MessageEventArgs e)
    {
        RawMessageReceived?.Invoke(e.Data);
        Log?.Invoke($"[{_username}] RAW WS DATA: {e.Data}");

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

            // Control event responses (WsEvent)
            if (root.TryGetProperty("Event", out var eventElement))
            {
                var wsEvent = (WsEvent)eventElement.GetInt32();

                switch (wsEvent)
                {
                    case WsEvent.KeyRegistered:
                        Log?.Invoke($"[{_username}] Key registered on server.");
                        return;

                    case WsEvent.Authenticated:
                        Log?.Invoke($"[{_username}] Authenticated with server.");
                        return;

                    case WsEvent.Error:
                        var detail = root.TryGetProperty("Detail", out var d) ? d.GetString() : null;
                        Log?.Invoke($"[{_username}] Server error: {detail}");
                        return;
                }

                return;
            }

            // Data messages (SignalType)
            if (!root.TryGetProperty("Type", out var typeElement))
                return;

            var type = (SignalType)typeElement.GetInt32();

            switch (type)
            {
                case SignalType.ChannelList:
                {
                    var channelList = JsonSerializer.Deserialize<SocketChannelList>(e.Data);
                    if (channelList is not null)
                        ChannelListReceived?.Invoke(channelList);

                    return;
                }

                case SignalType.ServerPublicKey:
                {
                    var serverKeyMessage = JsonSerializer.Deserialize<ServerPublicKeyMessage>(e.Data);
                    if (serverKeyMessage is not null)
                    {
                        ServerPublicKey = serverKeyMessage.PublicKey;
                        ServerPublicKeyReceived?.Invoke(serverKeyMessage.PublicKey);
                    }

                    return;
                }

                case SignalType.EncryptedSignal:
                {
                    var payload = JsonSerializer.Deserialize<SocketRtcSignalMessage>(e.Data);
                    if (payload is not null)
                        EncryptedRtcSignalReceived?.Invoke(payload);

                    return;
                }

                case SignalType.EncryptedChat:
                {
                    var payload = JsonSerializer.Deserialize<SocketEncryptedMessage>(e.Data);
                    if (payload is not null)
                        EncryptedChatReceived?.Invoke(payload);

                    return;
                }
            }
        }
        catch (Exception ex)
        {
            Log?.Invoke($"[{_username}] failed to process websocket message: {ex.Message}");
        }
    }
}