[Fix] Stability

omikron-connector/src/client.rs

@@ -27,7 +27,9 @@ pub enum OmikronStartupError {
     InitialConnectionTimeout {
         connection: std::sync::Arc<crate::omikron_connection::OmikronConnection>,
     },
-    Authentication,
+    Authentication {
+        connection: std::sync::Arc<crate::omikron_connection::OmikronConnection>,
+    },
 }
 
 #[async_trait]
@@ -39,5 +41,9 @@ pub trait OmikronClient: Send + Sync {
         timeout: Duration,
     ) -> Result<CommunicationValue, OmikronError>;
     async fn reconnect(&self) -> Result<(), OmikronError>;
+    /// Replace the local Iota identity and wait for the new identity to
+    /// register/authenticate. This is deliberately available while offline:
+    /// it is the recovery operation for an authentication failure.
+    async fn rotate_identity(&self) -> Result<(), OmikronError>;
     async fn is_connected(&self) -> bool;
 }

omikron-connector/src/omikron_connection.rs

@@ -13,7 +13,7 @@ use std::env;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, LazyLock};
-use std::time::{Duration, Instant};
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 use tokio::sync::{Mutex, RwLock, Semaphore, oneshot, watch};
 use tokio::task::JoinHandle;
 use tokio::time::sleep;
@@ -221,6 +221,11 @@ impl OmikronConnection {
         let _ = self.state_watch_tx.send(new_state);
     }
 
+    /// Subscribe to connection transitions for daemon health reporting.
+    pub fn connection_state(&self) -> watch::Receiver<ConnectionState> {
+        self.state_watch_tx.subscribe()
+    }
+
     // -------------------------------------------------------------------------
     // Connection Management
     // -------------------------------------------------------------------------
@@ -1880,10 +1885,11 @@ impl OmikronConnection {
                     let reason = response_cv
                         .get_data(DataType::Message)
                         .as_str()
+                        .or_else(|| response_cv.get_data(DataType::ErrorType).as_str())
                         .unwrap_or("connection error")
                         .to_string();
                     Err(format!(
-                        "Request failed due to disconnect (msg_id={}, reason={})",
+                        "Request rejected (msg_id={}, reason={})",
                         msg_id, reason
                     ))
                 } else {
@@ -1955,6 +1961,75 @@ impl OmikronConnection {
         self.stop().await;
         self.connect().await;
     }
+
+    /// Create a new local keyring and register it as a new Iota identity.
+    /// The existing keyring is retained as a timestamped backup so a failed
+    /// recovery does not silently destroy the user's previous identity.
+    pub async fn rotate_identity(self: &Arc<Self>) -> Result<(), OmikronError> {
+        log!("Iota identity rotation requested");
+        self.stop().await;
+
+        let path = identity_path();
+        if path.exists() {
+            let stamp = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap_or_default()
+                .as_millis();
+            let backup = path.with_extension(format!("mk.backup-{stamp}"));
+            std::fs::rename(path, &backup).map_err(|error| {
+                OmikronError::Internal(format!(
+                    "could not back up identity {}: {error}",
+                    path.display()
+                ))
+            })?;
+            log!("Existing Iota identity backed up to {}", backup.display());
+        }
+
+        let keyring = crypto_helper::generate_keyring();
+        if let Some(parent) = path.parent() {
+            std::fs::create_dir_all(parent).map_err(|error| {
+                OmikronError::Internal(format!(
+                    "could not create identity directory {}: {error}",
+                    parent.display()
+                ))
+            })?;
+        }
+        mtp::files::save_keyring_raw(&keyring, path).map_err(|error| {
+            OmikronError::Internal(format!(
+                "could not save new identity {}: {error}",
+                path.display()
+            ))
+        })?;
+        modify_config(|config| {
+            config.iota_id = None;
+            config.keyring = None;
+            config.public_key = None;
+            config.private_key = None;
+        });
+        log!("New Iota identity generated; registration started");
+
+        self.clear_auth_failure().await;
+        self.connect().await;
+        match self.await_connection(Some(CONNECTION_TIMEOUT)).await {
+            Ok(()) => {
+                let id = CONFIG.load().iota_id;
+                log!(
+                    "New Iota identity registered{}",
+                    id.map(|v| format!(" (Iota-ID: {v})")).unwrap_or_default()
+                );
+                Ok(())
+            }
+            Err(timeout) => {
+                if let Some(reason) = self.get_auth_failure().await {
+                    log!("Iota identity registration failed: {}", reason);
+                    Err(OmikronError::Authentication(reason))
+                } else {
+                    log!("Iota identity registration did not complete: {}", timeout);
+                    Err(OmikronError::Timeout(timeout))
+                }
+            }
+        }
+    }
 }
 
 // ============================================================================
