(fix): fix connection issues

.gitignore

@@ -20,3 +20,5 @@ target
 
 *.mk
 *.mpkb
+
+result

flake.lock

@@ -18,6 +18,22 @@
         "type": "github"
       }
     },
+    "mtp-type-maps": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1785165682,
+        "narHash": "sha256-QAtadTecsOlrsXUA6uGGwsG+yLrlWgd4m6bwwFv/lho=",
+        "ref": "refs/heads/main",
+        "rev": "594646ac39d986f0787aa614a99d580035a67318",
+        "revCount": 5,
+        "type": "git",
+        "url": "https://git.methanium.net/tensamin/mtp-type-maps"
+      },
+      "original": {
+        "type": "git",
+        "url": "https://git.methanium.net/tensamin/mtp-type-maps"
+      }
+    },
     "nixpkgs": {
       "locked": {
         "lastModified": 1780243769,
@@ -52,6 +68,7 @@
     "root": {
       "inputs": {
         "flake-parts": "flake-parts",
+        "mtp-type-maps": "mtp-type-maps",
         "nixpkgs": "nixpkgs",
         "rust-overlay": "rust-overlay"
       }

flake.nix

@@ -8,6 +8,10 @@
       url = "github:oxalica/rust-overlay";
       inputs.nixpkgs.follows = "nixpkgs";
     };
+    mtp-type-maps = {
+      url = "git+https://git.methanium.net/tensamin/mtp-type-maps";
+      flake = false;
+    };
   };
 
   outputs =
@@ -65,6 +69,7 @@
               ];
               buildInputs = with pkgs; [ openssl ];
               dontUseCmakeConfigure = true;
+              MTP_TYPE_MAPS = "${inputs.mtp-type-maps}/type-maps.yaml";
             };
           };
 

src/omega/omega_connection.rs

@@ -3,7 +3,7 @@ use crate::{
     rho::rho_manager::RhoManager, util::logger::PrintType,
 };
 use dashmap::DashMap;
