tensamin/iota
Watch
6
0
Fork
You've already forked iota
 0
Code Issues Pull requests 3 Projects Releases 1 Packages Wiki Activity Actions

Reconnection

Browse source
This commit is contained in:
Alex Emmet 2026-05-21 22:17:26 +02:00
commit e07ea9ea7e
3 changed files with 121 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

31
iota-cli/src/elements/console_card.rs
View file

@ -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;

12
iota-core/src/main.rs
View file

@ -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 {

86
omikron-connector/src/omikron_connection.rs
View file

@ -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,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<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;
}
}
// ============================================================================