[Fix] Iota connection

Cargo.lock

@@ -572,8 +572,9 @@ checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
 [[package]]
 name = "epsilon-core"
 version = "0.1.0"
-source = "git+https://github.com/Tensamin/Epsilon.git#05a7b22b2e05a9b023b0d976d65dddf06edccd7c"
+source = "git+https://github.com/Tensamin/Epsilon.git#809a32581843e663d35d4b3222f2cc59371d2786"
 dependencies = [
+ "base64 0.22.1",
  "byteorder",
  "rand 0.8.5",
  "strum",
@@ -583,7 +584,7 @@ dependencies = [
 [[package]]
 name = "epsilon-native"
 version = "0.1.0"
-source = "git+https://github.com/Tensamin/Epsilon.git#05a7b22b2e05a9b023b0d976d65dddf06edccd7c"
+source = "git+https://github.com/Tensamin/Epsilon.git#809a32581843e663d35d4b3222f2cc59371d2786"
 dependencies = [
  "epsilon-core",
  "quinn",
@@ -1799,9 +1800,9 @@ dependencies = [
 
 [[package]]
 name = "quinn-proto"
-version = "0.11.13"
+version = "0.11.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31"
+checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098"
 dependencies = [
  "aws-lc-rs",
  "bytes",
@@ -3267,18 +3268,18 @@ dependencies = [
 
 [[package]]
 name = "zerocopy"
-version = "0.8.40"
+version = "0.8.42"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a789c6e490b576db9f7e6b6d661bcc9799f7c0ac8352f56ea20193b2681532e5"
+checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3"
 dependencies = [
  "zerocopy-derive",
 ]
 
 [[package]]
 name = "zerocopy-derive"
-version = "0.8.40"
+version = "0.8.42"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f65c489a7071a749c849713807783f70672b28094011623e200cb86dcb835953"
+checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f"
 dependencies = [
  "proc-macro2",
  "quote",

src/main.rs

@@ -41,7 +41,9 @@ async fn main() {
 
     get_omega_connection();
     tokio::spawn(async move {
-        let _ = start(959).await;
+        if let Err(e) = start(959).await {
+            log_err!(0, util::logger::PrintType::General, "{}", e);
+        }
     });
     garbage_collect_calls();
 

src/omega/omega_connection.rs

@@ -457,7 +457,9 @@ impl OmegaConnection {
                 result = receiver.receive() => {
                     match result {
                         Ok(cv) => {
-                            log_cv_in!(&cv);
+                            if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
+                                log_cv_in!(PrintType::Omega, &cv);
+                            }
 
                             if cv.is_type(CommunicationType::pong) || cv.is_type(CommunicationType::ping) {
                                 self.handle_pong(&cv).await;
@@ -545,7 +547,9 @@ impl OmegaConnection {
     // -------------------------------------------------------------------------
 
     pub async fn send_message(&self, cv: &CommunicationValue) {
-        log_cv_out!(cv);
+        if !cv.is_type(CommunicationType::pong) && !cv.is_type(CommunicationType::ping) {
+            log_cv_out!(PrintType::Omega, &cv);
+        }
 
         let sender_guard = self.sender.read().await;
         if let Some(sender) = sender_guard.as_ref() {

src/rho/connection.rs

@@ -6,12 +6,16 @@ use tokio::sync::RwLock;
 
 use crate::{
     anonymous_clients::anonymous_client_connection::AnonymousClientConnection,
-    get_private_key, get_public_key,
+    get_private_key, get_public_key, log_err, log_in, log_out,
     omega::omega_connection::get_omega_connection,
-    rho::{client_connection::ClientConnection, iota_connection::IotaConnection},
+    rho::{
+        client_connection::ClientConnection, iota_connection::IotaConnection,
+        rho_connection::RhoConnection, rho_manager,
+    },
     util::{
         crypto_helper::{load_public_key, public_key_to_base64},
         crypto_util::{DataFormat, SecurePayload},
+        logger::PrintType,
     },
 };
 
@@ -52,10 +56,31 @@ impl GeneralConnection {
 }
 impl GeneralConnection {
     pub async fn handle(self: Arc<Self>) {
+        log_in!(0, PrintType::General, "General connection handler started");
+
         loop {
             let cv = match self.receiver.receive().await {
-                Ok(v) => v,
+                Ok(v) => {
-                Err(_) => break,
+                    log_in!(
+                        0,
+                        PrintType::General,
+                        "General connection received message type={:?} id={} identified={} challenged={}",
+                        v.get_type(),
+                        v.get_id(),
+                        *self.identified.read().await,
+                        *self.challenged.read().await
+                    );
+                    v
+                }
+                Err(e) => {
+                    log_err!(
+                        0,
+                        PrintType::General,
+                        "General connection receive error before upgrade completion: {:?}",
+                        e
+                    );
+                    break;
+                }
             };
 
             if !*self.identified.read().await {
@@ -65,33 +90,91 @@ impl GeneralConnection {
 
             if !*self.challenged.read().await {
                 self.handle_challenge_response(cv).await;
+                if *self.challenged.read().await {
+                    log_out!(
+                        0,
+                        PrintType::General,
+                        "General connection challenge flow completed, handler will stop after immediate migration"
+                    );
+                    break;
+                }
                 continue;
             }
 
+            log_in!(
+                0,
+                PrintType::General,
+                "General connection ready to migrate for id={} kind={:?}",
+                *self.id.read().await,
+                *self.connection_kind.read().await
+            );
+
             if self.migrate().await {
+                log_out!(
+                    0,
+                    PrintType::General,
+                    "General connection migration completed, handing over to specialized connection"
+                );
                 break;
             }
         }
+
+        log_out!(0, PrintType::General, "General connection handler stopped");
     }
     async fn handle_identification(self: &Arc<Self>, cv: CommunicationValue) {
         if !cv.is_type(CommunicationType::identification) {
+            log_in!(
+                0,
+                PrintType::General,
+                "Ignoring pre-identification message type={:?} id={}",
+                cv.get_type(),
+                cv.get_id()
+            );
             return;
         }
 
         if let DataValue::Number(iota_id) = cv.get_data(DataTypes::iota_id) {
+            log_in!(
+                *iota_id,
+                PrintType::Iota,
+                "Received Iota identification request message_id={}",
+                cv.get_id()
+            );
+
             *self.id.write().await = *iota_id as u64;
             *self.connection_kind.write().await = Some(ConnectionKind::Iota);
 
             let get_pub_key_msg = CommunicationValue::new(CommunicationType::get_iota_data)
                 .add_data(DataTypes::iota_id, DataValue::Number(*iota_id));
 
+            log_out!(
+                *iota_id,
+                PrintType::Iota,
+                "Requesting Iota public key from Omega"
+            );
+
             let response_cv = get_omega_connection()
                 .await_response(&get_pub_key_msg, Some(Duration::from_secs(20)))
                 .await;
 
             let response_cv = match response_cv {
-                Ok(r) => r,
+                Ok(r) => {
-                Err(_) => return,
+                    log_in!(
+                        *iota_id,
+                        PrintType::Iota,
+                        "Received Iota public key response from Omega"
+                    );
+                    r
+                }
+                Err(e) => {
+                    log_err!(
+                        *iota_id,
+                        PrintType::Iota,
+                        "Failed to load Iota public key from Omega: {:?}",
+                        e
+                    );
+                    return;
+                }
             };
 
             let base64_pub = response_cv
@@ -101,7 +184,14 @@ impl GeneralConnection {
 
             let pub_key = match load_public_key(base64_pub) {
                 Some(pk) => pk,
-                None => return,
+                None => {
+                    log_err!(
+                        *iota_id,
+                        PrintType::Iota,
+                        "Failed to decode Iota public key from Omega response"
+                    );
+                    return;
+                }
             };
 
             *self.pub_key.write().await = Some(pub_key.as_bytes().to_vec());
@@ -112,11 +202,18 @@ impl GeneralConnection {
                 .map(char::from)
                 .collect();
 
+            log_out!(
+                *iota_id,
+                PrintType::Iota,
+                "Generated challenge for Iota identification challenge_len={}",
+                challenge.len()
+            );
+
             *self.challenge.write().await = challenge.clone();
             *self.identified.write().await = true;
 
             let encrypted_challenge =
-                SecurePayload::new(&challenge, DataFormat::Base64, get_private_key())
+                SecurePayload::new(challenge.as_bytes(), DataFormat::Raw, get_private_key())
                     .unwrap()
                     .encrypt_x448(pub_key)
                     .unwrap()
@@ -129,44 +226,205 @@ impl GeneralConnection {
                 )
                 .add_data(DataTypes::challenge, DataValue::Str(encrypted_challenge));
 
-            let _ = self.sender.send(&response).await;
+            log_out!(
+                *iota_id,
+                PrintType::Iota,
+                "Sending encrypted identification challenge to Iota"
+            );
+
+            if let Err(e) = self.sender.send(&response).await {
+                log_err!(
+                    *iota_id,
+                    PrintType::Iota,
+                    "Failed to send challenge to Iota: {:?}",
+                    e
+                );
+            }
+        } else {
+            log_err!(
+                0,
+                PrintType::General,
+                "Identification message missing iota_id payload"
+            );
         }
     }
     async fn handle_challenge_response(self: &Arc<Self>, cv: CommunicationValue) {
+        let id = *self.id.read().await as i64;
+
         if !cv.is_type(CommunicationType::challenge_response) {
+            log_in!(
+                id,
+                PrintType::Iota,
+                "Ignoring pre-challenge-completion message type={:?} id={}",
+                cv.get_type(),
+                cv.get_id()
+            );
             return;
         }
 
         if let DataValue::Str(response) = cv.get_data(DataTypes::challenge) {
-            if *response == *self.challenge.read().await {
+            let expected = self.challenge.read().await.clone();
+
+            log_in!(
+                id,
+                PrintType::Iota,
+                "Received challenge response message_id={} response_len={} expected_len={}",
+                cv.get_id(),
+                response.len(),
+                expected.len()
+            );
+
+            if *response == expected {
+                log_in!(
+                    id,
+                    PrintType::Iota,
+                    "Challenge response validated successfully"
+                );
+
                 *self.challenged.write().await = true;
+
+                let response = CommunicationValue::new(CommunicationType::identification_response)
+                    .with_id(cv.get_id())
+                    .add_data(DataTypes::accepted, DataValue::Bool(true));
+
+                log_out!(
+                    id,
+                    PrintType::Iota,
+                    "Sending identification_response accepted=true"
+                );
+
+                if let Err(e) = self.sender.send(&response).await {
+                    log_err!(
+                        id,
+                        PrintType::Iota,
+                        "Failed to send identification_response: {:?}",
+                        e
+                    );
+                    return;
                 }
+
+                log_in!(
+                    id,
+                    PrintType::Iota,
+                    "Immediately migrating upgraded connection after successful challenge validation"
+                );
+
+                if self.migrate().await {
+                    log_out!(
+                        id,
+                        PrintType::Iota,
+                        "Immediate migration after challenge validation completed successfully"
+                    );
+                } else {
+                    log_err!(
+                        id,
+                        PrintType::Iota,
+                        "Immediate migration after challenge validation failed"
+                    );
+                }
+            } else {
+                log_err!(
+                    id,
+                    PrintType::Iota,
+                    "Challenge response mismatch expected={} actual={}",
+                    expected,
+                    response
+                );
+            }
+        } else {
+            log_err!(
+                id,
+                PrintType::Iota,
+                "Challenge response missing challenge payload"
+            );
         }
     }
 
     async fn migrate(self: &Arc<Self>) -> bool {
         let kind = match *self.connection_kind.read().await {
             Some(kind) => kind,
-            None => return false,
+            None => {
+                log_err!(
+                    0,
+                    PrintType::General,
+                    "Migration requested without a resolved connection kind"
+                );
+                return false;
+            }
         };
         let id = *self.id.read().await;
 
+        log_in!(
+            id as i64,
+            PrintType::General,
+            "Starting migration for kind={:?} id={}",
+            kind,
+            id
+        );
+
         match kind {
             ConnectionKind::Client => {
                 let client = ClientConnection::from_general(self.clone(), id).await;
                 client.start();
+                log_out!(
+                    id as i64,
+                    PrintType::Client,
+                    "Migrated general connection into ClientConnection"
+                );
             }
             ConnectionKind::Iota => {
                 let iota = IotaConnection::from_general(self.clone(), id).await;
+                log_in!(
+                    id as i64,
+                    PrintType::Iota,
+                    "Created upgraded IotaConnection from GeneralConnection"
+                );
+
+                let rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
+                log_in!(
+                    id as i64,
+                    PrintType::Iota,
+                    "Created RhoConnection for upgraded Iota connection"
+                );
+
+                iota.set_rho_connection(Arc::downgrade(&rho)).await;
+                log_in!(
+                    id as i64,
+                    PrintType::Iota,
+                    "Attached weak RhoConnection reference to IotaConnection"
+                );
+
+                rho_manager::add_rho(rho).await;
+                log_out!(
+                    id as i64,
+                    PrintType::Iota,
+                    "Registered upgraded Iota connection in rho_manager"
+                );
+
                 iota.start();
+                log_out!(
+                    id as i64,
+                    PrintType::Iota,
+                    "Started upgraded IotaConnection read loop"
+                );
             }
             ConnectionKind::AnonymousClient => {
                 let client = AnonymousClientConnection::from_general(self.clone(), id).await;
                 client.start();
+                log_out!(
+                    id as i64,
+                    PrintType::Client,
+                    "Migrated general connection into AnonymousClientConnection"
+                );
             }
             ConnectionKind::Phi => {
                 let iota = ClientConnection::from_general(self.clone(), id).await;
                 iota.start();
+                log_out!(
+                    id as i64,
+                    PrintType::General,
+                    "Migrated general connection into Phi/Client handler"
+                );
             }
         }
         true

src/rho/iota_connection.rs

@@ -3,6 +3,8 @@ use crate::calls::call_manager;
 use crate::log_cv_in;
 use crate::log_cv_out;
 use crate::log_err;
+use crate::log_in;
+use crate::log_out;
 use crate::omega::omega_connection::get_omega_connection;
 use crate::rho::connection::GeneralConnection;
 use crate::util::logger::PrintType;
@@ -54,9 +56,41 @@ impl IotaConnection {
     pub fn start(self: Arc<Self>) {
         let self_clone = self.clone();
         tokio::spawn(async move {
-            while let Ok(cv) = self_clone.receiver.receive().await {
+            log_in!(
+                self_clone.iota_id as i64,
+                PrintType::Iota,
+                "Upgraded IotaConnection read loop started"
+            );
+
+            loop {
+                match self_clone.receiver.receive().await {
+                    Ok(cv) => {
+                        log_in!(
+                            self_clone.iota_id as i64,
+                            PrintType::Iota,
+                            "Upgraded IotaConnection received message type={:?} id={}",
+                            cv.get_type(),
+                            cv.get_id()
+                        );
                         self_clone.clone().handle_message(cv).await;
                     }
+                    Err(e) => {
+                        log_err!(
+                            self_clone.iota_id as i64,
+                            PrintType::Iota,
+                            "Upgraded IotaConnection read loop stopped due to receive error: {:?}",
+                            e
+                        );
+                        break;
+                    }
+                }
+            }
+
+            log_out!(
+                self_clone.iota_id as i64,
+                PrintType::Iota,
+                "Upgraded IotaConnection read loop exited"
+            );
         });
     }
 
@@ -116,8 +150,23 @@ impl IotaConnection {
 
     /// Handle incoming message from Iota
     pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
+        log_in!(
+            self.iota_id as i64,
+            PrintType::Iota,
+            "Handling upgraded Iota message type={:?} id={} sender={} receiver={}",
+            cv.get_type(),
+            cv.get_id(),
+            cv.get_sender(),
+            cv.get_receiver()
+        );
+
         // Handle ping
         if cv.is_type(CommunicationType::ping) || cv.is_type(CommunicationType::pong) {
+            log_in!(
+                self.iota_id as i64,
+                PrintType::Iota,
+                "Routing upgraded Iota message to ping handler"
+            );
             self.handle_ping(cv).await;
             return;
         }
@@ -126,6 +175,11 @@ impl IotaConnection {
 
         // Handle GET_CHATS
         if cv.is_type(CommunicationType::get_chats) {
+            log_in!(
+                self.iota_id as i64,
+                PrintType::Iota,
+                "Routing upgraded Iota message to get_chats handler"
+            );
             self.handle_get_chats(cv).await;
             return;
         }
@@ -136,6 +190,12 @@ impl IotaConnection {
             || cv.is_type(CommunicationType::message_other_iota)
             || cv.is_type(CommunicationType::send_chat)
         {
+            log_in!(
+                self.iota_id as i64,
+                PrintType::Iota,
+                "Routing upgraded Iota message to forward_message handler receiver_id={}",
+                receiver_id
+            );
             self.handle_forward_message(cv).await;
             return;
         }
@@ -149,11 +209,22 @@ impl IotaConnection {
             || cv.is_type(CommunicationType::delete_iota)
         {
             let sender = self.get_iota_id().await;
+            log_in!(
+                self.iota_id as i64,
+                PrintType::Iota,
+                "Routing upgraded Iota message to Omega forwarder with sender={}",
+                sender
+            );
             self.handle_omega_forward(cv.with_sender(sender as u64))
                 .await;
             return;
         }
         // Forward to client
+        log_in!(
+            self.iota_id as i64,
+            PrintType::Iota,
+            "Routing upgraded Iota message to client forwarder"
+        );
         self.forward_to_client(cv).await;
     }
 
@@ -169,16 +240,45 @@ impl IotaConnection {
     async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
         let iota_for_closure = self.clone();
         tokio::spawn(async move {
+            log_out!(
+                self.iota_id as i64,
+                PrintType::Iota,
+                "Forwarding upgraded Iota message to Omega type={:?} id={}",
+                cv.get_type(),
+                cv.get_id()
+            );
+
             let response_cv = get_omega_connection()
                 .await_response(&cv.with_sender(self.iota_id), Some(Duration::from_secs(20)))
                 .await;
             if let Ok(response_cv) = response_cv {
+                log_in!(
+                    self.iota_id as i64,
+                    PrintType::Iota,
+                    "Received Omega response for upgraded Iota type={:?} id={}",
+                    response_cv.get_type(),
+                    response_cv.get_id()
+                );
                 iota_for_closure.send_message(&response_cv).await;
+            } else if let Err(e) = response_cv {
+                log_err!(
+                    self.iota_id as i64,
+                    PrintType::Iota,
+                    "Omega forward failed for upgraded Iota connection: {}",
+                    e
+                );
             }
         });
     }
     /// Handle ping message
     async fn handle_ping(&self, cv: CommunicationValue) {
+        log_in!(
+            self.iota_id as i64,
+            PrintType::Iota,
+            "Handling ping/pong for upgraded Iota connection message_id={}",
+            cv.get_id()
+        );
+
         if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
             if let Ok(ping_val) = last_ping.to_string().parse::<i64>() {
                 let mut ping_guard = self.ping.write().await;
@@ -199,6 +299,14 @@ impl IotaConnection {
         let response = CommunicationValue::new(CommunicationType::pong)
             .with_id(cv.get_id())
             .add_data(DataTypes::ping_clients, DataValue::Container(pings));
+
+        log_out!(
+            self.iota_id as i64,
+            PrintType::Iota,
+            "Sending pong from upgraded Iota connection message_id={}",
+            cv.get_id()
+        );
+
         self.send_message(&response).await;
     }
 
@@ -207,16 +315,43 @@ impl IotaConnection {
         let receiver_id = cv.get_receiver();
         let sender_id = cv.get_sender();
 
+        log_in!(
+            self.iota_id as i64,
+            PrintType::Iota,
+            "Handling cross-routing message sender={} receiver={} type={:?}",
+            sender_id,
+            receiver_id,
+            cv.get_type()
+        );
+
         if self.get_user_ids().await.contains(&(sender_id as u64)) {
             if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id as i64).await {
+                log_out!(
+                    self.iota_id as i64,
+                    PrintType::Iota,
+                    "Forwarding upgraded Iota message to target rho receiver={}",
+                    receiver_id
+                );
                 target_rho.message_to_iota(cv).await;
             } else {
+                log_err!(
+                    self.iota_id as i64,
+                    PrintType::Iota,
+                    "No target rho found for receiver={}",
+                    receiver_id
+                );
                 let error = CommunicationValue::new(CommunicationType::error_no_iota)
                     .with_id(cv.get_id())
                     .with_sender(cv.get_sender());
                 self.send_message(&error).await;
             }
         } else {
+            log_err!(
+                self.iota_id as i64,
+                PrintType::Iota,
+                "Unauthorized sender attempted cross-routing sender={}",
+                sender_id
+            );
             self.send_message(
                 &CommunicationValue::new(CommunicationType::error_invalid_user_id).add_data(
                     DataTypes::error_type,
@@ -334,19 +469,39 @@ impl IotaConnection {
     async fn forward_to_client(&self, cv: CommunicationValue) {
         if let Some(rho_conn) = self.get_rho_connection().await {
             let updated_cv = cv.with_sender(self.get_iota_id().await);
+            log_out!(
+                self.get_iota_id().await as i64,
+                PrintType::Iota,
+                "Forwarding upgraded Iota message to client type={:?} id={} receiver={}",
+                updated_cv.get_type(),
+                updated_cv.get_id(),
+                updated_cv.get_receiver()
+            );
             rho_conn.message_to_client(updated_cv).await;
         } else {
             log_err!(
                 self.get_iota_id().await as i64,
                 PrintType::General,
-                "Failed to forward message to client"
+                "Failed to forward message to client because rho connection is missing"
             );
         }
     }
 
     pub async fn handle_close(&self) {
+        log_out!(
+            self.get_iota_id().await as i64,
+            PrintType::Iota,
+            "Handling upgraded Iota connection close"
+        );
+
         if let Some(rho_conn) = self.get_rho_connection().await {
             rho_conn.close_iota_connection().await;
+        } else {
+            log_err!(
+                self.get_iota_id().await as i64,
+                PrintType::Iota,
+                "No rho connection available during upgraded Iota close"
+            );
         }
     }