From 8a48e7d46e9df911cd712b7c523d8393186ebdea Mon Sep 17 00:00:00 2001 From: Alois Date: Tue, 28 Jul 2026 01:02:05 +0200 Subject: [PATCH] (fix): fix connection issues --- .gitignore | 2 + flake.lock | 17 ++++ flake.nix | 5 ++ src/omega/omega_connection.rs | 15 ++-- src/rho/client_connection.rs | 4 +- src/rho/connection.rs | 153 ++++------------------------------ src/rho/iota_connection.rs | 32 +++++++ 7 files changed, 80 insertions(+), 148 deletions(-) diff --git a/.gitignore b/.gitignore index 11e3519..dbf5fe5 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,5 @@ target *.mk *.mpkb + +result \ No newline at end of file diff --git a/flake.lock b/flake.lock index 4c197a5..495af8b 100644 --- a/flake.lock +++ b/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" } diff --git a/flake.nix b/flake.nix index 21cf4f9..282d1f3 100644 --- a/flake.nix +++ b/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"; }; }; diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index 45d0fb3..a7806ea 100644 --- a/src/omega/omega_connection.rs +++ b/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_loop(&mut connection.receiver, sender_handle) - .await; + read_self.read_loop(connection, sender_handle).await; }); // Start heartbeat @@ -411,7 +410,7 @@ impl OmegaConnection { async fn read_loop( self: Arc, - receiver: &mut Receiver, + connection: Arc, sender_handle: Arc, ) { // 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) { diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index 0cd3e39..438abd7 100644 --- a/src/rho/client_connection.rs +++ b/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, diff --git a/src/rho/connection.rs b/src/rho/connection.rs index b909573..3a545f4 100755 --- a/src/rho/connection.rs +++ b/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 { - 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 - .request_initial_client_state(&rho_conn, user_id, handshake) - .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; - } + self.send_initial_client_state_request(&rho_conn, user_id) + .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( - &self, - rho: &Arc, - user_id: i64, - handshake: CommunicationValue, - ) -> Option { + async fn send_initial_client_state_request(&self, rho: &Arc, user_id: i64) { 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( - DataType::VersionNumber, - handshake.get_data(DataType::VersionNumber).clone(), - ) - .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 + .add_typed_default(DataType::VersionNumber, DataValue::SignedNumber(0)) + .add_typed_default(DataType::CacheValid, DataValue::Bool(false)) + .add_typed_default(DataType::CacheSchemaVersion, DataValue::SignedNumber(0)); + rho.get_iota_connection().send_message(&request).await; } async fn migrate_iota(self: &Arc) { @@ -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); - } -} diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index 5e26f80..7f6fe4c 100755 --- a/src/rho/iota_connection.rs +++ b/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,