@@ -1975,7 +2050,7 @@ pub async fn connect_initial(
     match conn.await_connection(Some(CONNECTION_TIMEOUT)).await {
         Ok(()) => Ok(conn),
         Err(_) if conn.has_auth_failure().await => {
-            Err(crate::client::OmikronStartupError::Authentication)
+            Err(crate::client::OmikronStartupError::Authentication { connection: conn })
         }
         Err(_) => {
             Err(crate::client::OmikronStartupError::InitialConnectionTimeout { connection: conn })
@@ -2027,6 +2102,8 @@ impl OmikronClient for OmikronConnection {
             .map_err(|error| {
                 if error.contains("timed out") {
                     OmikronError::Timeout(error)
+                } else if error.starts_with("Request rejected") {
+                    OmikronError::Internal(error)
                 } else {
                     OmikronError::Disconnected(error)
                 }
@@ -2057,6 +2134,29 @@ impl OmikronClient for OmikronConnection {
         Ok(())
     }
 
+    async fn rotate_identity(&self) -> Result<(), OmikronError> {
+        let this = Arc::new(Self {
+            state: self.state.clone(),
+            state_watch_tx: self.state_watch_tx.clone(),
+            sender: self.sender.clone(),
+            connection_loop_handle: self.connection_loop_handle.clone(),
+            last_ping: self.last_ping.clone(),
+            heartbeat_handle: self.heartbeat_handle.clone(),
+            connection_id: self.connection_id,
+            shutdown_tx: self.shutdown_tx.clone(),
+            reconnect_on_close: self.reconnect_on_close.clone(),
+            auth_failure: self.auth_failure.clone(),
+            app_challenges: self.app_challenges.clone(),
+            app_sessions: self.app_sessions.clone(),
+            missed_pongs: self.missed_pongs.clone(),
+            handler_semaphore: self.handler_semaphore.clone(),
+            cancellation: self.cancellation.clone(),
+            active_tasks: self.active_tasks.clone(),
+            app: self.app.clone(),
+        });
+        Self::rotate_identity(&this).await
+    }
+
     async fn is_connected(&self) -> bool {
         Self::is_connected(self).await
     }

omikron-connector/src/user_ops.rs

@@ -1,9 +1,9 @@
 use base64::{Engine as _, engine::general_purpose::STANDARD};
 use iota_logger::{PrintType, log, log_cv, log_t};
-use iota_storage::users::user_manager::{add_user, save_users};
+use iota_storage::users::user_manager::try_add_user;
 use iota_storage::users::user_profile::UserProfile;
 use iota_util::crypto_helper::{self, hex_hash, public_key_bundle_to_base64};
-use iota_util::file_util::save_file;
+use iota_util::file_util::try_save_file;
 use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
 use rand_core::{OsRng, RngCore};
 use std::time::Duration;
@@ -11,34 +11,75 @@ use std::time::Duration;
 use crate::OmikronClient;
 use crate::omega_discovery;
 
+#[derive(Debug)]
+pub enum CreateUserError {
+    InvalidUsername,
+    Transport(crate::OmikronError),
+    InvalidResponse,
+    RemoteRejected,
+    LocalPersistence(String),
+}
+
+fn valid_username(username: &str) -> bool {
+    !username.is_empty()
+        && username.chars().count() <= 15
+        && !username.chars().any(char::is_control)
+        && !username.contains(['/', '\\'])
+}
+
+async fn request_user_id(connection: &dyn OmikronClient) -> Result<(i64, String), CreateUserError> {
+    let request = CommunicationValue::new(CommunicationType::GetRegister);
+    let response = connection
+        .await_response(&request, Duration::from_secs(20))
+        .await
+        .map_err(CreateUserError::Transport)?;
+
+    if !response.is_type(CommunicationType::GetRegister) {
+        return Err(CreateUserError::InvalidResponse);
+    }
+
+    let user_id = response
+        .get_data(DataType::UserId)
+        .as_number()
+        .and_then(|id| i64::try_from(id).ok())
+        .filter(|id| (1..(1_i64 << 48)).contains(id))
+        .ok_or(CreateUserError::InvalidResponse)?;
+    let registration_token = response
+        .get_data(DataType::RegisterId)
+        .as_str()
+        .filter(|token| uuid::Uuid::parse_str(token).is_ok())
+        .map(str::to_owned)
+        .ok_or(CreateUserError::InvalidResponse)?;
+    Ok((user_id, registration_token))
+}
+
+/// A completion response can be lost after Omega commits the user.  Confirm
+/// the exact remote record before treating that transport failure as success.
+async fn registration_committed(connection: &dyn OmikronClient, profile: &UserProfile) -> bool {
+    let request = CommunicationValue::new(CommunicationType::GetUserData).add_typed_default(
+        DataType::UserId,
+        DataValue::SignedNumber(profile.user_id.into()),
+    );
+    let Ok(response) = connection
+        .await_response(&request, Duration::from_secs(5))
+        .await
+    else {
+        return false;
+    };
+    response.get_data(DataType::UserId).as_number() == Some(profile.user_id.into())
+        && response.get_data(DataType::Username).as_str() == Some(profile.username.as_str())
+        && response.get_data(DataType::PublicKey).as_str() == Some(profile.public_key.as_str())
+}
+
 pub async fn create_user(
     connection: &dyn OmikronClient,
     username: &str,
-) -> (Option<UserProfile>, Option<String>) {
-    let register_communication_value = CommunicationValue::new(CommunicationType::GetRegister);
-
-    let response_communication_value = match connection
-        .await_response(&register_communication_value, Duration::from_secs(20))
-        .await
-    {
-        Ok(communication_value) => communication_value,
-        Err(e) => {
-            log_t!("User creation: {}", e.to_string());
-            return (None, None);
-        }
-    };
-    log_cv!(PrintType::Omega, response_communication_value);
-
-    let user_id = match response_communication_value
-        .get_data(DataType::UserId)
-        .as_number()
-    {
-        Some(id) => id,
-        None => {
-            log_t!("User creation: Response returned none");
-            return (None, None);
-        }
-    };
+) -> Result<UserProfile, CreateUserError> {
+    if !valid_username(username) {
+        return Err(CreateUserError::InvalidUsername);
+    }
+    let (user_id, registration_token) = request_user_id(connection).await?;
+    log!("User creation: Omega allocated user ID {user_id}");
     let keyring = crypto_helper::generate_keyring();
     let pub_key_bundle = keyring.public_key_bundle();
     let keyring_b64 = crypto_helper::keyring_to_base64(&keyring);
@@ -50,7 +91,7 @@ pub async fn create_user(
     let reset_token = STANDARD.encode(&bytes);
 
     let user_profile = UserProfile::new(
-        user_id as i64,
+        user_id,
         username.to_string(),
         None,
         public_key_bundle_to_base64(&pub_key_bundle),
@@ -59,30 +100,43 @@ pub async fn create_user(
     );
 
     let communication_value = CommunicationValue::new(CommunicationType::CompleteRegisterUser)
-        .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id as i128))
+        .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
         .add_typed_default(DataType::Username, DataValue::Str(username.to_string()))
         .add_typed_default(
             DataType::PublicKey,
             DataValue::Str(public_key_bundle_to_base64(&pub_key_bundle)),
         )
-        .add_typed_default(DataType::IotaId, DataValue::SignedNumber(user_id as i128))
-        .add_typed_default(DataType::ResetToken, DataValue::Str(reset_token));
+        .add_typed_default(DataType::ResetToken, DataValue::Str(reset_token))
+        .add_typed_default(DataType::RegisterId, DataValue::Str(registration_token));
 
     let response_communication_value = connection
         .await_response(&communication_value, Duration::from_secs(20))
         .await;
 
-    if let Ok(response) = response_communication_value {
-        log_cv!(PrintType::Omega, response);
-        if !response.is_type(CommunicationType::Success) {
-            return (None, None);
+    match response_communication_value {
+        Ok(response) => {
+            log_cv!(PrintType::Omega, response);
+            if !response.is_type(CommunicationType::Success) {
+                return Err(CreateUserError::RemoteRejected);
+            }
+        }
+        Err(error) => {
+            if registration_committed(connection, &user_profile).await {
+                log!(
+                    "User creation: completion response was lost; verified user {} remotely",
+                    user_id
+                );
+            } else {
+                log_t!("User creation: {}", error.to_string());
+                return Err(match error {
+                    crate::OmikronError::Internal(_) => CreateUserError::RemoteRejected,
+                    error => CreateUserError::Transport(error),
+                });
+            }
         }
-    } else {
-        log_t!("User creation: Response returned none");
-        return (None, None);
     }
     log!("Created User");
-    save_file(
+    try_save_file(
         "",
         &format!("{}.tu", username),
         &format!(
@@ -91,9 +145,91 @@ pub async fn create_user(
             omega_discovery::omega_host(),
             keyring_b64
         ),
-    );
+    )
+    .map_err(|error| CreateUserError::LocalPersistence(error.to_string()))?;
 
-    add_user(user_profile.clone());
-    save_users();
-    (Some(user_profile), Some(keyring_b64))
+    try_add_user(user_profile.clone())
+        .map_err(|error| CreateUserError::LocalPersistence(error.to_string()))?;
+    Ok(user_profile)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{CreateUserError, request_user_id, valid_username};
+    use crate::{OmikronClient, OmikronError};
+    use async_trait::async_trait;
+    use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+    use std::time::Duration;
+
+    struct RegistrationClient {
+        response: CommunicationValue,
+    }
+
+    #[async_trait]
+    impl OmikronClient for RegistrationClient {
+        async fn send_message(&self, _: &CommunicationValue) -> Result<(), OmikronError> {
+            unreachable!()
+        }
+
+        async fn await_response(
+            &self,
+            request: &CommunicationValue,
+            _: Duration,
+        ) -> Result<CommunicationValue, OmikronError> {
+            assert!(request.is_type(CommunicationType::GetRegister));
+            Ok(self.response.clone().with_id(request.get_id()))
+        }
+
+        async fn reconnect(&self) -> Result<(), OmikronError> {
+            unreachable!()
+        }
+
+        async fn rotate_identity(&self) -> Result<(), OmikronError> {
+            unreachable!()
+        }
+
+        async fn is_connected(&self) -> bool {
+            true
+        }
+    }
+
+    #[test]
+    fn validates_usernames_before_remote_registration() {
+        assert!(valid_username("alice"));
+        assert!(valid_username("fifteen_char_ok"));
+        assert!(!valid_username(""));
+        assert!(!valid_username("sixteen_chars_bad"));
+        assert!(!valid_username("path/name"));
+        assert!(!valid_username("line\nbreak"));
+    }
+
+    #[tokio::test]
+    async fn uses_user_id_allocated_by_omega() {
+        let client = RegistrationClient {
+            response: CommunicationValue::new(CommunicationType::GetRegister)
+                .add_typed_default(DataType::UserId, DataValue::SignedNumber(4_294_967_311))
+                .add_typed_default(
+                    DataType::RegisterId,
+                    DataValue::Str("00000000-0000-4000-8000-000000000001".into()),
+                ),
+        };
+
+        assert_eq!(
+            request_user_id(&client).await.unwrap(),
+            (4_294_967_311, "00000000-0000-4000-8000-000000000001".into())
+        );
+    }
+
+    #[tokio::test]
+    async fn rejects_registration_response_without_a_positive_user_id() {
+        let client = RegistrationClient {
+            response: CommunicationValue::new(CommunicationType::GetRegister)
+                .add_typed_default(DataType::UserId, DataValue::SignedNumber(0)),
+        };
+
+        assert!(matches!(
+            request_user_id(&client).await,
+            Err(CreateUserError::InvalidResponse)
+        ));
+    }
 }