[Add] Proper User managment

Cargo.lock

@@ -2013,7 +2013,7 @@ name = "omikron"
 version = "0.1.0"
 dependencies = [
  "ansi_term",
- "base64 0.23.1",
+ "base64 0.22.1",
  "dashmap",
  "dotenv",
  "http",

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.23.0"
+base64 = "0.22.1"
 dashmap = "6.2.1"
 http = "1"
 once_cell = "1.21.4"

mtp-type-maps

@@ -1 +1 @@
-Subproject commit 4b82f4f8139ed9aa74fa86f73ba8f0d565703c86
+Subproject commit 909b3977cb6a233e74a925eb467412fb384844bc

src/main.rs

@@ -32,10 +32,9 @@ use crate::{
 };
 
 #[tokio::main]
-async fn main() {
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
     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<Keyring, String> {

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<CircuitBreaker>,
     keyring: mtp::crypto::Keyring,
     rho: Arc<RhoManager>,
+    authenticated_key_cache: DashMap<u64, mtp::crypto::PublicKeyBundle>,
     waiting_tasks: DashMap<u32, WaitingTask>,
     session_presence: DashMap<(i64, i64), SessionPresenceState>,
     peer_capabilities: Arc<RwLock<PeerCapabilities>>,
@@ -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<mtp::crypto::PublicKeyBundle> {
+        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)

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;
 

src/rho/relay_router.rs

@@ -156,7 +156,11 @@ pub async fn route_relay(
     source: RelaySource,
     frame: CommunicationValue,
 ) -> Result<CommunicationValue, RelayRouteError> {
-    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);

src/rho/server.rs

@@ -74,6 +74,10 @@ pub async fn get_by_connector_id(
 ) -> Option<PublicKeyBundle> {
     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<AppState>) -> Result<(), Box<dyn std::error::Error>> {
-    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,