diff --git a/Cargo.lock b/Cargo.lock index 988bd38..b30f3a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2030,6 +2030,7 @@ dependencies = [ "tokio", "trust-dns-resolver", "uuid", + "zeroize", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 8595400..b758cd8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,6 @@ mtp = { git = "https://git.methanium.net/Methanium/mtp.git", features = [ "client", "crypto", "files", - "raw", ] } mtp-transport = { git = "https://git.methanium.net/Methanium/mtp.git" } @@ -34,3 +33,12 @@ livekit-protocol = "=0.7.10" thiserror = "2.0.19" trust-dns-resolver = "0.23.2" serde_json = "1.0.151" +zeroize = "1.9" + +[features] +raw-migration = ["mtp/raw"] + +[[bin]] +name = "migrate_raw_keyring" +path = "src/bin/migrate_raw_keyring.rs" +required-features = ["raw-migration"] diff --git a/README.md b/README.md index 401bec9..1ca8ee3 100644 --- a/README.md +++ b/README.md @@ -4,3 +4,26 @@ It's primary purpose is to connect you're client to your Iota & hide your IP and The Omikron also host Voice-Calls. The Omikron is only used when the Iota is in Centralized and Hybrid mode, or when the Client uses the Tensamin Client with default configuration. + +## Configuration + +`OMIKRON_IDENTITY_SECRET` is required. Provision it through the deployment's secret environment before starting Omikron. Omikron uses it to load `omikron.mk` as a protected keyring and fails startup if the secret or existing identity cannot be loaded. + +Rho connection budgets can be configured with: + +- `RHO_MAX_CONNECTIONS`, default `256` +- `RHO_MAX_ANONYMOUS_CONNECTIONS`, default `128` +- `RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP`, default `16` + +`RHO_MAX_CONNECTIONS` is applied both to Omikron's application session +semaphore and to the MTP WebServer admission semaphore. This means the same +budget limits transport handshakes and authenticated Rho sessions instead of +only limiting connections after authentication. + +To migrate an existing raw `omikron.mk`, provision `OMIKRON_IDENTITY_SECRET` and run: + +```sh +cargo run --features raw-migration --bin migrate_raw_keyring +``` + +The normal Omikron binary does not enable raw keyring support. diff --git a/src/anonymous_clients/anonymous_client_connection.rs b/src/anonymous_clients/anonymous_client_connection.rs index d0ac89c..884936e 100644 --- a/src/anonymous_clients/anonymous_client_connection.rs +++ b/src/anonymous_clients/anonymous_client_connection.rs @@ -9,7 +9,7 @@ use crate::anonymous_clients::anonymous_manager::{self, generate_username}; use crate::app_state::AppState; use crate::calls::call_group::call_invite_secret_from_cv; use crate::rho::connection::{ - GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat, + GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields, }; use crate::util::data_type_id; use crate::util::logger::PrintType; @@ -109,10 +109,25 @@ impl AnonymousClientConnection { }; tokio::spawn(async move { let _permit = permit; + let message_id = match cv.require_id() { + Ok(message_id) => message_id, + Err(error) => { + let response = + CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id(); + log_out!( + self.user_id as i64, + PrintType::Client, + "Rejected malformed message: {}", + error + ); + self.send_message(&response).await; + return; + } + }; log_cv_in!(PrintType::Client, &cv); if cv.is_type(CommunicationType::Relay) { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNotAuthenticated) + self.send_error_response(message_id, CommunicationType::ErrorNotAuthenticated) .await; return; } @@ -126,18 +141,15 @@ impl AnonymousClientConnection { call } else { self.send_error_response( - &cv.get_id(), + message_id, CommunicationType::ErrorNotAuthenticated, ) .await; return; } } else { - self.send_error_response( - &cv.get_id(), - CommunicationType::ErrorNotAuthenticated, - ) - .await; + self.send_error_response(message_id, CommunicationType::ErrorNotAuthenticated) + .await; return; }; @@ -210,7 +222,7 @@ impl AnonymousClientConnection { self.clone() .send_message( &&CommunicationValue::new(CommunicationType::IdentificationResponse) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::UserId, DataValue::SignedNumber(self.user_id.into()), @@ -245,7 +257,7 @@ impl AnonymousClientConnection { // Presence is account-scoped and anonymous sessions have no // persisted account preference to change. if cv.is_type(CommunicationType::ClientChanged) { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNoUserId) + self.send_error_response(message_id, CommunicationType::ErrorNoUserId) .await; return; } @@ -291,7 +303,7 @@ impl AnonymousClientConnection { } } { let response = CommunicationValue::new(CommunicationType::GetUserData) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::Username, DataValue::Str(anonymous.get_user_name().await), @@ -344,9 +356,12 @@ impl AnonymousClientConnection { /// Handle call invite async fn handle_call_invite(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let receiver_id: i64 = cv.get_data(DataType::ReceiverId).as_number().unwrap_or(0) as i64; if receiver_id == 0 { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNoUserId) + self.send_error_response(message_id, CommunicationType::ErrorNoUserId) .await; return; } @@ -355,13 +370,13 @@ impl AnonymousClientConnection { Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) { Ok(id) => id, Err(_) => { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } }, _ => { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNoCallId) + self.send_error_response(message_id, CommunicationType::ErrorNoCallId) .await; return; } @@ -370,7 +385,7 @@ impl AnonymousClientConnection { let secret = match call_invite_secret_from_cv(&cv) { Some(secret) => secret, None => { - self.send_error_response(&cv.get_id(), CommunicationType::BadRequest) + self.send_error_response(message_id, CommunicationType::BadRequest) .await; return; } @@ -381,7 +396,7 @@ impl AnonymousClientConnection { .add_invite(call_id, self.user_id, receiver_id as u64, secret.clone()) .await; if !invited { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } @@ -391,7 +406,7 @@ impl AnonymousClientConnection { .call_manager .should_forward_invite(self.user_id, receiver_id as u64) { - let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id()); + let response = CommunicationValue::new(CommunicationType::Success).with_id(message_id); self.send_message(&response).await; return; } @@ -423,7 +438,7 @@ impl AnonymousClientConnection { }); let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::ReceiverId, DataValue::SignedNumber(receiver_id.into()), @@ -453,25 +468,28 @@ impl AnonymousClientConnection { target_rho.message_to_client(forward).await; - let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id()); + let response = CommunicationValue::new(CommunicationType::Success).with_id(message_id); self.send_message(&response).await; } /// Handle get call request async fn handle_get_call(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let user_id = self.get_user_id(); let call_id = match cv.get_data(DataType::CallId) { Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) { Ok(id) => id, Err(_) => { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } }, _ => { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNoCallId) + self.send_error_response(message_id, CommunicationType::ErrorNoCallId) .await; return; } @@ -485,7 +503,7 @@ impl AnonymousClientConnection { { Ok(token) => { let response = CommunicationValue::new(CommunicationType::CallToken) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(user_id) .add_typed_default(DataType::CallToken, DataValue::Str(token)); self.send_message(&response).await; @@ -497,16 +515,19 @@ impl AnonymousClientConnection { error ); let error_cv = CommunicationValue::new(CommunicationType::ErrorNoCallId) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string())); self.send_message(&error_cv).await; } } } async fn handle_call_timeout_user(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or("")) else { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; }; @@ -520,12 +541,12 @@ impl AnonymousClientConnection { .unwrap_or(0); let Some(call) = self.state.call_manager.get_call(call_id).await else { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNotFound) + self.send_error_response(message_id, CommunicationType::ErrorNotFound) .await; return; }; let Some(caller) = call.get_caller(self.get_user_id()).await else { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidUserId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId) .await; return; }; @@ -536,9 +557,12 @@ impl AnonymousClientConnection { } } async fn handle_call_disconnect_user(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or("")) else { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; }; @@ -548,12 +572,12 @@ impl AnonymousClientConnection { .unwrap_or(0); let Some(call) = self.state.call_manager.get_call(call_id).await else { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorNotFound) + self.send_error_response(message_id, CommunicationType::ErrorNotFound) .await; return; }; let Some(caller) = call.get_caller(self.get_user_id()).await else { - self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidUserId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId) .await; return; }; @@ -563,8 +587,8 @@ impl AnonymousClientConnection { } /// Send error response - async fn send_error_response(self: Arc, message_id: &u32, error_type: CommunicationType) { - let error = CommunicationValue::new(error_type).with_id(*message_id); + async fn send_error_response(self: Arc, message_id: u32, error_type: CommunicationType) { + let error = CommunicationValue::new(error_type).with_id(message_id); self.send_message(&error).await; } /// Close the connection diff --git a/src/app_state.rs b/src/app_state.rs index 750354a..bb986fe 100644 --- a/src/app_state.rs +++ b/src/app_state.rs @@ -1,6 +1,8 @@ -use std::sync::Arc; +use std::{net::IpAddr, sync::Arc}; +use dashmap::DashMap; use mtp::crypto::Keyring; +use tokio::sync::Semaphore; use crate::{ calls::{call_manager::CallManager, call_util::LiveKitService}, @@ -10,6 +12,58 @@ use crate::{ services::call_state::CallStateAggregator, }; +pub struct RhoConnectionLimits { + pub all: Arc, + pub anonymous: Arc, + anonymous_by_ip: Arc>, + max_anonymous_per_ip: usize, +} + +impl RhoConnectionLimits { + pub fn new(max_all: usize, max_anonymous: usize, max_anonymous_per_ip: usize) -> Self { + Self { + all: Arc::new(Semaphore::new(max_all)), + anonymous: Arc::new(Semaphore::new(max_anonymous)), + anonymous_by_ip: Arc::new(DashMap::new()), + max_anonymous_per_ip, + } + } + + pub fn try_acquire_anonymous_per_ip(&self, ip: IpAddr) -> Option { + let mut count = self.anonymous_by_ip.entry(ip).or_insert(0); + if *count >= self.max_anonymous_per_ip { + return None; + } + *count += 1; + drop(count); + + Some(RhoIpConnectionPermit { + ip, + active_by_ip: self.anonymous_by_ip.clone(), + }) + } + + #[cfg(test)] + fn anonymous_count_for_ip(&self, ip: IpAddr) -> usize { + self.anonymous_by_ip.get(&ip).map_or(0, |count| *count) + } +} + +pub struct RhoIpConnectionPermit { + ip: IpAddr, + active_by_ip: Arc>, +} + +impl Drop for RhoIpConnectionPermit { + fn drop(&mut self) { + if let Some(mut count) = self.active_by_ip.get_mut(&self.ip) { + *count = count.saturating_sub(1); + } + self.active_by_ip + .remove_if(&self.ip, |_, count| *count == 0); + } +} + /* * Holds startup dependencies so listener and connection constructors can take * one explicit handle while the remaining manager migrations are completed. @@ -22,6 +76,7 @@ pub struct AppState { pub call_manager: Arc, pub call_state_aggregator: Arc, pub livekit: Arc, + pub rho_connection_limits: Arc, } impl AppState { @@ -33,6 +88,11 @@ impl AppState { call_manager: Arc, livekit: Arc, ) -> Arc { + let rho_connection_limits = Arc::new(RhoConnectionLimits::new( + config.rho_max_connections, + config.rho_max_anonymous_connections, + config.rho_max_anonymous_connections_per_ip, + )); Arc::new(Self { config, keyring, @@ -41,6 +101,7 @@ impl AppState { call_state_aggregator: Arc::new(CallStateAggregator::new(call_manager.clone())), call_manager, livekit, + rho_connection_limits, }) } diff --git a/src/bin/migrate_raw_keyring.rs b/src/bin/migrate_raw_keyring.rs new file mode 100644 index 0000000..bdd82d4 --- /dev/null +++ b/src/bin/migrate_raw_keyring.rs @@ -0,0 +1,23 @@ +#[path = "../identity.rs"] +mod identity; + +use identity::{ + KEYRING_PATH, PUBLIC_KEY_PATH, identity_secret_from_environment, migrate_raw_keyring, +}; + +fn main() { + let secret = match identity_secret_from_environment() { + Ok(secret) => secret, + Err(error) => { + eprintln!("Unable to load Omikron identity secret: {error}"); + std::process::exit(1); + } + }; + + if let Err(error) = migrate_raw_keyring(secret.as_slice(), KEYRING_PATH, PUBLIC_KEY_PATH) { + eprintln!("Unable to migrate Omikron identity: {error}"); + std::process::exit(1); + } + + println!("Migrated {KEYRING_PATH} to protected keyring storage"); +} diff --git a/src/calls/call_group.rs b/src/calls/call_group.rs index 62f773e..4626f7e 100755 --- a/src/calls/call_group.rs +++ b/src/calls/call_group.rs @@ -15,6 +15,7 @@ use crate::{ pub struct CallGroup { pub call_id: Uuid, pub members: RwLock>>, + #[allow(dead_code)] pub show: RwLock, pub anonymous_joining: RwLock, pub short_link: RwLock>, @@ -100,6 +101,7 @@ pub fn call_invite_secret_from_cv(cv: &CommunicationValue) -> Option) -> Self { Self::new_with_service(call_id, user, Arc::new(LiveKitService::new(None))) } diff --git a/src/calls/call_util.rs b/src/calls/call_util.rs index 8476c8c..edd2b9b 100644 --- a/src/calls/call_util.rs +++ b/src/calls/call_util.rs @@ -134,6 +134,7 @@ impl LiveKitService { Ok(()) } + #[allow(dead_code)] pub fn garbage_collect_calls(self: Arc, manager: Arc) { tokio::spawn(async move { loop { @@ -148,6 +149,7 @@ impl LiveKitService { } } +#[allow(dead_code)] pub async fn clean_calls(manager: &CallManager, room_service: RoomClient) { let rooms = match tokio::time::timeout(LIVEKIT_REQUEST_TIMEOUT, room_service.list_rooms(Vec::new())) diff --git a/src/calls/caller.rs b/src/calls/caller.rs index 6a69566..ee5c480 100644 --- a/src/calls/caller.rs +++ b/src/calls/caller.rs @@ -39,6 +39,7 @@ impl Caller { pub async fn set_timeout(&self, timeout: i64) { *self.timeout.write().await = timeout; } + #[allow(dead_code)] pub fn create_token(&self, livekit: &LiveKitService) -> Result { livekit.create_token(self.user_id, self.call_id, self.has_admin()) } diff --git a/src/config.rs b/src/config.rs index fce6f3c..e69f383 100644 --- a/src/config.rs +++ b/src/config.rs @@ -8,6 +8,9 @@ const DEFAULT_OMEGA_HOST: &str = "tensamin.net"; const DEFAULT_OMEGA_PORT: u16 = 9187; const DEFAULT_OMEGA_SYNC_TIMEOUT_SECONDS: u64 = 20; const DEFAULT_OMEGA_SYNC_RETRIES: u32 = 3; +const DEFAULT_RHO_MAX_CONNECTIONS: usize = 256; +const DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS: usize = 128; +const DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP: usize = 16; #[derive(Clone, Debug, Eq, PartialEq)] pub struct LiveKitConfig { @@ -28,6 +31,12 @@ pub struct Config { /// Number of synchronization requests before the transport is closed and /// the normal reconnect loop starts. pub omega_sync_retries: u32, + /// Maximum number of application sessions accepted by the Rho listener. + pub rho_max_connections: usize, + /// Maximum number of anonymous application sessions. + pub rho_max_anonymous_connections: usize, + /// Maximum number of anonymous sessions from one peer IP address. + pub rho_max_anonymous_connections_per_ip: usize, pub livekit: Option, } @@ -57,6 +66,16 @@ impl Config { )?); let omega_sync_retries = parse_or_default("OMEGA_SYNC_RETRIES", DEFAULT_OMEGA_SYNC_RETRIES)?.max(1); + let rho_max_connections = + parse_positive_or_default("RHO_MAX_CONNECTIONS", DEFAULT_RHO_MAX_CONNECTIONS)?; + let rho_max_anonymous_connections = parse_positive_or_default( + "RHO_MAX_ANONYMOUS_CONNECTIONS", + DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS, + )?; + let rho_max_anonymous_connections_per_ip = parse_positive_or_default( + "RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP", + DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP, + )?; let omega_host = env::var("OMEGA_HOST") .unwrap_or_else(|_| DEFAULT_OMEGA_HOST.to_string()) .trim() @@ -77,6 +96,9 @@ impl Config { omikron_id, omega_sync_timeout, omega_sync_retries, + rho_max_connections, + rho_max_anonymous_connections, + rho_max_anonymous_connections_per_ip, livekit: livekit_from_environment()?, }) } diff --git a/src/identity.rs b/src/identity.rs new file mode 100644 index 0000000..c0459c2 --- /dev/null +++ b/src/identity.rs @@ -0,0 +1,181 @@ +use std::{env, io::ErrorKind, path::Path}; + +use mtp::crypto::Keyring; +use mtp::files::{FileError, load_keyring, save_keyring, save_public_key_bundle}; +use zeroize::Zeroizing; + +pub const KEYRING_PATH: &str = "./omikron.mk"; +pub const PUBLIC_KEY_PATH: &str = "./omikron.mpkb"; +const IDENTITY_SECRET_ENV: &str = "OMIKRON_IDENTITY_SECRET"; + +pub fn identity_secret_from_environment() -> Result>, String> { + let secret = env::var(IDENTITY_SECRET_ENV).map_err(|_| { + format!("{IDENTITY_SECRET_ENV} must be set before Omikron networking starts") + })?; + if secret.trim().is_empty() { + return Err(format!("{IDENTITY_SECRET_ENV} must not be empty")); + } + Ok(Zeroizing::new(secret.into_bytes())) +} + +#[allow(dead_code)] +pub fn load_or_create_keyring( + passphrase: &[u8], + keyring_path: impl AsRef, + public_key_path: impl AsRef, +) -> Result { + let keyring_path = keyring_path.as_ref(); + let public_key_path = public_key_path.as_ref(); + + match load_keyring(keyring_path, passphrase) { + Ok(keyring) => Ok(keyring), + Err(FileError::Io(error)) if error.kind() == ErrorKind::NotFound => { + let keyring = Keyring::generate(); + save_keyring(&keyring, keyring_path, passphrase).map_err(|error| { + format!("unable to persist Omikron keyring at {keyring_path:?}: {error}") + })?; + save_public_key_bundle(&keyring.public_key_bundle(), public_key_path).map_err( + |error| { + format!( + "unable to persist Omikron public key bundle at {public_key_path:?}: {error}" + ) + }, + )?; + eprintln!( + "Generated new protected keyring at {}", + keyring_path.display() + ); + Ok(keyring) + } + Err(error) => Err(format!( + "unable to load existing Omikron identity from {}: {error}", + keyring_path.display() + )), + } +} + +#[cfg(feature = "raw-migration")] +#[allow(dead_code)] +pub fn migrate_raw_keyring( + passphrase: &[u8], + keyring_path: impl AsRef, + public_key_path: impl AsRef, +) -> Result { + use mtp::files::load_keyring_raw; + + let keyring = load_keyring_raw(keyring_path.as_ref()) + .map_err(|error| format!("raw keyring migration failed: {error}"))?; + save_keyring(&keyring, keyring_path.as_ref(), passphrase) + .map_err(|error| format!("protected keyring write failed: {error}"))?; + save_public_key_bundle(&keyring.public_key_bundle(), public_key_path.as_ref()) + .map_err(|error| format!("public key bundle write failed: {error}"))?; + Ok(keyring) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{ + fs, + time::{SystemTime, UNIX_EPOCH}, + }; + + struct TestPaths { + directory: std::path::PathBuf, + keyring: std::path::PathBuf, + public_key: std::path::PathBuf, + } + + impl TestPaths { + fn new() -> Self { + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let directory = env::temp_dir().join(format!("omikron-identity-{suffix}")); + fs::create_dir_all(&directory).unwrap(); + Self { + keyring: directory.join("omikron.mk"), + public_key: directory.join("omikron.mpkb"), + directory, + } + } + } + + impl Drop for TestPaths { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.directory); + } + } + + #[test] + fn missing_identity_is_created_and_can_be_loaded_again() { + let paths = TestPaths::new(); + let passphrase = b"test identity secret"; + let first = load_or_create_keyring(passphrase, &paths.keyring, &paths.public_key).unwrap(); + let first_public = first.public_key_bundle().try_as_bytes().unwrap(); + let second = load_or_create_keyring(passphrase, &paths.keyring, &paths.public_key).unwrap(); + + assert_eq!( + second.public_key_bundle().try_as_bytes().unwrap(), + first_public + ); + } + + #[test] + fn corrupt_identity_fails_closed() { + let paths = TestPaths::new(); + fs::write(&paths.keyring, b"not a keyring").unwrap(); + + let error = + load_or_create_keyring(b"test identity secret", &paths.keyring, &paths.public_key) + .unwrap_err(); + assert!(error.contains("unable to load existing Omikron identity")); + } + + #[test] + fn unreadable_identity_fails_closed() { + let paths = TestPaths::new(); + fs::create_dir(&paths.keyring).unwrap(); + + let error = + load_or_create_keyring(b"test identity secret", &paths.keyring, &paths.public_key) + .unwrap_err(); + assert!(error.contains("unable to load existing Omikron identity")); + } + + #[test] + fn public_key_persistence_failure_fails_closed() { + let paths = TestPaths::new(); + fs::create_dir(&paths.public_key).unwrap(); + + let error = + load_or_create_keyring(b"test identity secret", &paths.keyring, &paths.public_key) + .unwrap_err(); + assert!(error.contains("unable to persist Omikron public key bundle")); + } + + #[cfg(feature = "raw-migration")] + #[test] + fn raw_identity_migration_rewrites_protected_storage() { + use mtp::files::{load_keyring, save_keyring_raw}; + + let paths = TestPaths::new(); + let original = Keyring::generate(); + save_keyring_raw(&original, &paths.keyring).unwrap(); + + let migrated = + migrate_raw_keyring(b"test identity secret", &paths.keyring, &paths.public_key) + .unwrap(); + let loaded = load_keyring(&paths.keyring, b"test identity secret").unwrap(); + + assert_eq!( + migrated.public_key_bundle().try_as_bytes().unwrap(), + original.public_key_bundle().try_as_bytes().unwrap() + ); + assert_eq!( + loaded.public_key_bundle().try_as_bytes().unwrap(), + original.public_key_bundle().try_as_bytes().unwrap() + ); + } +} diff --git a/src/main.rs b/src/main.rs index 25bf29e..7202ca2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod app_state; mod calls; mod config; mod data; +mod identity; mod omega; mod rho; mod services; @@ -18,35 +19,20 @@ pub static WORKING_DIR: Lazy = use rustls::crypto::aws_lc_rs::default_provider; use mtp::crypto::Keyring; -use mtp::files::{load_keyring_raw, save_keyring_raw, save_public_key_bundle}; use crate::{ app_state::AppState, calls::{call_manager::CallManager, call_util::LiveKitService}, config::Config, + identity::{ + KEYRING_PATH, PUBLIC_KEY_PATH, identity_secret_from_environment, load_or_create_keyring, + }, omega::omega_connection::{OmegaConnection, start_task_cleanup_loop}, rho::rho_manager::RhoManager, rho::server::start, util::logger::{PrintType, startup}, }; -const KEYRING_PATH: &str = "./omikron.mk"; -const PUBLIC_KEY_PATH: &str = "./omikron.mpkb"; - -fn load_keyring() -> Result { - match load_keyring_raw(KEYRING_PATH) { - Ok(keyring) => Ok(keyring), - Err(_) => { - let kr = Keyring::generate(); - save_keyring_raw(&kr, KEYRING_PATH).map_err(|error| error.to_string())?; - save_public_key_bundle(&kr.public_key_bundle(), PUBLIC_KEY_PATH) - .map_err(|error| error.to_string())?; - eprintln!("Generated new keyring at {}", KEYRING_PATH); - Ok(kr) - } - } -} - #[tokio::main] async fn main() { if let Err(_) = default_provider().install_default() { @@ -64,10 +50,10 @@ async fn main() { } }; - let keyring = match load_keyring() { - Ok(keyring) => keyring, + let identity_secret = match identity_secret_from_environment() { + Ok(secret) => secret, Err(error) => { - eprintln!("Unable to load keyring: {error}"); + eprintln!("Unable to load Omikron identity secret: {error}"); return; } }; diff --git a/src/omega/capabilities.rs b/src/omega/capabilities.rs index b658068..cff09a8 100644 --- a/src/omega/capabilities.rs +++ b/src/omega/capabilities.rs @@ -42,6 +42,7 @@ impl PeerCapabilities { format!("{OMIKRON_PREFIX}{}", names.join(",")) } + #[allow(dead_code)] pub fn from_identification_description(description: Option<&str>) -> Result { parse_capabilities(description, OMIKRON_PREFIX) } diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index f8d505f..3e7755c 100644 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -2,7 +2,7 @@ use super::capabilities::PeerCapabilities; use crate::{ config::Config, log_cv_in, log_cv_out, log_err, log_in, log_out, - rho::connection::{MtpValueCompat, OptionalDataValueCompat}, + rho::connection::{OptionalDataValueCompat, RequiredMtpFields}, rho::relay_router, rho::rho_manager::RhoManager, util::{data_type_id, logger::PrintType}, @@ -16,6 +16,7 @@ use mtp::{ host::{Policy, SendMode}, }; use mtp_transport::ConnectionHandle; +use rand::RngExt; use std::{sync::Arc, time::Duration}; use tokio::{ sync::{Mutex, RwLock, mpsc, oneshot, watch}, @@ -27,6 +28,7 @@ use uuid::Uuid; const RECONNECT_DELAY: Duration = Duration::from_secs(5); const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(300); +const RECONNECT_JITTER: Duration = Duration::from_secs(1); const CONNECTION_TIMEOUT: Duration = Duration::from_secs(10); const CAPABILITY_NEGOTIATION_TIMEOUT: Duration = Duration::from_secs(1); const PING_INTERVAL: Duration = Duration::from_secs(5); @@ -36,13 +38,29 @@ const MAX_CONCURRENT_REQUESTS: usize = 128; const CIRCUIT_BREAKER_FAILURE_THRESHOLD: u32 = 3; const CIRCUIT_BREAKER_COOLDOWN: Duration = Duration::from_secs(30); +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ConnectionOutcome { + HealthySessionEnded, + FailedBeforeHealthy, +} + +fn reconnect_base_after_outcome(current: Duration, outcome: ConnectionOutcome) -> Duration { + match outcome { + ConnectionOutcome::HealthySessionEnded => RECONNECT_DELAY, + ConnectionOutcome::FailedBeforeHealthy => current, + } +} + +fn reconnect_delay_with_jitter(delay: Duration) -> Duration { + let max_jitter_ms = std::cmp::min(delay, RECONNECT_JITTER).as_millis() as u64; + delay + Duration::from_millis(rand::rng().random_range(0..=max_jitter_ms)) +} + fn client_changed_target(value: &CommunicationValue) -> Option<(i64, i64)> { if !value.is_type(CommunicationType::ClientChanged) { return None; } - let receiver = i64::try_from(value.get_receiver()) - .ok() - .filter(|id| *id > 0)?; + let receiver = i64::try_from(value.receiver()?).ok().filter(|id| *id > 0)?; let session_id = value .get_data(DataType::SessionId) .as_signed_number() @@ -284,33 +302,45 @@ impl OmegaConnection { } match self.clone().connect_once().await { - Ok(()) => { - // Connection closed gracefully, check if we should reconnect - if *self.reconnect_on_close.read().await { - log_err!( - 0, - PrintType::Omega, - "Connection lost, reconnecting in {:?}...", - reconnect_delay - ); - } else { - log_in!(0, PrintType::Omega, "Connection closed, not reconnecting"); - break; + Ok(outcome) => { + reconnect_delay = reconnect_base_after_outcome(reconnect_delay, outcome); + match outcome { + ConnectionOutcome::HealthySessionEnded => { + log_err!( + 0, + PrintType::Omega, + "Healthy Omega connection ended; reconnecting" + ); + } + ConnectionOutcome::FailedBeforeHealthy => { + log_err!( + 0, + PrintType::Omega, + "Omega connection failed before synchronization" + ); + } } } - Err(e) => { - log_err!( - 0, - PrintType::Omega, - "Connection failed: {}, retrying in {:?}...", - e, - reconnect_delay - ); + Err(error) => { + log_err!(0, PrintType::Omega, "Connection failed: {}", error); } } + if !*self.reconnect_on_close.read().await { + log_in!(0, PrintType::Omega, "Connection closed, not reconnecting"); + break; + } + + let retry_delay = reconnect_delay_with_jitter(reconnect_delay); + log_err!( + 0, + PrintType::Omega, + "Retrying Omega connection in {:?}...", + retry_delay + ); + tokio::select! { - _ = sleep(reconnect_delay) => {} + _ = sleep(retry_delay) => {} _ = shutdown_rx.changed() => { if *shutdown_rx.borrow() { break; @@ -322,7 +352,7 @@ impl OmegaConnection { } } - async fn connect_once(self: Arc) -> Result<(), String> { + async fn connect_once(self: Arc) -> Result { *self.state.write().await = ConnectionState::Connecting; let client_config = ClientConfig::new(format!("https://{}:{}", self.host, self.port)) @@ -331,7 +361,6 @@ impl OmegaConnection { .with_policy( Policy::default() .with_send_mode(SendMode::SingleStreamPerMessage) - .with_max_message_size(1_000_000_000) .with_timeouts( Duration::from_millis(5_000), Duration::from_millis(5_000), @@ -343,8 +372,7 @@ impl OmegaConnection { .with_max_concurrent_stream_tasks(64) .with_persistent_stream_retries(5, Duration::from_secs(5)), ) - .with_ping_interval(PING_INTERVAL) - .with_max_missed_pings(0); + .with_ping_interval(PING_INTERVAL); let host_public_key = load_public_key_bundle("./omega.mpkb") .map_err(|e| format!("Failed to load omega.mpkb: {}", e))?; @@ -389,7 +417,7 @@ impl OmegaConnection { self.fail_synchronization("capability negotiation", error) .await; let _ = read_handle.await; - return Err("Omega capability negotiation failed".to_string()); + return Ok(ConnectionOutcome::FailedBeforeHealthy); } Ok(Err(_)) | Err(_) => { // An Omega from before capability negotiation sends no second @@ -403,7 +431,7 @@ impl OmegaConnection { // Tell omega our current state now that we're actually connected - // doing this after teardown (as before) sent into a sender that had // already been cleared, silently dropping the sync every time. - self.clone().sync_client_iota_status().await; + let reached_healthy_state = self.clone().sync_client_iota_status().await; // Wait for read loop to complete (connection closed) let result = read_handle.await; @@ -428,19 +456,13 @@ impl OmegaConnection { *self.state.write().await = ConnectionState::Disconnected; match result { - Ok(()) => { - // Check if we should reconnect - if *self.reconnect_on_close.read().await { - Err("Connection closed, will reconnect".to_string()) - } else { - Ok(()) - } - } + Ok(()) if reached_healthy_state => Ok(ConnectionOutcome::HealthySessionEnded), + Ok(()) => Ok(ConnectionOutcome::FailedBeforeHealthy), Err(e) => Err(format!("Read loop error: {}", e)), } } - async fn sync_client_iota_status(&self) { + async fn sync_client_iota_status(&self) -> bool { *self.state.write().await = ConnectionState::SynchronizingRoutes; let mut connected_iota_ids: Vec = Vec::new(); let mut connected_sessions: Vec = Vec::new(); @@ -497,16 +519,16 @@ impl OmegaConnection { if let Err(error) = self.send_message_result(&sync_msg).await { self.fail_synchronization("legacy route synchronization", error) .await; - return; + return false; } *self.state.write().await = ConnectionState::SynchronizingSubscriptions; if let Err(error) = self.restore_presence_subscriptions().await { self.fail_synchronization("subscription restoration", error) .await; - return; + return false; } *self.state.write().await = ConnectionState::Ready; - return; + return true; } let mut response = Err("route synchronization did not start".to_string()); @@ -535,7 +557,7 @@ impl OmegaConnection { ); self.fail_synchronization("subscription restoration", error) .await; - return; + return false; } *self.state.write().await = ConnectionState::Ready; } @@ -546,14 +568,15 @@ impl OmegaConnection { ); self.fail_synchronization("route synchronization", error) .await; - return; + return false; } Err(error) => { self.fail_synchronization("route synchronization", error) .await; - return; + return false; } } + true } async fn fail_synchronization(&self, phase: &str, error: String) { @@ -600,8 +623,39 @@ impl OmegaConnection { } if cv.is_type(CommunicationType::Relay) { - let destination_iota = cv.receiver().unwrap_or_default(); - let request_id = cv.get_id(); + let request_id = match cv.require_id() { + Ok(request_id) => request_id, + Err(error) => { + log_err!( + self.omikron_id as i64, + PrintType::Omega, + "Rejected malformed relay: {}", + error + ); + let response = + CommunicationValue::new(CommunicationType::ErrorInvalidData) + .without_id(); + let _ = self.send_message_result(&response).await; + continue; + } + }; + let destination_iota = match cv.require_receiver() { + Ok(destination_iota) => destination_iota, + Err(error) => { + log_err!( + self.omikron_id as i64, + PrintType::Omega, + "Rejected malformed relay: {}", + error + ); + let response = CommunicationValue::new( + CommunicationType::ErrorInvalidData, + ) + .with_id(request_id); + let _ = self.send_message_result(&response).await; + continue; + } + }; let response = match relay_router::route_from_omega(&self.rho, cv).await { Ok(()) => CommunicationValue::new(CommunicationType::Success) .with_id(request_id), @@ -630,7 +684,18 @@ impl OmegaConnection { continue; } - let msg_id = cv.get_id(); + let msg_id = match cv.require_id() { + Ok(msg_id) => msg_id, + Err(error) => { + log_err!( + self.omikron_id as i64, + PrintType::Omega, + "Discarded response without an MTP id: {}", + error + ); + continue; + } + }; if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) { if (task.task)(self.clone(), cv.clone()) { continue; @@ -749,6 +814,7 @@ impl OmegaConnection { self.try_send_message(cv).await } + #[allow(dead_code)] pub async fn supports_set_user_state(&self) -> bool { self.peer_capabilities.read().await.set_user_state_v1 } @@ -807,10 +873,12 @@ impl OmegaConnection { cv: &CommunicationValue, timeout_duration: Option, ) -> Result { + let msg_id = cv + .require_id() + .map_err(|error| format!("request is missing correlation id: {error}"))?; self.await_connection(timeout_duration).await?; let (tx, mut rx) = mpsc::channel(1); - let msg_id = cv.get_id(); self.waiting_tasks.insert( msg_id, @@ -819,9 +887,9 @@ impl OmegaConnection { log_in!( 0, PrintType::Omega, - "Matched Omega response (request_id={}, response_id={}, type={})", + "Matched Omega response (request_id={}, response_id={:?}, type={})", msg_id, - response_cv.get_id(), + response_cv.id(), response_cv .get_comm_type_enum() .map(|kind| kind.to_string()) @@ -880,7 +948,7 @@ impl OmegaConnection { pub async fn close_iota(&self, iota_id: i64) -> Result<(), String> { let cv = CommunicationValue::new(CommunicationType::IotaDisconnected) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())); - let request_id = cv.get_id(); + let request_id = cv.require_id().map_err(|error| error.to_string())?; let result = self.lifecycle_request(&cv).await; if let Err(error) = &result { self.log_lifecycle_failure("IotaDisconnected", iota_id, None, request_id, error); @@ -897,7 +965,7 @@ impl OmegaConnection { DataType::SessionId, DataValue::SignedNumber(session_id.into()), ); - let request_id = cv.get_id(); + let request_id = cv.require_id().map_err(|error| error.to_string())?; let result = self.lifecycle_request(&cv).await; if let Err(error) = &result { self.log_lifecycle_failure( @@ -938,7 +1006,7 @@ impl OmegaConnection { .as_millis() as i128, ), ); - let request_id = request.get_id(); + let request_id = request.require_id().map_err(|error| error.to_string())?; let result = self.lifecycle_request(&request).await; if let Err(error) = &result { self.log_lifecycle_failure( @@ -956,7 +1024,7 @@ impl OmegaConnection { pub async fn iota_connected(&self, iota_id: i64) -> Result<(), String> { let request = CommunicationValue::new(CommunicationType::IotaConnected) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())); - let request_id = request.get_id(); + let request_id = request.require_id().map_err(|error| error.to_string())?; let result = self.lifecycle_request(&request).await; if let Err(error) = &result { self.log_lifecycle_failure("IotaConnected", iota_id, None, request_id, error); @@ -965,6 +1033,7 @@ impl OmegaConnection { result } + #[allow(dead_code)] pub async fn reconcile_routes(&self) { self.sync_client_iota_status().await; } @@ -1174,7 +1243,11 @@ impl OmegaConnection { #[cfg(test)] mod tests { - use super::{PeerCapabilities, client_changed_target, parse_omega_capability_response}; + use super::{ + ConnectionOutcome, MAX_RECONNECT_DELAY, PeerCapabilities, RECONNECT_DELAY, + client_changed_target, parse_omega_capability_response, reconnect_base_after_outcome, + reconnect_delay_with_jitter, + }; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; fn notification() -> CommunicationValue { @@ -1196,6 +1269,9 @@ mod tests { let mut missing_state = notification(); missing_state.remove_data(DataType::UserState); assert_eq!(client_changed_target(&missing_state), None); + + let missing_receiver = notification().without_receiver(); + assert_eq!(client_changed_target(&missing_receiver), None); } #[test] @@ -1208,6 +1284,35 @@ mod tests { assert_eq!(client_changed_target(&invalid), None); } + #[test] + fn reconnect_jitter_is_bounded_without_changing_the_base_delay() { + for _ in 0..32 { + let delay = reconnect_delay_with_jitter(RECONNECT_DELAY); + assert!(delay >= RECONNECT_DELAY); + assert!(delay <= RECONNECT_DELAY + std::time::Duration::from_secs(1)); + } + + assert!(MAX_RECONNECT_DELAY > RECONNECT_DELAY); + } + + #[test] + fn healthy_session_resets_accumulated_reconnect_backoff() { + let mut delay = RECONNECT_DELAY; + for _ in 0..4 { + delay = std::cmp::min(delay * 2, MAX_RECONNECT_DELAY); + } + assert!(delay > RECONNECT_DELAY); + + assert_eq!( + reconnect_base_after_outcome(delay, ConnectionOutcome::FailedBeforeHealthy), + delay + ); + assert_eq!( + reconnect_base_after_outcome(delay, ConnectionOutcome::HealthySessionEnded), + RECONNECT_DELAY + ); + } + #[test] fn capability_response_negotiates_new_omega() { let response = CommunicationValue::new(CommunicationType::IdentificationResponse) diff --git a/src/rho/app_connection.rs b/src/rho/app_connection.rs index 7e8778a..c6af5e9 100644 --- a/src/rho/app_connection.rs +++ b/src/rho/app_connection.rs @@ -1,7 +1,7 @@ use crate::anonymous_clients::anonymous_manager; use crate::app_state::AppState; use crate::rho::connection::{ - GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat, + GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields, }; use crate::rho::relay_router::{self, RelaySource}; use crate::rho::rho_connection::RhoConnection; @@ -89,12 +89,38 @@ impl AppConnection { /// Handle incoming message from app pub async fn handle_message(self: Arc, cv: CommunicationValue) { tokio::spawn(async move { + let message_id = match cv.require_id() { + Ok(message_id) => message_id, + Err(error) => { + log_err!( + self.user_id as i64, + PrintType::App, + "Rejected malformed message: {}", + error + ); + let response = + CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id(); + self.send_message(&response).await; + return; + } + }; log_cv_in!(PrintType::App, cv); if cv.is_type(CommunicationType::Relay) { - let cv = relay_router::ensure_relay_frame_id(cv); - let request_id = cv.get_id(); - let next_hop = cv.receiver().unwrap_or_default(); + let next_hop = match cv.require_receiver() { + Ok(next_hop) => next_hop, + Err(error) => { + log_err!( + self.user_id as i64, + PrintType::App, + "Rejected malformed relay: {}", + error + ); + self.send_error_response(message_id, CommunicationType::ErrorInvalidData) + .await; + return; + } + }; let result = match self.get_rho_connection().await { Some(rho) => { relay_router::route_relay( @@ -102,7 +128,7 @@ impl AppConnection { RelaySource::Client { iota_id: rho.get_iota_id().await, }, - cv, + relay_router::ensure_relay_frame_id(cv), ) .await } @@ -110,7 +136,7 @@ impl AppConnection { }; let response = match result { Ok(()) => { - CommunicationValue::new(CommunicationType::Success).with_id(request_id) + CommunicationValue::new(CommunicationType::Success).with_id(message_id) } Err(error) => { log_err!( @@ -121,7 +147,7 @@ impl AppConnection { error ); CommunicationValue::new(relay_router::error_response_type(&error)) - .with_id(request_id) + .with_id(message_id) } }; self.send_message(&response).await; @@ -130,7 +156,7 @@ impl AppConnection { if cv.is_type(CommunicationType::Success) { if let Some(rho) = self.get_rho_connection().await { - rho.forward_relay_ack(self.user_id, cv.get_id()).await; + rho.forward_relay_ack(self.user_id, message_id).await; } return; } @@ -139,7 +165,7 @@ impl AppConnection { relay_router::message_security_class(&cv), relay_router::MessageSecurityClass::RelayOnly ) { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData) + self.send_error_response(message_id, CommunicationType::ErrorInvalidData) .await; return; } @@ -155,7 +181,7 @@ impl AppConnection { } } { let response = CommunicationValue::new(CommunicationType::GetUserData) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::Username, DataValue::Str(anonymous.get_user_name().await), @@ -197,7 +223,7 @@ impl AppConnection { "Rejected unsupported communication type {}", cv.get_type() ); - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData) + self.send_error_response(message_id, CommunicationType::ErrorInvalidData) .await; }); } diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index ba83728..5a25ec0 100644 --- a/src/rho/client_connection.rs +++ b/src/rho/client_connection.rs @@ -3,7 +3,7 @@ use crate::app_state::AppState; use crate::calls::call_group::call_invite_secret_from_cv; use crate::data::user::UserStatus; use crate::rho::connection::{ - GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat, + GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields, }; use crate::rho::relay_router::{self, RelaySource}; use crate::rho::rho_connection::RhoConnection; @@ -117,14 +117,41 @@ impl ClientConnection { }; tokio::spawn(async move { let _permit = permit; + let message_id = match cv.require_id() { + Ok(message_id) => message_id, + Err(error) => { + log_err!( + self.user_id as i64, + PrintType::Client, + "Rejected malformed message: {}", + error + ); + let response = + CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id(); + self.send_message(&response).await; + return; + } + }; log_cv_in!(PrintType::Client, cv); let mut cv = cv; if cv.is_type(CommunicationType::Relay) { + let next_hop = match cv.require_receiver() { + Ok(next_hop) => next_hop, + Err(error) => { + log_err!( + self.user_id as i64, + PrintType::Client, + "Rejected malformed relay: {}", + error + ); + self.send_error_response(message_id, CommunicationType::ErrorInvalidData) + .await; + return; + } + }; cv = relay_router::ensure_relay_frame_id(cv); - let request_id = cv.get_id(); - let next_hop = cv.receiver().unwrap_or_default(); let result = match self.get_rho_connection().await { Some(rho) => { relay_router::route_relay( @@ -140,7 +167,7 @@ impl ClientConnection { }; let response = match result { Ok(()) => { - CommunicationValue::new(CommunicationType::Success).with_id(request_id) + CommunicationValue::new(CommunicationType::Success).with_id(message_id) } Err(error) => { log_err!( @@ -151,7 +178,7 @@ impl ClientConnection { error ); CommunicationValue::new(relay_router::error_response_type(&error)) - .with_id(request_id) + .with_id(message_id) } }; self.send_message(&response).await; @@ -160,7 +187,7 @@ impl ClientConnection { if cv.is_type(CommunicationType::Success) { if let Some(rho) = self.get_rho_connection().await { - rho.forward_relay_ack(self.user_id, cv.get_id()).await; + rho.forward_relay_ack(self.user_id, message_id).await; } return; } @@ -169,7 +196,7 @@ impl ClientConnection { relay_router::message_security_class(&cv), relay_router::MessageSecurityClass::RelayOnly ) { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData) + self.send_error_response(message_id, CommunicationType::ErrorInvalidData) .await; return; } @@ -178,11 +205,11 @@ impl ClientConnection { // user fields, if present, are deliberately ignored: an // authenticated connection may only change its own state. if cv.is_type(CommunicationType::ClientChanged) - && cv.get_data_opt(DataType::UserState).is_some() + && cv.get_data(DataType::UserState).is_some() { self.handle_set_user_state( CommunicationValue::new(CommunicationType::ClientChanged) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::UserState, cv.get_data(DataType::UserState) @@ -240,7 +267,7 @@ impl ClientConnection { } } { let response = CommunicationValue::new(CommunicationType::GetUserData) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::Username, DataValue::Str(anonymous.get_user_name().await), @@ -276,19 +303,19 @@ impl ClientConnection { || cv.is_type(CommunicationType::DeleteUser) { if cv.is_type(CommunicationType::ChangeUserData) - && cv.get_data_opt(DataType::OnlineStatus).is_some() + && cv.get_data(DataType::OnlineStatus).is_some() { let mut profile_request = cv.clone(); let preference = profile_request .remove_data(DataType::OnlineStatus) .unwrap_or(DataValue::Null); let state_request = CommunicationValue::new(CommunicationType::ClientChanged) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::UserState, preference); let state_response = match self.request_set_user_state(state_request).await { Ok(response) => response, Err(error_type) => { - self.send_error_response(cv.get_id(), error_type).await; + self.send_error_response(message_id, error_type).await; return; } }; @@ -311,7 +338,7 @@ impl ClientConnection { } Ok(response) => self.send_message(&response).await, Err(_) => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInternal) + self.send_error_response(message_id, CommunicationType::ErrorInternal) .await; } } @@ -336,7 +363,7 @@ impl ClientConnection { if is_per_device_settings { let Some(session_id) = session_id else { let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(self.user_id) .add_typed_default( DataType::Message, @@ -352,7 +379,7 @@ impl ClientConnection { if session_id != expected_session_id { let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(self.user_id) .add_typed_default( DataType::Message, @@ -368,7 +395,7 @@ impl ClientConnection { } else if let Some(session_id) = session_id { if session_id != expected_session_id { let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(self.user_id) .add_typed_default( DataType::Message, @@ -396,7 +423,7 @@ impl ClientConnection { if let Some(session_id) = cv.get_data(DataType::SessionId).as_signed_number() { if session_id != expected_session_id { let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::SessionId, DataValue::SignedNumber(expected_session_id), @@ -417,12 +444,14 @@ impl ClientConnection { "Rejected unsupported communication type {}", cv.get_type() ); - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData) + self.send_error_response(message_id, CommunicationType::ErrorInvalidData) .await; }); } async fn handle_omega_forward(self: Arc, cv: CommunicationValue) { - let request_id = cv.get_id(); + let Ok(request_id) = cv.require_id() else { + return; + }; match self.await_omega_response(cv).await { Ok(response_cv) => self.send_message(&response_cv).await, Err(_) => { @@ -445,6 +474,9 @@ impl ClientConnection { &self, cv: CommunicationValue, ) -> Result { + let message_id = cv + .require_id() + .map_err(|_| CommunicationType::ErrorInvalidData)?; if !self.state.omega.is_ready().await { return Err(CommunicationType::ErrorInternal); } @@ -459,7 +491,7 @@ impl ClientConnection { return Err(CommunicationType::ErrorNoIota); }; let request = CommunicationValue::new(CommunicationType::ClientChanged) - .with_id(cv.get_id()) + .with_id(message_id) .with_sender(self.user_id) .add_typed_default( DataType::UserId, @@ -472,22 +504,28 @@ impl ClientConnection { .await .map_err(|_| CommunicationType::ErrorInternal)?; return Ok(CommunicationValue::new(CommunicationType::Success) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::UserState, DataValue::Str(state.to_string()))); } async fn handle_set_user_state(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; match self.request_set_user_state(cv.clone()).await { Ok(response) => self.send_message(&response).await, - Err(error_type) => self.send_error_response(cv.get_id(), error_type).await, + Err(error_type) => self.send_error_response(message_id, error_type).await, } } /// Handle call invite async fn handle_call_invite(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let receiver_id: i128 = cv.get_data(DataType::ReceiverId).as_number().unwrap_or(0); if receiver_id == 0 { - self.send_error_response(cv.get_id(), CommunicationType::ErrorNoUserId) + self.send_error_response(message_id, CommunicationType::ErrorNoUserId) .await; return; } @@ -496,13 +534,13 @@ impl ClientConnection { Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) { Ok(id) => id, Err(_) => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } }, _ => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorNoCallId) + self.send_error_response(message_id, CommunicationType::ErrorNoCallId) .await; return; } @@ -511,7 +549,7 @@ impl ClientConnection { let secret = match call_invite_secret_from_cv(&cv) { Some(secret) => secret, None => { - self.send_error_response(cv.get_id(), CommunicationType::BadRequest) + self.send_error_response(message_id, CommunicationType::BadRequest) .await; return; } @@ -522,7 +560,7 @@ impl ClientConnection { .add_invite(call_id, self.user_id, receiver_id as u64, secret.clone()) .await; if !invited { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } @@ -532,7 +570,7 @@ impl ClientConnection { .call_manager .should_forward_invite(self.user_id, receiver_id as u64) { - let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id()); + let response = CommunicationValue::new(CommunicationType::Success).with_id(message_id); self.send_message(&response).await; return; } @@ -564,7 +602,7 @@ impl ClientConnection { }); let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default( DataType::ReceiverId, DataValue::SignedNumber(receiver_id.into()), @@ -594,25 +632,28 @@ impl ClientConnection { target_rho.message_to_client(forward).await; - let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id()); + let response = CommunicationValue::new(CommunicationType::Success).with_id(message_id); self.send_message(&response).await; } /// Handle get call request async fn handle_get_call(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let user_id = self.get_user_id().await; let call_id = match cv.get_data(DataType::CallId) { Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) { Ok(id) => id, Err(_) => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } }, _ => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorNoCallId) + self.send_error_response(message_id, CommunicationType::ErrorNoCallId) .await; return; } @@ -626,7 +667,7 @@ impl ClientConnection { { Ok(token) => { let response = CommunicationValue::new(CommunicationType::CallToken) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(user_id as u64) .add_typed_default(DataType::CallToken, DataValue::Str(token)); self.send_message(&response).await; @@ -634,26 +675,29 @@ impl ClientConnection { Err(error) => { log::warn!("Unable to create call token for {}: {}", call_id, error); let error_cv = CommunicationValue::new(CommunicationType::ErrorNoCallId) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string())); self.send_message(&error_cv).await; } } } async fn handle_get_call_data(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let user_id = self.get_user_id().await; let call_id = match cv.get_data(DataType::CallId) { Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) { Ok(id) => id, Err(_) => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; } }, _ => { - self.send_error_response(cv.get_id(), CommunicationType::ErrorNoCallId) + self.send_error_response(message_id, CommunicationType::ErrorNoCallId) .await; return; } @@ -670,20 +714,20 @@ impl ClientConnection { } let response = CommunicationValue::new(CommunicationType::CallData) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(user_id as u64) .add_typed_default(DataType::UserIds, DataValue::Array(user_ids)); self.send_message(&response).await; } else { let error_cv = CommunicationValue::new(CommunicationType::ErrorInvalidUserId) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())); self.send_message(&error_cv).await; return; } } else { let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string())); self.send_message(&error_cv).await; return; @@ -691,9 +735,12 @@ impl ClientConnection { } async fn handle_call_timeout_user(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or("")) else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; }; @@ -707,13 +754,13 @@ impl ClientConnection { .unwrap_or(0); let Some(call) = self.state.call_manager.get_call(call_id).await else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound) + self.send_error_response(message_id, CommunicationType::ErrorNotFound) .await; return; }; let Some(caller) = call.get_caller(self.get_user_id().await).await else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidUserId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId) .await; return; }; @@ -729,9 +776,12 @@ impl ClientConnection { } } async fn handle_call_disconnect_user(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or("")) else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; }; @@ -741,12 +791,12 @@ impl ClientConnection { .unwrap_or(0); let Some(call) = self.state.call_manager.get_call(call_id).await else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound) + self.send_error_response(message_id, CommunicationType::ErrorNotFound) .await; return; }; let Some(caller) = call.get_caller(self.get_user_id().await).await else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidUserId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId) .await; return; }; @@ -755,9 +805,12 @@ impl ClientConnection { } } async fn handle_call_set_anonymous_joining(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or("")) else { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId) + self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId) .await; return; }; @@ -780,7 +833,7 @@ impl ClientConnection { short_link = call.get_short_link().await; } let mut response_cv = CommunicationValue::new(CommunicationType::CallSetAnonymousJoining) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string())) .add_typed_default(DataType::Enabled, DataValue::Bool(enable)); if let Some(short_link) = short_link { @@ -790,6 +843,9 @@ impl ClientConnection { } async fn handle_load_txt_record(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; if let Some(path) = cv.get_data(DataType::Path).as_str() { let resolver = match TokioAsyncResolver::tokio_from_system_conf() { Ok(r) => r, @@ -799,7 +855,7 @@ impl ClientConnection { .cloned() .unwrap_or(DataValue::Null); let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::Path, path_data); self.send_message(&error_cv).await; return; @@ -818,7 +874,7 @@ impl ClientConnection { Ok(text) => text, Err(_) => { self.send_error_response( - cv.get_id(), + message_id, CommunicationType::ErrorInvalidData, ) .await; @@ -827,8 +883,8 @@ impl ClientConnection { }; let response = CommunicationValue::new(CommunicationType::LoadTxtRecord) - .with_id(cv.get_id()) - .add_typed_default(DataType::Content, DataValue::Str(record_text)); + .with_id(message_id) + .add_typed_default(DataType::AppContent, DataValue::Str(record_text)); self.send_message(&response).await; return; } @@ -838,7 +894,7 @@ impl ClientConnection { .cloned() .unwrap_or(DataValue::Null); let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::Path, path_data); self.send_message(&error_cv).await; } @@ -848,7 +904,7 @@ impl ClientConnection { .cloned() .unwrap_or(DataValue::Null); let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::Path, path_data); self.send_message(&error_cv).await; } @@ -862,7 +918,7 @@ impl ClientConnection { .cloned() .unwrap_or(DataValue::Null); let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) + .with_id(message_id) .add_typed_default(DataType::Path, path_data); self.send_message(&error_cv).await; } diff --git a/src/rho/connection.rs b/src/rho/connection.rs index 8109f36..fed5169 100755 --- a/src/rho/connection.rs +++ b/src/rho/connection.rs @@ -16,37 +16,38 @@ use crate::{ use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use mtp::host::AuthState; use mtp::webserver::{WebMTPConnection, WebMtpReceiver, WebMtpSender}; +use thiserror::Error; pub type MtpSender = WebMtpSender; pub type MtpReceiver = WebMtpReceiver; -/* - * MTP 0.3 exposes absent frame fields and data entries as Options. These - * adapters keep legacy control handlers explicit while Relay code uses the - * native optional accessors directly. - */ -pub(crate) trait MtpValueCompat { - fn get_id(&self) -> u32; - fn get_sender(&self) -> u64; - fn get_receiver(&self) -> u64; - fn get_data_opt(&self, data_type: DataType) -> Option<&DataValue>; +#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] +pub enum FrameValidationError { + #[error("message is missing an MTP id")] + MissingId, + #[error("message is missing an MTP sender")] + MissingSender, + #[error("message is missing an MTP receiver")] + MissingReceiver, } -impl MtpValueCompat for CommunicationValue { - fn get_id(&self) -> u32 { - self.id().unwrap_or_default() +pub trait RequiredMtpFields { + fn require_id(&self) -> Result; + fn require_sender(&self) -> Result; + fn require_receiver(&self) -> Result; +} + +impl RequiredMtpFields for CommunicationValue { + fn require_id(&self) -> Result { + self.id().ok_or(FrameValidationError::MissingId) } - fn get_sender(&self) -> u64 { - self.sender().unwrap_or_default() + fn require_sender(&self) -> Result { + self.sender().ok_or(FrameValidationError::MissingSender) } - fn get_receiver(&self) -> u64 { - self.receiver().unwrap_or_default() - } - - fn get_data_opt(&self, data_type: DataType) -> Option<&DataValue> { - self.get_data(data_type) + fn require_receiver(&self) -> Result { + self.receiver().ok_or(FrameValidationError::MissingReceiver) } } @@ -55,6 +56,7 @@ pub(crate) trait OptionalDataValueCompat { fn as_number(&self) -> Option; fn as_signed_number(&self) -> Option; fn as_str(&self) -> Option<&str>; + #[allow(dead_code)] fn as_bytes(&self) -> Option>; } @@ -80,6 +82,42 @@ impl OptionalDataValueCompat for Option<&DataValue> { } } +#[cfg(test)] +mod tests { + use super::{FrameValidationError, RequiredMtpFields}; + use mtp::codec::{CommunicationType, CommunicationValue}; + + #[test] + fn required_fields_preserve_missing_field_errors() { + let frame = CommunicationValue::new(CommunicationType::Success) + .without_id() + .without_sender() + .without_receiver(); + + assert_eq!(frame.require_id(), Err(FrameValidationError::MissingId)); + assert_eq!( + frame.require_sender(), + Err(FrameValidationError::MissingSender) + ); + assert_eq!( + frame.require_receiver(), + Err(FrameValidationError::MissingReceiver) + ); + } + + #[test] + fn zero_is_a_present_routing_value() { + let frame = CommunicationValue::new(CommunicationType::Success) + .with_id(0) + .with_sender(0) + .with_receiver(0); + + assert_eq!(frame.require_id(), Ok(0)); + assert_eq!(frame.require_sender(), Ok(0)); + assert_eq!(frame.require_receiver(), Ok(0)); + } +} + /* * How a connection identified itself during the mtp handshake driven by * `server.rs` ("iota" / "client" authenticated logins, "anonymous" @@ -149,6 +187,10 @@ impl GeneralConnection { })) } + pub fn connection_kind(&self) -> ConnectionKind { + self.connection_kind + } + pub async fn handle(self: Arc) { log_in!(0, PrintType::General, "General connection handler started"); if self.migrate().await { diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index de4f02a..bd822cb 100755 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -6,7 +6,7 @@ use crate::log_err; use crate::log_in; use crate::log_out; use crate::rho::connection::{ - GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat, + GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields, }; use crate::rho::relay_router::{self, RelaySource}; use crate::util::data_type_id; @@ -33,8 +33,9 @@ fn contact_snapshot(value: &CommunicationValue) -> Option<(i64, i64, Vec)> return None; } - let user_id = i64::try_from(value.get_receiver()) - .ok() + let user_id = value + .receiver() + .and_then(|id| i64::try_from(id).ok()) .filter(|id| *id > 0)?; let session_id = value .get_data(DataType::SessionId) @@ -222,6 +223,7 @@ impl IotaConnection { } } + #[allow(dead_code)] pub async fn send_relay(&self, cv: &CommunicationValue) -> Result<(), String> { self.sender .send(cv) @@ -235,20 +237,46 @@ impl IotaConnection { return; }; let _permit = permit; + let message_id = match cv.require_id() { + Ok(message_id) => message_id, + Err(error) => { + log_err!( + self.iota_id as i64, + PrintType::Iota, + "Rejected malformed message: {}", + error + ); + let response = + CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id(); + self.send_message(&response).await; + return; + } + }; if cv.is_type(CommunicationType::Relay) { - let cv = relay_router::ensure_relay_frame_id(cv); - let request_id = cv.get_id(); - let next_hop = cv.receiver().unwrap_or_default(); + let next_hop = match cv.require_receiver() { + Ok(next_hop) => next_hop, + Err(error) => { + log_err!( + self.iota_id as i64, + PrintType::Iota, + "Rejected malformed relay: {}", + error + ); + self.send_error_response(message_id, CommunicationType::ErrorInvalidData, None) + .await; + return; + } + }; let response = match relay_router::route_relay( &self.state, RelaySource::Iota { iota_id: self.iota_id, }, - cv, + relay_router::ensure_relay_frame_id(cv), ) .await { - Ok(()) => CommunicationValue::new(CommunicationType::Success).with_id(request_id), + Ok(()) => CommunicationValue::new(CommunicationType::Success).with_id(message_id), Err(error) => { log_err!( self.iota_id as i64, @@ -258,7 +286,7 @@ impl IotaConnection { error ); CommunicationValue::new(relay_router::error_response_type(&error)) - .with_id(request_id) + .with_id(message_id) } }; self.send_message(&response).await; @@ -269,12 +297,12 @@ impl IotaConnection { crate::rho::relay_router::message_security_class(&cv), crate::rho::relay_router::MessageSecurityClass::RelayOnly ) { - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData, None) + self.send_error_response(message_id, CommunicationType::ErrorInvalidData, None) .await; return; } - let msg_id = cv.get_id(); + let msg_id = message_id; if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) { if (task)(self.clone(), cv.clone()) { return; @@ -309,7 +337,7 @@ impl IotaConnection { if cv.is_type(CommunicationType::StateSubscribe) { self.send_error_response( - cv.get_id(), + message_id, CommunicationType::ErrorInvalidData, Some("StateSubscribe must come from an authoritative contact snapshot"), ) @@ -335,7 +363,7 @@ impl IotaConnection { self.iota_id as i64, PrintType::Omega, "Forwarding CompleteRegisterUser to Omega (request_id={})", - request.get_id() + message_id ); let mut response_cv = self .state @@ -348,7 +376,7 @@ impl IotaConnection { self.iota_id as i64, PrintType::Omega, "CompleteRegisterUser request_id={} failed: {}; retrying once", - request.get_id(), + message_id, error ); response_cv = self @@ -363,9 +391,9 @@ impl IotaConnection { log_in!( self.iota_id as i64, PrintType::Omega, - "Omega completed registration (request_id={}, response_id={}, type={})", - request.get_id(), - response_cv.get_id(), + "Omega completed registration (request_id={}, response_id={:?}, type={})", + message_id, + response_cv.id(), response_cv .get_comm_type_enum() .map(|kind| kind.to_string()) @@ -408,7 +436,7 @@ impl IotaConnection { self.add_user_id(user_id as u64).await; self.send_message( &CommunicationValue::new(CommunicationType::Success) - .with_id(cv.get_id()), + .with_id(message_id), ) .await; return; @@ -416,15 +444,15 @@ impl IotaConnection { Ok(verified) => log_err!( self.iota_id as i64, PrintType::Omega, - "Registration verification returned an unexpected user (request_id={}, response_id={})", - verification.get_id(), - verified.get_id() + "Registration verification returned an unexpected user (request_id={:?}, response_id={:?})", + verification.id(), + verified.id() ), Err(verify_error) => log_err!( self.iota_id as i64, PrintType::Omega, - "Registration verification failed after request_id={}: {}", - verification.get_id(), + "Registration verification failed after request_id={:?}: {}", + verification.id(), verify_error ), } @@ -436,7 +464,7 @@ impl IotaConnection { error ); self.send_error_response( - cv.get_id(), + message_id, CommunicationType::ErrorInternal, Some(&format!("Omega forwarding failed: {error}")), ) @@ -477,7 +505,7 @@ impl IotaConnection { "Rejected unsupported communication type {}", cv.get_type() ); - self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData, None) + self.send_error_response(message_id, CommunicationType::ErrorInvalidData, None) .await; } @@ -514,6 +542,9 @@ impl IotaConnection { } async fn handle_omega_forward_without_sender(self: Arc, cv: CommunicationValue) { + let Ok(message_id) = cv.require_id() else { + return; + }; let iota_for_closure = self.clone(); let request = cv.clone().add_typed_default( DataType::IotaId, @@ -530,7 +561,7 @@ impl IotaConnection { self.iota_id as i64, PrintType::Omega, "GetRegister request_id={} failed: {}; retrying once", - request.get_id(), + message_id, error ); response_cv = self @@ -550,7 +581,7 @@ impl IotaConnection { error ); self.send_error_response( - cv.get_id(), + message_id, CommunicationType::ErrorInternal, Some(&format!("Omega forwarding failed: {error}")), ) @@ -560,7 +591,17 @@ impl IotaConnection { } /// Handle GET_CHATS message async fn handle_get_chats(&self, cv: CommunicationValue) { - let user_id = cv.get_sender(); + let Ok(message_id) = cv.require_id() else { + return; + }; + let Ok(user_id) = cv.require_sender() else { + log_err!( + self.iota_id as i64, + PrintType::Iota, + "Rejected get_chats without an MTP sender" + ); + return; + }; // Authority check: user must be linked to this Iota if !self.get_user_ids().await.contains(&user_id) { @@ -582,7 +623,7 @@ impl IotaConnection { else { self.forward_to_client( CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) + .with_id(message_id) .with_receiver(user_id), ) .await; @@ -763,7 +804,14 @@ impl IotaConnection { async fn add_call_state(&self, response: CommunicationValue) -> CommunicationValue { let mut output = response.clone(); - let user_id = response.get_receiver(); + let Some(user_id) = response.receiver() else { + log_err!( + self.iota_id as i64, + PrintType::Iota, + "Discarded response without an MTP receiver" + ); + return output; + }; let typed_data: Vec<_> = response.iter_typed_data().collect(); for (key, value) in typed_data { @@ -807,7 +855,9 @@ impl IotaConnection { timeout_duration: Option, ) -> Result { let (tx, mut rx) = mpsc::channel(1); - let msg_id = cv.get_id(); + let msg_id = cv + .require_id() + .map_err(|error| format!("request is missing correlation id: {error}"))?; let task_tx = tx.clone(); self.waiting_tasks.insert( diff --git a/src/rho/relay_router.rs b/src/rho/relay_router.rs index 0e0ffaa..bd0f0ab 100644 --- a/src/rho/relay_router.rs +++ b/src/rho/relay_router.rs @@ -26,6 +26,7 @@ pub enum RouteTarget { } impl RouteTarget { + #[allow(dead_code)] pub fn wire_id(self) -> Option { let (kind, id) = match self { Self::User(id) => (USER_TARGET_KIND, id), @@ -251,7 +252,7 @@ fn route_response(response: CommunicationValue) -> Result<(), RelayRouteError> { } pub fn ensure_relay_frame_id(frame: CommunicationValue) -> CommunicationValue { - if frame.id().is_some_and(|id| id != 0) { + if frame.id().is_some() { return frame; } let id = NEXT_RELAY_FRAME_ID.fetch_add(1, Ordering::Relaxed).max(1); diff --git a/src/rho/rho_connection.rs b/src/rho/rho_connection.rs index 094bb1b..f4aa658 100644 --- a/src/rho/rho_connection.rs +++ b/src/rho/rho_connection.rs @@ -3,10 +3,7 @@ use super::{client_connection::ClientConnection, iota_connection::IotaConnection use super::relay_router::RouteTarget; use crate::{ log_err, - rho::{ - app_connection::AppConnection, - connection::{MtpValueCompat, OptionalDataValueCompat}, - }, + rho::{app_connection::AppConnection, connection::OptionalDataValueCompat}, }; use dashmap::DashMap; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; @@ -257,7 +254,14 @@ impl RhoConnection { /// Send message from Iota to specific client pub async fn message_to_client(&self, cv: CommunicationValue) { let connections = self.get_client_connections().await; - let receiver_id = cv.get_receiver(); + let Some(receiver_id) = cv.receiver() else { + log_err!( + 0, + crate::util::logger::PrintType::General, + "Discarded message without an MTP receiver" + ); + return; + }; let session_id = cv.get_data(DataType::SessionId).as_number(); for connection in connections.iter() { @@ -273,6 +277,7 @@ impl RhoConnection { } } + #[allow(dead_code)] pub async fn message_to_iota(&self, cv: CommunicationValue) { self.iota_connection.send_message(&cv).await; } @@ -308,6 +313,7 @@ impl RhoConnection { send_error.map_or(Ok(()), Err) } + #[allow(dead_code)] pub async fn send_relay_to_iota(&self, cv: &CommunicationValue) -> Result<(), String> { self.iota_connection.send_relay(cv).await } diff --git a/src/rho/rho_manager.rs b/src/rho/rho_manager.rs index 2eb4632..18003bf 100644 --- a/src/rho/rho_manager.rs +++ b/src/rho/rho_manager.rs @@ -24,6 +24,7 @@ impl RhoManager { self.users.get(&user_id).map(|entry| entry.value().clone()) } + #[allow(dead_code)] pub async fn contains_iota(&self, iota_id: i64) -> bool { self.connections.contains_key(&iota_id) } diff --git a/src/rho/server.rs b/src/rho/server.rs index c3e73ab..b94885b 100644 --- a/src/rho/server.rs +++ b/src/rho/server.rs @@ -8,7 +8,7 @@ use crate::{ app_state::AppState, log, log_err, omega::omega_connection::OmegaConnection, - rho::connection::{GeneralConnection, OptionalDataValueCompat}, + rho::connection::{ConnectionKind, GeneralConnection, OptionalDataValueCompat}, util::{file_util::load_file_vec, logger::PrintType}, }; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; @@ -16,6 +16,27 @@ use mtp::crypto::PublicKeyBundle; use mtp::host::{AuthenticationPolicy, HostConfig, Policy, SendMode}; use mtp::webserver::{MTPWebServer, WebServerConfig}; +fn web_config(max_connections: usize) -> Result { + WebServerConfig::new() + .max_connections(max_connections) + .route("/", |_request, response| async move { response.body("OK") }) +} + +fn rho_policy() -> Policy { + Policy::default() + .with_send_mode(SendMode::SingleStreamPerMessage) + .with_timeouts( + Duration::from_millis(2_000), + Duration::from_millis(2_000), + Duration::from_millis(30_000), + ) + .with_keep_alive(Some(Duration::from_secs(6))) + .with_max_idle_timeout(Some(Duration::from_secs(30))) + .with_receiver_queue_capacity(1000) + .with_max_concurrent_stream_tasks(10) + .with_persistent_stream_retries(5, Duration::from_secs(5)) +} + /* * Resolves the PublicKeyBundle mtp needs to verify a login's signed * challenge response. "iota"/"client" ids are looked up through Omega, the @@ -126,21 +147,7 @@ pub async fn start(state: Arc) -> Result<(), Box) -> Result<(), Box) -> Result<(), Box permit, + Err(_) => { + log_err!( + 0, + PrintType::General, + "Rejected connection: global Rho connection limit reached" + ); + continue; + } + }; + let peer_ip = conn.remote_addr.map(|address| address.ip()); let state = state.clone(); tokio::spawn(async move { - let Some(conn) = GeneralConnection::new(conn, state) else { + let Some(conn) = GeneralConnection::new(conn, state.clone()) else { log_err!( 0, PrintType::General, @@ -199,6 +217,58 @@ pub async fn start(state: Arc) -> Result<(), Box Some(permit), + Err(_) => { + log_err!( + 0, + PrintType::General, + "Rejected anonymous connection: anonymous limit reached" + ); + return; + } + } + } else { + None + }; + + let anonymous_ip_permit = if conn.connection_kind() == ConnectionKind::AnonymousClient { + let Some(peer_ip) = peer_ip else { + log_err!( + 0, + PrintType::General, + "Rejected anonymous connection: peer address unavailable" + ); + return; + }; + match state + .rho_connection_limits + .try_acquire_anonymous_per_ip(peer_ip) + { + Some(permit) => Some(permit), + None => { + log_err!( + 0, + PrintType::General, + "Rejected anonymous connection: per-IP limit reached" + ); + return; + } + } + } else { + None + }; + + let _global_permit = global_permit; + let _anonymous_permit = anonymous_permit; + let _anonymous_ip_permit = anonymous_ip_permit; conn.handle().await; }); } @@ -206,3 +276,22 @@ pub async fn start(state: Arc) -> Result<(), Box, } +#[allow(dead_code)] impl CallService { pub fn new(state: Arc) -> Self { Self { state } diff --git a/src/services/routing_service.rs b/src/services/routing_service.rs index 86f034f..5410644 100644 --- a/src/services/routing_service.rs +++ b/src/services/routing_service.rs @@ -4,10 +4,12 @@ use mtp::codec::CommunicationValue; use crate::{app_state::AppState, rho::rho_connection::RhoConnection}; +#[allow(dead_code)] pub struct RoutingService { state: Arc, } +#[allow(dead_code)] impl RoutingService { pub fn new(state: Arc) -> Self { Self { state } diff --git a/src/services/user_service.rs b/src/services/user_service.rs index efc72a5..12e6e59 100644 --- a/src/services/user_service.rs +++ b/src/services/user_service.rs @@ -2,10 +2,12 @@ use std::sync::Arc; use crate::{app_state::AppState, rho::rho_connection::RhoConnection}; +#[allow(dead_code)] pub struct UserService { state: Arc, } +#[allow(dead_code)] impl UserService { pub fn new(state: Arc) -> Self { Self { state } diff --git a/src/util/logger.rs b/src/util/logger.rs index b83714f..35a3f89 100644 --- a/src/util/logger.rs +++ b/src/util/logger.rs @@ -9,8 +9,6 @@ use std::{ use ansi_term::Color; use mtp::codec::{CommunicationType, CommunicationValue, DataTypeId, DataValue, Version}; -use crate::rho::connection::MtpValueCompat; - static LOGGER: OnceLock> = OnceLock::new(); #[allow(dead_code)] @@ -153,7 +151,9 @@ pub fn log_cv_internal( let formatted = format_cv(cv); log_internal( - cv.get_sender() as i64, + cv.sender() + .and_then(|sender| i64::try_from(sender).ok()) + .unwrap_or(0), print_type.unwrap_or(PrintType::General), prefix, false, @@ -164,14 +164,14 @@ pub fn log_cv_internal( pub fn format_cv(cv: &CommunicationValue) -> String { let mut parts = Vec::new(); - let sender = cv.get_sender(); - let receiver = cv.get_receiver(); + let sender = cv.sender(); + let receiver = cv.receiver(); - if sender > 0 && receiver > 0 { + if let (Some(sender), Some(receiver)) = (sender, receiver) { parts.push(format!("{} > {}", sender, receiver)); - } else if sender > 0 { + } else if let Some(sender) = sender { parts.push(format!("{}", sender)); - } else if receiver > 0 { + } else if let Some(receiver) = receiver { parts.push(format!("> {}", receiver)); } @@ -179,7 +179,7 @@ pub fn format_cv(cv: &CommunicationValue) -> String { .get_comm_type_enum() .map(|kind| kind.to_string()) .unwrap_or_else(|| cv.get_type().to_string()); - parts.push(format!("{} (id={})", comm_type, cv.get_id())); + parts.push(format!("{} (id={:?})", comm_type, cv.id())); if cv.is_type(CommunicationType::Relay) { parts.push("".to_string());