From e07ea9ea7ea93b594b49e0e8fcec0f64fe7470cf Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Thu, 21 May 2026 22:17:26 +0200 Subject: [PATCH] Reconnection --- iota-cli/src/elements/console_card.rs | 31 +++++++- iota-core/src/main.rs | 12 ++- omikron-connector/src/omikron_connection.rs | 86 ++++++++++++++++++++- 3 files changed, 121 insertions(+), 8 deletions(-) diff --git a/iota-cli/src/elements/console_card.rs b/iota-cli/src/elements/console_card.rs index a394e96..d911b05 100644 --- a/iota-cli/src/elements/console_card.rs +++ b/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 | user remove | 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; diff --git a/iota-core/src/main.rs b/iota-core/src/main.rs index ea57b40..d95f268 100644 --- a/iota-core/src/main.rs +++ b/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 { diff --git a/omikron-connector/src/omikron_connection.rs b/omikron-connector/src/omikron_connection.rs index c967a57..558e628 100755 --- a/omikron-connector/src/omikron_connection.rs +++ b/omikron-connector/src/omikron_connection.rs @@ -91,6 +91,7 @@ pub struct OmikronConnection { pub connection_id: Uuid, shutdown_tx: Arc>>>, reconnect_on_close: Arc>, + auth_failure: Arc>>, pub app_challenges: Arc>>, pub app_sessions: Arc>>, } @@ -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,11 +723,21 @@ impl OmikronConnection { } if cv.is_type(CommunicationType::identification_response) { - if let Some(_accepted) = cv.get_data(DataTypes::accepted).as_bool() { - let mut state = self.state.write().await; - if let ConnectionState::Connected { identified: _ } = *state { - *state = ConnectionState::Connected { identified: true }; + 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 { + 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.clear_auth_failure().await; + *self.reconnect_on_close.write().await = true; + self.stop().await; + self.connect().await; + } } // ============================================================================