From 2ff706e2934147aa99039be288d4b34248f8ce22 Mon Sep 17 00:00:00 2001 From: Rasensprenger Date: Wed, 2 Sep 2026 00:00:42 +0300 Subject: [PATCH 1/2] Update Rust crate base64 to 0.23.0 --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b234e94..7819992 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2013,7 +2013,7 @@ name = "omikron" version = "0.1.0" dependencies = [ "ansi_term", - "base64 0.22.1", + "base64 0.23.1", "dashmap", "dotenv", "http", diff --git a/Cargo.toml b/Cargo.toml index d59e492..1328e6b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ mtp-transport = { git = "https://git.methanium.net/Methanium/mtp.git" } ansi_term = "0.12.1" uuid = { version = "1.24.0", features = ["v4"] } -base64 = "0.22.1" +base64 = "0.23.0" dashmap = "6.2.1" http = "1" once_cell = "1.21.4" From f6ffb58fbb765e4a69429bf678891e857211823b Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:42:24 +0200 Subject: [PATCH 2/2] [Add] Proper User managment --- mtp-type-maps | 2 +- src/main.rs | 40 +++++++++++++----------- src/omega/omega_connection.rs | 59 ++++++++++++++++++++++++++++++++++- src/rho/iota_connection.rs | 1 + src/rho/relay_router.rs | 52 ++++++++++++++++++++++++++++-- src/rho/server.rs | 16 +++++++--- 6 files changed, 144 insertions(+), 26 deletions(-) diff --git a/mtp-type-maps b/mtp-type-maps index 4b82f4f..909b397 160000 --- a/mtp-type-maps +++ b/mtp-type-maps @@ -1 +1 @@ -Subproject commit 4b82f4f8139ed9aa74fa86f73ba8f0d565703c86 +Subproject commit 909b3977cb6a233e74a925eb467412fb384844bc diff --git a/src/main.rs b/src/main.rs index f7769fe..b04de2b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -32,10 +32,9 @@ use crate::{ }; #[tokio::main] -async fn main() { +async fn main() -> Result<(), Box> { if let Err(_) = default_provider().install_default() { - println!("Error loading Provider"); - return; + return Err("Error loading Provider".into()); } dotenv().ok(); startup(); @@ -43,23 +42,20 @@ async fn main() { let config = match Config::from_environment() { Ok(config) => config, Err(error) => { - eprintln!("Invalid configuration: {error}"); - return; + return Err(format!("Invalid configuration: {error}").into()); } }; let keyring = match load_or_create_keyring(KEYRING_PATH, PUBLIC_KEY_PATH) { Ok(keyring) => keyring, Err(error) => { - eprintln!("Unable to load Omikron keyring: {error}"); - return; + return Err(format!("Unable to load Omikron keyring: {error}").into()); } }; let public_key = match omega_database_public_key(&keyring) { Ok(public_key) => public_key, Err(error) => { - eprintln!("Unable to serialize public key for Omega: {error}"); - return; + return Err(format!("Unable to serialize public key for Omega: {error}").into()); } }; log!( @@ -71,8 +67,7 @@ async fn main() { let omega_keyring = match keyring_for_omega(&keyring) { Ok(keyring) => keyring, Err(error) => { - eprintln!("Unable to copy keyring for Omega: {error}"); - return; + return Err(format!("Unable to copy keyring for Omega: {error}").into()); } }; let rho = std::sync::Arc::new(RhoManager::new()); @@ -83,18 +78,27 @@ async fn main() { )); omega.clone().start().await; start_task_cleanup_loop(omega.clone()); + tokio::select! { + _ = omega.await_ready() => {} + result = tokio::signal::ctrl_c() => { + result.map_err(|error| format!("Unable to wait for shutdown signal: {error}"))?; + omega.stop().await; + return Ok(()); + } + } let livekit = std::sync::Arc::new(LiveKitService::new(config.livekit.clone())); let call_manager = std::sync::Arc::new(CallManager::new(livekit.clone())); let state = AppState::new(config.clone(), keyring, omega, rho, call_manager, livekit); - tokio::spawn(async move { - if let Err(e) = start(state).await { - log_err!(0, util::logger::PrintType::General, "{}", e); + tokio::select! { + result = start(state) => { + result.map_err(|error| format!("Rho listener terminated: {error}"))?; + return Err("Rho listener stopped unexpectedly".into()); + } + result = tokio::signal::ctrl_c() => { + result.map_err(|error| format!("Unable to wait for shutdown signal: {error}"))?; } - }); - - if let Err(error) = tokio::signal::ctrl_c().await { - eprintln!("Unable to wait for shutdown signal: {error}"); } + Ok(()) } fn keyring_for_omega(keyring: &Keyring) -> Result { diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index 5f28fc9..90e6bc9 100644 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -12,7 +12,7 @@ use mtp::client::{Client, MTPConnection, Sender}; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use mtp::{ client::ClientConfig, - files::load_public_key_bundle, + files::{load_public_key_bundle, save_public_key_bundle}, host::{Policy, SendMode}, }; use mtp_transport::ConnectionHandle; @@ -38,6 +38,12 @@ const MAX_CONCURRENT_REQUESTS: usize = 128; const CIRCUIT_BREAKER_FAILURE_THRESHOLD: u32 = 3; const CIRCUIT_BREAKER_COOLDOWN: Duration = Duration::from_secs(30); +fn authenticated_key_path(client_id: u64) -> std::path::PathBuf { + crate::WORKING_DIR + .join("iota-public-keys") + .join(format!("{client_id}.mpkb")) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ConnectionOutcome { HealthySessionEnded, @@ -218,6 +224,7 @@ pub struct OmegaConnection { circuit_breaker: Mutex, keyring: mtp::crypto::Keyring, rho: Arc, + authenticated_key_cache: DashMap, waiting_tasks: DashMap, session_presence: DashMap<(i64, i64), SessionPresenceState>, peer_capabilities: Arc>, @@ -260,6 +267,7 @@ impl OmegaConnection { }), keyring, rho, + authenticated_key_cache: DashMap::new(), waiting_tasks: DashMap::new(), session_presence: DashMap::new(), peer_capabilities: Arc::new(RwLock::new(PeerCapabilities::default())), @@ -959,6 +967,55 @@ impl OmegaConnection { self.state.read().await.is_ready() } + pub async fn await_ready(&self) { + while !self.is_ready().await { + sleep(Duration::from_millis(100)).await; + } + } + + pub fn cached_authenticated_key(&self, client_id: u64) -> Option { + if let Some(key) = self + .authenticated_key_cache + .get(&client_id) + .map(|key| key.clone()) + { + return Some(key); + } + + let key = load_public_key_bundle(&authenticated_key_path(client_id)).ok()?; + self.authenticated_key_cache.insert(client_id, key.clone()); + Some(key) + } + + pub fn cache_authenticated_key( + &self, + client_id: u64, + public_key: mtp::crypto::PublicKeyBundle, + ) { + let path = authenticated_key_path(client_id); + if let Some(parent) = path.parent() { + if let Err(error) = std::fs::create_dir_all(parent) { + log_err!( + client_id as i64, + PrintType::General, + "Create authenticated key cache directory failed: {}", + error + ); + return; + } + } + if let Err(error) = save_public_key_bundle(&public_key, &path) { + log_err!( + client_id as i64, + PrintType::General, + "Persist authenticated key cache failed: {}", + error + ); + return; + } + self.authenticated_key_cache.insert(client_id, public_key); + } + #[allow(dead_code)] pub async fn close_iota(&self, iota_id: i64) -> Result<(), String> { let cv = CommunicationValue::new(CommunicationType::IotaDisconnected) diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index 0130218..9311cde 100755 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -513,6 +513,7 @@ impl IotaConnection { || cv.is_type(CommunicationType::DeleteUserCredentialBegin) || cv.is_type(CommunicationType::DeleteUserCredentialComplete) || cv.is_type(CommunicationType::EraseHostedUserDataAck) + || cv.is_type(CommunicationType::AcknowledgeIotaUserProvision) { let sender = self.get_iota_id().await; diff --git a/src/rho/relay_router.rs b/src/rho/relay_router.rs index 5117db9..0f4a092 100644 --- a/src/rho/relay_router.rs +++ b/src/rho/relay_router.rs @@ -156,7 +156,11 @@ pub async fn route_relay( source: RelaySource, frame: CommunicationValue, ) -> Result { - let (next_hop, frame) = prepare_frame(ensure_relay_frame_id(frame))?; + let frame = ensure_relay_frame_id(frame); + let (next_hop, frame) = match source { + RelaySource::Client { iota_id } => prepare_client_frame(frame, iota_id)?, + RelaySource::Iota { .. } => prepare_frame(frame)?, + }; validate_source_next_hop(source, next_hop)?; match source { @@ -288,6 +292,31 @@ fn prepare_frame( Ok((target, frame)) } +/* + * Older clients sent their authenticated Iota's raw ID as the first relay + * hop. Accept only that exact value, then restore the namespace tag before + * forwarding so legacy clients cannot select a different route. + */ +fn prepare_client_frame( + frame: CommunicationValue, + source_iota_id: u64, +) -> Result<(RouteTarget, CommunicationValue), RelayRouteError> { + let receiver = frame.receiver().ok_or(RelayRouteError::MissingReceiver)?; + if let Some(target) = RouteTarget::from_wire_id(receiver) { + let frame = forward_relay_frame(&frame, receiver)?; + return Ok((target, frame)); + } + if receiver != source_iota_id { + return Err(RelayRouteError::InvalidRouteTarget(receiver)); + } + let target = RouteTarget::Iota(source_iota_id); + let wire_id = target + .wire_id() + .ok_or(RelayRouteError::InvalidRouteTarget(source_iota_id))?; + let frame = forward_relay_frame(&frame, wire_id)?; + Ok((target, frame)) +} + fn validate_source_next_hop( source: RelaySource, next_hop: RouteTarget, @@ -316,7 +345,7 @@ mod tests { use super::{ MessageSecurityClass, RelayRouteError, RelaySource, RouteTarget, message_security_class, - prepare_frame, validate_source_next_hop, + prepare_client_frame, prepare_frame, validate_source_next_hop, }; fn wire(target: RouteTarget) -> u64 { @@ -354,6 +383,25 @@ mod tests { assert_ne!(RouteTarget::from_wire_id(42), Some(RouteTarget::Iota(42))); } + #[test] + fn client_raw_iota_receiver_is_rewritten_to_a_typed_route() { + let (target, forwarded) = match prepare_client_frame(relay(Some(42)), 42) { + Ok(value) => value, + Err(error) => panic!("legacy client relay was rejected: {error}"), + }; + + assert_eq!(target, RouteTarget::Iota(42)); + assert_eq!(forwarded.receiver(), Some(wire(RouteTarget::Iota(42)))); + } + + #[test] + fn client_raw_receiver_cannot_select_another_iota() { + assert!(matches!( + prepare_client_frame(relay(Some(43)), 42), + Err(RelayRouteError::InvalidRouteTarget(43)) + )); + } + #[test] fn relay_rejects_an_outer_sender() { let frame = relay(Some(wire(RouteTarget::Iota(7)))).with_sender(9); diff --git a/src/rho/server.rs b/src/rho/server.rs index 7290190..2f9ee5a 100644 --- a/src/rho/server.rs +++ b/src/rho/server.rs @@ -74,6 +74,10 @@ pub async fn get_by_connector_id( ) -> Option { let request = public_key_lookup_request(omega.omikron_id(), client_id, description.as_deref())?; + if !omega.is_ready().await { + return omega.cached_authenticated_key(client_id); + } + let response = match omega .await_response(&request, Some(Duration::from_secs(20))) .await @@ -87,14 +91,16 @@ pub async fn get_by_connector_id( description, e ); - return None; + return omega.cached_authenticated_key(client_id); } }; let bytes = BASE64_STD .decode(response.get_data(DataType::PublicKey).as_str()?) .ok()?; - PublicKeyBundle::from_bytes(&bytes).ok() + let public_key = PublicKeyBundle::from_bytes(&bytes).ok()?; + omega.cache_authenticated_key(client_id, public_key.clone()); + Some(public_key) } fn public_key_lookup_request( @@ -170,9 +176,11 @@ pub async fn complete_register( } pub async fn start(state: Arc) -> Result<(), Box> { - let cert_pem = load_file_vec("certs", "cert.pem").expect("Error loading Pemfile"); + let cert_pem = load_file_vec("certs", "cert.pem") + .map_err(|error| format!("load certs/cert.pem: {error}"))?; - let key_pem = load_file_vec("certs", "key.pem").expect("Error loading Keyfile"); + let key_pem = load_file_vec("certs", "key.pem") + .map_err(|error| format!("load certs/key.pem: {error}"))?; let host_config = HostConfig::new( state.config.bind_address,