diff --git a/.gitignore b/.gitignore index dbf5fe5..11e3519 100644 --- a/.gitignore +++ b/.gitignore @@ -20,5 +20,3 @@ target *.mk *.mpkb - -result \ No newline at end of file diff --git a/flake.lock b/flake.lock index 495af8b..4c197a5 100644 --- a/flake.lock +++ b/flake.lock @@ -18,22 +18,6 @@ "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, @@ -68,7 +52,6 @@ "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 282d1f3..21cf4f9 100644 --- a/flake.nix +++ b/flake.nix @@ -8,10 +8,6 @@ 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 = @@ -69,7 +65,6 @@ ]; buildInputs = with pkgs; [ openssl ]; dontUseCmakeConfigure = true; - MTP_TYPE_MAPS = "${inputs.mtp-type-maps}/type-maps.yaml"; }; }; diff --git a/mtp-type-maps b/mtp-type-maps index 9776f43..594646a 160000 --- a/mtp-type-maps +++ b/mtp-type-maps @@ -1 +1 @@ -Subproject commit 9776f430510d9aaeaa163f662ef5184108b6eec8 +Subproject commit 594646ac39d986f0787aa614a99d580035a67318 diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index a7806ea..45d0fb3 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, MTPConnection, Sender}; +use mtp::client::{Client, Receiver, 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 connection = Client::auth_connect(client_config, &self.keyring, &host_public_key) + let mut connection = Client::auth_connect(client_config, &self.keyring, &host_public_key) .await .map_err(|e| format!("Connection failed: {}", e))?; @@ -309,18 +309,19 @@ impl OmegaConnection { ); // Store sender - let sender_arc = Arc::new(connection.sender.clone()); + let sender_arc = Arc::new(connection.sender); *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(connection, sender_handle).await; + read_self + .read_loop(&mut connection.receiver, sender_handle) + .await; }); // Start heartbeat @@ -410,7 +411,7 @@ impl OmegaConnection { async fn read_loop( self: Arc, - connection: Arc, + receiver: &mut Receiver, sender_handle: Arc, ) { // Monitor both receiver and sender handle for close @@ -418,7 +419,7 @@ impl OmegaConnection { loop { tokio::select! { - result = connection.receive() => { + result = receiver.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 438abd7..0cd3e39 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::ClientPing) { + if cv.is_type(CommunicationType::Ping) { self.handle_ping(cv).await; return; } @@ -331,7 +331,7 @@ impl ClientConnection { }; // Send pong response - let response = CommunicationValue::new(CommunicationType::ClientPing) + let response = CommunicationValue::new(CommunicationType::Pong) .with_id(cv.get_id()) .add_typed_default( DataType::PingIota, diff --git a/src/rho/connection.rs b/src/rho/connection.rs index 3a545f4..b909573 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_err, log_in, log_out, + log_cv_out, log_err, log_in, log_out, rho::{ app_connection::AppConnection, client_connection::ClientConnection, iota_connection::IotaConnection, rho_connection::RhoConnection, @@ -20,6 +20,14 @@ 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" @@ -77,12 +85,7 @@ impl GeneralConnection { connection_kind: kind, id: conn.client_id, rho_connection: Arc::new(RwLock::new(None)), - 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, - })), + session_id: Arc::new(RwLock::new(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())), @@ -129,6 +132,49 @@ 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(); @@ -138,8 +184,24 @@ 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; - self.send_initial_client_state_request(&rho_conn, user_id) - .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; + } } else { log_err!( user_id, @@ -147,7 +209,7 @@ impl GeneralConnection { "No RhoConnection found for user {}, client not attached to iota", id ); - let error = CommunicationValue::new(CommunicationType::ErrorNoIota); + let error = CommunicationValue::new(CommunicationType::ErrorNoIota).with_id(request_id); let _ = self.sender.send(&error).await; } @@ -209,19 +271,67 @@ impl GeneralConnection { Some(rho) } - async fn send_initial_client_state_request(&self, rho: &Arc, user_id: i64) { + async fn request_initial_client_state( + &self, + rho: &Arc, + user_id: i64, + handshake: CommunicationValue, + ) -> Option { 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, 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; + .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 } async fn migrate_iota(self: &Arc) { @@ -281,3 +391,16 @@ 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 7f6fe4c..5e26f80 100755 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -208,12 +208,6 @@ 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; @@ -741,32 +735,6 @@ 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,