Finished. Have at it.
This commit is contained in:
449
RelayServer/Services/Chat/ChatSocketBehavior.cs
Normal file
449
RelayServer/Services/Chat/ChatSocketBehavior.cs
Normal file
@@ -0,0 +1,449 @@
|
||||
using System.Text.Json;
|
||||
using RelayServer.Models;
|
||||
using RelayServer.Services.Crypto;
|
||||
using RelayServer.Services.Data;
|
||||
using WebSocketSharp;
|
||||
using WebSocketSharp.Server;
|
||||
|
||||
namespace RelayServer.Services.Chat;
|
||||
|
||||
/// <summary>
|
||||
/// Handles websocket-based chat operations including client key registration,
|
||||
/// server key retrieval, channel listing, channel history loading, and encrypted
|
||||
/// channel message relay.
|
||||
/// </summary>
|
||||
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; }
|
||||
|
||||
/// <summary>
|
||||
/// Routes incoming websocket messages to the appropriate chat handler.
|
||||
/// </summary>
|
||||
/// <param name="e">The websocket message event arguments.</param>
|
||||
protected override void OnMessage(MessageEventArgs e)
|
||||
{
|
||||
var msg = e.Data;
|
||||
Console.WriteLine(msg);
|
||||
|
||||
if (msg.StartsWith("REGISTER_KEY|"))
|
||||
{
|
||||
HandleRegisterKey(msg);
|
||||
return;
|
||||
}
|
||||
|
||||
if (msg == "GET_SERVER_KEY")
|
||||
{
|
||||
HandleGetServerKey();
|
||||
return;
|
||||
}
|
||||
|
||||
if (msg == "GET_CHANNELS")
|
||||
{
|
||||
HandleGetChannels();
|
||||
return;
|
||||
}
|
||||
|
||||
if (msg.StartsWith("GET_HISTORY|"))
|
||||
{
|
||||
HandleGetHistory(msg);
|
||||
return;
|
||||
}
|
||||
|
||||
HandleEncryptedChatMessage(msg);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Extracts a display username from a stored user record id value.
|
||||
/// </summary>
|
||||
/// <param name="senderUserId">The stored sender user id.</param>
|
||||
/// <returns>
|
||||
/// The extracted username when possible; otherwise, a fallback value.
|
||||
/// </returns>
|
||||
private static string ExtractUsernameFromUserId(string senderUserId)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(senderUserId))
|
||||
return "Unknown";
|
||||
|
||||
var parts = senderUserId.Split(':', 2);
|
||||
return parts.Length == 2 ? parts[1] : senderUserId;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Registers or updates a client's public key from a websocket registration payload.
|
||||
/// </summary>
|
||||
/// <param name="msg">The raw websocket registration message.</param>
|
||||
private void HandleRegisterKey(string msg)
|
||||
{
|
||||
var parts = msg.Split('|', 3);
|
||||
|
||||
if (parts.Length < 3)
|
||||
{
|
||||
Console.WriteLine("Invalid REGISTER_KEY payload.");
|
||||
return;
|
||||
}
|
||||
|
||||
var username = parts[1];
|
||||
var publicKey = parts[2];
|
||||
|
||||
if (ClientKeyService is null)
|
||||
{
|
||||
Console.WriteLine("ClientKeyService is not initialized.");
|
||||
return;
|
||||
}
|
||||
|
||||
RegisterOrUpdateClientKeySync(username, publicKey);
|
||||
|
||||
Send($"SERVER:REGISTERED_KEY:{username}");
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sends the current list of channels to the connected websocket client.
|
||||
/// </summary>
|
||||
private void HandleGetChannels()
|
||||
{
|
||||
if (Db is null)
|
||||
{
|
||||
Console.WriteLine("Db is not initialized.");
|
||||
return;
|
||||
}
|
||||
|
||||
var channels = GetChannelsSync()
|
||||
.OrderBy(c => c.CreatedAt)
|
||||
.Select(c => new SocketChannelInfo
|
||||
{
|
||||
ChannelId = GetRecordId(c.Id),
|
||||
Name = c.Name,
|
||||
CreatedAt = c.CreatedAt
|
||||
})
|
||||
.ToList();
|
||||
|
||||
var payload = new SocketChannelList
|
||||
{
|
||||
Type = "channel_list",
|
||||
Channels = channels
|
||||
};
|
||||
|
||||
Send(JsonSerializer.Serialize(payload));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sends the server's public key to the connected websocket client.
|
||||
/// </summary>
|
||||
private void HandleGetServerKey()
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(ServerPublicKey))
|
||||
{
|
||||
Console.WriteLine("Server public key is not initialized.");
|
||||
return;
|
||||
}
|
||||
|
||||
var payload = new ServerPublicKeyMessage
|
||||
{
|
||||
Type = "server_public_key",
|
||||
PublicKey = ServerPublicKey
|
||||
};
|
||||
|
||||
Send(JsonSerializer.Serialize(payload));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Decrypts an incoming encrypted chat payload, stores it in the database,
|
||||
/// and rebroadcasts it to connected clients encrypted with each client's public key.
|
||||
/// </summary>
|
||||
/// <param name="msg">The raw encrypted chat websocket message.</param>
|
||||
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 != "client_encrypted_chat")
|
||||
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 allKeys = GetAllClientPublicKeysSync();
|
||||
|
||||
foreach (var client in allKeys)
|
||||
{
|
||||
var encrypted = E2EeHelper.EncryptForRecipient(plainText, client.PublicKey);
|
||||
|
||||
Console.WriteLine($"Encrypting outbound message from {clientPayload.SenderUsername} for {client.Username}");
|
||||
|
||||
var outbound = new SocketEncryptedMessage
|
||||
{
|
||||
Type = "encrypted_chat",
|
||||
SenderUsername = clientPayload.SenderUsername,
|
||||
RecipientUsername = client.Username,
|
||||
ChannelId = clientPayload.ChannelId,
|
||||
CipherText = encrypted.CipherText,
|
||||
Nonce = encrypted.Nonce,
|
||||
Tag = encrypted.Tag,
|
||||
EncryptedKey = encrypted.EncryptedKey
|
||||
};
|
||||
|
||||
Sessions.Broadcast(JsonSerializer.Serialize(outbound));
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Loads stored channel history for a specific user and channel, decrypts it from
|
||||
/// database storage format, and sends it back encrypted for the requesting client.
|
||||
/// </summary>
|
||||
/// <param name="msg">The raw history request websocket message.</param>
|
||||
private void HandleGetHistory(string msg)
|
||||
{
|
||||
var parts = msg.Split('|', 3);
|
||||
|
||||
if (parts.Length < 3)
|
||||
{
|
||||
Console.WriteLine("Invalid GET_HISTORY payload.");
|
||||
return;
|
||||
}
|
||||
|
||||
var username = parts[1];
|
||||
var channelId = parts[2];
|
||||
|
||||
if (!EnsureCoreReady() || ChannelCryptoService is null || string.IsNullOrWhiteSpace(ChannelDbKey))
|
||||
{
|
||||
Console.WriteLine("History dependencies are not initialized.");
|
||||
return;
|
||||
}
|
||||
|
||||
var targetClient = GetClientPublicKeyByUsernameSync(username);
|
||||
|
||||
if (targetClient is null)
|
||||
{
|
||||
Console.WriteLine($"No public key found for history request user {username}");
|
||||
return;
|
||||
}
|
||||
|
||||
var allMessages = GetChannelMessagesSync();
|
||||
|
||||
var channelMessages = allMessages
|
||||
.Where(m => m.ChannelId == channelId)
|
||||
.OrderBy(m => m.CreatedAt)
|
||||
.ToList();
|
||||
|
||||
Console.WriteLine($"Sending {channelMessages.Count} history messages to {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 = "encrypted_chat",
|
||||
SenderUsername = ExtractUsernameFromUserId(dbMessage.SenderUserId),
|
||||
RecipientUsername = username,
|
||||
ChannelId = channelId,
|
||||
CipherText = encrypted.CipherText,
|
||||
Nonce = encrypted.Nonce,
|
||||
Tag = encrypted.Tag,
|
||||
EncryptedKey = encrypted.EncryptedKey
|
||||
};
|
||||
|
||||
Send(JsonSerializer.Serialize(outbound));
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Converts a SurrealDB record id object into a table:id string representation.
|
||||
/// </summary>
|
||||
/// <param name="id">The raw record id object.</param>
|
||||
/// <returns>
|
||||
/// A formatted record id string, or an empty string if the input is null.
|
||||
/// </returns>
|
||||
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}";
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Synchronously registers or updates a stored client public key using the async key service.
|
||||
/// </summary>
|
||||
/// <param name="username">The client username.</param>
|
||||
/// <param name="publicKey">The client's public key.</param>
|
||||
private void RegisterOrUpdateClientKeySync(string username, string publicKey)
|
||||
{
|
||||
Task.Run(async () => await ClientKeyService!.RegisterOrUpdateKeyAsync(username, publicKey))
|
||||
.GetAwaiter()
|
||||
.GetResult();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Synchronously loads all channels from the database.
|
||||
/// </summary>
|
||||
/// <returns>A list of channel records.</returns>
|
||||
private List<Channels> GetChannelsSync()
|
||||
{
|
||||
return Task.Run(async () => await Db!.Select<Channels>("channels"))
|
||||
.GetAwaiter()
|
||||
.GetResult()
|
||||
.ToList();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Synchronously gets the stored public key record for the specified user.
|
||||
/// </summary>
|
||||
/// <param name="username">The username to look up.</param>
|
||||
/// <returns>
|
||||
/// The matching client public key record, or null if none exists.
|
||||
/// </returns>
|
||||
private ClientPublicKeys? GetClientPublicKeyByUsernameSync(string username)
|
||||
{
|
||||
return Task.Run(async () => await ClientKeyService!.GetByUsernameAsync(username))
|
||||
.GetAwaiter()
|
||||
.GetResult();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Synchronously loads all stored client public key records.
|
||||
/// </summary>
|
||||
/// <returns>A list of all client public key records.</returns>
|
||||
private List<ClientPublicKeys> GetAllClientPublicKeysSync()
|
||||
{
|
||||
return Task.Run(async () => await ClientKeyService!.GetAllAsync())
|
||||
.GetAwaiter()
|
||||
.GetResult();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Synchronously loads all stored channel messages from the database.
|
||||
/// </summary>
|
||||
/// <returns>A list of channel message records.</returns>
|
||||
private List<ChannelMessages> GetChannelMessagesSync()
|
||||
{
|
||||
return Task.Run(async () => await Db!.Select<ChannelMessages>("channel_messages"))
|
||||
.GetAwaiter()
|
||||
.GetResult()
|
||||
.ToList();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Synchronously creates a new channel message record in the database.
|
||||
/// </summary>
|
||||
/// <param name="message">The message record to create.</param>
|
||||
/// <returns>The created channel message record.</returns>
|
||||
private ChannelMessages CreateChannelMessageSync(ChannelMessages message)
|
||||
{
|
||||
return Task.Run(async () => await Db!.Create("channel_messages", message))
|
||||
.GetAwaiter()
|
||||
.GetResult();
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user