-use mtp::client::{Client, Receiver, Sender};
+use mtp::client::{Client, MTPConnection, Sender};
 use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
 use mtp::{
     client::ClientConfig,
@@ -297,7 +297,7 @@ impl OmegaConnection {
         let host_public_key = load_public_key_bundle("./omega.mpkb")
             .map_err(|e| format!("Failed to load omega.mpkb: {}", e))?;
 
-        let mut connection = Client::auth_connect(client_config, &self.keyring, &host_public_key)
+        let connection = Client::auth_connect(client_config, &self.keyring, &host_public_key)
             .await
             .map_err(|e| format!("Connection failed: {}", e))?;
 
@@ -309,19 +309,18 @@ impl OmegaConnection {
         );
 
         // Store sender
-        let sender_arc = Arc::new(connection.sender);
+        let sender_arc = Arc::new(connection.sender.clone());
         *self.sender.write().await = Some(sender_arc.clone());
         *self.state.write().await = ConnectionState::Connected { identified: false };
 
         // Get handle for close monitoring
         let sender_handle = sender_arc.handle().clone();
+        let connection = Arc::new(connection);
 
         // Start read loop
         let read_self = self.clone();
         let read_handle = tokio::spawn(async move {
-            read_self
+            read_self.read_loop(connection, sender_handle).await;
-                .read_loop(&mut connection.receiver, sender_handle)
-                .await;
         });
 
         // Start heartbeat
@@ -411,7 +410,7 @@ impl OmegaConnection {
 
     async fn read_loop(
         self: Arc<Self>,
-        receiver: &mut Receiver,
+        connection: Arc<MTPConnection>,
         sender_handle: Arc<ConnectionHandle>,
     ) {
         // Monitor both receiver and sender handle for close
@@ -419,7 +418,7 @@ impl OmegaConnection {
 
         loop {
             tokio::select! {
-                result = receiver.receive() => {
+                result = connection.receive() => {
                     match result {
                         Ok(cv) => {
                             if !cv.is_type(CommunicationType::Pong) && !cv.is_type(CommunicationType::Ping) {

src/rho/client_connection.rs

@@ -105,7 +105,7 @@ impl ClientConnection {
         };
         tokio::spawn(async move {
             let _permit = permit;
-            if cv.is_type(CommunicationType::Ping) {
+            if cv.is_type(CommunicationType::ClientPing) {
                 self.handle_ping(cv).await;
                 return;
             }
@@ -331,7 +331,7 @@ impl ClientConnection {
         };
 
         // Send pong response
-        let response = CommunicationValue::new(CommunicationType::Pong)
+        let response = CommunicationValue::new(CommunicationType::ClientPing)
             .with_id(cv.get_id())
             .add_typed_default(
                 DataType::PingIota,

src/rho/connection.rs

@@ -6,7 +6,7 @@ use uuid::Uuid;
 use crate::{
     anonymous_clients::anonymous_client_connection::AnonymousClientConnection,
     app_state::AppState,
-    log_cv_out, log_err, log_in, log_out,
+    log_err, log_in, log_out,
     rho::{
         app_connection::AppConnection, client_connection::ClientConnection,
         iota_connection::IotaConnection, rho_connection::RhoConnection,
@@ -20,14 +20,6 @@ use mtp::webserver::{WebMTPConnection, WebMtpReceiver, WebMtpSender};
 pub type MtpSender = WebMtpSender;
 pub type MtpReceiver = WebMtpReceiver;
 
-fn data_i64(value: &DataValue) -> Option<i64> {
-    match value {
-        DataValue::SignedNumber(number) => i64::try_from(*number).ok(),
-        DataValue::UnsignedNumber(number) => i64::try_from(*number).ok(),
-        _ => None,
-    }
-}
-
 /*
  * How a connection identified itself during the mtp handshake driven by
  * `server.rs` ("iota" / "client" authenticated logins, "anonymous"
@@ -85,7 +77,12 @@ impl GeneralConnection {
             connection_kind: kind,
             id: conn.client_id,
             rho_connection: Arc::new(RwLock::new(None)),
-            session_id: Arc::new(RwLock::new(conn.client_id)),
+            session_id: Arc::new(RwLock::new(match kind {
+                ConnectionKind::Client => {
+                    ((Uuid::new_v4().as_u128() as u64) & ((1_u64 << 53) - 1)).max(1)
+                }
+                _ => conn.client_id,
+            })),
             app_identifier: Arc::new(RwLock::new(None)),
             app_session: Arc::new(RwLock::new(None)),
             client_version: Arc::new(RwLock::new(conn.version.to_string())),
@@ -132,49 +129,6 @@ impl GeneralConnection {
         let id = self.id;
         let user_id = id as i64;
 
-        let Ok(handshake) = self.receiver.receive().await else {
-            return false;
-        };
-        if !handshake.is_type(CommunicationType::ClientConnected) {
-            let error = CommunicationValue::new(CommunicationType::ErrorInvalidData)
-                .with_id(handshake.get_id());
-            let _ = self.sender.send(&error).await;
-            return false;
-        }
-        let Some(session_id) = data_i64(handshake.get_data(DataType::SessionId)) else {
-            let _ = self
-                .sender
-                .send(
-                    &CommunicationValue::new(CommunicationType::ErrorInvalidData)
-                        .with_id(handshake.get_id()),
-                )
-                .await;
-            return false;
-        };
-        if session_id <= 0 {
-            let _ = self
-                .sender
-                .send(
-                    &CommunicationValue::new(CommunicationType::ErrorInvalidData)
-                        .with_id(handshake.get_id()),
-                )
-                .await;
-            return false;
-        }
-        let version = data_i64(handshake.get_data(DataType::VersionNumber));
-        if !matches!(version, Some(version) if version >= 0) {
-            let _ = self
-                .sender
-                .send(
-                    &CommunicationValue::new(CommunicationType::ErrorInvalidData)
-                        .with_id(handshake.get_id()),
-                )
-                .await;
-            return false;
-        }
-        *self.session_id.write().await = session_id as u64;
-        let request_id = handshake.get_id();
-
         let client = ClientConnection::from_general(self.clone(), id).await;
         let rho = self.find_user_rho(user_id).await;
         *self.rho_connection.write().await = rho.clone();
@@ -184,24 +138,8 @@ impl GeneralConnection {
             rho_conn.add_client_connection(client.clone()).await;
             self.notify_user_connected(user_id, rho_conn.get_iota_id().await as i64)
                 .await;
-            if let Some(response) = self
+            self.send_initial_client_state_request(&rho_conn, user_id)
-                .request_initial_client_state(&rho_conn, user_id, handshake)
+                .await;
-                .await
-            {
-                let response = self.add_call_state(response, user_id).await;
-                log_cv_out!(response);
-                let _ = self.sender.send(&response).await;
-            } else {
-                let error =
-                    CommunicationValue::new(CommunicationType::ErrorNoIota).with_id(request_id);
-                log_err!(
-                    user_id,
-                    PrintType::Client,
-                    "Initial state request failed for user {}",
-                    id
-                );
-                let _ = self.sender.send(&error).await;
-            }
         } else {
             log_err!(
                 user_id,
@@ -209,7 +147,7 @@ impl GeneralConnection {
                 "No RhoConnection found for user {}, client not attached to iota",
                 id
             );
-            let error = CommunicationValue::new(CommunicationType::ErrorNoIota).with_id(request_id);
+            let error = CommunicationValue::new(CommunicationType::ErrorNoIota);
             let _ = self.sender.send(&error).await;
         }
 
@@ -271,67 +209,19 @@ impl GeneralConnection {
         Some(rho)
     }
 
-    async fn request_initial_client_state(
+    async fn send_initial_client_state_request(&self, rho: &Arc<RhoConnection>, user_id: i64) {
-        &self,
-        rho: &Arc<RhoConnection>,
-        user_id: i64,
-        handshake: CommunicationValue,
-    ) -> Option<CommunicationValue> {
         let session_id = *self.session_id.read().await as i64;
         let request = CommunicationValue::new(CommunicationType::ClientConnected)
-            .with_id(handshake.get_id())
             .with_sender(user_id as u64)
             .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
             .add_typed_default(
                 DataType::SessionId,
                 DataValue::SignedNumber(session_id.into()),
             )
-            .add_typed_default(
+            .add_typed_default(DataType::VersionNumber, DataValue::SignedNumber(0))
-                DataType::VersionNumber,
+            .add_typed_default(DataType::CacheValid, DataValue::Bool(false))
-                handshake.get_data(DataType::VersionNumber).clone(),
+            .add_typed_default(DataType::CacheSchemaVersion, DataValue::SignedNumber(0));
-            )
+        rho.get_iota_connection().send_message(&request).await;
-            .add_typed_default(
-                DataType::CacheValid,
-                handshake.get_data(DataType::CacheValid).clone(),
-            )
-            .add_typed_default(
-                DataType::CacheSchemaVersion,
-                handshake.get_data(DataType::CacheSchemaVersion).clone(),
-            );
-        rho.get_iota_connection()
-            .clone()
-            .await_response(&request, Some(Duration::from_secs(20)))
-            .await
-            .ok()
-    }
-
-    async fn add_call_state(
-        &self,
-        response: CommunicationValue,
-        user_id: i64,
-    ) -> CommunicationValue {
-        let mut output = response.clone();
-
-        for (key, value) in response.iter_typed_data() {
-            if key == Some(DataType::Contacts) {
-                if let Some(contacts) = value.as_array() {
-                    let (contacts, global_calls) = self
-                        .state
-                        .call_state_aggregator
-                        .augment_contacts(user_id as u64, contacts.clone())
-                        .await;
-                    output =
-                        output.add_typed_default(DataType::Contacts, DataValue::Array(contacts));
-                    output =
-                        output.add_typed_default(DataType::Calls, DataValue::Array(global_calls));
-                    continue;
-                }
-            }
-            if let Some(data_type) = key {
-                output = output.add_typed_default(data_type, value.clone());
-            }
-        }
-        output
     }
 
     async fn migrate_iota(self: &Arc<Self>) {
@@ -391,16 +281,3 @@ impl GeneralConnection {
         app_conn.start();
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::data_i64;
-    use mtp::codec::DataValue;
-
-    #[test]
-    fn data_i64_accepts_signed_and_unsigned_values() {
-        assert_eq!(data_i64(&DataValue::SignedNumber(42)), Some(42));
-        assert_eq!(data_i64(&DataValue::UnsignedNumber(42)), Some(42));
-        assert_eq!(data_i64(&DataValue::UnsignedNumber(u128::MAX)), None);
-    }
-}

src/rho/iota_connection.rs

@@ -208,6 +208,12 @@ impl IotaConnection {
 
         log_cv_in!(PrintType::Iota, cv);
 
+        let cv = if cv.is_type(CommunicationType::ClientStateSync) {
+            self.add_call_state(cv).await
+        } else {
+            cv
+        };
+
         // Handle GET_CHATS
         if cv.is_type(CommunicationType::GetChats) {
             self.handle_get_chats(cv).await;
@@ -735,6 +741,32 @@ impl IotaConnection {
         }
     }
 
+    async fn add_call_state(&self, response: CommunicationValue) -> CommunicationValue {
+        let mut output = response.clone();
+        let user_id = response.get_receiver();
+
+        for (key, value) in response.iter_typed_data() {
+            if key == Some(DataType::Contacts) {
+                if let Some(contacts) = value.as_array() {
+                    let (contacts, global_calls) = self
+                        .state
+                        .call_state_aggregator
+                        .augment_contacts(user_id, contacts.clone())
+                        .await;
+                    output =
+                        output.add_typed_default(DataType::Contacts, DataValue::Array(contacts));
+                    output =
+                        output.add_typed_default(DataType::Calls, DataValue::Array(global_calls));
+                    continue;
+                }
+            }
+            if let Some(data_type) = key {
+                output = output.add_typed_default(data_type, value.clone());
+            }
+        }
+        output
+    }
+
     pub async fn handle_close(&self) {
         log_out!(
             self.iota_id as i64,