Rewritten server code

src/crypto.rs

@@ -142,7 +142,7 @@ pub async fn crypto_handshake(
 ) -> anyhow::Result<EnclaveWebSocket> {
     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)?));
 

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)

src/protocol/initialize.rs

@@ -5,16 +5,18 @@ 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(
-    pub async fn initialize(
     server: &Arc<Server>,
     socket: &EnclaveWebSocket,
-    ) -> anyhow::Result<(VerifyingKey, ClientMeta)> {
+) -> anyhow::Result<(VerifyingKey, ClientMeta)> {
     let Some(ServerMethod::Initialize {
         public_key: public_key_string,
         signature,
@@ -88,8 +90,8 @@ impl UserConnections {
     {
         socket
             .send(&ClientMethod::Initialized {
-                    public_key: crate::crypto::to_string(&server.key.verifying_key()),
+                public_key: crate::crypto::to_string(&server.identity.key.verifying_key()),
-                    signature: crate::crypto::to_string_sig(&server.key.sign(
+                signature: crate::crypto::to_string_sig(&server.identity.key.sign(
                     format!("{server_timestamp}@{hostname}@{public_key_string}").as_bytes(),
                 )),
 
@@ -111,15 +113,14 @@ impl UserConnections {
         ));
     };
 
-        for (pubkey, channel_id) in server.voice_pins.lock().await.values() {
+    for pin in server.voice.pins.lock().await.values() {
         socket
             .send(&ClientMethod::UserJoinedVoice {
-                    channel_id: channel_id.clone(),
+                channel_id: pin.channel_id.clone(),
-                    pubkey: crate::crypto::to_string(pubkey),
+                pubkey: crate::crypto::to_string(&pin.pubkey),
             })
             .await?;
     }
 
     Ok((public_key, meta))
-    }
 }

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,9 +42,10 @@ pub async fn send_message(
         data: message,
     };
 
-    server.message_store.insert_message(&channel_id, &stored)?;
+    server.store.messages.insert_message(&channel_id, &stored)?;
 
-    server
+server
+    .sessions
     .broadcast(&ClientMethod::Messages {
             messages: HashMap::from([(channel_id, vec![stored])]),
         })
@@ -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,7 +125,8 @@ pub async fn edit_message(
         },
     };
 
-    server
+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,10 +154,12 @@ pub async fn delete_message(
     }
 
     server
-        .message_store
+        .store
+        .messages
         .delete_message(&channel_id, &message_id)?;
 
-    server
+server
+    .sessions
     .broadcast(&ClientMethod::MessageDeleted {
             channel_id: channel_id.clone(),
             message_id,

src/protocol/user.rs

@@ -10,7 +10,7 @@ pub async fn get_users(
     socket: &Arc<crate::ws::EnclaveWebSocket>,
     pubkeys: Vec<String>,
 ) -> 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?;
 

src/protocol/voice.rs

@@ -10,14 +10,13 @@ pub async fn join(
     socket: &Arc<crate::ws::EnclaveWebSocket>,
     channel_id: String,
 ) -> anyhow::Result<()> {
-    {
+    let user = server
-        let pin = rand::random::<u64>() % (1 << 53);
+        .sessions
-
+        .get(&verifying_key)
-        server
-            .voice_pins
-            .lock()
         .await
-            .insert(pin, (verifying_key, channel_id.clone()));
+        .ok_or_else(|| anyhow::anyhow!("Not connected"))?;
+
+    let pin = server.voice.join(verifying_key, user, &channel_id).await;
 
     socket
         .send(&ClientMethod::JoinVoice {
@@ -25,9 +24,9 @@ pub async fn join(
             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<Server>, verifying_key: VerifyingKey) -> anyhow::Result<()> {
-    let Some(channel_id) = ({
+    let Some(channel_id) = server.voice.remove(verifying_key).await else {
-        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 {
         return Ok(());
     };
 
     server
-        .voice_pins
+        .sessions
-        .lock()
-        .await
-        .retain(|_, v| v.0 != verifying_key);
-
-    server
         .broadcast(&ClientMethod::UserLeftVoice {
             channel_id,
             pubkey: crate::crypto::to_string(&verifying_key),

src/server.rs

@@ -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<HashMap<u16, Arc<crate::ws::EnclaveWebSocket>>>,
-    pub cihper: Arc<Mutex<SessionCipher>>,
-    pub voice: Mutex<Option<VoiceConnection>>,
-}
-
-pub struct Server {
-    pub key: SigningKey,
-    pub x_keypair: (X25519Public, X25519Secret),
-    pub config: Config,
-    pub clients: Mutex<HashMap<VerifyingKey, Arc<UserConnections>>>,
-    pub voice_pins: Mutex<HashMap<u64, (VerifyingKey, String)>>,
-    pub message_store: MessageStore,
-    pub user_store: UserMetaStore,
-    pub voice_socket: OnceCell<UdpSocket>,
-}
-
-impl Server {
-    pub async fn new() -> anyhow::Result<Arc<Self>> {
-        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<Self>, 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<Self>, 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<bool> {
-        if let Some(conn) = self.connections.lock().await.get(&id) {
-            conn.send(message).await?;
-
-            Ok(true)
-        } else {
-            Ok(false)
-        }
-    }
-}

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<Self> {
+        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 })
+    }
+}

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<VoiceServer>,
+    pub store: DataStore,
+}
+
+impl Server {
+    pub async fn new() -> anyhow::Result<Arc<Self>> {
+        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<Self>, 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}")
+                }
+            }
+        })
+    }
+}

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<HashMap<u16, Arc<EnclaveWebSocket>>>,
+    pub cipher: Arc<Mutex<SessionCipher>>,
+}
+
+pub struct SessionRegistry {
+    pub clients: Mutex<HashMap<VerifyingKey, Arc<UserConnections>>>,
+}
+
+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<Arc<UserConnections>> {
+        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<EnclaveWebSocket>, 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<bool> {
+        if let Some(conn) = self.connections.lock().await.get(&id) {
+            conn.send(message).await?;
+
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+}

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<Self> {
+        Ok(Self {
+            messages: MessageStore::new(PathBuf::from("messages"))?,
+            users: UserMetaStore::new(PathBuf::from("users.db"))?,
+        })
+    }
+}

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<UserConnections>,
+    pub channel_id: String,
+    pub addr: SocketAddr,
+    pub last_speaking_sent: Instant,
+}
+
+pub struct VoiceServer {
+    pub pins: Mutex<HashMap<u64, VoicePin>>,
+    pub socket: OnceCell<UdpSocket>,
+    pub participants: Mutex<HashMap<VerifyingKey, VoiceConnection>>,
+}
+
+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<UserConnections>,
+        channel_id: &str,
+    ) -> u64 {
+        let pin = rand::random::<u64>() % (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<String> {
+        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<Self>, 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<Mutex<SessionCipher>>,
+        addr: &SocketAddr,
+        payload: &[u8],
+    ) -> anyhow::Result<()> {
+        let socket = self.get_voice_socket()?;
+
+        socket
+            .send_to(&cipher.lock().await.encrypt(payload)?, addr)
+            .await?;
+
+        Ok(())
+    }
+}

src/vc_server.rs

@@ -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<Self>) -> 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<Mutex<SessionCipher>>,
-        addr: &SocketAddr,
-        payload: &[u8],
-    ) -> anyhow::Result<()> {
-        let socket = self.get_voice_socket()?;
-
-        socket
-            .send_to(&cipher.lock().await.encrypt(payload)?, addr)
-            .await?;
-
-        Ok(())
-    }
-}