diff --git a/src/main.rs b/src/main.rs index 1dc0627..68a08b9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ pub mod data; pub mod protocol; pub mod server; pub mod types; +pub mod vc_server; use std::{ net::{IpAddr, Ipv4Addr, SocketAddr}, @@ -26,6 +27,8 @@ 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 cors = CorsLayer::new() .allow_origin(Any) .allow_methods(Any) @@ -46,6 +49,8 @@ async fn main() -> anyhow::Result<()> { axum::serve(listener, app).await?; + udp_server.abort(); + Ok(()) } diff --git a/src/server.rs b/src/server.rs index 2dce567..0584278 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,5 +1,6 @@ use std::{ collections::HashMap, + net::SocketAddr, path::PathBuf, sync::{Arc, atomic::AtomicU16}, }; @@ -9,7 +10,11 @@ use axum::{ response::Response, }; use ed25519_dalek::{SigningKey, VerifyingKey}; -use tokio::{sync::Mutex, task::JoinSet}; +use tokio::{ + net::UdpSocket, + sync::{Mutex, OnceCell}, + task::JoinSet, +}; use crate::{ data::{config::Config, messages::MessageStore, users::UserMetaStore}, @@ -22,6 +27,7 @@ pub struct UserConnections { pub counter: AtomicU16, pub public_key: VerifyingKey, pub connections: Mutex>>>, + pub voice: Mutex>, } pub struct Server { @@ -31,6 +37,7 @@ pub struct Server { pub voice_pins: Mutex>, pub message_store: MessageStore, pub user_store: UserMetaStore, + pub voice_socket: OnceCell, } impl Server { @@ -42,6 +49,7 @@ impl Server { 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(), })) } } @@ -73,6 +81,7 @@ impl Server { public_key: public_key, counter: AtomicU16::new(0), connections: Mutex::new(HashMap::new()), + voice: Mutex::new(None), }) }) .clone(); diff --git a/src/vc_server.rs b/src/vc_server.rs new file mode 100644 index 0000000..b8ad59c --- /dev/null +++ b/src/vc_server.rs @@ -0,0 +1,118 @@ +use std::{net::SocketAddr, sync::Arc}; + +use anyhow::Context; +use ed25519_dalek::VerifyingKey; +use tokio::net::UdpSocket; + +use crate::server::Server; + +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]; + + println!("UDP server started"); + + loop { + let (len, addr) = self.get_voice_socket()?.recv_from(&mut buf).await?; + + if len < 8 { + continue; // too short to even contain a pincode, drop silently + } + + let pin_bytes: [u8; 8] = buf[..8].try_into().unwrap(); + let pin = u64::from_be_bytes(pin_bytes); + let payload = &buf[8..len]; + + // Resolve the sender's identity + channel for this pin. + // First packet for a pin consumes it (single-use) and binds the address. + let sender = { + let mut pins = self.voice_pins.lock().await; + + if let Some((pubkey, channel_id)) = pins.remove(&pin) { + Some((pubkey, channel_id)) + } else { + None + } + }; + + let (sender_pubkey, channel_id) = match sender { + Some(v) => v, + + // Not a first-time pin — check if this addr is already a known + // voice participant, so we know who's speaking and where to relay. + None => match self.find_voice_sender(&addr).await { + Some(v) => v, + None => continue, // unknown pin, unknown addr — drop + }, + }; + + let clients = self.clients.lock().await; + let Some(user) = clients.get(&sender_pubkey).cloned() else { + continue; // pin referenced a user that's since disconnected + }; + drop(clients); + + // Record/refresh this user's known voice address + channel. + *user.voice.lock().await = Some((addr, channel_id.clone())); + + self.relay_voice(&sender_pubkey, &channel_id, payload).await; + } + } + + /// 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((user_addr, channel)) = &*user.voice.lock().await { + if *user_addr == *addr { + return Some((*pubkey, channel.clone())); + } + } + } + + None + } + + /// Sends `payload` to every other voice participant currently in `channel_id`. + async fn relay_voice(&self, sender: &VerifyingKey, channel_id: &str, payload: &[u8]) { + let clients = self.clients.lock().await; + + for (pubkey, user) in clients.iter() { + // if pubkey == sender { + // continue; + // } + + let Some((addr, channel)) = &*user.voice.lock().await else { + continue; + }; + + if channel_id != channel { + continue; + } + + let _ = self.udp_send_to(addr, payload).await; + } + } + + 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, addr: &SocketAddr, payload: &[u8]) -> anyhow::Result<()> { + let socket = self.get_voice_socket()?; + + socket.send_to(payload, addr).await?; + + Ok(()) + } +}