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<SocketEncryptedMessage>? MessageEdited;
    public event Action<SocketMessageDeletedEvent>? MessageDeleted;
    public event Action<SocketTypingEvent>? TypingReceived;
    public event Action<SocketEditHistoryResponse>? EditHistoryReceived;
    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 Disconnect()
    {
        _socket.OnMessage -= OnMessage;
        if (_socket.ReadyState == WebSocketState.Open)
            _socket.Close();
    }

    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 SendTyping(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.SendTyping, Username = _username, ChannelId = channelId });

    public void SendGetEditHistory(string messageId, string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.GetEditHistory, Username = _username, MessageId = messageId, ChannelId = channelId });

    public void SendCreateChannel(string name, ChannelType type, string group = "") =>
        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.CreateChannel,
            ChannelName = name,
            ChannelType = (int)type,
            ChannelGroup = group
        });

    public void SendDeleteChannel(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.DeleteChannel, ChannelId = channelId });

    public void SendEditMessage(string messageId, string channelId, EncryptedPayload encrypted) =>
        SendJson(new SocketEncryptedMessage
        {
            Type = SignalType.ClientEditMessage, MessageId = messageId,
            SenderUsername = _username, ChannelId = channelId,
            CipherText = encrypted.CipherText, Nonce = encrypted.Nonce,
            Tag = encrypted.Tag, EncryptedKey = encrypted.EncryptedKey
        });

    public void SendDeleteMessage(string messageId, string channelId) =>
        SendJson(new SocketEncryptedMessage
        {
            Type = SignalType.ClientDeleteMessage, MessageId = messageId,
            SenderUsername = _username, ChannelId = channelId
        });

    public void SendRaw(string message)
    {
        if (_socket.ReadyState != WebSocketState.Open)
        {
            Log?.Invoke($"[{_username}] Drop: socket not open ({_socket.ReadyState}), {message.Length} bytes.");
            return;
        }

        try
        {
            _socket.Send(message);
        }
        catch (Exception ex)
        {
            Log?.Invoke($"[{_username}] Send failed ({message.Length} bytes): {ex.Message}");
            throw;
        }
    }

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

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

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

            if (root.TryGetProperty("Event", out var evEl))
            {
                var wsEvent = (WsEvent)evEl.GetInt32();
                switch (wsEvent)
                {
                    case WsEvent.KeyRegistered: Log?.Invoke($"[{_username}] Key registered."); return;
                    case WsEvent.Authenticated: Log?.Invoke($"[{_username}] Authenticated.");  return;
                    case WsEvent.Error:
                        var det = root.TryGetProperty("Detail", out var d) ? d.GetString() : null;
                        Log?.Invoke($"[{_username}] Server error: {det}");
                        return;
                }
                return;
            }

            if (!root.TryGetProperty("Type", out var typeEl)) return;
            var type = (SignalType)typeEl.GetInt32();

            switch (type)
            {
                case SignalType.ChannelList:
                {
                    var p = JsonSerializer.Deserialize<SocketChannelList>(e.Data);
                    if (p is not null) ChannelListReceived?.Invoke(p);
                    return;
                }
                case SignalType.ServerPublicKey:
                {
                    var p = JsonSerializer.Deserialize<ServerPublicKeyMessage>(e.Data);
                    if (p is not null) { ServerPublicKey = p.PublicKey; ServerPublicKeyReceived?.Invoke(p.PublicKey); }
                    return;
                }
                case SignalType.EncryptedSignal:
                {
                    var p = JsonSerializer.Deserialize<SocketRtcSignalMessage>(e.Data);
                    if (p is not null) EncryptedRtcSignalReceived?.Invoke(p);
                    return;
                }
                case SignalType.EncryptedChat:
                {
                    var p = JsonSerializer.Deserialize<SocketEncryptedMessage>(e.Data);
                    if (p is not null) EncryptedChatReceived?.Invoke(p);
                    return;
                }
                case SignalType.MessageEdited:
                {
                    var p = JsonSerializer.Deserialize<SocketEncryptedMessage>(e.Data);
                    if (p is not null) MessageEdited?.Invoke(p);
                    return;
                }
                case SignalType.MessageDeleted:
                {
                    var p = JsonSerializer.Deserialize<SocketMessageDeletedEvent>(e.Data);
                    if (p is not null) MessageDeleted?.Invoke(p);
                    return;
                }
                case SignalType.TypingIndicator:
                {
                    var p = JsonSerializer.Deserialize<SocketTypingEvent>(e.Data);
                    if (p is not null) TypingReceived?.Invoke(p);
                    return;
                }
                case SignalType.EditHistory:
                {
                    var p = JsonSerializer.Deserialize<SocketEditHistoryResponse>(e.Data);
                    if (p is not null) EditHistoryReceived?.Invoke(p);
                    return;
                }
            }
        }
        catch (Exception ex)
        {
            Log?.Invoke($"[{_username}] WS parse error: {ex.Message}");
        }
    }
}