Reconnection

iota-cli/src/elements/console_card.rs

@@ -2,7 +2,8 @@ use crossterm::event::{KeyCode, KeyEvent};
 use iota_logger::{log, log_command, log_cv};
 use iota_state::{ACTIVE_TASKS, RELOAD, SHUTDOWN};
 use iota_storage::users::{user_manager, user_profile::UserProfile};
-use iota_util::file_util;
+use iota_storage::util::config_util::CONFIG;
+use iota_util::{crypto_helper, file_util};
 use omikron_connector::omikron_connection::OMIKRON_CONNECTION;
 use ratatui::{
     Frame,
@@ -400,7 +401,7 @@ pub async fn run_command(command: &str) {
         }
 
         ["help"] => {
-            log!("Available commands: tasks, fps, ping, user");
+            log!("Available commands: tasks, fps, ping, user, reconnect, regenerate");
         }
 
         ["help", "tasks"] => {
@@ -415,6 +416,12 @@ pub async fn run_command(command: &str) {
         ["help", "user"] => {
             log!("User command usage: user add <username> | user remove <username> | user list");
         }
+        ["help", "reconnect"] => {
+            log!("Reconnect command usage: reconnect — retry connecting to the Omikron server");
+        }
+        ["help", "regenerate"] => {
+            log!("Regenerate command usage: regenerate private-key — generate a new Iota key pair and reconnect");
+        }
 
         ["ping"] => {
             ping(20).await;
@@ -460,6 +467,26 @@ pub async fn run_command(command: &str) {
                 log!("User info: Username doesn't exist");
             }
         }
+        ["reconnect"] => {
+            log!("Reconnecting to Omikron server...");
+            OMIKRON_CONNECTION.reconnect().await;
+            log!("Reconnected to Omikron server");
+        }
+        ["regenerate", "private-key"] => {
+            log!("Regenerating Iota key pair...");
+            let key_pair = crypto_helper::generate_keypair();
+            let public_key_base64 = crypto_helper::public_key_to_base64(&key_pair.public);
+            let private_key_base64 = crypto_helper::secret_key_to_base64(&key_pair.secret);
+            {
+                let mut conf = CONFIG.write().await;
+                conf.change("public_key", json::JsonValue::from(public_key_base64));
+                conf.change("private_key", json::JsonValue::from(private_key_base64));
+                conf.update();
+            }
+            log!("Key pair regenerated. Reconnecting to Omikron server...");
+            OMIKRON_CONNECTION.reconnect().await;
+            log!("Reconnected with new key pair");
+        }
         ["reload"] | ["restart"] => {
             log!("Restarting");
             *RELOAD.write().await = true;

iota-core/src/main.rs

@@ -13,7 +13,7 @@ use iota_storage::users::user_manager;
 use iota_storage::util::config_util::CONFIG;
 use iota_util::file_util::{download_and_extract_zip, has_dir};
 use omikron_connector as omikron;
-use ttp_core::{CommunicationType, DataTypes};
+use omikron_connector::omikron_connection::OMIKRON_CONNECTION;
 
 #[tokio::main(flavor = "multi_thread", worker_threads = 16)]
 #[allow(unused_must_use, dead_code, unused_assignments)]
@@ -159,7 +159,15 @@ async fn main() {
                 break;
             }
 
-            sleep(Duration::from_millis(100)).await;
+            if OMIKRON_CONNECTION.has_auth_failure().await {
+                if let Some(reason) = OMIKRON_CONNECTION.get_auth_failure().await {
+                    log!("Authentication failed: {}", reason);
+                    log!("Use /reconnect to try again or /regenerate private-key to create a new key pair");
+                    OMIKRON_CONNECTION.clear_auth_failure().await;
+                }
+            }
+
+            sleep(Duration::from_millis(500)).await;
         }
         if *RELOAD.read().await {
             loop {

omikron-connector/src/omikron_connection.rs

@@ -91,6 +91,7 @@ pub struct OmikronConnection {
     pub connection_id: Uuid,
     shutdown_tx: Arc<Mutex<Option<watch::Sender<bool>>>>,
     reconnect_on_close: Arc<RwLock<bool>>,
+    auth_failure: Arc<RwLock<Option<String>>>,
     pub app_challenges: Arc<RwLock<HashMap<u64, String>>>,
     pub app_sessions: Arc<RwLock<HashMap<u64, (i64, String)>>>,
 }
@@ -114,6 +115,7 @@ impl OmikronConnection {
             connection_id: Uuid::new_v4(),
             shutdown_tx: Arc::new(Mutex::new(Some(shutdown_tx))),
             reconnect_on_close: Arc::new(RwLock::new(true)),
+            auth_failure: Arc::new(RwLock::new(None)),
             app_challenges: Arc::new(RwLock::new(HashMap::new())),
             app_sessions: Arc::new(RwLock::new(HashMap::new())),
         }
@@ -134,6 +136,11 @@ impl OmikronConnection {
             handle.abort();
         }
 
+        if self.shutdown_tx.lock().await.is_none() {
+            let (shutdown_tx, _) = watch::channel(false);
+            *self.shutdown_tx.lock().await = Some(shutdown_tx);
+        }
+
         *self.reconnect_on_close.write().await = true;
 
         let self_clone = self.clone();
@@ -191,6 +198,10 @@ impl OmikronConnection {
                     }
                 }
                 Err(e) => {
+                    if self.auth_failure.read().await.is_some() {
+                        log!("Authentication failed, stopping reconnection: {}", e);
+                        break;
+                    }
                     log!(
                         "Connection failed: {}, retrying in {:?}...",
                         e,
@@ -256,6 +267,22 @@ impl OmikronConnection {
         // Handle registration/identification
         self.handle_authentication().await;
 
+        // Wait for identification to complete
+        if !self.await_identification(Duration::from_secs(30)).await {
+            *self.reconnect_on_close.write().await = false;
+            let reason = "Authentication failed: server did not accept the challenge. Your Iota keys may be invalid or the private key has changed on the server."
+                .to_string();
+            *self.auth_failure.write().await = Some(reason.clone());
+
+            if let Some(sender) = self.sender.write().await.take() {
+                sender.close();
+            }
+            *self.state.write().await = ConnectionState::Disconnected;
+            return Err(reason);
+        }
+
+        log_t!("omikron_authenticated");
+
         // Start heartbeat
         let heartbeat_self = self.clone();
         let heartbeat_handle = tokio::spawn(async move {
@@ -696,12 +723,22 @@ impl OmikronConnection {
         }
 
         if cv.is_type(CommunicationType::identification_response) {
-            if let Some(_accepted) = cv.get_data(DataTypes::accepted).as_bool() {
+            match cv.get_data(DataTypes::accepted).as_bool() {
+                Some(true) => {
                     let mut state = self.state.write().await;
                     if let ConnectionState::Connected { identified: _ } = *state {
                         *state = ConnectionState::Connected { identified: true };
                     }
                 }
+                Some(false) => {
+                    *self.auth_failure.write().await = Some(
+                        "Server rejected the challenge response — your Iota keys may be invalid."
+                            .to_string(),
+                    );
+                    log_t!("omikron_auth_rejected");
+                }
+                None => {}
+            }
             return;
         }
 
@@ -1360,6 +1397,28 @@ impl OmikronConnection {
                 .add_data(DataTypes::challenge, DataValue::Str(solved));
 
             self.send_message(&response).await;
+        } else {
+            log_t!("omikron_challenge_decryption_failed");
+            *self.auth_failure.write().await = Some(
+                "Challenge decryption failed — your Iota private key may not match the registered key on the server."
+                    .to_string(),
+            );
+        }
+    }
+
+    async fn await_identification(&self, timeout: Duration) -> bool {
+        let start = Instant::now();
+        loop {
+            if self.state.read().await.is_identified() {
+                return true;
+            }
+            if self.auth_failure.read().await.is_some() {
+                return false;
+            }
+            if start.elapsed() >= timeout {
+                return false;
+            }
+            sleep(Duration::from_millis(100)).await;
         }
     }
 
@@ -1521,6 +1580,25 @@ impl OmikronConnection {
             sleep(Duration::from_millis(100)).await;
         }
     }
+
+    pub async fn has_auth_failure(&self) -> bool {
+        self.auth_failure.read().await.is_some()
+    }
+
+    pub async fn get_auth_failure(&self) -> Option<String> {
+        self.auth_failure.read().await.clone()
+    }
+
+    pub async fn clear_auth_failure(&self) {
+        *self.auth_failure.write().await = None;
+    }
+
+    pub async fn reconnect(self: &Arc<Self>) {
+        self.clear_auth_failure().await;
+        *self.reconnect_on_close.write().await = true;
+        self.stop().await;
+        self.connect().await;
+    }
 }
 
 // ============================================================================