From 647eeec8a8f7be4057cfcb47c7bfaad0a76ee977 Mon Sep 17 00:00:00 2001 From: Klesti Selimaj Date: Sat, 29 Aug 2026 13:27:17 +0200 Subject: [PATCH] Rewritten server code --- src/crypto.rs | 4 +- src/main.rs | 3 +- src/protocol/initialize.rs | 217 +++++++++++++++---------------- src/protocol/message.rs | 36 ++++-- src/protocol/user.rs | 2 +- src/protocol/voice.rs | 46 +++---- src/server.rs | 205 ------------------------------ src/server/identity.rs | 26 ++++ src/server/mod.rs | 90 +++++++++++++ src/server/session.rs | 133 +++++++++++++++++++ src/server/store.rs | 17 +++ src/server/vc_server.rs | 254 +++++++++++++++++++++++++++++++++++++ src/vc_server.rs | 179 -------------------------- 13 files changed, 671 insertions(+), 541 deletions(-) delete mode 100644 src/server.rs create mode 100644 src/server/identity.rs create mode 100644 src/server/mod.rs create mode 100644 src/server/session.rs create mode 100644 src/server/store.rs create mode 100644 src/server/vc_server.rs delete mode 100644 src/vc_server.rs diff --git a/src/crypto.rs b/src/crypto.rs index a848622..d0caaa2 100644 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -142,7 +142,7 @@ pub async fn crypto_handshake( ) -> anyhow::Result { socket .send(axum::extract::ws::Message::Binary( - server.x_keypair.0.to_bytes().to_vec().into(), + server.identity.x25519.public.to_bytes().to_vec().into(), )) .await?; @@ -159,7 +159,7 @@ pub async fn crypto_handshake( "Failed to get proper length of client x key" ))?); - let shared_secret = server.x_keypair.1.diffie_hellman(&client_pubkey); + let shared_secret = server.identity.x25519.secret.diffie_hellman(&client_pubkey); let cipher = Arc::new(Mutex::new(SessionCipher::new(&shared_secret)?)); diff --git a/src/main.rs b/src/main.rs index a4524bf..7f5b833 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,7 +3,6 @@ pub mod data; pub mod protocol; pub mod server; pub mod types; -pub mod vc_server; pub mod ws; use std::{ @@ -28,7 +27,7 @@ use crate::server::Server; async fn main() -> anyhow::Result<()> { let server = Server::new().await?; - let udp_server = tokio::spawn(server.clone().start_udp_server()); + let udp_server = tokio::spawn(server.voice.clone().run(server.config.port)); let cors = CorsLayer::new() .allow_origin(Any) diff --git a/src/protocol/initialize.rs b/src/protocol/initialize.rs index b8c7148..59452c7 100644 --- a/src/protocol/initialize.rs +++ b/src/protocol/initialize.rs @@ -5,121 +5,122 @@ use std::{ use ed25519_dalek::{Signer, VerifyingKey}; -use crate::{server::Server, ws::EnclaveWebSocket}; +use crate::{ + server::Server, + types::ClientMeta, + ws::EnclaveWebSocket, +}; use super::*; -use crate::server::UserConnections; -impl UserConnections { - pub async fn initialize( - server: &Arc, - socket: &EnclaveWebSocket, - ) -> anyhow::Result<(VerifyingKey, ClientMeta)> { - let Some(ServerMethod::Initialize { - public_key: public_key_string, - signature, +pub async fn initialize( + server: &Arc, + socket: &EnclaveWebSocket, +) -> anyhow::Result<(VerifyingKey, ClientMeta)> { + let Some(ServerMethod::Initialize { + public_key: public_key_string, + signature, - timestamp, - hostname, - }) = socket.read().await? - else { - socket - .send(&ClientMethod::Error { - error: Cow::Borrowed("Initialization required"), - }) - .await?; - - return Err(anyhow::anyhow!( - "Failed to initialize: Client sent the wrong method" - )); - }; - - let server_timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64; - - if server_timestamp.saturating_sub(timestamp) > 2000 { - socket - .send(&ClientMethod::Error { - error: Cow::Borrowed( - "Timestamp doesn't match, make sure it's in secs and is (<= 2secs)", - ), - }) - .await?; - - return Err(anyhow::anyhow!("Client tampstamp wasn't correct")); - } - - if !server.config.hostnames.contains(&hostname) { - socket.send( - &ClientMethod::Error { - error: Cow::Owned(format!("Invalid Hostname, to avoid man-in-the-middle attacks, please use the correct hostname(s): {}", server.config.hostnames.clone().into_iter().collect::>().join(", "))), - }, - ) + timestamp, + hostname, + }) = socket.read().await? + else { + socket + .send(&ClientMethod::Error { + error: Cow::Borrowed("Initialization required"), + }) .await?; - return Err(anyhow::anyhow!("Client's hostname wasn't correct")); - } + return Err(anyhow::anyhow!( + "Failed to initialize: Client sent the wrong method" + )); + }; - let Ok(public_key) = crate::crypto::from_string(&public_key_string) else { - socket - .send(&ClientMethod::Error { - error: Cow::Borrowed("Invalid public key"), - }) - .await?; + let server_timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64; - return Err(anyhow::anyhow!("Invalid public key")); - }; + if server_timestamp.saturating_sub(timestamp) > 2000 { + socket + .send(&ClientMethod::Error { + error: Cow::Borrowed( + "Timestamp doesn't match, make sure it's in secs and is (<= 2secs)", + ), + }) + .await?; - if public_key - .verify_strict( - format!("{timestamp}@{hostname}").as_bytes(), - &crate::crypto::from_string_sig(&signature)?, - ) - .is_err() - { - socket - .send(&ClientMethod::Error { - error: Cow::Borrowed("Invalid signature"), - }) - .await?; - - return Err(anyhow::anyhow!("Invalid signature")); - } - - { - socket - .send(&ClientMethod::Initialized { - public_key: crate::crypto::to_string(&server.key.verifying_key()), - signature: crate::crypto::to_string_sig(&server.key.sign( - format!("{server_timestamp}@{hostname}@{public_key_string}").as_bytes(), - )), - - timestamp: server_timestamp, - hostname, - }) - .await?; - } - - let Some(ServerMethod::Meta(meta)) = socket.read().await? else { - socket - .send(&ClientMethod::Error { - error: Cow::Borrowed("Expected meta"), - }) - .await?; - - return Err(anyhow::anyhow!( - "Expected meta, client called another method" - )); - }; - - for (pubkey, channel_id) in server.voice_pins.lock().await.values() { - socket - .send(&ClientMethod::UserJoinedVoice { - channel_id: channel_id.clone(), - pubkey: crate::crypto::to_string(pubkey), - }) - .await?; - } - - Ok((public_key, meta)) + return Err(anyhow::anyhow!("Client tampstamp wasn't correct")); } -} + + if !server.config.hostnames.contains(&hostname) { + socket.send( + &ClientMethod::Error { + error: Cow::Owned(format!("Invalid Hostname, to avoid man-in-the-middle attacks, please use the correct hostname(s): {}", server.config.hostnames.clone().into_iter().collect::>().join(", "))), + }, + ) + .await?; + + return Err(anyhow::anyhow!("Client's hostname wasn't correct")); + } + + let Ok(public_key) = crate::crypto::from_string(&public_key_string) else { + socket + .send(&ClientMethod::Error { + error: Cow::Borrowed("Invalid public key"), + }) + .await?; + + return Err(anyhow::anyhow!("Invalid public key")); + }; + + if public_key + .verify_strict( + format!("{timestamp}@{hostname}").as_bytes(), + &crate::crypto::from_string_sig(&signature)?, + ) + .is_err() + { + socket + .send(&ClientMethod::Error { + error: Cow::Borrowed("Invalid signature"), + }) + .await?; + + return Err(anyhow::anyhow!("Invalid signature")); + } + + { + socket + .send(&ClientMethod::Initialized { + public_key: crate::crypto::to_string(&server.identity.key.verifying_key()), + signature: crate::crypto::to_string_sig(&server.identity.key.sign( + format!("{server_timestamp}@{hostname}@{public_key_string}").as_bytes(), + )), + + timestamp: server_timestamp, + hostname, + }) + .await?; + } + + let Some(ServerMethod::Meta(meta)) = socket.read().await? else { + socket + .send(&ClientMethod::Error { + error: Cow::Borrowed("Expected meta"), + }) + .await?; + + return Err(anyhow::anyhow!( + "Expected meta, client called another method" + )); + }; + + for pin in server.voice.pins.lock().await.values() { + socket + .send(&ClientMethod::UserJoinedVoice { + channel_id: pin.channel_id.clone(), + pubkey: crate::crypto::to_string(&pin.pubkey), + }) + .await?; + } + + Ok((public_key, meta)) +} \ No newline at end of file diff --git a/src/protocol/message.rs b/src/protocol/message.rs index 91d7a1e..c6776e7 100644 --- a/src/protocol/message.rs +++ b/src/protocol/message.rs @@ -22,7 +22,7 @@ pub async fn send_message( ); } - let server_pubkey_string = crate::crypto::to_string(&server.key.verifying_key()); + let server_pubkey_string = crate::crypto::to_string(&server.identity.key.verifying_key()); let signed_string = format!( "{}@{}@{}", message.timestamp, server_pubkey_string, message.content @@ -42,10 +42,11 @@ pub async fn send_message( data: message, }; - server.message_store.insert_message(&channel_id, &stored)?; + server.store.messages.insert_message(&channel_id, &stored)?; - server - .broadcast(&ClientMethod::Messages { +server + .sessions + .broadcast(&ClientMethod::Messages { messages: HashMap::from([(channel_id, vec![stored])]), }) .await?; @@ -63,7 +64,8 @@ pub async fn get_messages( const CHUNK_SIZE: u32 = 16; let messages = server - .message_store + .store + .messages .get_recent_messages(&channel_id, CHUNK_SIZE, chunk)?; socket @@ -84,7 +86,8 @@ pub async fn edit_message( new_signature: String, ) -> anyhow::Result<()> { let existing = server - .message_store + .store + .messages .get_message(&channel_id, &message_id)? .ok_or_else(|| anyhow::anyhow!("Message not found"))?; @@ -93,7 +96,7 @@ pub async fn edit_message( anyhow::bail!("Not authorized to edit this message"); } - let server_pubkey_string = crate::crypto::to_string(&server.key.verifying_key()); + let server_pubkey_string = crate::crypto::to_string(&server.identity.key.verifying_key()); let signed_string = format!( "{}@{}@{}", existing.data.timestamp, server_pubkey_string, new_content @@ -107,7 +110,8 @@ pub async fn edit_message( .map_err(|_| anyhow::anyhow!("Signature verification failed"))?; server - .message_store + .store + .messages .update_message(&channel_id, &message_id, &new_content, &new_signature)?; let updated = StoredMessage { @@ -121,8 +125,9 @@ pub async fn edit_message( }, }; - server - .broadcast(&ClientMethod::MessageEdited { +server + .sessions + .broadcast(&ClientMethod::MessageEdited { channel_id: channel_id.clone(), message: updated, }) @@ -138,7 +143,8 @@ pub async fn delete_message( channel_id: String, ) -> anyhow::Result<()> { let existing = server - .message_store + .store + .messages .get_message(&channel_id, &message_id)? .ok_or_else(|| anyhow::anyhow!("Message not found"))?; @@ -148,11 +154,13 @@ pub async fn delete_message( } server - .message_store + .store + .messages .delete_message(&channel_id, &message_id)?; - server - .broadcast(&ClientMethod::MessageDeleted { +server + .sessions + .broadcast(&ClientMethod::MessageDeleted { channel_id: channel_id.clone(), message_id, }) diff --git a/src/protocol/user.rs b/src/protocol/user.rs index c56e76f..1295aa8 100644 --- a/src/protocol/user.rs +++ b/src/protocol/user.rs @@ -10,7 +10,7 @@ pub async fn get_users( socket: &Arc, pubkeys: Vec, ) -> anyhow::Result<()> { - let users = server.user_store.get_users(&pubkeys).await?; + let users = server.store.users.get_users(&pubkeys).await?; socket.send(&ClientMethod::Users { users }).await?; diff --git a/src/protocol/voice.rs b/src/protocol/voice.rs index 563225b..d4f60c4 100644 --- a/src/protocol/voice.rs +++ b/src/protocol/voice.rs @@ -10,24 +10,23 @@ pub async fn join( socket: &Arc, channel_id: String, ) -> anyhow::Result<()> { - { - let pin = rand::random::() % (1 << 53); + let user = server + .sessions + .get(&verifying_key) + .await + .ok_or_else(|| anyhow::anyhow!("Not connected"))?; - server - .voice_pins - .lock() - .await - .insert(pin, (verifying_key, channel_id.clone())); + let pin = server.voice.join(verifying_key, user, &channel_id).await; - socket - .send(&ClientMethod::JoinVoice { - channel_id: channel_id.clone(), - pin, - }) - .await?; - } + socket + .send(&ClientMethod::JoinVoice { + channel_id: channel_id.clone(), + pin, + }) + .await?; server + .sessions .broadcast(&ClientMethod::UserJoinedVoice { channel_id, pubkey: crate::crypto::to_string(&verifying_key), @@ -38,25 +37,12 @@ pub async fn join( } pub async fn leave(server: &Arc, verifying_key: VerifyingKey) -> anyhow::Result<()> { - let Some(channel_id) = ({ - let clients = server.clients.lock().await; - - let Some(client) = clients.get(&verifying_key) else { - return Ok(()); - }; - - client.voice.lock().await.take().map(|v| v.channel_id) - }) else { + let Some(channel_id) = server.voice.remove(verifying_key).await else { return Ok(()); }; server - .voice_pins - .lock() - .await - .retain(|_, v| v.0 != verifying_key); - - server + .sessions .broadcast(&ClientMethod::UserLeftVoice { channel_id, pubkey: crate::crypto::to_string(&verifying_key), @@ -64,4 +50,4 @@ pub async fn leave(server: &Arc, verifying_key: VerifyingKey) -> anyhow: .await?; Ok(()) -} +} \ No newline at end of file diff --git a/src/server.rs b/src/server.rs deleted file mode 100644 index 6c908ee..0000000 --- a/src/server.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::{ - collections::HashMap, - net::SocketAddr, - path::PathBuf, - sync::{Arc, atomic::AtomicU16}, -}; - -use axum::{ - extract::{WebSocketUpgrade, ws::WebSocket}, - response::Response, -}; -use ed25519_dalek::{SigningKey, VerifyingKey}; -use tokio::{ - net::UdpSocket, - sync::{Mutex, OnceCell}, - task::JoinSet, - time::Instant, -}; - -use crate::{ - crypto::SessionCipher, - data::{config::Config, messages::MessageStore, users::UserMetaStore}, - protocol::{ClientMethod, read_loop}, - types::ClientMeta, -}; -use x25519_dalek::{PublicKey as X25519Public, StaticSecret as X25519Secret}; - -pub struct VoiceConnection { - pub addr: SocketAddr, - pub channel_id: String, - pub last_speaking_sent: Instant, -} - -pub struct UserConnections { - pub meta: ClientMeta, - pub counter: AtomicU16, - pub public_key: VerifyingKey, - pub connections: Mutex>>, - pub cihper: Arc>, - pub voice: Mutex>, -} - -pub struct Server { - pub key: SigningKey, - pub x_keypair: (X25519Public, X25519Secret), - pub config: Config, - pub clients: Mutex>>, - pub voice_pins: Mutex>, - pub message_store: MessageStore, - pub user_store: UserMetaStore, - pub voice_socket: OnceCell, -} - -impl Server { - pub async fn new() -> anyhow::Result> { - let key = crate::crypto::get().await?; - Ok(Arc::new(Self { - x_keypair: ( - crate::crypto::ed25519_verifying_key_to_x25519(&key.verifying_key()) - .ok_or(anyhow::anyhow!("Failed to convert ed pubkey to x"))?, - crate::crypto::ed25519_signing_key_to_x25519(&key), - ), - key, - config: Config::get().await?, - clients: Mutex::new(HashMap::new()), - voice_pins: Mutex::new(HashMap::new()), - message_store: MessageStore::new(PathBuf::from("messages"))?, - user_store: UserMetaStore::new(PathBuf::from("users.db"))?, - voice_socket: OnceCell::new(), - })) - } -} - -impl Server { - pub async fn ws_handler(self: &Arc, ws: WebSocketUpgrade) -> Response { - let s = self.clone(); - - ws.on_upgrade(move |socket: WebSocket| async move { - let mut client = match crate::crypto::crypto_handshake(&s, socket).await { - Ok(client) => client, - Err(err) => { - eprintln!("Failed to initialize crypto: {err}"); - return; - } - }; - - match UserConnections::initialize(&s, &client).await { - Ok((public_key, meta)) => { - if let Err(e) = s - .user_store - .upsert_user(&crate::crypto::to_string(&public_key), &meta) - .await - { - eprintln!("Failed to upsert client: {e}"); - } - - let mut clients_meta = s.clients.lock().await; - - let clients = clients_meta - .entry(public_key) - .or_insert_with(|| { - Arc::new(UserConnections { - meta, - public_key: public_key, - counter: AtomicU16::new(0), - connections: Mutex::new(HashMap::new()), - voice: Mutex::new(None), - cihper: client.cipher.clone(), - }) - }) - .clone(); - - client.cipher = clients.cihper.clone(); - - let client = Arc::new(client); - - let conid = clients - .counter - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - - clients - .connections - .lock() - .await - .insert(conid, client.clone()); - - drop(clients_meta); - - if let Err(e) = read_loop(&s, public_key, &client).await { - eprintln!("Failed to handle client: {e}"); - } - - let mut clients_meta = s.clients.lock().await; - - let mut connections = clients.connections.lock().await; - - connections.remove(&conid); - - if connections.is_empty() { - clients_meta.remove(&public_key); - - if let Some(voice) = clients.voice.lock().await.take() { - s.voice_pins - .lock() - .await - .retain(|_, v| v.0 != public_key); - - s.broadcast(&crate::protocol::ClientMethod::UserLeftVoice { - channel_id: voice.channel_id, - pubkey: crate::crypto::to_string(&public_key), - }) - .await - .ok(); - } - } - } - - Err(e) => { - eprintln!("Failed to initialize client: {e}") - } - } - }) - } - - pub async fn broadcast(self: &Arc, message: &ClientMethod) -> anyhow::Result<()> { - let mut set = JoinSet::new(); - - for (_, client) in self.clients.lock().await.iter() { - let msg = message.clone(); - let client = client.clone(); - - set.spawn(async move { client.send(&msg).await }); - } - - // Await all spawned tasks to finish - while let Some(res) = set.join_next().await { - // handle task panic or errors if necessary - if let Ok(Err(e)) = res { - eprintln!("Failed to send to a client: {:?}", e); - } - } - - Ok(()) - } -} - -impl UserConnections { - pub async fn send(&self, message: &ClientMethod) -> anyhow::Result<()> { - for (_, conn) in self.connections.lock().await.iter() { - conn.send(message).await?; - } - - Ok(()) - } - - pub async fn send_to(&self, id: u16, message: &ClientMethod) -> anyhow::Result { - if let Some(conn) = self.connections.lock().await.get(&id) { - conn.send(message).await?; - - Ok(true) - } else { - Ok(false) - } - } -} diff --git a/src/server/identity.rs b/src/server/identity.rs new file mode 100644 index 0000000..3d76747 --- /dev/null +++ b/src/server/identity.rs @@ -0,0 +1,26 @@ +use ed25519_dalek::SigningKey; +use x25519_dalek::{PublicKey as X25519Public, StaticSecret as X25519Secret}; + +pub struct ServerIdentity { + pub key: SigningKey, + pub x25519: X25519KeyPair, +} + +pub struct X25519KeyPair { + pub public: X25519Public, + pub secret: X25519Secret, +} + +impl ServerIdentity { + pub async fn load() -> anyhow::Result { + let key = crate::crypto::get().await?; + + let x25519 = X25519KeyPair { + public: crate::crypto::ed25519_verifying_key_to_x25519(&key.verifying_key()) + .ok_or(anyhow::anyhow!("Failed to convert ed pubkey to x"))?, + secret: crate::crypto::ed25519_signing_key_to_x25519(&key), + }; + + Ok(Self { key, x25519 }) + } +} \ No newline at end of file diff --git a/src/server/mod.rs b/src/server/mod.rs new file mode 100644 index 0000000..fe9bfbb --- /dev/null +++ b/src/server/mod.rs @@ -0,0 +1,90 @@ +pub mod identity; +pub mod session; +pub mod store; +pub mod vc_server; + +use std::sync::Arc; + +use axum::{ + extract::{WebSocketUpgrade, ws::WebSocket}, + response::Response, +}; + +use crate::{ + data::config::Config, + protocol::{read_loop, ClientMethod}, +}; + +pub use identity::{ServerIdentity, X25519KeyPair}; +pub use session::{SessionRegistry, UserConnections}; +pub use store::DataStore; +pub use vc_server::{VoiceConnection, VoicePin, VoiceServer}; + +pub struct Server { + pub identity: ServerIdentity, + pub config: Config, + pub sessions: SessionRegistry, + pub voice: Arc, + pub store: DataStore, +} + +impl Server { + pub async fn new() -> anyhow::Result> { + Ok(Arc::new(Self { + identity: ServerIdentity::load().await?, + config: Config::get().await?, + sessions: SessionRegistry::new(), + voice: Arc::new(VoiceServer::new()), + store: DataStore::new()?, + })) + } + + pub async fn ws_handler(self: &Arc, ws: WebSocketUpgrade) -> Response { + let s = self.clone(); + + ws.on_upgrade(move |socket: WebSocket| async move { + let client = match crate::crypto::crypto_handshake(&s, socket).await { + Ok(client) => client, + Err(err) => { + eprintln!("Failed to initialize crypto: {err}"); + return; + } + }; + + match crate::protocol::initialize::initialize(&s, &client).await { + Ok((public_key, meta)) => { + if let Err(e) = s + .store + .users + .upsert_user(&crate::crypto::to_string(&public_key), &meta) + .await + { + eprintln!("Failed to upsert client: {e}"); + } + + let (client, conid) = s.sessions.register(public_key, meta, client).await; + + if let Err(e) = read_loop(&s, public_key, &client).await { + eprintln!("Failed to handle client: {e}"); + } + + if s.sessions.deregister(public_key, conid).await + && let Some(channel_id) = s.voice.remove(public_key).await + { + s.sessions + .broadcast(&ClientMethod::UserLeftVoice { + channel_id, + pubkey: crate::crypto::to_string(&public_key), + }) + .await + .ok(); + } + } + + Err(e) => { + eprintln!("Failed to initialize client: {e}") + } + } + }) + } +} \ No newline at end of file diff --git a/src/server/session.rs b/src/server/session.rs new file mode 100644 index 0000000..da0d153 --- /dev/null +++ b/src/server/session.rs @@ -0,0 +1,133 @@ +use std::{ + collections::HashMap, + sync::{Arc, atomic::AtomicU16}, +}; + +use ed25519_dalek::VerifyingKey; +use tokio::sync::Mutex; + +use crate::{ + crypto::SessionCipher, + protocol::ClientMethod, + types::ClientMeta, + ws::EnclaveWebSocket, +}; + +pub struct UserConnections { + pub meta: ClientMeta, + pub counter: AtomicU16, + pub public_key: VerifyingKey, + pub connections: Mutex>>, + pub cipher: Arc>, +} + +pub struct SessionRegistry { + pub clients: Mutex>>, +} + +impl Default for SessionRegistry { + fn default() -> Self { + Self::new() + } +} + +impl SessionRegistry { + pub fn new() -> Self { + Self { + clients: Mutex::new(HashMap::new()), + } + } + + pub async fn get(&self, public_key: &VerifyingKey) -> Option> { + self.clients.lock().await.get(public_key).cloned() + } + + /// Registers a new websocket connection for a user, returning the + /// connection and its id. The websocket is given the shared cipher + /// of the user's existing connections so voice traffic stays keyed + /// consistently across devices. + pub async fn register( + &self, + public_key: VerifyingKey, + meta: ClientMeta, + mut client: EnclaveWebSocket, + ) -> (Arc, u16) { + let mut clients = self.clients.lock().await; + + let user = clients + .entry(public_key) + .or_insert_with(|| { + Arc::new(UserConnections { + meta, + public_key, + counter: AtomicU16::new(0), + connections: Mutex::new(HashMap::new()), + cipher: client.cipher.clone(), + }) + }) + .clone(); + + client.cipher = user.cipher.clone(); + + let client = Arc::new(client); + + let conid = user.counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + user.connections.lock().await.insert(conid, client.clone()); + + (client, conid) + } + + /// Removes a connection from a user's session. Returns `true` when + /// that was the user's last connection and they have been dropped + /// from the registry entirely. + pub async fn deregister(&self, public_key: VerifyingKey, conid: u16) -> bool { + let mut clients = self.clients.lock().await; + + let Some(user) = clients.get(&public_key).cloned() else { + return false; + }; + + let mut connections = user.connections.lock().await; + connections.remove(&conid); + + if connections.is_empty() { + clients.remove(&public_key); + true + } else { + false + } + } + + pub async fn broadcast(&self, message: &ClientMethod) -> anyhow::Result<()> { + let users: Vec<_> = self.clients.lock().await.values().cloned().collect(); + + for user in users { + if let Err(e) = user.send(message).await { + eprintln!("Failed to send to a client: {:?}", e); + } + } + + Ok(()) + } +} + +impl UserConnections { + pub async fn send(&self, message: &ClientMethod) -> anyhow::Result<()> { + for (_, conn) in self.connections.lock().await.iter() { + conn.send(message).await?; + } + + Ok(()) + } + + pub async fn send_to(&self, id: u16, message: &ClientMethod) -> anyhow::Result { + if let Some(conn) = self.connections.lock().await.get(&id) { + conn.send(message).await?; + + Ok(true) + } else { + Ok(false) + } + } +} \ No newline at end of file diff --git a/src/server/store.rs b/src/server/store.rs new file mode 100644 index 0000000..d3465ab --- /dev/null +++ b/src/server/store.rs @@ -0,0 +1,17 @@ +use std::path::PathBuf; + +use crate::data::{messages::MessageStore, users::UserMetaStore}; + +pub struct DataStore { + pub messages: MessageStore, + pub users: UserMetaStore, +} + +impl DataStore { + pub fn new() -> anyhow::Result { + Ok(Self { + messages: MessageStore::new(PathBuf::from("messages"))?, + users: UserMetaStore::new(PathBuf::from("users.db"))?, + }) + } +} \ No newline at end of file diff --git a/src/server/vc_server.rs b/src/server/vc_server.rs new file mode 100644 index 0000000..ed43c23 --- /dev/null +++ b/src/server/vc_server.rs @@ -0,0 +1,254 @@ +use std::{ + collections::HashMap, + net::{IpAddr, Ipv4Addr, SocketAddr}, + sync::Arc, +}; + +use anyhow::Context; +use ed25519_dalek::VerifyingKey; +use tokio::{ + net::UdpSocket, + sync::{Mutex, OnceCell}, + time::Instant, +}; + +use crate::{ + crypto::SessionCipher, + protocol::ClientMethod, + server::UserConnections, +}; + +pub struct VoicePin { + pub pubkey: VerifyingKey, + pub channel_id: String, +} + +pub struct VoiceConnection { + pub user: Arc, + pub channel_id: String, + pub addr: SocketAddr, + pub last_speaking_sent: Instant, +} + +pub struct VoiceServer { + pub pins: Mutex>, + pub socket: OnceCell, + pub participants: Mutex>, +} + +impl Default for VoiceServer { + fn default() -> Self { + Self::new() + } +} + +impl VoiceServer { + pub fn new() -> Self { + Self { + pins: Mutex::new(HashMap::new()), + socket: OnceCell::new(), + participants: Mutex::new(HashMap::new()), + } + } + + /// Adds a user to a voice channel and allocates a one-time pin that + /// their next UDP packet must carry to bind their address. + pub async fn join( + &self, + public_key: VerifyingKey, + user: Arc, + channel_id: &str, + ) -> u64 { + let pin = rand::random::() % (1 << 53); + + self.participants.lock().await.insert( + public_key, + VoiceConnection { + user, + channel_id: channel_id.to_string(), + addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0), + last_speaking_sent: Instant::now(), + }, + ); + + self.pins.lock().await.insert( + pin, + VoicePin { + pubkey: public_key, + channel_id: channel_id.to_string(), + }, + ); + + pin + } + + /// Removes a user from every voice channel/pin they hold, returning + /// the channel they were in, if any. + pub async fn remove(&self, public_key: VerifyingKey) -> Option { + let channel_id = self + .participants + .lock() + .await + .remove(&public_key)? + .channel_id; + + self.pins.lock().await.retain(|_, pin| pin.pubkey != public_key); + + Some(channel_id) + } + + /// Runs the UDP loop that terminates voice audio: binds addresses on + /// the first (pin-bearing) packet and relays encrypted audio between + /// the participants of each channel. + pub async fn run(self: Arc, port: u16) -> anyhow::Result<()> { + let socket = UdpSocket::bind(("0.0.0.0", port)).await?; + + self.socket + .set(socket) + .map_err(|_| anyhow::anyhow!("UDP server already started"))?; + + let mut buf = [0u8; 4096]; + + eprintln!("[vc] UDP server listening on port {port}"); + + loop { + let (len, addr) = self.get_voice_socket()?.recv_from(&mut buf).await?; + + if len < 8 { + eprintln!("[vc] dropping packet too short for pin"); + continue; + } + + let pin_bytes: [u8; 8] = buf[..8].try_into().unwrap(); + let pin = u64::from_be_bytes(pin_bytes); + + // is_first_packet distinguishes the plaintext pin-bootstrap packet + // from subsequent encrypted audio packets. + let (sender_pubkey, channel_id, payload, is_first_packet) = { + let mut pins = self.pins.lock().await; + + if let Some(pin) = pins.remove(&pin) { + (pin.pubkey, pin.channel_id, &buf[8..len], true) + } else { + drop(pins); + + match self.find_sender(&addr).await { + Some((pubkey, channel_id)) => (pubkey, channel_id, &buf[..len], false), + None => continue, + } + } + }; + + if let Some(participant) = self.participants.lock().await.get_mut(&sender_pubkey) { + participant.addr = addr; + } + + // The pin-bearing bootstrap packet carries no payload to decrypt — + // it's purely "here's my pin, bind my address." + if is_first_packet { + continue; + } + + let Some(cipher) = ({ + let participants = self.participants.lock().await; + participants.get(&sender_pubkey).map(|p| p.user.cipher.clone()) + }) else { + eprintln!("[vc] sender pubkey not found in participants, dropping"); + continue; + }; + + let decrypted_payload = match cipher.lock().await.decrypt(payload) { + Ok(pt) => pt, + Err(e) => { + eprintln!("[vc] dropping packet: decryption failed: {e}"); + continue; + } + }; + + let s = self.clone(); + + tokio::spawn(async move { + if let Err(e) = s + .relay_voice(&sender_pubkey, &channel_id, &decrypted_payload) + .await + { + eprintln!("{e}"); + } + }); + } + } + + /// Looks up which known voice participant a UDP address belongs to, + /// for packets arriving after the initial pin-bearing packet. + async fn find_sender(&self, addr: &SocketAddr) -> Option<(VerifyingKey, String)> { + let participants = self.participants.lock().await; + + for (pubkey, participant) in participants.iter() { + if *addr == participant.addr { + return Some((*pubkey, participant.channel_id.clone())); + } + } + + None + } + + /// Sends `payload` to every voice participant currently in `channel_id`. + async fn relay_voice( + &self, + sender: &VerifyingKey, + channel_id: &str, + payload: &[u8], + ) -> anyhow::Result<()> { + let mut participants = self.participants.lock().await; + + for (_pubkey, participant) in participants.iter_mut() { + if channel_id != participant.channel_id { + continue; + } + + let now = Instant::now(); + + if now.duration_since(participant.last_speaking_sent).as_millis() >= 600 { + participant + .user + .send(&ClientMethod::Speaking { + pubkey: crate::crypto::to_string(sender), + }) + .await?; + + participant.last_speaking_sent = now; + } + + if *sender == participant.user.public_key { + continue; + } + + let _ = self + .udp_send_to(&participant.user.cipher, &participant.addr, payload) + .await; + } + + Ok(()) + } + + pub fn get_voice_socket(&self) -> anyhow::Result<&UdpSocket> { + self.socket + .get() + .context("Failed to get voice socket") + } + + pub async fn udp_send_to( + &self, + cipher: &Arc>, + addr: &SocketAddr, + payload: &[u8], + ) -> anyhow::Result<()> { + let socket = self.get_voice_socket()?; + + socket + .send_to(&cipher.lock().await.encrypt(payload)?, addr) + .await?; + + Ok(()) + } +} \ No newline at end of file diff --git a/src/vc_server.rs b/src/vc_server.rs deleted file mode 100644 index 8da599f..0000000 --- a/src/vc_server.rs +++ /dev/null @@ -1,179 +0,0 @@ -use std::{net::SocketAddr, sync::Arc}; - -use anyhow::Context; -use ed25519_dalek::VerifyingKey; -use tokio::{net::UdpSocket, sync::Mutex}; - -use crate::{crypto::SessionCipher, protocol::ClientMethod, server::Server}; - -use tokio::time::Instant; - -impl Server { - pub async fn start_udp_server(self: Arc) -> anyhow::Result<()> { - let socket = UdpSocket::bind(("0.0.0.0", self.config.port)).await?; - - self.voice_socket - .set(socket) - .map_err(|_| anyhow::anyhow!("UDP server already started"))?; - - let mut buf = [0u8; 4096]; - - eprintln!("[vc] UDP server listening on port {}", self.config.port); - - loop { - let (len, addr) = self.get_voice_socket()?.recv_from(&mut buf).await?; - - if len < 8 { - eprintln!("[vc] dropping packet too short for pin"); - continue; - } - - let pin_bytes: [u8; 8] = buf[..8].try_into().unwrap(); - let pin = u64::from_be_bytes(pin_bytes); - - // is_first_packet distinguishes the plaintext pin-bootstrap packet - // from subsequent encrypted audio packets. - let (sender_pubkey, channel_id, payload, is_first_packet) = { - let mut pins = self.voice_pins.lock().await; - - if let Some((pubkey, channel_id)) = pins.remove(&pin) { - (pubkey, channel_id, &buf[8..len], true) - } else { - drop(pins); - - match self.find_voice_sender(&addr).await { - Some((pubkey, channel_id)) => (pubkey, channel_id, &buf[..len], false), - None => { - continue; - } - } - } - }; - - let clients = self.clients.lock().await; - let Some(user) = clients.get(&sender_pubkey).cloned() else { - eprintln!("[vc] sender pubkey not found in clients, dropping"); - continue; - }; - drop(clients); - - let mut voice = user.voice.lock().await; - - if let Some(voice) = &mut *voice { - voice.addr = addr; - } else { - *voice = Some(crate::server::VoiceConnection { - addr, - channel_id: channel_id.clone(), - last_speaking_sent: Instant::now(), - }); - } - - drop(voice); - - // The pin-bearing bootstrap packet carries no payload to decrypt — - // it's purely "here's my pin, bind my address." Everything after - // this first packet is the real, encrypted audio stream. - if is_first_packet { - continue; - } - - let decrypted_payload = match user.cihper.lock().await.decrypt(payload) { - Ok(pt) => pt, - Err(e) => { - eprintln!("[vc] dropping packet: decryption failed: {e}"); - continue; - } - }; - - let s = self.clone(); - - tokio::spawn(async move { - if let Err(e) = s - .relay_voice(&sender_pubkey, &channel_id, &decrypted_payload) - .await - { - eprintln!("{e}"); - } - }); - } - } - - /// Looks up which known voice participant a UDP address belongs to, - /// for packets arriving after the initial pin-bearing packet. - async fn find_voice_sender(&self, addr: &SocketAddr) -> Option<(VerifyingKey, String)> { - let clients = self.clients.lock().await; - - for (pubkey, user) in clients.iter() { - if let Some(voice) = &*user.voice.lock().await { - if *addr == voice.addr { - return Some((*pubkey, voice.channel_id.clone())); - } - } - } - - None - } - - /// Sends `payload` to every voice participant currently in `channel_id`. - async fn relay_voice( - &self, - sender: &VerifyingKey, - channel_id: &str, - payload: &[u8], - ) -> anyhow::Result<()> { - let clients = self.clients.lock().await; - - for (_pubkey, user) in clients.iter() { - let Some(voice) = &mut *user.voice.lock().await else { - continue; - }; - - if channel_id != voice.channel_id { - continue; - } - - let now = Instant::now(); - - if now.duration_since(voice.last_speaking_sent).as_millis() >= 600 { - for conn in user.connections.lock().await.values() { - conn.send(&ClientMethod::Speaking { - pubkey: crate::crypto::to_string(sender), - }) - .await?; - } - - voice.last_speaking_sent = now; - } - - if *sender == user.public_key { - continue; - } - - let _ = self.udp_send_to(&user.cihper, &voice.addr, payload).await; - } - - Ok(()) - } - - pub fn get_voice_socket(&self) -> anyhow::Result<&UdpSocket> { - self.voice_socket - .get() - .context("Failed to get voice socket") - } - - pub async fn udp_send_to( - &self, - cipher: &Arc>, - addr: &SocketAddr, - payload: &[u8], - ) -> anyhow::Result<()> { - let socket = self.get_voice_socket()?; - - socket - .send_to(&cipher.lock().await.encrypt(payload)?, addr) - .await?; - - Ok(()) - } -}