diff --git a/Cargo.lock b/Cargo.lock index cf4ad0b..cc8ec44 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -568,9 +568,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytestring" @@ -2179,6 +2179,7 @@ dependencies = [ "json", "mtp", "once_cell", + "r2d2", "rand 0.8.6", "rand_core 0.6.4", "ratatui", @@ -2186,8 +2187,10 @@ dependencies = [ "rusqlite", "serde", "serde_json", + "serde_yaml", "sha2 0.10.9", "sysinfo", + "thiserror 2.0.18", "tokio", "uuid", "walkdir", @@ -2548,9 +2551,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.2" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "memmem" @@ -3535,6 +3538,17 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "rand" version = "0.8.6" @@ -4016,6 +4030,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + [[package]] name = "scoped-tls" version = "1.0.1" @@ -5628,18 +5651,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.53" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75726053136156d419e285b9b7eddaaea9e3fea6ce32eed44a89901f0bd98de1" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.53" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4714fd92cf900833d49538023a9b3915155210801d1c1169eba513b2addefd71" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" dependencies = [ "proc-macro2", "quote", diff --git a/client/src/client_connection.rs b/client/src/client_connection.rs index 774119f..f1474ec 100644 --- a/client/src/client_connection.rs +++ b/client/src/client_connection.rs @@ -512,24 +512,18 @@ impl ClientConnection { let messages = chat_files::get_messages(user_id, contact.user_id, 0, amount); let mut msg_array = Vec::new(); - for m in messages.members() { - let message_time = m["message_time"].as_i64().unwrap_or(0); - let content = m["content"].as_str().unwrap_or("").to_string(); - let sent_by_self = m["sent_by_self"].as_bool().unwrap_or(false); - let height = m["height"].as_i64().unwrap_or(0); - let message_state = m["message_state"].as_str().unwrap_or("").to_string(); - + for m in &messages { let mut msg_container = Vec::new(); msg_container.push(( DataType::SendTime, - DataValue::SignedNumber(message_time as i128), + DataValue::SignedNumber(m.message_time as i128), )); - msg_container.push((DataType::Content, DataValue::Str(content.clone()))); - msg_container.push((DataType::MessageState, DataValue::Str(message_state))); - msg_container.push((DataType::Height, DataValue::SignedNumber(height as i128))); + msg_container.push((DataType::Content, DataValue::Str(m.content.clone()))); + msg_container.push((DataType::MessageState, DataValue::Str(m.message_state.clone()))); + msg_container.push((DataType::Height, DataValue::SignedNumber(m.height as i128))); msg_container.push(( DataType::SenderId, - DataValue::UnsignedNumber(if sent_by_self { + DataValue::UnsignedNumber(if m.sent_by_self { user_id as u128 } else { contact.user_id as u128 @@ -538,13 +532,13 @@ impl ClientConnection { msg_array.push(typed_container(msg_container)); if msg_array.len() == 1 { - let sender_id = if sent_by_self { + let sender_id = if m.sent_by_self { user_id } else { contact.user_id }; let mut last_msg = Vec::new(); - last_msg.push((DataType::Content, DataValue::Str(content))); + last_msg.push((DataType::Content, DataValue::Str(m.content.clone()))); last_msg.push(( DataType::SenderId, DataValue::SignedNumber(sender_id as i128), @@ -749,12 +743,8 @@ impl ClientConnection { amount as i64, ); let mut msg_array: Vec = Vec::new(); - for m in messages.members() { - let message_time: i64 = m["message_time"].as_i64().unwrap_or(0); - let content: String = m["content"].as_str().unwrap_or("").to_string(); - let sent_by_self: bool = m["sent_by_self"].as_bool().unwrap_or(false); - let height: i64 = m["height"].as_i64().unwrap_or(0); - let sender_id: i64 = if sent_by_self { + for m in &messages { + let sender_id: i64 = if m.sent_by_self { my_id as i64 } else { if let Some(n) = cv.get_data(DataType::ChatPartnerId).as_number() { @@ -765,23 +755,22 @@ impl ClientConnection { partner_id as i64 } }; - let message_state: String = m["message_state"].as_str().unwrap_or("").to_string(); let mut container = Vec::new(); container.push(( DataType::SendTime, - DataValue::SignedNumber(message_time as i128), + DataValue::SignedNumber(m.message_time as i128), )); - container.push((DataType::Content, DataValue::Str(content))); + container.push((DataType::Content, DataValue::Str(m.content.clone()))); container.push(( DataType::SenderId, DataValue::SignedNumber(sender_id as i128), )); - container.push((DataType::MessageState, DataValue::Str(message_state))); - container.push((DataType::Height, DataValue::SignedNumber(height as i128))); + container.push((DataType::MessageState, DataValue::Str(m.message_state.clone()))); + container.push((DataType::Height, DataValue::SignedNumber(m.height as i128))); container.push(( DataType::SenderId, - DataValue::UnsignedNumber(if sent_by_self { + DataValue::UnsignedNumber(if m.sent_by_self { my_id as u128 } else { partner_id as u128 @@ -883,18 +872,12 @@ impl ClientConnection { let mut comm_array = Vec::new(); for c in CommunitiesUtil::get_communities(cv.get_sender() as i64) { let mut container: Vec<(DataType, DataValue)> = Vec::new(); - if let Some(address) = c["address"].as_str() { - container.push(( - DataType::CommunityAddress, - DataValue::Str(address.to_string()), - )); - } - if let Some(title) = c["title"].as_str() { - container.push((DataType::CommunityTitle, DataValue::Str(title.to_string()))); - } - if let Some(position) = c["position"].as_str() { - container.push((DataType::Position, DataValue::Str(position.to_string()))); - } + container.push(( + DataType::CommunityAddress, + DataValue::Str(c.address.clone()), + )); + container.push((DataType::CommunityTitle, DataValue::Str(c.title.clone()))); + container.push((DataType::Position, DataValue::Str(c.position.clone()))); comm_array.push(typed_container(container)); } diff --git a/iota-cli/src/elements/console_card.rs b/iota-cli/src/elements/console_card.rs index 6e9a7a0..820ae1e 100644 --- a/iota-cli/src/elements/console_card.rs +++ b/iota-cli/src/elements/console_card.rs @@ -444,7 +444,7 @@ pub async fn run_command(command: &str) { if let Some(user) = user_manager::get_user_by_username(username) { let msg = CommunicationValue::new(CommunicationType::DeleteUser) .with_sender(user.user_id as u64); - OMIKRON_CONNECTION.send_message(&msg).await; + let _ = OMIKRON_CONNECTION.send_message(&msg).await; user_manager::remove_user(user.user_id); log!("Removed user {}", user.user_id); } else { diff --git a/iota-storage/Cargo.toml b/iota-storage/Cargo.toml index 78359be..24d7603 100644 --- a/iota-storage/Cargo.toml +++ b/iota-storage/Cargo.toml @@ -17,8 +17,11 @@ hkdf = "0.12.4" json = "*" arc-swap = "1" once_cell = "1.21.3" +r2d2 = "0.8" serde = { version = "1", features = ["derive"] } serde_json = "1" +serde_yaml = "0.9" +thiserror = "2" rand = "0.8" rand_core = { version = "0.6", features = ["getrandom", "std"] } ratatui = "0.30.0" diff --git a/iota-storage/src/lib.rs b/iota-storage/src/lib.rs index 7f0e8ef..a0e5061 100644 --- a/iota-storage/src/lib.rs +++ b/iota-storage/src/lib.rs @@ -1,2 +1,3 @@ +pub mod storage_error; pub mod users; pub mod util; diff --git a/iota-storage/src/storage_error.rs b/iota-storage/src/storage_error.rs new file mode 100644 index 0000000..100bf08 --- /dev/null +++ b/iota-storage/src/storage_error.rs @@ -0,0 +1,13 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum StorageError { + #[error("Database error: {0}")] + Db(#[from] rusqlite::Error), + #[error("Connection pool error: {0}")] + Pool(String), + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + #[error("{0}")] + Other(String), +} diff --git a/iota-storage/src/users/contact.rs b/iota-storage/src/users/contact.rs index 632cea3..50d224d 100644 --- a/iota-storage/src/users/contact.rs +++ b/iota-storage/src/users/contact.rs @@ -1,4 +1,3 @@ -use json::{self, JsonValue, number::Number}; use std::time::{SystemTime, UNIX_EPOCH}; #[derive(Debug, Clone)] @@ -33,29 +32,4 @@ impl Contact { pub fn set_last_message_at(&mut self, p0: i64) { self.last_message_at = Option::from(p0); } - - pub fn to_json(&self) -> JsonValue { - let mut obj = JsonValue::new_object(); - obj["user_id"] = JsonValue::Number(Number::from(self.user_id)); - if let Some(name) = &self.user_name { - obj["user_name"] = JsonValue::from(name.as_str()); - } - if let Some(ts) = &self.last_message_at { - obj["last_message_at"] = JsonValue::Number(Number::from(*ts)); - } - obj - } - pub fn from_json(o: &JsonValue) -> Contact { - let user_id = o["user_id"].as_i64().unwrap_or(0); - - let user_name = o["user_name"].as_str().map(|s| s.to_string()); - - let last_message_at = o["last_message_at"].as_i64(); - - Contact { - user_id, - user_name, - last_message_at, - } - } } diff --git a/iota-storage/src/users/mod.rs b/iota-storage/src/users/mod.rs index aef5a02..cd4fa45 100644 --- a/iota-storage/src/users/mod.rs +++ b/iota-storage/src/users/mod.rs @@ -1,4 +1,3 @@ pub mod contact; -pub mod user_community_util; pub mod user_manager; pub mod user_profile; diff --git a/iota-storage/src/users/user_community_util.rs b/iota-storage/src/users/user_community_util.rs deleted file mode 100644 index 89e0fb6..0000000 --- a/iota-storage/src/users/user_community_util.rs +++ /dev/null @@ -1,58 +0,0 @@ -use iota_util::file_util::{load_file, save_file}; -use json::{self, Array, JsonValue}; - -pub struct UserCommunityUtil; - -impl UserCommunityUtil { - pub fn add_community(storage_owner: i64, address: String, title: String, position: String) { - let file_path = format!("users/{}/", storage_owner); - let mut communities = Self::load_array(&file_path, "communities.json"); - - let mut community = JsonValue::new_object(); - community["title"] = JsonValue::String(title); - community["address"] = JsonValue::String(address); - community["position"] = JsonValue::String(position); - - communities.push(community); - - save_file( - &file_path, - "communities.json", - &JsonValue::Array(communities).to_string(), - ); - } - - pub fn remove_community(storage_owner: i64, community_address: String) { - let file_path = format!("users/{}/", storage_owner); - let communities = Self::load_array(&file_path, "communities.json"); - - let filtered: Array = communities - .iter() - .filter(|entry| entry["address"].as_str() != Some(&community_address)) - .cloned() - .collect(); - save_file( - &file_path, - "communities.json", - &JsonValue::Array(filtered).to_string(), - ); - } - - pub fn get_communities(storage_owner: i64) -> Array { - let file_path = format!("users/{}/", storage_owner); - Self::load_array(&file_path, "communities.json") - } - - fn load_array(dir: &str, name: &str) -> Array { - let content = load_file(dir, name); - if content.is_empty() { - return Array::new(); - } - - let parsed = json::parse(&content); - match parsed { - Ok(JsonValue::Array(arr)) => arr, - _ => Array::new(), - } - } -} diff --git a/iota-storage/src/users/user_manager.rs b/iota-storage/src/users/user_manager.rs index bf1dc00..68afe5e 100644 --- a/iota-storage/src/users/user_manager.rs +++ b/iota-storage/src/users/user_manager.rs @@ -1,15 +1,212 @@ use crate::users::user_profile::UserProfile; +use crate::util::db; use base64::{Engine as _, engine::general_purpose::STANDARD}; use iota_util::crypto_helper::{self, hex_hash, keyring_from_base64, public_key_bundle_to_base64}; use iota_util::file_util::{load_file, save_file}; -use json::JsonValue; -use once_cell::sync::Lazy; +use rusqlite::params; use rand_core::{OsRng, RngCore}; -use std::io::{self}; -use std::sync::Mutex; -static USERS: Lazy>> = Lazy::new(|| Mutex::new(Vec::new())); -static UNIQUE: Lazy> = Lazy::new(|| Mutex::new(false)); +pub fn add_user(user: UserProfile) { + if let Err(e) = db::with_db(|conn| { + conn.execute( + r#" + INSERT INTO users (user_id, username, public_key, private_key_hash, reset_token, created_at, display_name) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) + ON CONFLICT(user_id) DO UPDATE SET + username = excluded.username, + public_key = excluded.public_key, + private_key_hash = excluded.private_key_hash, + reset_token = excluded.reset_token, + display_name = excluded.display_name + "#, + params![ + user.user_id, + user.username, + user.public_key, + user.private_key_hash, + user.reset_token, + user.created_at, + user.display_name, + ], + )?; + + for (app_id, app_secret) in &user.trusted_apps { + conn.execute( + r#" + INSERT OR REPLACE INTO trusted_apps (user_id, app_id, app_secret) + VALUES (?1, ?2, ?3) + "#, + params![user.user_id, app_id, app_secret], + )?; + } + Ok(()) + }) { + eprintln!("Failed to add_user: {}", e); + } +} + +pub fn update_user(user: UserProfile) { + add_user(user); +} + +pub fn get_user_by_username(username: &str) -> Option { + match db::with_db(|conn| { + match conn.query_row( + "SELECT user_id, username, public_key, private_key_hash, reset_token, created_at, display_name FROM users WHERE username = ?1 LIMIT 1", + params![username], + |r| { + let user_id: i64 = r.get(0)?; + Ok(UserProfile { + user_id, + username: r.get(1)?, + display_name: r.get(6)?, + public_key: r.get(2)?, + private_key_hash: r.get(3)?, + created_at: r.get(5)?, + reset_token: r.get(4)?, + trusted_apps: load_trusted_apps(user_id), + }) + }, + ) { + Ok(user) => Ok(Some(user)), + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(e.into()), + } + }) { + Ok(opt) => opt, + Err(e) => { + eprintln!("Error querying user by username: {}", e); + None + } + } +} + +pub fn get_user(user_id: i64) -> Option { + match db::with_db(|conn| { + match conn.query_row( + "SELECT user_id, username, public_key, private_key_hash, reset_token, created_at, display_name FROM users WHERE user_id = ?1 LIMIT 1", + params![user_id], + |r| { + let user_id: i64 = r.get(0)?; + Ok(UserProfile { + user_id, + username: r.get(1)?, + display_name: r.get(6)?, + public_key: r.get(2)?, + private_key_hash: r.get(3)?, + created_at: r.get(5)?, + reset_token: r.get(4)?, + trusted_apps: load_trusted_apps(user_id), + }) + }, + ) { + Ok(user) => Ok(Some(user)), + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(e.into()), + } + }) { + Ok(opt) => opt, + Err(e) => { + eprintln!("Error querying user: {}", e); + None + } + } +} + +pub fn get_users() -> Vec { + match db::with_db(|conn| { + let mut stmt = conn.prepare( + r#" + SELECT user_id, username, public_key, private_key_hash, reset_token, created_at, display_name + FROM users + ORDER BY username + "#, + )?; + + let rows = stmt.query_map([], |r| { + let user_id: i64 = r.get(0)?; + let username: String = r.get(1)?; + let public_key: String = r.get(2)?; + let private_key_hash: String = r.get(3)?; + let reset_token: String = r.get(4)?; + let created_at: i64 = r.get(5)?; + let display_name: Option = r.get(6)?; + + Ok(UserProfile { + user_id, + username, + display_name, + public_key, + private_key_hash, + created_at, + reset_token, + trusted_apps: std::collections::HashMap::new(), + }) + })?; + + let mut out = Vec::new(); + for row in rows { + match row { + Ok(mut user) => { + user.trusted_apps = load_trusted_apps(user.user_id); + out.push(user); + } + Err(e) => eprintln!("Failed to read user row: {}", e), + } + } + Ok(out) + }) { + Ok(v) => v, + Err(e) => { + eprintln!("Failed to query users: {}", e); + Vec::new() + } + } +} + +fn load_trusted_apps(user_id: i64) -> std::collections::HashMap { + match db::with_db(|conn| { + let mut stmt = conn.prepare( + "SELECT app_id, app_secret FROM trusted_apps WHERE user_id = ?1", + )?; + let rows = stmt.query_map(params![user_id], |r| { + Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)) + })?; + + let mut map = std::collections::HashMap::new(); + for row in rows { + if let Ok((k, v)) = row { + map.insert(k, v); + } + } + Ok(map) + }) { + Ok(m) => m, + Err(e) => { + eprintln!("Failed to load trusted apps: {}", e); + std::collections::HashMap::new() + } + } +} + +pub fn remove_user(user_id: i64) { + if let Err(e) = db::with_db(|conn| { + conn.execute("DELETE FROM trusted_apps WHERE user_id = ?1", params![user_id])?; + conn.execute("DELETE FROM users WHERE user_id = ?1", params![user_id])?; + Ok(()) + }) { + eprintln!("Failed to remove_user: {}", e); + } +} + +pub fn clear() { + if let Err(e) = db::with_db(|conn| { + conn.execute_batch("DELETE FROM trusted_apps; DELETE FROM users;")?; + Ok(()) + }) { + eprintln!("Failed to clear users: {}", e); + } +} #[allow(dead_code)] pub async fn load_from_tu(username: &str) -> Result<(), ()> { @@ -35,91 +232,38 @@ pub async fn load_from_tu(username: &str) -> Result<(), ()> { hex_hash(&keyring_b64), reset_token, ); - USERS.lock().unwrap().push(user_profile); + add_user(user_profile); Ok(()) } -pub fn add_user(user: UserProfile) { - USERS.lock().unwrap().push(user); -} - -pub fn update_user(user: UserProfile) { - let mut users = USERS.lock().unwrap(); - if let Some(pos) = users.iter().position(|u| u.user_id == user.user_id) { - users[pos] = user; - } - *UNIQUE.lock().unwrap() = true; -} -pub fn get_user_by_username(username: &str) -> Option { - USERS - .lock() - .unwrap() - .iter() - .cloned() - .find(|u| u.username == username) -} - -pub fn get_user(user_id: i64) -> Option { - USERS - .lock() - .unwrap() - .iter() - .cloned() - .find(|u| u.user_id == user_id) -} - -pub fn get_users() -> Vec { - USERS.lock().unwrap().clone() -} - -pub fn remove_user(user_id: i64) { - let mut users = USERS.lock().unwrap(); - users.retain(|u| u.user_id != user_id); - *UNIQUE.lock().unwrap() = true; -} - pub fn save_users() { - *UNIQUE.lock().unwrap() = false; - let users = USERS.lock().unwrap(); - let arr: Vec = users.iter().map(|u| u.to_json()).collect(); - let json_str = JsonValue::Array(arr).dump(); - - save_file("", "users.json", &json_str); + // No-op: users are auto-saved via SQLite. } -pub fn clear() { - let mut users = USERS.lock().unwrap(); - users.clear(); - *UNIQUE.lock().unwrap() = true; -} - -pub async fn load_users() -> io::Result<()> { +pub async fn load_users() -> std::io::Result<()> { + // Users are loaded from SQLite on demand. This function is kept for API compat. + // If we need to migrate from a legacy users.json file, we can do so here. let content = load_file("", "users.json"); if content.trim().is_empty() { return Ok(()); } - - let parsed = - json::parse(&content).map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?; - if let JsonValue::Array(arr) = parsed { - let mut users = USERS.lock().unwrap(); - for j in arr.iter() { - if let Some(up) = UserProfile::from_json(j).await { - users.push(up); + if let Ok(parsed) = json::parse(&content) { + if let json::JsonValue::Array(arr) = parsed { + for j in arr.iter() { + if let Some(up) = UserProfile::from_json(j) { + add_user(up); + } } } } - if *UNIQUE.lock().unwrap() { - save_users(); - } + // Rename the old file so we don't re-import + let _ = std::fs::rename( + std::path::PathBuf::from(iota_util::file_util::get_directory()).join("users.json"), + std::path::PathBuf::from(iota_util::file_util::get_directory()).join("users.json.imported"), + ); Ok(()) } -#[allow(dead_code)] -pub fn set_unique(val: bool) { - *UNIQUE.lock().unwrap() = val; -} - pub fn save_app_data(user_id: i64, app_identifier: &str, data: &str) { let path = format!("users/{}/apps", user_id); let name = format!("{}.json", app_identifier); diff --git a/iota-storage/src/users/user_profile.rs b/iota-storage/src/users/user_profile.rs index c0ae20e..2ecb9cc 100644 --- a/iota-storage/src/users/user_profile.rs +++ b/iota-storage/src/users/user_profile.rs @@ -5,9 +5,9 @@ use iota_util::file_util::{has_file, load_file, used_dir_space}; use json::{JsonValue, object}; use rand::Rng; use rand::rngs::OsRng; +use serde::{Deserialize, Serialize}; -// --- UserProfile --- -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct UserProfile { pub user_id: i64, pub username: String, @@ -43,26 +43,6 @@ impl UserProfile { } } - pub fn to_json(&self) -> JsonValue { - let mut trusted_apps_obj = json::JsonValue::new_object(); - for (k, v) in &self.trusted_apps { - trusted_apps_obj[k] = v.clone().into(); - } - - let mut obj = object! { - "uuid" => self.user_id, - "username" => self.username.clone(), - "public_key" => self.public_key.clone(), - "private_key_hash" => self.private_key_hash.clone(), - "created_at" => self.created_at, - "reset_token" => self.reset_token.clone(), - "trusted_apps" => trusted_apps_obj, - }; - if let Some(d) = &self.display_name { - obj["display_name"] = d.clone().into(); - } - obj - } pub fn frontend(&self) -> JsonValue { let mut obj = object! { "uuid" => self.user_id, @@ -78,10 +58,11 @@ impl UserProfile { if has_file("", &format!("{}.tu", self.username.clone())) { obj["tu"] = load_file("", &format!("{}.tu", self.username.clone())).into(); } - obj } - pub async fn from_json(j: &JsonValue) -> Option { + + /// Legacy JSON import - used when migrating from users.json to SQLite. + pub fn from_json(j: &JsonValue) -> Option { let user_id = j["uuid"].as_i64()?; let username = j["username"].as_str()?.to_string(); let public_key = j["public_key"].as_str()?.to_string(); @@ -99,7 +80,7 @@ impl UserProfile { } } - let up = UserProfile { + Some(UserProfile { user_id, username, display_name, @@ -108,22 +89,15 @@ impl UserProfile { created_at, reset_token, trusted_apps, - }; + }) + } - // TODO: Migrate to Omikron / Wss - /* if j.has_key("migrate") - || j.has_key("migrating") - || j.has_key("changing") - || j.has_key("move") - || j.has_key("moving") - { - if auth_connector::migrate_user(&mut up).await { - log_message(format!("[INFO] Migration triggered for {}", up.username)); - user_manager::set_unique(true); - } - } */ + pub fn from_yaml(s: &str) -> Result { + serde_yaml::from_str(s) + } - Some(up) + pub fn to_yaml(&self) -> Result { + serde_yaml::to_string(self) } #[allow(dead_code)] diff --git a/iota-storage/src/util/chat_files.rs b/iota-storage/src/util/chat_files.rs index a5b50fe..4aba87d 100644 --- a/iota-storage/src/util/chat_files.rs +++ b/iota-storage/src/util/chat_files.rs @@ -1,9 +1,7 @@ use crate::util::db; use iota_logger::log; -use json::{JsonValue, array, object}; use rusqlite::params; -use std::io; -use std::sync::{Arc, LazyLock, Mutex}; +use crate::storage_error::StorageError; #[derive(PartialEq, Debug, Clone)] pub enum MessageState { @@ -45,11 +43,202 @@ impl MessageState { } } -// Shared DB created via helper. -// The db helper constructs the messages sqlite file and ensures PRAGMAs and schema exist. -static MESSAGES_DB: LazyLock>> = LazyLock::new(|| { - db::create_general_messages_db().expect("Failed to create or initialize general messages DB") -}); +#[derive(Debug, Clone)] +pub struct StoredMessage { + pub id: i64, + pub message_time: i64, + pub content: String, + pub edited: bool, + pub sent_by_self: bool, + pub message_state: String, + pub height: i64, + pub reply_to: Option, + pub reactions: Vec, +} + +/* + * Each edit is recorded in message_edits with the before/after content and a + * timestamp. Only the original sender (sent_by_self = 1) may edit. + */ +pub fn edit_message( + storage_owner: i64, + external_user: i64, + message_time: i64, + editor_id: i64, + new_content: &str, +) -> Result<(), StorageError> { + db::with_db(|conn| { + let msg = conn.query_row( + r#" + SELECT id, content, sent_by_self + FROM messages + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + ORDER BY id DESC LIMIT 1 + "#, + params![storage_owner, external_user, message_time], + |row| { + Ok(( + row.get::<_, i64>(0)?, + row.get::<_, String>(1)?, + row.get::<_, i64>(2)?, + )) + }, + )?; + + let (msg_id, old_content, sent_by_self) = msg; + if sent_by_self != 1 { + return Err(StorageError::Other( + "Only the original sender can edit this message".into(), + )); + } + + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as i64; + + conn.execute( + r#" + INSERT INTO message_edits (message_id, content_before, content_after, edited_at, edited_by) + VALUES (?1, ?2, ?3, ?4, ?5) + "#, + params![msg_id, old_content, new_content, now, editor_id], + )?; + + conn.execute( + r#" + UPDATE messages + SET content = ?1, edited_count = edited_count + 1 + WHERE id = ?2 + "#, + params![new_content, msg_id], + )?; + + Ok(()) + }) +} + +pub fn hard_delete_message(storage_owner: i64, external_user: i64, message_time: i64) -> Result<(), StorageError> { + db::with_db(|conn| { + let msg_id: i64 = conn.query_row( + r#" + SELECT id FROM messages + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + ORDER BY id DESC LIMIT 1 + "#, + params![storage_owner, external_user, message_time], + |row| row.get(0), + )?; + + conn.execute("DELETE FROM message_edits WHERE message_id = ?1", params![msg_id])?; + conn.execute("DELETE FROM reactions WHERE message_id = ?1", params![msg_id])?; + conn.execute("DELETE FROM messages WHERE id = ?1", params![msg_id])?; + Ok(()) + }) +} + +/* + * Marks a message as deleted by the external user rather than removing the row, + * so the storage owner still sees a tombstone in the UI. + */ +pub fn flag_deleted_by_external(storage_owner: i64, external_user: i64, message_time: i64) -> Result<(), StorageError> { + db::with_db(|conn| { + let affected = conn.execute( + r#" + UPDATE messages + SET deleted_by_external = 1 + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + "#, + params![storage_owner, external_user, message_time], + )?; + if affected == 0 { + return Err(StorageError::Other("Message not found".into())); + } + Ok(()) + }) +} + +/* + * Removes the edit trail but keeps the message with edited_count > 0 so + * the UI still shows the "edited" indicator. Only the own user should + * call this. + */ +pub fn delete_edit_history(storage_owner: i64, external_user: i64, message_time: i64) -> Result<(), StorageError> { + db::with_db(|conn| { + let msg_id: i64 = conn.query_row( + r#" + SELECT id FROM messages + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + ORDER BY id DESC LIMIT 1 + "#, + params![storage_owner, external_user, message_time], + |row| row.get(0), + )?; + + conn.execute("DELETE FROM message_edits WHERE message_id = ?1", params![msg_id])?; + Ok(()) + }) +} + +pub fn add_reaction( + storage_owner: i64, + external_user: i64, + message_time: i64, + user_id: i64, + reaction: &str, +) -> Result<(), StorageError> { + db::with_db(|conn| { + let msg_id: i64 = conn.query_row( + r#" + SELECT id FROM messages + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + ORDER BY id DESC LIMIT 1 + "#, + params![storage_owner, external_user, message_time], + |row| row.get(0), + )?; + + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as i64; + + conn.execute( + r#" + INSERT OR IGNORE INTO reactions (message_id, user_id, reaction, created_at) + VALUES (?1, ?2, ?3, ?4) + "#, + params![msg_id, user_id, reaction, now], + )?; + Ok(()) + }) +} + +pub fn remove_reaction( + storage_owner: i64, + external_user: i64, + message_time: i64, + user_id: i64, + reaction: &str, +) -> Result<(), StorageError> { + db::with_db(|conn| { + let msg_id: i64 = conn.query_row( + r#" + SELECT id FROM messages + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + ORDER BY id DESC LIMIT 1 + "#, + params![storage_owner, external_user, message_time], + |row| row.get(0), + )?; + + conn.execute( + "DELETE FROM reactions WHERE message_id = ?1 AND user_id = ?2 AND reaction = ?3", + params![msg_id, user_id, reaction], + )?; + Ok(()) + }) +} pub fn add_message( send_time: u128, @@ -68,19 +257,12 @@ pub fn add_message( } }; - // Insert the message into the DB - let insert_result = db::with_conn(&MESSAGES_DB, |conn| { + if let Err(e) = db::with_db(|conn| { conn.execute( r#" INSERT INTO messages ( - storage_owner, - external_user, - message_time, - content, - sent_by_self, - message_state, - height, - reply_to + storage_owner, external_user, message_time, content, + sent_by_self, message_state, height, reply_to ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) "#, params![ @@ -99,18 +281,13 @@ pub fn add_message( ], )?; Ok(()) - }); - - if let Err(e) = insert_result { + }) { log!("Failed to insert message into sqlite: {}", e); return; } - // Update contacts table to reflect that this conversation exists and has a recent message. - // Use the Contact helper to set last_message_at to the message timestamp. let mut contact = crate::users::contact::Contact::new(external_user); contact.set_last_message_at(message_time); - // This will insert or update the contact for the storage owner. crate::util::chats_util::mod_user(storage_owner, &contact); } @@ -119,25 +296,21 @@ pub fn change_message_state( storage_owner: i64, external_user: i64, new_state: MessageState, -) -> io::Result<()> { - // Run the SELECT and UPDATE inside with_conn to centralize connection access. - let res: Result<(), String> = db::with_conn(&MESSAGES_DB, |conn| { +) -> std::io::Result<()> { + db::with_db(|conn| { let current: Option = match conn.query_row( r#" SELECT message_state FROM messages - WHERE storage_owner = ?1 - AND external_user = ?2 - AND message_time = ?3 - ORDER BY id DESC - LIMIT 1 + WHERE storage_owner = ?1 AND external_user = ?2 AND message_time = ?3 + ORDER BY id DESC LIMIT 1 "#, params![storage_owner, external_user, timestamp], |row| row.get(0), ) { Ok(state) => Some(state), Err(rusqlite::Error::QueryReturnedNoRows) => None, - Err(e) => return Err(e), + Err(e) => return Err(e.into()), }; let Some(current_state_raw) = current else { @@ -154,24 +327,45 @@ pub fn change_message_state( UPDATE messages SET message_state = ?1 WHERE id = ( - SELECT id - FROM messages - WHERE storage_owner = ?2 - AND external_user = ?3 - AND message_time = ?4 - ORDER BY id DESC - LIMIT 1 + SELECT id FROM messages + WHERE storage_owner = ?2 AND external_user = ?3 AND message_time = ?4 + ORDER BY id DESC LIMIT 1 ) "#, params![upgraded, storage_owner, external_user, timestamp], )?; Ok(()) - }); + }) + .map_err(|e: StorageError| std::io::Error::new(std::io::ErrorKind::Other, e.to_string())) +} - match res { - Ok(_) => Ok(()), - Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)), +fn load_reactions(conn: &rusqlite::Connection, msg_ids: &[i64]) -> std::collections::HashMap> { + if msg_ids.is_empty() { + return std::collections::HashMap::new(); } + + let placeholders: Vec = msg_ids.iter().enumerate() + .map(|(i, _)| format!("?{}", i + 1)) + .collect(); + let query = format!( + "SELECT message_id, reaction || ':' || COUNT(*) FROM reactions WHERE message_id IN ({}) GROUP BY message_id, reaction", + placeholders.join(", ") + ); + + let mut map: std::collections::HashMap> = std::collections::HashMap::new(); + if let Ok(mut stmt) = conn.prepare(&query) { + let params: Vec<&dyn rusqlite::types::ToSql> = msg_ids.iter() + .map(|id| id as &dyn rusqlite::types::ToSql) + .collect(); + if let Ok(rows) = stmt.query_map(params.as_slice(), |row| { + Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)) + }) { + for row in rows.flatten() { + map.entry(row.0).or_default().push(row.1); + } + } + } + map } pub fn get_messages( @@ -179,26 +373,17 @@ pub fn get_messages( external_user: i64, loaded_messages: i64, amount: i64, -) -> JsonValue { - let messages = array![]; - +) -> Vec { if amount <= 0 || loaded_messages < 0 { - return messages; + return Vec::new(); } - let res: Result = db::with_conn(&MESSAGES_DB, |conn| { + match db::with_db(|conn| { let mut stmt = conn.prepare( r#" - SELECT - message_time, - content, - sent_by_self, - message_state, - height, - reply_to + SELECT id, message_time, content, sent_by_self, message_state, height, reply_to, edited_count FROM messages - WHERE storage_owner = ?1 - AND external_user = ?2 + WHERE storage_owner = ?1 AND external_user = ?2 AND deleted_by_external = 0 ORDER BY message_time DESC, id DESC LIMIT ?3 OFFSET ?4 "#, @@ -207,49 +392,40 @@ pub fn get_messages( let rows = stmt.query_map( params![storage_owner, external_user, amount, loaded_messages], |row| { - let message_time: i64 = row.get(0)?; - let content: String = row.get(1)?; - let sent_by_self: i64 = row.get(2)?; - let message_state: String = row.get(3)?; - let height: i64 = row.get(4).unwrap_or(0); - let reply_to: Option = row.get(5).ok().flatten(); - Ok((message_time, content, sent_by_self, message_state, height, reply_to)) + Ok(StoredMessage { + id: row.get(0)?, + message_time: row.get(1)?, + content: row.get(2)?, + sent_by_self: row.get::<_, i64>(3)? != 0, + message_state: row.get(4)?, + height: row.get(5).unwrap_or(0), + reply_to: row.get(6).ok().flatten(), + edited: row.get::<_, i64>(7).unwrap_or(0) > 0, + reactions: Vec::new(), + }) }, )?; - let mut out = array![]; + let mut out = Vec::new(); for row in rows { match row { - Ok((message_time, content, sent_by_self, message_state, height, reply_to)) => { - let mut msg = object! { - "message_time" => message_time, - "content" => content, - "sent_by_self" => (sent_by_self != 0), - "message_state" => message_state, - "height" => height - }; - if let Some(rt) = reply_to { - let _ = msg.insert("reply_to", rt); - } - if let Err(e) = out.push(msg) { - // out.push returns a JsonError; log it instead of using `?` to avoid - // incompatible error conversions inside the DB closure. - log!("Failed to append message to output array: {:?}", e); - } - } - Err(e) => { - log!("Failed to read row from sqlite: {}", e); - } + Ok(msg) => out.push(msg), + Err(e) => log!("Failed to read row from sqlite: {}", e), } } - Ok(out) - }); - match res { + let msg_ids: Vec = out.iter().map(|m| m.id).collect(); + let reaction_map = load_reactions(conn, &msg_ids); + for msg in &mut out { + msg.reactions = reaction_map.get(&msg.id).cloned().unwrap_or_default(); + } + + Ok(out) + }) { Ok(v) => v, Err(e) => { log!("Failed to query messages: {}", e); - messages + Vec::new() } } } diff --git a/iota-storage/src/util/chats_util.rs b/iota-storage/src/util/chats_util.rs index b8e0ad2..eb8f069 100644 --- a/iota-storage/src/util/chats_util.rs +++ b/iota-storage/src/util/chats_util.rs @@ -1,24 +1,13 @@ use crate::users::contact::Contact; use crate::util::db; use rusqlite::params; -use std::sync::{Arc, LazyLock, Mutex}; -/// Shared DB connection for contacts/messages (created by db helper). -static MESSAGES_DB: LazyLock>> = LazyLock::new(|| { - db::create_general_messages_db().expect("Failed to create or initialize general messages DB") -}); - -/// Insert or update a contact for the given storage owner. pub fn mod_user(storage_owner: i64, contact: &Contact) { - if let Err(e) = db::with_conn(&MESSAGES_DB, |conn| { + if let Err(e) = db::with_db(|conn| { conn.execute( r#" - INSERT INTO contacts ( - storage_owner, - user_id, - user_name, - last_message_at - ) VALUES (?1, ?2, ?3, ?4) + INSERT INTO contacts (storage_owner, user_id, user_name, last_message_at) + VALUES (?1, ?2, ?3, ?4) ON CONFLICT(storage_owner, user_id) DO UPDATE SET user_name = excluded.user_name, last_message_at = excluded.last_message_at @@ -27,7 +16,7 @@ pub fn mod_user(storage_owner: i64, contact: &Contact) { storage_owner, contact.user_id, contact.user_name.clone(), - contact.last_message_at + contact.last_message_at, ], )?; Ok(()) @@ -36,9 +25,8 @@ pub fn mod_user(storage_owner: i64, contact: &Contact) { } } -/// Retrieve a single contact for storage_owner/user_id. pub fn get_user(storage_owner: i64, user_id: i64) -> Option { - let res: Result, String> = db::with_conn(&MESSAGES_DB, |conn| { + match db::with_db(|conn| { match conn.query_row( r#" SELECT user_id, user_name, last_message_at @@ -48,23 +36,18 @@ pub fn get_user(storage_owner: i64, user_id: i64) -> Option { "#, params![storage_owner, user_id], |r| { - let user_id: i64 = r.get(0)?; - let user_name: Option = r.get(1)?; - let last_message_at: Option = r.get(2)?; Ok(Contact { - user_id, - user_name, - last_message_at, + user_id: r.get(0)?, + user_name: r.get(1)?, + last_message_at: r.get(2)?, }) }, ) { Ok(c) => Ok(Some(c)), Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), - Err(e) => Err(e), + Err(e) => Err(e.into()), } - }); - - match res { + }) { Ok(opt) => opt, Err(e) => { eprintln!("Error querying user in get_user: {}", e); @@ -73,11 +56,8 @@ pub fn get_user(storage_owner: i64, user_id: i64) -> Option { } } -/// Retrieve all contacts for a storage owner, ordered by last_message_at desc / user_id asc. pub fn get_users(storage_owner: i64) -> Vec { - let contacts_out = Vec::new(); - - let res: Result, String> = db::with_conn(&MESSAGES_DB, |conn| { + match db::with_db(|conn| { let mut stmt = conn.prepare( r#" SELECT user_id, user_name, last_message_at @@ -91,13 +71,10 @@ pub fn get_users(storage_owner: i64) -> Vec { )?; let rows = stmt.query_map(params![storage_owner], |r| { - let user_id: i64 = r.get(0)?; - let user_name: Option = r.get(1)?; - let last_message_at: Option = r.get(2)?; Ok(Contact { - user_id, - user_name, - last_message_at, + user_id: r.get(0)?, + user_name: r.get(1)?, + last_message_at: r.get(2)?, }) })?; @@ -109,13 +86,11 @@ pub fn get_users(storage_owner: i64) -> Vec { } } Ok(out) - }); - - match res { + }) { Ok(v) => v, Err(e) => { eprintln!("Failed to query contacts in get_users: {}", e); - contacts_out + Vec::new() } } } diff --git a/iota-storage/src/util/communities_util.rs b/iota-storage/src/util/communities_util.rs index ed8b7ae..c560056 100644 --- a/iota-storage/src/util/communities_util.rs +++ b/iota-storage/src/util/communities_util.rs @@ -1,25 +1,22 @@ use crate::util::db; -use json::Array; use rusqlite::params; -use std::sync::{Arc, LazyLock, Mutex}; -static MESSAGES_DB: LazyLock>> = LazyLock::new(|| { - db::create_general_messages_db().expect("Failed to create or initialize general messages DB") -}); +#[derive(Debug, Clone)] +pub struct StoredCommunity { + pub address: String, + pub title: String, + pub position: String, +} pub struct CommunitiesUtil; impl CommunitiesUtil { pub fn add_community(storage_owner: i64, address: String, title: String, position: String) { - if let Err(e) = db::with_conn(&MESSAGES_DB, |conn| { + if let Err(e) = db::with_db(|conn| { conn.execute( r#" - INSERT INTO communities ( - storage_owner, - address, - title, - position - ) VALUES (?1, ?2, ?3, ?4) + INSERT INTO communities (storage_owner, address, title, position) + VALUES (?1, ?2, ?3, ?4) ON CONFLICT(storage_owner, address) DO UPDATE SET title = excluded.title, position = excluded.position @@ -33,7 +30,7 @@ impl CommunitiesUtil { } pub fn remove_community(storage_owner: i64, community_address: String) { - if let Err(e) = db::with_conn(&MESSAGES_DB, |conn| { + if let Err(e) = db::with_db(|conn| { conn.execute( "DELETE FROM communities WHERE storage_owner = ?1 AND address = ?2", params![storage_owner, community_address], @@ -44,10 +41,8 @@ impl CommunitiesUtil { } } - pub fn get_communities(storage_owner: i64) -> Array { - let communities_out = Array::new(); - - let res: Result = db::with_conn(&MESSAGES_DB, |conn| { + pub fn get_communities(storage_owner: i64) -> Vec { + match db::with_db(|conn| { let mut stmt = conn.prepare( r#" SELECT address, title, position @@ -57,33 +52,26 @@ impl CommunitiesUtil { )?; let rows = stmt.query_map(params![storage_owner], |r| { - let address: String = r.get(0)?; - let title: String = r.get(1)?; - let position: String = r.get(2)?; - Ok((address, title, position)) + Ok(StoredCommunity { + address: r.get(0)?, + title: r.get(1)?, + position: r.get(2)?, + }) })?; - let mut out = Array::new(); + let mut out = Vec::new(); for row in rows { match row { - Ok((address, title, position)) => { - let mut community = json::JsonValue::new_object(); - community["title"] = json::JsonValue::String(title); - community["address"] = json::JsonValue::String(address); - community["position"] = json::JsonValue::String(position); - out.push(community); - } + Ok(community) => out.push(community), Err(e) => eprintln!("Failed to read community row: {}", e), } } Ok(out) - }); - - match res { - Ok(arr) => arr, + }) { + Ok(v) => v, Err(e) => { eprintln!("Failed to query communities in get_communities: {}", e); - communities_out + Vec::new() } } } diff --git a/iota-storage/src/util/config_util.rs b/iota-storage/src/util/config_util.rs index 450ffb4..82fbbaa 100644 --- a/iota-storage/src/util/config_util.rs +++ b/iota-storage/src/util/config_util.rs @@ -51,17 +51,17 @@ impl Default for IotaConfig { } pub fn load_config() { - let s = load_file("", "config.json"); + let s = load_file("", "config.yaml"); if s.is_empty() { return; } - match serde_json::from_str::(&s) { + match serde_yaml::from_str::(&s) { Ok(parsed) => { CONFIG.store(Arc::new(parsed)); } Err(e) => { - eprintln!("Failed to parse config.json: {}. Content: '{}'", e, s); + eprintln!("Failed to parse config.yaml: {}. Content: '{}'", e, s); } } } @@ -72,8 +72,8 @@ pub fn clear_config() { } pub fn save_config() { - if let Ok(json) = serde_json::to_string(&**CONFIG.load()) { - save_file("", "config.json", &json); + if let Ok(yaml) = serde_yaml::to_string(&**CONFIG.load()) { + save_file("", "config.yaml", &yaml); } } diff --git a/iota-storage/src/util/db.rs b/iota-storage/src/util/db.rs index c39a4bf..3263f68 100644 --- a/iota-storage/src/util/db.rs +++ b/iota-storage/src/util/db.rs @@ -1,88 +1,207 @@ -//! Database helper utilities. -//! -//! This module provides small helpers to open/init sqlite databases and to -//! create a shared (Arc>) connection wrapper callers can -//! reuse. The goal is to centralize the "open and initialize" logic and -//! provide small convenience helpers used by other util modules. - use iota_util::file_util::get_directory; -use rusqlite::{Connection, Error as RusqliteError}; +use once_cell::sync::Lazy; +use r2d2::ManageConnection; +use rusqlite::Connection; use std::path::PathBuf; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::time::Duration; -/// Returns the file path for a named DB inside the application's data directory. -/// -/// Arguments: -/// - `db_name` : name of the DB (without extension). Example: `"messages"`. -pub fn db_file_path(db_name: &str) -> String { +use crate::storage_error::StorageError; + +const DB_NAME: &str = "messages"; + +/// A simple r2d2 manager for rusqlite connections. +pub struct SqliteManager; + +impl ManageConnection for SqliteManager { + type Connection = Connection; + type Error = rusqlite::Error; + + fn connect(&self) -> Result { + let path = db_file_path(DB_NAME); + let conn = Connection::open(path)?; + conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL;")?; + conn.busy_timeout(Duration::from_millis(250))?; + Ok(conn) + } + + fn is_valid(&self, conn: &mut Connection) -> Result<(), rusqlite::Error> { + conn.execute_batch("SELECT 1") + } + + fn has_broken(&self, _conn: &mut Connection) -> bool { + false + } +} + +static POOL: Lazy>> = Lazy::new(|| { + let manager = SqliteManager; + let pool = r2d2::Pool::builder() + .max_size(8) + .build(manager) + .expect("Failed to create database connection pool"); + run_migrations(&pool).expect("Failed to run database migrations"); + Arc::new(pool) +}); + +pub fn pool() -> Arc> { + POOL.clone() +} + +pub fn with_db(f: F) -> Result +where + F: FnOnce(&Connection) -> Result, +{ + let conn = POOL.get().map_err(|e| StorageError::Pool(e.to_string()))?; + f(&conn) +} + +fn db_file_path(db_name: &str) -> String { let mut p = PathBuf::from(get_directory()); p.push(format!("{db_name}.sqlite3")); p.to_string_lossy().to_string() } -/// Open a sqlite connection to the named DB file (no initialization). -/// -/// Arguments: -/// - `db_name`: name of the DB (without extension). -pub fn open_connection(db_name: &str) -> Result { +fn run_migrations(pool: &r2d2::Pool) -> Result<(), StorageError> { + let conn = pool.get().map_err(|e| StorageError::Pool(e.to_string()))?; + let current_version: i64 = conn + .pragma_query_value(None, "user_version", |r| r.get(0)) + .unwrap_or(0); + + if current_version < 1 { + conn.execute_batch( + r#" + CREATE TABLE IF NOT EXISTS messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + storage_owner INTEGER NOT NULL, + external_user INTEGER NOT NULL, + message_time INTEGER NOT NULL, + content TEXT NOT NULL, + sent_by_self INTEGER NOT NULL, + message_state TEXT NOT NULL, + height INTEGER NOT NULL DEFAULT 0 + ); + CREATE INDEX IF NOT EXISTS idx_messages_lookup + ON messages (storage_owner, external_user, message_time DESC); + + CREATE TABLE IF NOT EXISTS contacts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + storage_owner INTEGER NOT NULL, + user_id INTEGER NOT NULL, + user_name TEXT, + last_message_at INTEGER, + UNIQUE(storage_owner, user_id) + ); + CREATE INDEX IF NOT EXISTS idx_contacts_owner + ON contacts (storage_owner, last_message_at DESC, user_id ASC); + + CREATE TABLE IF NOT EXISTS communities ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + storage_owner INTEGER NOT NULL, + address TEXT NOT NULL, + title TEXT NOT NULL, + position TEXT NOT NULL, + UNIQUE(storage_owner, address) + ); + CREATE INDEX IF NOT EXISTS idx_communities_owner + ON communities (storage_owner); + + CREATE TABLE IF NOT EXISTS users ( + user_id INTEGER PRIMARY KEY, + username TEXT NOT NULL UNIQUE, + public_key TEXT NOT NULL, + private_key_hash TEXT NOT NULL, + reset_token TEXT NOT NULL, + created_at INTEGER NOT NULL, + display_name TEXT + ); + + CREATE TABLE IF NOT EXISTS trusted_apps ( + user_id INTEGER NOT NULL, + app_id TEXT NOT NULL, + app_secret TEXT NOT NULL, + PRIMARY KEY (user_id, app_id) + ); + + PRAGMA user_version = 1; + "#, + )?; + } + + if current_version < 2 { + conn.execute_batch( + r#" + ALTER TABLE messages ADD COLUMN height INTEGER NOT NULL DEFAULT 0; + PRAGMA user_version = 2; + "#, + )?; + } + + if current_version < 3 { + conn.execute_batch( + r#" + ALTER TABLE messages ADD COLUMN reply_to INTEGER; + PRAGMA user_version = 3; + "#, + )?; + } + + if current_version < 4 { + conn.execute_batch( + r#" + ALTER TABLE messages ADD COLUMN edited_count INTEGER NOT NULL DEFAULT 0; + ALTER TABLE messages ADD COLUMN deleted_by_external INTEGER NOT NULL DEFAULT 0; + + CREATE TABLE IF NOT EXISTS message_edits ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + message_id INTEGER NOT NULL REFERENCES messages(id), + content_before TEXT NOT NULL, + content_after TEXT NOT NULL, + edited_at INTEGER NOT NULL, + edited_by INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_message_edits_msg + ON message_edits (message_id, edited_at DESC); + + CREATE TABLE IF NOT EXISTS reactions ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + message_id INTEGER NOT NULL REFERENCES messages(id), + user_id INTEGER NOT NULL, + reaction TEXT NOT NULL, + created_at INTEGER NOT NULL, + UNIQUE(message_id, user_id, reaction) + ); + CREATE INDEX IF NOT EXISTS idx_reactions_msg + ON reactions (message_id, reaction); + + PRAGMA user_version = 4; + "#, + )?; + } + + Ok(()) +} + +pub fn open_connection(db_name: &str) -> Result { let path = db_file_path(db_name); Connection::open(path) } -/// Open a connection and immediately run `init_sql` via `execute_batch`. -/// -/// Arguments: -/// - `db_name`: name of the DB (without extension). -/// - `init_sql`: SQL statements to initialize schema & PRAGMAs (can be multiple). -pub fn open_and_init(db_name: &str, init_sql: &str) -> Result { - let conn = open_connection(db_name)?; - conn.execute_batch(init_sql)?; - Ok(conn) -} - -/// Create a shared, Arc> initialized with the given SQL. -/// -/// This is a convenience wrapper that returns an owned Arc> -/// so caller modules can store it in a `static` or pass it around. -/// -/// Arguments: -/// - `db_name`: DB name (without extension). -/// - `init_sql`: init SQL (eg PRAGMA + CREATE TABLE statements). pub fn create_shared_connection( db_name: &str, init_sql: &str, -) -> Result>, String> { - match open_and_init(db_name, init_sql) { - Ok(conn) => { - // Configure some sensible defaults for concurrency - // Attempt to set a busy timeout to reduce SQLITE_BUSY failures. - let _ = conn.busy_timeout(Duration::from_millis(250)); - Ok(Arc::new(Mutex::new(conn))) - } - Err(e) => Err(format!("Failed to open/init DB '{}': {}", db_name, e)), - } +) -> Result>, String> { + let path = db_file_path(db_name); + let conn = Connection::open(path).map_err(|e| e.to_string())?; + conn.execute_batch(init_sql).map_err(|e| e.to_string())?; + let _ = conn.busy_timeout(Duration::from_millis(250)); + Ok(Arc::new(std::sync::Mutex::new(conn))) } -/// Acquire the Connection from an Arc> and run the provided -/// closure. Converts rusqlite::Error into a String on error. -/// -/// Arguments: -/// - `shared`: Arc> -/// - `f`: closure that receives &Connection and returns Result -/// -/// Returns Ok(T) or Err(String). -pub fn with_conn(shared: &Arc>, f: F) -> Result +pub fn with_conn(shared: &Arc>, f: F) -> Result where - F: FnOnce(&Connection) -> Result, + F: FnOnce(&Connection) -> Result, { - // When invoked from within an async runtime (such as Tokio), taking a blocking - // std::sync::Mutex lock on the runtime thread can cause deadlocks or permanent - // awaits. Detect whether we're running inside a Tokio runtime and, if so, - // execute the blocking lock + database closure using Tokio's blocking helper. - // - // The blocking section returns Result so we can propagate errors - // in the same form as before. if tokio::runtime::Handle::try_current().is_ok() { tokio::task::block_in_place(|| { let guard = shared @@ -98,97 +217,7 @@ where } } -/// Initialize a general-purpose messages+contacts DB and return a shared -/// connection. This helper creates a single DB file that can contain multiple -/// tables (messages, contacts, ...). The SQL here is conservative and intended -/// to be safe if called multiple times. -/// -/// Callers may prefer to call `create_shared_connection("messages", INIT_SQL)` -/// directly, but this convenience is useful for code that expects both tables. -pub fn create_general_messages_db() -> Result>, String> { - // Keep PRAGMA and schema in one multi-statement string so callers only - // need to call a single execute_batch. - const INIT_SQL: &str = r#" - PRAGMA journal_mode = WAL; - PRAGMA synchronous = NORMAL; - - CREATE TABLE IF NOT EXISTS messages ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - storage_owner INTEGER NOT NULL, - external_user INTEGER NOT NULL, - message_time INTEGER NOT NULL, - content TEXT NOT NULL, - sent_by_self INTEGER NOT NULL, - message_state TEXT NOT NULL, - height INTEGER NOT NULL DEFAULT 0 - ); - - CREATE INDEX IF NOT EXISTS idx_messages_lookup - ON messages (storage_owner, external_user, message_time DESC); - - CREATE TABLE IF NOT EXISTS contacts ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - storage_owner INTEGER NOT NULL, - user_id INTEGER NOT NULL, - user_name TEXT, - last_message_at INTEGER, - UNIQUE(storage_owner, user_id) - ); - - CREATE INDEX IF NOT EXISTS idx_contacts_owner - ON contacts (storage_owner, last_message_at DESC, user_id ASC); - - CREATE TABLE IF NOT EXISTS communities ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - storage_owner INTEGER NOT NULL, - address TEXT NOT NULL, - title TEXT NOT NULL, - position TEXT NOT NULL, - UNIQUE(storage_owner, address) - ); - - CREATE INDEX IF NOT EXISTS idx_communities_owner - ON communities (storage_owner); - "#; - - match create_shared_connection("messages", INIT_SQL) { - Ok(shared_conn) => { - // Attempt to add the height column for backwards compatibility. - // This will fail if the column already exists, which is expected. - let _ = with_conn(&shared_conn, |conn| { - let _ = conn.execute( - "ALTER TABLE messages ADD COLUMN height INTEGER NOT NULL DEFAULT 0", - [], - ); - Ok(()) - }); - // Attempt to add the reply_to column for backwards compatibility. - // This will fail if the column already exists, which is expected. - let _ = with_conn(&shared_conn, |conn| { - let _ = conn.execute( - "ALTER TABLE messages ADD COLUMN reply_to INTEGER", - [], - ); - Ok(()) - }); - Ok(shared_conn) - } - Err(e) => Err(e), - } +/// Legacy - kept for e2ee_storage which uses its own DB. +pub fn create_general_messages_db() -> Result>, String> { + create_shared_connection(DB_NAME, "") } - -/* -Example usage: - -// In some util module (at init time, e.g. lazy_static or LazyLock) -static MESSAGES_DB: LazyLock>> = LazyLock::new(|| { - create_general_messages_db().expect("failed to create messages DB") -}); - -// Later, to run a query: -let res: Result, String> = with_conn(&MESSAGES_DB, |conn| { - let mut stmt = conn.prepare("SELECT ...")?; - let rows = stmt.query_map(...)?; - // collect and return Ok(...) -}); -*/ diff --git a/omikron-connector/src/omikron_connection.rs b/omikron-connector/src/omikron_connection.rs index 19ed8da..984c9dd 100755 --- a/omikron-connector/src/omikron_connection.rs +++ b/omikron-connector/src/omikron_connection.rs @@ -16,11 +16,11 @@ use mtp::client::{Client, ClientConfig, Policy, Receiver, SendMode, Sender}; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use mtp::crypto::{Keyring, PublicKeyBundle}; use mtp::type_map::TypeMap; -use std::collections::HashMap; use std::env; +use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, LazyLock}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use tokio::sync::{Mutex, RwLock, mpsc, watch}; +use tokio::sync::{Mutex, RwLock, oneshot, watch, Semaphore}; use tokio::task::JoinHandle; use tokio::time::sleep; use uuid::Uuid; @@ -235,13 +235,15 @@ const CONNECTION_TIMEOUT: Duration = Duration::from_secs(10); const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); const TASK_CLEANUP_INTERVAL: Duration = Duration::from_secs(60); const TASK_MAX_AGE: Duration = Duration::from_secs(60); +const MAX_MISSED_PONGS: u32 = 3; +const MAX_CONCURRENT_HANDLERS: usize = 20; // ============================================================================ // Waiting Task System // ============================================================================ pub struct WaitingTask { - pub task: Box, CommunicationValue) -> bool + Send + Sync>, + pub task: Box, CommunicationValue) -> bool + Send + Sync>, pub inserted_at: Instant, } @@ -284,6 +286,7 @@ impl ConnectionState { #[allow(dead_code)] // message_send_times is unused. pub struct OmikronConnection { state: Arc>, + state_watch_tx: watch::Sender, sender: Arc>>>, connection_loop_handle: Arc>>>, pub last_ping: Arc>, @@ -292,16 +295,20 @@ pub struct OmikronConnection { shutdown_tx: Arc>>>, reconnect_on_close: Arc>, auth_failure: Arc>>, - pub app_challenges: Arc>>, - pub app_sessions: Arc>>, + pub app_challenges: Arc>, + pub app_sessions: Arc>, + pub(crate) missed_pongs: Arc, + handler_semaphore: Arc, } impl OmikronConnection { pub fn new() -> Self { let (shutdown_tx, _) = watch::channel(false); + let (state_watch_tx, _) = watch::channel(ConnectionState::Disconnected); OmikronConnection { state: Arc::new(RwLock::new(ConnectionState::Disconnected)), + state_watch_tx, sender: Arc::new(RwLock::new(None)), connection_loop_handle: Arc::new(Mutex::new(None)), last_ping: Arc::new(Mutex::new(-1)), @@ -310,11 +317,18 @@ impl OmikronConnection { 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())), + app_challenges: Arc::new(DashMap::new()), + app_sessions: Arc::new(DashMap::new()), + missed_pongs: Arc::new(AtomicU32::new(0)), + handler_semaphore: Arc::new(Semaphore::new(MAX_CONCURRENT_HANDLERS)), } } + async fn set_state(&self, new_state: ConnectionState) { + *self.state.write().await = new_state; + let _ = self.state_watch_tx.send(new_state); + } + // ------------------------------------------------------------------------- // Connection Management // ------------------------------------------------------------------------- @@ -364,7 +378,7 @@ impl OmikronConnection { sender.close(); } - *self.state.write().await = ConnectionState::Disconnected; + self.set_state(ConnectionState::Disconnected).await; *self.sender.write().await = None; } @@ -418,7 +432,7 @@ impl OmikronConnection { } async fn connect_once(self: Arc) -> Result<(), String> { - *self.state.write().await = ConnectionState::Connecting; + self.set_state(ConnectionState::Connecting).await; log_t!("omikron_connecting"); let keyring = self.load_or_migrate_keyring().await; @@ -467,7 +481,7 @@ impl OmikronConnection { ); *self.reconnect_on_close.write().await = false; *self.auth_failure.write().await = Some(reason.clone()); - *self.state.write().await = ConnectionState::Disconnected; + self.set_state(ConnectionState::Disconnected).await; return Err(reason); } Err(e) => return Err(format!("Connection failed: {}", e)), @@ -482,7 +496,7 @@ impl OmikronConnection { let sender_arc = Arc::new(connection.sender); *self.sender.write().await = Some(sender_arc.clone()); - *self.state.write().await = ConnectionState::Connected { identified: true }; + self.set_state(ConnectionState::Connected { identified: true }).await; // Start read loop let mut receiver = connection.receiver; @@ -507,7 +521,7 @@ impl OmikronConnection { // Wait for read loop to complete let result = read_handle.await; *self.sender.write().await = None; - *self.state.write().await = ConnectionState::Disconnected; + self.set_state(ConnectionState::Disconnected).await; { ACTIVE_TASKS.remove("Omikron Listener"); } @@ -547,7 +561,15 @@ impl OmikronConnection { let legacy = CONFIG.load().keyring.clone(); let keyring = legacy .and_then(|b64| keyring_from_base64(&b64)) - .unwrap_or_else(crypto_helper::generate_keyring); + .unwrap_or_else(|| { + log!( + "WARNING: No existing keyring found. Neither {} nor config.json \ + contain a keyring; generating a new identity. If you already had \ + an Iota identity, restore {} from a backup to avoid losing access.", + IOTA_KEYRING_PATH, IOTA_KEYRING_PATH + ); + crypto_helper::generate_keyring() + }); if let Err(e) = mtp::files::save_keyring(&keyring, IOTA_KEYRING_PATH) { log!("Failed to persist {}: {}", IOTA_KEYRING_PATH, e); @@ -672,7 +694,23 @@ impl OmikronConnection { let result = receiver.receive().await; match result { Ok(cv) => { - self.clone().handle_message(cv).await; + let msg_id = cv.get_id(); + if let Some((_, task)) = WAITING_TASKS.remove(&msg_id) { + if (task.task)(self.clone(), cv.clone()) { + continue; + } + } + if cv.is_type(CommunicationType::Pong) { + self.handle_pong(&cv).await; + continue; + } + + let permit = self.handler_semaphore.clone().acquire_owned().await; + let self_clone = self.clone(); + tokio::spawn(async move { + let _permit = permit; + self_clone.handle_message_impl(cv).await; + }); } Err(e) => { self.fail_all_waiting_tasks(format!( @@ -710,6 +748,17 @@ impl OmikronConnection { break; } + if self.missed_pongs.load(Ordering::Relaxed) > MAX_MISSED_PONGS { + log!( + "Connection appears dead ({} consecutive missed pongs), closing sender", + self.missed_pongs.load(Ordering::Relaxed) + ); + if let Some(sender) = self.sender.read().await.as_ref() { + sender.close(); + } + break; + } + self.flush_pending_chat_secret_forwards().await; self.send_ping().await; } @@ -754,8 +803,109 @@ impl OmikronConnection { } } + async fn forward_message_live( + &self, + message_id: u32, + receiver_id: u64, + sender_id: i64, + timestamp: i64, + content: &str, + height: i64, + reply_to: Option, + ) -> Option { + let mut msg_fields = vec![ + (DataType::Content, DataValue::Str(content.to_string())), + (DataType::SendTime, DataValue::SignedNumber(timestamp as i128)), + (DataType::Height, DataValue::SignedNumber(height as i128)), + ]; + if let Some(rt) = reply_to { + msg_fields.push((DataType::ReplyId, DataValue::UnsignedNumber(rt as u64 as u128))); + } + + let user_forward = CommunicationValue::new(CommunicationType::MessageLive) + .with_id(message_id) + .with_receiver(receiver_id) + .add_typed_default(DataType::SenderId, DataValue::SignedNumber(sender_id as i128)) + .add_typed_default(DataType::Message, typed_container(msg_fields)); + + match self.await_response(&user_forward, Some(Duration::from_secs(3))).await { + Ok(user_resp) => { + let ms_raw = user_resp + .get_data(DataType::MessageState) + .as_string() + .unwrap_or_else(|| "".to_string()); + Some(MessageState::from_str(&ms_raw).upgrade(MessageState::Received)) + } + Err(_) => None, + } + } + + async fn forward_to_remote_iota( + &self, + cv: &CommunicationValue, + sender_id: i64, + receiver_id: i64, + timestamp: i64, + content: &str, + height: i64, + reply_to: Option, + ) { + let mut fw_msg = CommunicationValue::new(CommunicationType::MessageOtherIota) + .with_id(cv.get_id()) + .with_receiver(receiver_id as u64) + .with_sender(sender_id as u64) + .add_typed_default(DataType::Height, DataValue::SignedNumber(height as i128)) + .add_typed_default(DataType::Content, DataValue::Str(content.to_string())) + .add_typed_default(DataType::SendTime, DataValue::SignedNumber(timestamp as i128)); + if let Some(rt) = reply_to { + fw_msg = fw_msg.add_typed_default( + DataType::ReplyId, + DataValue::UnsignedNumber(rt as u64 as u128), + ); + } + + match self.await_response(&fw_msg, Some(Duration::from_secs(10))).await { + Ok(resp) => { + let ms_raw = resp + .get_data(DataType::MessageState) + .as_string() + .unwrap_or_else(|| "".to_string()); + let ms = MessageState::from_str(&ms_raw).upgrade(MessageState::Received); + + let _ = chat_files::change_message_state( + timestamp, sender_id, receiver_id, ms.clone(), + ); + + let _ = self.send_message( + &CommunicationValue::new(CommunicationType::MessageState) + .with_id(cv.get_id()) + .with_receiver(sender_id as u64) + .with_sender(receiver_id as u64) + .add_typed_default(DataType::ChatPartnerId, DataValue::SignedNumber(receiver_id as i128)) + .add_typed_default(DataType::SendTime, DataValue::SignedNumber(timestamp as i128)) + .add_typed_default(DataType::MessageState, DataValue::Str(ms.as_str().to_string())), + ).await; + } + Err(_) => { + let _ = chat_files::change_message_state( + timestamp, sender_id, receiver_id, MessageState::Sent, + ); + + let _ = self.send_message( + &CommunicationValue::new(CommunicationType::MessageState) + .with_id(cv.get_id()) + .with_receiver(sender_id as u64) + .with_sender(receiver_id as u64) + .add_typed_default(DataType::ChatPartnerId, DataValue::SignedNumber(receiver_id as i128)) + .add_typed_default(DataType::SendTime, DataValue::SignedNumber(timestamp as i128)) + .add_typed_default(DataType::MessageState, DataValue::Str(MessageState::Sent.as_str().to_string())), + ).await; + } + } + } + // ------------------------------------------------------------------------- - // Message Handling (Preserved from original) + // Message Handling — Dispatch // ------------------------------------------------------------------------- pub async fn handle_message(self: Arc, cv: CommunicationValue) { @@ -765,7 +915,6 @@ impl OmikronConnection { let msg_id = cv.get_id(); - // Dispatch waiting task for this message id if let Some((_, task)) = WAITING_TASKS.remove(&msg_id) { if (task.task)(self.clone(), cv.clone()) { return; @@ -777,677 +926,579 @@ impl OmikronConnection { return; } - if cv.is_type(CommunicationType::SetChatSecret) { - let sender_id = cv.get_sender().to_string(); - let recipients = match chat_secret_recipients(&cv) { - Some(recipients) => recipients, - None => { - self.send_message(&error_response(&cv, CommunicationType::ErrorInvalidData)) - .await; + self.clone().handle_message_impl(cv).await; + } + + async fn handle_message_impl(self: Arc, cv: CommunicationValue) { + macro_rules! dispatch { + ($ty:ident, $method:ident) => { + if cv.is_type(CommunicationType::$ty) { + self.clone().$method(&cv).await; return; } }; - let now = now_millis_i64(); - let chat_id = data_string(&cv, DataType::ChatId); - let secret_id = data_string(&cv, DataType::SecretId); - let version = data_i64(&cv, DataType::VersionNumber); - let wrapping_scheme = data_string(&cv, DataType::WrappingScheme); - let created_at = data_i64(&cv, DataType::CreatedAt).unwrap_or(now); - - let Some((((chat_id, secret_id), version), wrapping_scheme)) = - chat_id.zip(secret_id).zip(version).zip(wrapping_scheme) - else { - self.send_message(&error_response(&cv, CommunicationType::ErrorInvalidData)) - .await; - return; - }; - - for recipient in &recipients { - let recipient_id = recipient.user_id.parse::().unwrap_or(0); - let is_local = iota_storage::users::user_manager::get_user(recipient_id).is_some(); - - if is_local { - if e2ee_storage::put_chat_secret(StoredChatSecret { - user_id: recipient.user_id.clone(), - chat_id: chat_id.clone(), - secret_id: secret_id.clone(), - version, - encrypted_secret: recipient.encrypted_secret.clone(), - kem_ciphertext: recipient.kem_ciphertext.clone(), - wrapping_scheme: wrapping_scheme.clone(), - created_at, - updated_at: now, - }) - .is_err() - { - self.send_message(&error_response( - &cv, - CommunicationType::ErrorInvalidData, - )) - .await; - return; - } - continue; - } - - if recipient.user_id != sender_id { - let forward = set_chat_secret_cv_for_recipient(&cv, recipient); - if !self.forward_chat_secret(&forward).await { - self.store_pending_chat_secret_forward(&forward).await; - } - } - } - - self.send_message(&error_response(&cv, CommunicationType::Success)) - .await; - return; } - if cv.is_type(CommunicationType::GetChatSecret) { - let Some(user_id) = data_string(&cv, DataType::UserId) else { - self.send_message(&error_response(&cv, CommunicationType::ErrorInvalidData)) - .await; - return; - }; - if user_id != cv.get_sender().to_string() { - self.send_message(&error_response(&cv, CommunicationType::ErrorNotFound)) - .await; - return; - } - let Some(chat_id) = data_string(&cv, DataType::ChatId) else { - self.send_message(&error_response(&cv, CommunicationType::ErrorInvalidData)) - .await; - return; - }; + dispatch!(SetChatSecret, handle_set_chat_secret); + dispatch!(GetChatSecret, handle_get_chat_secret); + dispatch!(ChatSecretForward, handle_chat_secret_forward); + dispatch!(AppIdentification, handle_app_identification); + dispatch!(AppChallengeResponse, handle_app_challenge_response); + dispatch!(SaveAppData, handle_save_app_data); + dispatch!(LoadAppData, handle_load_app_data); + dispatch!(CreateApp, handle_create_app); + dispatch!(DeleteApp, handle_delete_app); + dispatch!(ClientConnected, handle_client_connected); + dispatch!(MessageState, handle_message_state); + dispatch!(MessageSend, handle_message_send); + dispatch!(MessageOtherIota, handle_message_other_iota); + dispatch!(MessagesGet, handle_messages_get); + dispatch!(GetChats, handle_get_chats); + dispatch!(AddConversation, handle_add_conversation); + dispatch!(AddCommunity, handle_add_community); + dispatch!(GetCommunities, handle_get_communities); + dispatch!(RemoveCommunity, handle_remove_community); + dispatch!(GlobalSettingsSave, handle_global_settings_save); + dispatch!(GlobalSettingsLoad, handle_global_settings_load); + dispatch!(SettingsSave, handle_settings_save); + dispatch!(SettingsLoad, handle_settings_load); + dispatch!(SettingsList, handle_settings_list); + } - match e2ee_storage::get_chat_secret(ChatSecretQuery { - user_id, - chat_id, - secret_id: data_string(&cv, DataType::SecretId), - }) { - Ok(Some(record)) => { - let response = CommunicationValue::new(CommunicationType::ChatSecretResponse) - .with_id(cv.get_id()) - .with_receiver(cv.get_sender()) - .add_typed_default(DataType::UserId, DataValue::Str(record.user_id)) - .add_typed_default(DataType::ChatId, DataValue::Str(record.chat_id)) - .add_typed_default(DataType::SecretId, DataValue::Str(record.secret_id)) - .add_typed_default( - DataType::VersionNumber, - DataValue::SignedNumber(record.version as i128), - ) - .add_typed_default( - DataType::EncryptedSecret, - DataValue::Bytes(record.encrypted_secret), - ) - .add_typed_default( - DataType::KemCiphertext, - DataValue::Bytes(record.kem_ciphertext), - ) - .add_typed_default( - DataType::WrappingScheme, - DataValue::Str(record.wrapping_scheme), - ) - .add_typed_default( - DataType::CreatedAt, - DataValue::SignedNumber(record.created_at as i128), - ) - .add_typed_default( - DataType::UpdatedAt, - DataValue::SignedNumber(record.updated_at as i128), - ); - self.send_message(&response).await; - } - Ok(None) => { - self.send_message(&error_response(&cv, CommunicationType::ErrorNotSet)) - .await - } - Err(_) => { - self.send_message(&error_response(&cv, CommunicationType::ErrorInvalidData)) - .await - } + // ------------------------------------------------------------------------- + // Message Handlers + // ------------------------------------------------------------------------- + + async fn handle_set_chat_secret(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender().to_string(); + let recipients = match chat_secret_recipients(cv) { + Some(recipients) => recipients, + None => { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorInvalidData)).await; + return; } + }; + let now = now_millis_i64(); + let chat_id = data_string(cv, DataType::ChatId); + let secret_id = data_string(cv, DataType::SecretId); + let version = data_i64(cv, DataType::VersionNumber); + let wrapping_scheme = data_string(cv, DataType::WrappingScheme); + let created_at = data_i64(cv, DataType::CreatedAt).unwrap_or(now); + + let Some((((chat_id, secret_id), version), wrapping_scheme)) = + chat_id.zip(secret_id).zip(version).zip(wrapping_scheme) + else { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorInvalidData)).await; return; - } + }; - if cv.is_type(CommunicationType::ChatSecretForward) { - let sender_id = cv.get_sender().to_string(); - let recipient_user_id = data_string(&cv, DataType::RecipientUserId).unwrap_or_default(); - if data_string(&cv, DataType::SenderUserId).as_deref() != Some(sender_id.as_str()) - || recipient_user_id.is_empty() - || pending_chat_secret_forward_from_cv(&cv).is_none() - { - self.send_message(&error_response(&cv, CommunicationType::ErrorInvalidData)) - .await; - return; - } + let mut non_local_forwards: Vec = Vec::new(); - let forward = cv - .clone() - .with_receiver(recipient_user_id.parse::().unwrap_or(0)); - if self.forward_chat_secret(&forward).await { - self.send_message(&error_response(&cv, CommunicationType::Success)) - .await; - } else { - self.store_pending_chat_secret_forward(&cv).await; - self.send_message(&error_response(&cv, CommunicationType::Success)) - .await; - } - return; - } + for recipient in &recipients { + let recipient_id = recipient.user_id.parse::().unwrap_or(0); + let is_local = iota_storage::users::user_manager::get_user(recipient_id).is_some(); - if cv.is_type(CommunicationType::AppIdentification) { - let sender_id = cv.get_sender(); - let app_identifier = cv - .get_data(DataType::AppIdentifier) - .as_str() - .unwrap_or("") - .to_string(); - let app_public_key = cv - .get_data(DataType::AppPublicKey) - .as_str() - .unwrap_or("") - .to_string(); - let user_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0) as i64; - - let mut trusted = false; - if let Some(user) = iota_storage::users::user_manager::get_user(user_id) { - if let Some(pub_k) = user.trusted_apps.get(&app_identifier) { - if pub_k == &app_public_key { - trusted = true; - } - } - } - - if trusted { - let challenge = Uuid::new_v4().to_string(); - - self.app_challenges - .write() - .await - .insert(sender_id, challenge.clone()); - self.app_sessions - .write() - .await - .insert(sender_id, (user_id, app_identifier.clone())); - - if let Some(app_pub_bundle) = - iota_util::crypto_helper::public_key_bundle_from_base64(&app_public_key) + if is_local { + if e2ee_storage::put_chat_secret(StoredChatSecret { + user_id: recipient.user_id.clone(), + chat_id: chat_id.clone(), + secret_id: secret_id.clone(), + version, + encrypted_secret: recipient.encrypted_secret.clone(), + kem_ciphertext: recipient.kem_ciphertext.clone(), + wrapping_scheme: wrapping_scheme.clone(), + created_at, + updated_at: now, + }) + .is_err() { - let kr_str = CONFIG.load().keyring.clone().unwrap_or_default(); - - if let Some(keyring) = keyring_from_base64(&kr_str) { - if let Ok(encrypted_challenge) = - crypto_util::encrypt_challenge(&challenge, &app_pub_bundle) - { - let bundle = keyring.public_key_bundle(); - let pub_k_b64 = crypto_helper::public_key_bundle_to_base64(&bundle); - - let res = CommunicationValue::new(CommunicationType::AppChallenge) - .with_id(cv.get_id()) - .with_receiver(sender_id) - .add_typed_default(DataType::PublicKey, DataValue::Str(pub_k_b64)) - .add_typed_default( - DataType::Challenge, - DataValue::Str(encrypted_challenge), - ); - - self.send_message(&res).await; - return; - } - } + let _ = self.send_message(&error_response( + cv, + CommunicationType::ErrorInvalidData, + )).await; + return; } + continue; } - let res = CommunicationValue::new(CommunicationType::ErrorInvalidChallenge) - .with_id(cv.get_id()) - .with_receiver(sender_id); - self.send_message(&res).await; + if recipient.user_id != sender_id { + non_local_forwards + .push(set_chat_secret_cv_for_recipient(cv, recipient)); + } + } + + if !non_local_forwards.is_empty() { + let mut handles = Vec::new(); + for forward in &non_local_forwards { + let self_clone = self.clone(); + let fwd = forward.clone(); + handles.push(tokio::spawn(async move { + self_clone.forward_chat_secret(&fwd).await + })); + } + + for (forward, handle) in non_local_forwards.into_iter().zip(handles) { + match handle.await { + Ok(true) => {} + _ => self.store_pending_chat_secret_forward(&forward).await, + } + } + } + + let _ = self.send_message(&error_response(cv, CommunicationType::Success)).await; + } + + async fn handle_get_chat_secret(self: Arc, cv: &CommunicationValue) { + let Some(user_id) = data_string(cv, DataType::UserId) else { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorInvalidData)).await; + return; + }; + if user_id != cv.get_sender().to_string() { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorNotFound)).await; + return; + } + let Some(chat_id) = data_string(cv, DataType::ChatId) else { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorInvalidData)).await; + return; + }; + + match e2ee_storage::get_chat_secret(ChatSecretQuery { + user_id, + chat_id, + secret_id: data_string(cv, DataType::SecretId), + }) { + Ok(Some(record)) => { + let response = CommunicationValue::new(CommunicationType::ChatSecretResponse) + .with_id(cv.get_id()) + .with_receiver(cv.get_sender()) + .add_typed_default(DataType::UserId, DataValue::Str(record.user_id)) + .add_typed_default(DataType::ChatId, DataValue::Str(record.chat_id)) + .add_typed_default(DataType::SecretId, DataValue::Str(record.secret_id)) + .add_typed_default( + DataType::VersionNumber, + DataValue::SignedNumber(record.version as i128), + ) + .add_typed_default( + DataType::EncryptedSecret, + DataValue::Bytes(record.encrypted_secret), + ) + .add_typed_default( + DataType::KemCiphertext, + DataValue::Bytes(record.kem_ciphertext), + ) + .add_typed_default( + DataType::WrappingScheme, + DataValue::Str(record.wrapping_scheme), + ) + .add_typed_default( + DataType::CreatedAt, + DataValue::SignedNumber(record.created_at as i128), + ) + .add_typed_default( + DataType::UpdatedAt, + DataValue::SignedNumber(record.updated_at as i128), + ); + let _ = self.send_message(&response).await; + } + Ok(None) => { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorNotSet)).await; + } + Err(_) => { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorInvalidData)).await; + } + } + } + + async fn handle_chat_secret_forward(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender().to_string(); + let recipient_user_id = data_string(cv, DataType::RecipientUserId).unwrap_or_default(); + if data_string(cv, DataType::SenderUserId).as_deref() != Some(sender_id.as_str()) + || recipient_user_id.is_empty() + || pending_chat_secret_forward_from_cv(cv).is_none() + { + let _ = self.send_message(&error_response(cv, CommunicationType::ErrorInvalidData)).await; return; } - if cv.is_type(CommunicationType::AppChallengeResponse) { - let sender_id = cv.get_sender(); - let mut challenges = self.app_challenges.write().await; - if let Some(expected) = challenges.remove(&sender_id) { - if let DataValue::Str(response) = cv.get_data(DataType::Challenge) { - if expected == *response { - let res = - CommunicationValue::new(CommunicationType::AppIdentificationResponse) - .with_id(cv.get_id()) - .with_receiver(sender_id); - self.send_message(&res).await; + let forward = cv + .clone() + .with_receiver(recipient_user_id.parse::().unwrap_or(0)); + if self.forward_chat_secret(&forward).await { + let _ = self.send_message(&error_response(cv, CommunicationType::Success)).await; + } else { + self.store_pending_chat_secret_forward(cv).await; + let _ = self.send_message(&error_response(cv, CommunicationType::Success)).await; + } + } + + async fn handle_app_identification(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender(); + let app_identifier = cv + .get_data(DataType::AppIdentifier) + .as_str() + .unwrap_or("") + .to_string(); + let app_public_key = cv + .get_data(DataType::AppPublicKey) + .as_str() + .unwrap_or("") + .to_string(); + let user_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0) as i64; + + let mut trusted = false; + if let Some(user) = iota_storage::users::user_manager::get_user(user_id) { + if let Some(pub_k) = user.trusted_apps.get(&app_identifier) { + if pub_k == &app_public_key { + trusted = true; + } + } + } + + if trusted { + let challenge = Uuid::new_v4().to_string(); + + self.app_challenges.insert(sender_id, challenge.clone()); + self.app_sessions.insert(sender_id, (user_id, app_identifier.clone())); + + if let Some(app_pub_bundle) = + iota_util::crypto_helper::public_key_bundle_from_base64(&app_public_key) + { + let kr_str = CONFIG.load().keyring.clone().unwrap_or_default(); + + if let Some(keyring) = keyring_from_base64(&kr_str) { + if let Ok(encrypted_challenge) = + crypto_util::encrypt_challenge(&challenge, &app_pub_bundle) + { + let bundle = keyring.public_key_bundle(); + let pub_k_b64 = crypto_helper::public_key_bundle_to_base64(&bundle); + + let res = CommunicationValue::new(CommunicationType::AppChallenge) + .with_id(cv.get_id()) + .with_receiver(sender_id) + .add_typed_default(DataType::PublicKey, DataValue::Str(pub_k_b64)) + .add_typed_default( + DataType::Challenge, + DataValue::Str(encrypted_challenge), + ); + + let _ = self.send_message(&res).await; return; } } } - let res = CommunicationValue::new(CommunicationType::ErrorInvalidChallenge) - .with_id(cv.get_id()) - .with_receiver(sender_id); - self.send_message(&res).await; - return; } - if cv.is_type(CommunicationType::SaveAppData) { - let sender_id = cv.get_sender(); - let app_data = cv - .get_data(DataType::AppData) - .as_str() - .unwrap_or("") - .to_string(); + let res = CommunicationValue::new(CommunicationType::ErrorInvalidChallenge) + .with_id(cv.get_id()) + .with_receiver(sender_id); + let _ = self.send_message(&res).await; + } - let sessions = self.app_sessions.read().await; - if let Some((user_id, app_identifier)) = sessions.get(&sender_id) { - iota_storage::users::user_manager::save_app_data( - *user_id, - app_identifier, - &app_data, - ); - } - - let res = CommunicationValue::new(CommunicationType::SaveAppData) - .with_id(cv.get_id()) - .with_receiver(sender_id); - self.send_message(&res).await; - return; - } - - if cv.is_type(CommunicationType::LoadAppData) { - let sender_id = cv.get_sender(); - let mut app_data = String::new(); - - let sessions = self.app_sessions.read().await; - if let Some((user_id, app_identifier)) = sessions.get(&sender_id) { - app_data = - iota_storage::users::user_manager::load_app_data(*user_id, app_identifier); - } - - let res = CommunicationValue::new(CommunicationType::LoadAppData) - .with_id(cv.get_id()) - .with_receiver(sender_id) - .add_typed_default(DataType::AppData, DataValue::Str(app_data)); - self.send_message(&res).await; - return; - } - - if cv.is_type(CommunicationType::CreateApp) { - let sender_id = cv.get_sender() as i64; - let app_identifier = cv - .get_data(DataType::AppIdentifier) - .as_str() - .unwrap_or("") - .to_string(); - let app_public_key = cv - .get_data(DataType::AppPublicKey) - .as_str() - .unwrap_or("") - .to_string(); - - if !app_identifier.is_empty() && !app_public_key.is_empty() { - if let Some(mut user) = iota_storage::users::user_manager::get_user(sender_id) { - if !user.trusted_apps.contains_key(&app_identifier) { - user.trusted_apps.insert(app_identifier, app_public_key); - iota_storage::users::user_manager::update_user(user); - } + async fn handle_app_challenge_response(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender(); + if let Some((_, expected_challenge)) = self.app_challenges.remove(&sender_id) { + if let DataValue::Str(response) = cv.get_data(DataType::Challenge) { + if expected_challenge == *response { + let res = + CommunicationValue::new(CommunicationType::AppIdentificationResponse) + .with_id(cv.get_id()) + .with_receiver(sender_id); + let _ = self.send_message(&res).await; + return; } } - - let res = CommunicationValue::new(CommunicationType::CreateApp) - .with_id(cv.get_id()) - .with_receiver(sender_id as u64); - self.send_message(&res).await; - return; } + let res = CommunicationValue::new(CommunicationType::ErrorInvalidChallenge) + .with_id(cv.get_id()) + .with_receiver(sender_id); + let _ = self.send_message(&res).await; + } - if cv.is_type(CommunicationType::DeleteApp) { - let sender_id = cv.get_sender() as i64; - let app_identifier = cv - .get_data(DataType::AppIdentifier) - .as_str() - .unwrap_or("") - .to_string(); + async fn handle_save_app_data(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender(); + let app_data = cv + .get_data(DataType::AppData) + .as_str() + .unwrap_or("") + .to_string(); - if !app_identifier.is_empty() { - if let Some(mut user) = iota_storage::users::user_manager::get_user(sender_id) { - if user.trusted_apps.contains_key(&app_identifier) { - user.trusted_apps.remove(&app_identifier); - iota_storage::users::user_manager::update_user(user); - } - } - } - - let res = CommunicationValue::new(CommunicationType::DeleteApp) - .with_id(cv.get_id()) - .with_receiver(sender_id as u64); - self.send_message(&res).await; - return; - } - - if cv.is_type(CommunicationType::ClientConnected) { - let user_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0) as i64; - let _session_id = cv.get_data(DataType::SessionId).as_number().unwrap_or(0) as i64; - - let contacts = chats_util::get_users(user_id); - let mut contacts_array = Vec::new(); - - for (i, contact) in contacts.iter().enumerate() { - let mut contact_container = Vec::new(); - contact_container.push(( - DataType::UserId, - DataValue::SignedNumber(contact.user_id as i128), - )); - contact_container.push(( - DataType::LastMessageAt, - DataValue::SignedNumber(contact.last_message_at.unwrap_or(0) as i128), - )); - - if let Some(ref name) = contact.user_name { - contact_container.push((DataType::Username, DataValue::Str(name.clone()))); - } - - let amount = if i < 10 { 20 } else { 1 }; - let messages = chat_files::get_messages(user_id, contact.user_id, 0, amount); - - let mut msg_array = Vec::new(); - for m in messages.members() { - let message_time = m["message_time"].as_i64().unwrap_or(0); - let content = m["content"].as_str().unwrap_or("").to_string(); - let sent_by_self = m["sent_by_self"].as_bool().unwrap_or(false); - let height = m["height"].as_i64().unwrap_or(0); - let message_state = m["message_state"].as_str().unwrap_or("").to_string(); - - let mut msg_container = Vec::new(); - msg_container.push(( - DataType::SendTime, - DataValue::SignedNumber(message_time as i128), - )); - msg_container.push((DataType::Content, DataValue::Str(content.clone()))); - msg_container.push((DataType::MessageState, DataValue::Str(message_state))); - msg_container.push((DataType::Height, DataValue::SignedNumber(height as i128))); - msg_container.push(( - DataType::SenderId, - DataValue::UnsignedNumber(if sent_by_self { - user_id as u128 - } else { - contact.user_id as u128 - }), - )); - msg_array.push(typed_container(msg_container)); - - if msg_array.len() == 1 { - let sender_id = if sent_by_self { - user_id - } else { - contact.user_id - }; - let mut last_msg = Vec::new(); - last_msg.push((DataType::Content, DataValue::Str(content))); - last_msg.push(( - DataType::SenderId, - DataValue::SignedNumber(sender_id as i128), - )); - contact_container.push((DataType::LastMessage, typed_container(last_msg))); - } - } - contact_container.push((DataType::Messages, DataValue::Array(msg_array))); - contacts_array.push(typed_container(contact_container)); - } - - let resp = CommunicationValue::new(CommunicationType::ClientConnected) - .with_id(cv.get_id()) - .add_typed_default(DataType::Contacts, DataValue::Array(contacts_array)); - self.send_message(&resp).await; - return; - } - - // ************************************************ // - // Direct messages // - // ************************************************ // - - if cv.is_type(CommunicationType::MessageState) { - let sender_id = &cv.get_sender(); - let receiver_id = match cv.get_data(DataType::ChatPartnerId).as_number() { - Some(id) => id, - _ => return, - }; - - // Parse send_time robustly: accept numeric or string, fallback to current time - let send_time_val = cv.get_data(DataType::SendTime); - let now_i64 = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as i64; - let timestamp_i64 = if let Some(n) = send_time_val.as_number() { - n as i64 - } else if let Some(s) = send_time_val.as_str() { - s.parse::().unwrap_or(now_i64) - } else { - now_i64 - }; - - let _ = chat_files::change_message_state( - timestamp_i64, - receiver_id as i64, - *sender_id as i64, - MessageState::from_str(cv.get_data(DataType::MessageState).as_str().unwrap_or("")), + if let Some(session) = self.app_sessions.get(&sender_id) { + let (user_id, app_identifier) = session.value(); + iota_storage::users::user_manager::save_app_data( + *user_id, + app_identifier, + &app_data, ); } - // Incoming storsed message: store for the recipient, attempt local delivery, notify sender. - if cv.is_type(CommunicationType::MessageSend) { - let sender_id: u64 = cv.get_sender(); + let res = CommunicationValue::new(CommunicationType::SaveAppData) + .with_id(cv.get_id()) + .with_receiver(sender_id); + let _ = self.send_message(&res).await; + } - // parse receiver_id (the storage owner for this incoming message) - let receiver_id: i64 = if let Some(n) = cv.get_data(DataType::ReceiverId).as_number() { - n as i64 - } else if let Some(s) = cv.get_data(DataType::ReceiverId).as_str() { - s.parse::().unwrap_or(0) - } else { - 0 - }; + async fn handle_load_app_data(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender(); + let mut app_data = String::new(); - // parse send_time robustly (number or string), fallback to now - let send_time_val = cv.get_data(DataType::SendTime); - let now_i64 = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as i64; - let timestamp_i64 = if let Some(n) = send_time_val.as_number() { - n as i64 - } else if let Some(s) = send_time_val.as_str() { - s.parse::().unwrap_or(now_i64) - } else { - now_i64 - }; - let timestamp_u128 = timestamp_i64 as u128; + if let Some(session) = self.app_sessions.get(&sender_id) { + let (user_id, app_identifier) = session.value(); + app_data = + iota_storage::users::user_manager::load_app_data(*user_id, app_identifier); + } - // content may be missing; default to empty string - let content = cv - .get_data(DataType::Content) - .as_str() - .unwrap_or("") - .to_string(); + let res = CommunicationValue::new(CommunicationType::LoadAppData) + .with_id(cv.get_id()) + .with_receiver(sender_id) + .add_typed_default(DataType::AppData, DataValue::Str(app_data)); + let _ = self.send_message(&res).await; + } - let height = cv.get_data(DataType::Height).as_number().unwrap_or(0) as i64; - let reply_to = cv.get_data(DataType::ReplyId).as_number().map(|n| n as i64); + async fn handle_create_app(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender() as i64; + let app_identifier = cv + .get_data(DataType::AppIdentifier) + .as_str() + .unwrap_or("") + .to_string(); + let app_public_key = cv + .get_data(DataType::AppPublicKey) + .as_str() + .unwrap_or("") + .to_string(); - let is_local = iota_storage::users::user_manager::get_user(receiver_id).is_some(); + if !app_identifier.is_empty() && !app_public_key.is_empty() { + if let Some(mut user) = iota_storage::users::user_manager::get_user(sender_id) { + if !user.trusted_apps.contains_key(&app_identifier) { + user.trusted_apps.insert(app_identifier, app_public_key); + iota_storage::users::user_manager::update_user(user); + } + } + } - if is_local { - // persist message for the receiver (storage_owner = receiver_id) - chat_files::add_message( - timestamp_u128, - false, - receiver_id as i64, - sender_id as i64, - &content, - height, - reply_to, - ); + let res = CommunicationValue::new(CommunicationType::CreateApp) + .with_id(cv.get_id()) + .with_receiver(sender_id as u64); + let _ = self.send_message(&res).await; + } + + async fn handle_delete_app(self: Arc, cv: &CommunicationValue) { + let sender_id = cv.get_sender() as i64; + let app_identifier = cv + .get_data(DataType::AppIdentifier) + .as_str() + .unwrap_or("") + .to_string(); + + if !app_identifier.is_empty() { + if let Some(mut user) = iota_storage::users::user_manager::get_user(sender_id) { + if user.trusted_apps.contains_key(&app_identifier) { + user.trusted_apps.remove(&app_identifier); + iota_storage::users::user_manager::update_user(user); + } + } + } + + let res = CommunicationValue::new(CommunicationType::DeleteApp) + .with_id(cv.get_id()) + .with_receiver(sender_id as u64); + let _ = self.send_message(&res).await; + } + + async fn handle_client_connected(self: Arc, cv: &CommunicationValue) { + let user_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0) as i64; + let _session_id = cv.get_data(DataType::SessionId).as_number().unwrap_or(0) as i64; + + let contacts = chats_util::get_users(user_id); + let mut contacts_array = Vec::new(); + + for (i, contact) in contacts.iter().enumerate() { + let mut contact_container = Vec::new(); + contact_container.push(( + DataType::UserId, + DataValue::SignedNumber(contact.user_id as i128), + )); + contact_container.push(( + DataType::LastMessageAt, + DataValue::SignedNumber(contact.last_message_at.unwrap_or(0) as i128), + )); + + if let Some(ref name) = contact.user_name { + contact_container.push((DataType::Username, DataValue::Str(name.clone()))); } - // persist message for the sender (storage_owner = sender_id) + let amount = if i < 10 { 20 } else { 1 }; + let messages = chat_files::get_messages(user_id, contact.user_id, 0, amount); + + let mut msg_array = Vec::new(); + for m in &messages { + let mut msg_container = Vec::new(); + msg_container.push(( + DataType::SendTime, + DataValue::SignedNumber(m.message_time as i128), + )); + msg_container.push((DataType::Content, DataValue::Str(m.content.clone()))); + msg_container.push((DataType::MessageState, DataValue::Str(m.message_state.clone()))); + msg_container.push((DataType::Height, DataValue::SignedNumber(m.height as i128))); + msg_container.push(( + DataType::SenderId, + DataValue::UnsignedNumber(if m.sent_by_self { + user_id as u128 + } else { + contact.user_id as u128 + }), + )); + msg_array.push(typed_container(msg_container)); + + if msg_array.len() == 1 { + let sender_id = if m.sent_by_self { + user_id + } else { + contact.user_id + }; + let mut last_msg = Vec::new(); + last_msg.push((DataType::Content, DataValue::Str(m.content.clone()))); + last_msg.push(( + DataType::SenderId, + DataValue::SignedNumber(sender_id as i128), + )); + contact_container.push((DataType::LastMessage, typed_container(last_msg))); + } + } + contact_container.push((DataType::Messages, DataValue::Array(msg_array))); + contacts_array.push(typed_container(contact_container)); + } + + let resp = CommunicationValue::new(CommunicationType::ClientConnected) + .with_id(cv.get_id()) + .add_typed_default(DataType::Contacts, DataValue::Array(contacts_array)); + let _ = self.send_message(&resp).await; + } + + async fn handle_message_state(self: Arc, cv: &CommunicationValue) { + let sender_id = &cv.get_sender(); + let receiver_id = match cv.get_data(DataType::ChatPartnerId).as_number() { + Some(id) => id, + _ => return, + }; + + let timestamp_i64 = if let Some(n) = cv.get_data(DataType::SendTime).as_number() { + n as i64 + } else if let Some(s) = cv.get_data(DataType::SendTime).as_str() { + s.parse::().unwrap_or_else(|_| now_millis_i64()) + } else { + now_millis_i64() + }; + + let _ = chat_files::change_message_state( + timestamp_i64, + receiver_id as i64, + *sender_id as i64, + MessageState::from_str(cv.get_data(DataType::MessageState).as_str().unwrap_or("")), + ); + } + + async fn handle_message_send(self: Arc, cv: &CommunicationValue) { + let sender_id: u64 = cv.get_sender(); + + let receiver_id: i64 = if let Some(n) = cv.get_data(DataType::ReceiverId).as_number() { + n as i64 + } else if let Some(s) = cv.get_data(DataType::ReceiverId).as_str() { + s.parse::().unwrap_or(0) + } else { + 0 + }; + + let timestamp_i64 = if let Some(n) = cv.get_data(DataType::SendTime).as_number() { + n as i64 + } else if let Some(s) = cv.get_data(DataType::SendTime).as_str() { + s.parse::().unwrap_or_else(|_| now_millis_i64()) + } else { + now_millis_i64() + }; + let timestamp_u128 = timestamp_i64 as u128; + + let content = cv + .get_data(DataType::Content) + .as_str() + .unwrap_or("") + .to_string(); + + let height = cv.get_data(DataType::Height).as_number().unwrap_or(0) as i64; + let reply_to = cv.get_data(DataType::ReplyId).as_number().map(|n| n as i64); + + let is_local = iota_storage::users::user_manager::get_user(receiver_id).is_some(); + + if is_local { chat_files::add_message( timestamp_u128, - true, - sender_id as i64, + false, receiver_id as i64, + sender_id as i64, &content, height, reply_to, ); + } - // send confirmation back to sender - let conf_msg = CommunicationValue::new(CommunicationType::MessageSend) - .with_id(cv.get_id()) - .with_receiver(sender_id as u64); - self.send_message(&conf_msg).await; + chat_files::add_message( + timestamp_u128, + true, + sender_id as i64, + receiver_id as i64, + &content, + height, + reply_to, + ); - if !is_local { - let mut fw_msg = CommunicationValue::new(CommunicationType::MessageOtherIota) - .with_id(cv.get_id()) - .with_receiver(receiver_id as u64) - .with_sender(sender_id as u64) - .add_typed_default(DataType::Height, DataValue::SignedNumber(height as i128)) - .add_typed_default(DataType::Content, DataValue::Str(content)) - .add_typed_default( - DataType::SendTime, - DataValue::SignedNumber(timestamp_i64 as i128), - ); - if let Some(rt) = reply_to { - fw_msg = fw_msg.add_typed_default( - DataType::ReplyId, - DataValue::UnsignedNumber(rt as u64 as u128), - ); - } - - let other_iota_resp = self - .clone() - .await_response(&fw_msg, Some(Duration::from_secs(10))) - .await; - - if let Ok(resp) = other_iota_resp { - let ms_raw = resp - .get_data(DataType::MessageState) - .as_string() - .unwrap_or_else(|| "".to_string()); - let ms = MessageState::from_str(&ms_raw).upgrade(MessageState::Received); + let conf_msg = CommunicationValue::new(CommunicationType::MessageSend) + .with_id(cv.get_id()) + .with_receiver(sender_id as u64); + let _ = self.send_message(&conf_msg).await; + if !is_local { + self.forward_to_remote_iota( + cv, + sender_id as i64, + receiver_id, + timestamp_i64, + &content, + height, + reply_to, + ) + .await; + } else { + match self + .forward_message_live( + cv.get_id(), + receiver_id as u64, + sender_id as i64, + timestamp_i64, + &content, + height, + reply_to, + ) + .await + { + Some(ms) => { let _ = chat_files::change_message_state( timestamp_i64, - sender_id as i64, - receiver_id as i64, - ms.clone(), - ); - - self.send_message( - &CommunicationValue::new(CommunicationType::MessageState) - .with_id(cv.get_id()) - .with_receiver(sender_id as u64) - .with_sender(receiver_id as u64) - .add_typed_default( - DataType::ChatPartnerId, - DataValue::SignedNumber(receiver_id as i128), - ) - .add_typed_default( - DataType::SendTime, - DataValue::SignedNumber(timestamp_i64 as i128), - ) - .add_typed_default( - DataType::MessageState, - DataValue::Str(ms.as_str().to_string()), - ), - ) - .await; - } else { - let _ = chat_files::change_message_state( - timestamp_i64, - sender_id as i64, - receiver_id as i64, - MessageState::Sent, - ); - - self.send_message( - &CommunicationValue::new(CommunicationType::MessageState) - .with_id(cv.get_id()) - .with_receiver(sender_id as u64) - .with_sender(receiver_id as u64) - .add_typed_default( - DataType::ChatPartnerId, - DataValue::SignedNumber(receiver_id as i128), - ) - .add_typed_default( - DataType::SendTime, - DataValue::SignedNumber(timestamp_i64 as i128), - ) - .add_typed_default( - DataType::MessageState, - DataValue::Str(MessageState::Sent.as_str().to_string()), - ), - ) - .await; - } - return; - } else { - // Build a live-delivery message for the local client (recipient) - let user_forward = CommunicationValue::new(CommunicationType::MessageLive) - .with_id(cv.get_id()) - .with_receiver(receiver_id as u64) - .add_typed_default( - DataType::SenderId, - DataValue::SignedNumber(sender_id as i128), - ) - .add_typed_default(DataType::Message, { - let mut msg_fields = vec![ - (DataType::Content, DataValue::Str(content.clone())), - ( - DataType::SendTime, - DataValue::SignedNumber(timestamp_i64 as i128), - ), - (DataType::Height, DataValue::SignedNumber(height as i128)), - ]; - if let Some(rt) = reply_to { - msg_fields.push(( - DataType::ReplyId, - DataValue::UnsignedNumber(rt as u64 as u128), - )); - } - typed_container(msg_fields) - }); - - // Attempt delivery and await a response from the local client - let user_resp = self - .clone() - .await_response(&user_forward, Some(Duration::from_secs(3))) - .await; - - if let Ok(user_resp) = user_resp { - let ms_raw = user_resp - .get_data(DataType::MessageState) - .as_string() - .unwrap_or_else(|| "".to_string()); - let ms = MessageState::from_str(&ms_raw).upgrade(MessageState::Received); - - // update stored message state for receiver - let _ = chat_files::change_message_state( - timestamp_i64, - receiver_id as i64, + receiver_id, sender_id as i64, ms.clone(), ); - - // update stored message state for sender let _ = chat_files::change_message_state( timestamp_i64, sender_id as i64, - receiver_id as i64, + receiver_id, ms.clone(), ); - - // notify original sender about the delivered/read state (if read receipts are enabled) if is_read_receipts_enabled().await { - self.send_message( - &CommunicationValue::new(CommunicationType::MessageState) + let _ = self + .send_message( + &CommunicationValue::new( + CommunicationType::MessageState, + ) .with_id(cv.get_id()) .with_receiver(sender_id as u64) .with_sender(receiver_id as u64) @@ -1463,37 +1514,36 @@ impl OmikronConnection { DataType::MessageState, DataValue::Str(ms.as_str().to_string()), ), - ) - .await; + ) + .await; } - } else { - // Delivery failed or timed out; mark as Sent + } + None => { let _ = chat_files::change_message_state( timestamp_i64, - receiver_id as i64, + receiver_id, sender_id as i64, MessageState::Sent, ); - let _ = chat_files::change_message_state( timestamp_i64, sender_id as i64, - receiver_id as i64, + receiver_id, MessageState::Sent, ); - - // Send push notification to Omega since user is offline - let push_msg = CommunicationValue::new(CommunicationType::PushNotification) - .with_receiver(receiver_id as u64) - .add_typed_default( - DataType::SenderId, - DataValue::SignedNumber(sender_id as i128), - ); - self.send_message(&push_msg).await; - - // notify sender - self.send_message( - &CommunicationValue::new(CommunicationType::MessageState) + let push_msg = + CommunicationValue::new(CommunicationType::PushNotification) + .with_receiver(receiver_id as u64) + .add_typed_default( + DataType::SenderId, + DataValue::SignedNumber(sender_id as i128), + ); + let _ = self.send_message(&push_msg).await; + let _ = self + .send_message( + &CommunicationValue::new( + CommunicationType::MessageState, + ) .with_id(cv.get_id()) .with_receiver(sender_id as u64) .with_sender(receiver_id as u64) @@ -1507,91 +1557,61 @@ impl OmikronConnection { ) .add_typed_default( DataType::MessageState, - DataValue::Str(MessageState::Sent.as_str().to_string()), + DataValue::Str( + MessageState::Sent.as_str().to_string(), + ), ), - ) - .await; + ) + .await; } - return; } } + } - if cv.is_type(CommunicationType::MessageOtherIota) { - let sender_id = &cv.get_sender(); - let receiver_id = &cv.get_receiver(); + async fn handle_message_other_iota(self: Arc, cv: &CommunicationValue) { + let sender_id = &cv.get_sender(); + let receiver_id = &cv.get_receiver(); - // parse send_time safely (number or string), fallback to now - let send_time_val = cv.get_data(DataType::SendTime); - let now_i64 = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as i64; - let timestamp = if let Some(n) = send_time_val.as_number() { - n as i64 - } else if let Some(s) = send_time_val.as_str() { - s.parse::().unwrap_or(now_i64) - } else { - now_i64 - }; + let timestamp = if let Some(n) = cv.get_data(DataType::SendTime).as_number() { + n as i64 + } else if let Some(s) = cv.get_data(DataType::SendTime).as_str() { + s.parse::().unwrap_or_else(|_| now_millis_i64()) + } else { + now_millis_i64() + }; - // content may be missing or non-string; default to empty string - let content = cv - .get_data(DataType::Content) - .as_str() - .unwrap_or("") - .to_string(); + let content = cv + .get_data(DataType::Content) + .as_str() + .unwrap_or("") + .to_string(); - let height = cv.get_data(DataType::Height).as_number().unwrap_or(0) as i64; - let reply_to = cv.get_data(DataType::ReplyId).as_number().map(|n| n as i64); + let height = cv.get_data(DataType::Height).as_number().unwrap_or(0) as i64; + let reply_to = cv.get_data(DataType::ReplyId).as_number().map(|n| n as i64); - chat_files::add_message( - timestamp as u128, - false, - *receiver_id as i64, + chat_files::add_message( + timestamp as u128, + false, + *receiver_id as i64, + *sender_id as i64, + &content, + height, + reply_to, + ); + + match self + .forward_message_live( + cv.get_id(), + *receiver_id, *sender_id as i64, + timestamp, &content, height, reply_to, - ); - - // Build user_forward using the parsed numeric timestamp and safe content string - let user_forward = CommunicationValue::new(CommunicationType::MessageLive) - .with_id(cv.get_id()) - .with_receiver(*receiver_id) - .add_typed_default( - DataType::SenderId, - DataValue::SignedNumber(*sender_id as i128), - ) - .add_typed_default(DataType::Message, { - let mut msg_fields = vec![ - (DataType::Content, DataValue::Str(content.clone())), - ( - DataType::SendTime, - DataValue::SignedNumber(timestamp as i128), - ), - (DataType::Height, DataValue::SignedNumber(height as i128)), - ]; - if let Some(rt) = reply_to { - msg_fields.push(( - DataType::ReplyId, - DataValue::UnsignedNumber(rt as u64 as u128), - )); - } - typed_container(msg_fields) - }); - - let user_resp = self - .clone() - .await_response(&user_forward, Some(Duration::from_secs(3))) - .await; - - if let Ok(user_resp) = user_resp { - let ms_raw = user_resp - .get_data(DataType::MessageState) - .as_string() - .unwrap_or_else(|| "".to_string()); - let ms = MessageState::from_str(&ms_raw).upgrade(MessageState::Received); - + ) + .await + { + Some(ms) => { let _ = change_message_state( timestamp, *receiver_id as i64, @@ -1599,9 +1619,8 @@ impl OmikronConnection { ms.clone(), ); - // notify original sender about the delivered/read state (if read receipts are enabled) if is_read_receipts_enabled().await { - self.send_message( + let _ = self.send_message( &CommunicationValue::new(CommunicationType::MessageState) .with_id(cv.get_id()) .with_receiver(*sender_id) @@ -1618,11 +1637,10 @@ impl OmikronConnection { DataType::MessageState, DataValue::Str(ms.as_str().to_string()), ), - ) - .await; + ).await; } - } else { - // Delivery timed out/failed — update stored state and notify sender with numeric timestamp + } + None => { let _ = chat_files::change_message_state( timestamp, *receiver_id as i64, @@ -1630,16 +1648,16 @@ impl OmikronConnection { MessageState::Sent, ); - // Send push notification to Omega since user is offline - let push_msg = CommunicationValue::new(CommunicationType::PushNotification) - .with_receiver(*receiver_id) - .add_typed_default( - DataType::SenderId, - DataValue::SignedNumber(*sender_id as i128), - ); - self.send_message(&push_msg).await; + let push_msg = + CommunicationValue::new(CommunicationType::PushNotification) + .with_receiver(*receiver_id) + .add_typed_default( + DataType::SenderId, + DataValue::SignedNumber(*sender_id as i128), + ); + let _ = self.send_message(&push_msg).await; - self.send_message( + let _ = self.send_message( &CommunicationValue::new(CommunicationType::MessageState) .with_id(cv.get_id()) .with_receiver(*sender_id) @@ -1656,225 +1674,223 @@ impl OmikronConnection { DataType::MessageState, DataValue::Str(MessageState::Sent.as_str().to_string()), ), - ) - .await; + ).await; } - return; } + } - if cv.is_type(CommunicationType::MessagesGet) { - let my_id = cv.get_sender(); - let partner_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0); - let offset = cv.get_data(DataType::Offset).as_number().unwrap_or(0); - let amount = cv.get_data(DataType::Amount).as_number().unwrap_or(0); - let messages = chat_files::get_messages( - my_id as i64, - partner_id as i64, - offset as i64, - amount as i64, - ); - let mut msg_array: Vec = Vec::new(); - for m in messages.members() { - let message_time: i64 = m["message_time"].as_i64().unwrap_or(0); - let content: String = m["content"].as_str().unwrap_or("").to_string(); - let sent_by_self: bool = m["sent_by_self"].as_bool().unwrap_or(false); - let height: i64 = m["height"].as_i64().unwrap_or(0); - let sender_id: i64 = if sent_by_self { - my_id as i64 + async fn handle_messages_get(self: Arc, cv: &CommunicationValue) { + let my_id = cv.get_sender(); + let partner_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0); + let offset = cv.get_data(DataType::Offset).as_number().unwrap_or(0); + let amount = cv.get_data(DataType::Amount).as_number().unwrap_or(0); + let messages = chat_files::get_messages( + my_id as i64, + partner_id as i64, + offset as i64, + amount as i64, + ); + let mut msg_array: Vec = Vec::new(); + for m in &messages { + let sender_id: i64 = if m.sent_by_self { + my_id as i64 + } else { + if let Some(n) = cv.get_data(DataType::ChatPartnerId).as_number() { + n as i64 + } else if let Some(s) = cv.get_data(DataType::ChatPartnerId).as_str() { + s.parse::().unwrap_or(partner_id as i64) } else { - if let Some(n) = cv.get_data(DataType::ChatPartnerId).as_number() { - n as i64 - } else if let Some(s) = cv.get_data(DataType::ChatPartnerId).as_str() { - s.parse::().unwrap_or(partner_id as i64) - } else { - partner_id as i64 - } - }; - let message_state: String = m["message_state"].as_str().unwrap_or("").to_string(); - - let mut container = Vec::new(); - container.push(( - DataType::SendTime, - DataValue::SignedNumber(message_time as i128), - )); - container.push((DataType::Content, DataValue::Str(content))); - container.push(( - DataType::SenderId, - DataValue::SignedNumber(sender_id as i128), - )); - container.push((DataType::MessageState, DataValue::Str(message_state))); - container.push((DataType::Height, DataValue::SignedNumber(height as i128))); - container.push(( - DataType::SenderId, - DataValue::UnsignedNumber(if sent_by_self { - my_id as u128 - } else { - partner_id as u128 - }), - )); - if let Some(rt) = m["reply_to"].as_i64() { - container.push(( - DataType::ReplyId, - DataValue::UnsignedNumber(rt as u64 as u128), - )); + partner_id as i64 } - msg_array.push(typed_container(container)); - } - - let resp = CommunicationValue::new(CommunicationType::MessagesGet) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default(DataType::Messages, DataValue::Array(msg_array)); - - self.send_message(&resp).await; - return; - } - - if cv.is_type(CommunicationType::GetChats) { - let user_id = cv.get_sender(); - let users = chats_util::get_users(user_id as i64); - let mut user_array = Vec::new(); - for user in users { - let mut container = Vec::new(); - container.push(( - DataType::UserId, - DataValue::SignedNumber(user.user_id as i128), - )); - if let Some(name) = user.user_name { - container.push((DataType::Username, DataValue::Str(name))); - } - if let Some(ts) = user.last_message_at { - container.push((DataType::LastMessageAt, DataValue::SignedNumber(ts as i128))); - } - user_array.push(typed_container(container)); - } - let resp = CommunicationValue::new(CommunicationType::GetChats) - .with_id(cv.get_id()) - .with_receiver(user_id) - .add_typed_default(DataType::UserIds, DataValue::Array(user_array)); - self.send_message(&resp).await; - return; - } - - if cv.is_type(CommunicationType::AddConversation) { - let user_id = cv.get_sender(); - let other_id = match cv.get_data(DataType::ChatPartnerId).as_number() { - Some(n) => n as i64, - None => cv - .get_data(DataType::ChatPartnerId) - .as_str() - .unwrap_or("0") - .parse() - .unwrap_or(0), - }; - let mut contact = get_user(user_id as i64, other_id).unwrap_or(Contact::new(other_id)); - - if let Some(name) = cv.get_data(DataType::ChatPartnerName).as_str() { - contact.user_name = Some(name.to_string()); - } - - contact.set_last_message_at( - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_millis() as i64, - ); - mod_user(user_id as i64, &contact); - let resp = CommunicationValue::new(CommunicationType::AddConversation) - .with_id(cv.get_id()) - .with_receiver(user_id); - self.send_message(&resp).await; - return; - } - - if cv.is_type(CommunicationType::AddCommunity) { - CommunitiesUtil::add_community( - cv.get_sender() as i64, - cv.get_data(DataType::CommunityAddress) - .as_str() - .unwrap() - .to_string(), - cv.get_data(DataType::CommunityTitle) - .as_str() - .unwrap() - .to_string(), - cv.get_data(DataType::Position) - .as_str() - .unwrap() - .to_string(), - ); - let resp = CommunicationValue::new(CommunicationType::AddCommunity) - .with_id(cv.get_id()) - .with_receiver(cv.get_sender()); - self.send_message(&resp).await; - return; - } - - if cv.is_type(CommunicationType::GetCommunities) { - let mut comm_array = Vec::new(); - for c in CommunitiesUtil::get_communities(cv.get_sender() as i64) { - let mut container: Vec<(DataType, DataValue)> = Vec::new(); - if let Some(address) = c["address"].as_str() { - container.push(( - DataType::CommunityAddress, - DataValue::Str(address.to_string()), - )); - } - if let Some(title) = c["title"].as_str() { - container.push((DataType::CommunityTitle, DataValue::Str(title.to_string()))); - } - if let Some(position) = c["position"].as_str() { - container.push((DataType::Position, DataValue::Str(position.to_string()))); - } - comm_array.push(typed_container(container)); - } - - let resp = CommunicationValue::new(CommunicationType::GetCommunities) - .with_id(cv.get_id()) - .with_receiver(cv.get_sender()) - .add_typed_default(DataType::Communities, DataValue::Array(comm_array)); - self.send_message(&resp).await; - return; - } - - if cv.is_type(CommunicationType::RemoveCommunity) { - CommunitiesUtil::remove_community( - cv.get_sender() as i64, - cv.get_data(DataType::CommunityAddress) - .as_str() - .unwrap() - .to_string(), - ); - let resp = CommunicationValue::new(CommunicationType::RemoveCommunity) - .with_id(cv.get_id()) - .with_receiver(cv.get_sender()); - self.send_message(&resp).await; - return; - } - - if cv.is_type(CommunicationType::GlobalSettingsSave) { - let my_id = cv.get_sender(); - let Some(settings_value) = cv.get_data(DataType::Payload).as_str() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing settings payload".to_string()), - ); - self.send_message(&response).await; - return; }; - save_file( - &format!("users/{}", my_id), - "global.settings", - settings_value, - ); + let mut container = Vec::new(); + container.push(( + DataType::SendTime, + DataValue::SignedNumber(m.message_time as i128), + )); + container.push((DataType::Content, DataValue::Str(m.content.clone()))); + container.push(( + DataType::SenderId, + DataValue::SignedNumber(sender_id as i128), + )); + container.push((DataType::MessageState, DataValue::Str(m.message_state.clone()))); + container.push((DataType::Height, DataValue::SignedNumber(m.height as i128))); + container.push(( + DataType::SenderId, + DataValue::UnsignedNumber(if m.sent_by_self { + my_id as u128 + } else { + partner_id as u128 + }), + )); + if let Some(rt) = m.reply_to { + container.push(( + DataType::ReplyId, + DataValue::UnsignedNumber(rt as u64 as u128), + )); + } + msg_array.push(typed_container(container)); + } - let mut response = CommunicationValue::new(CommunicationType::GlobalSettingsSave) + let resp = CommunicationValue::new(CommunicationType::MessagesGet) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default(DataType::Messages, DataValue::Array(msg_array)); + + let _ = self.send_message(&resp).await; + } + + async fn handle_get_chats(self: Arc, cv: &CommunicationValue) { + let user_id = cv.get_sender(); + let users = chats_util::get_users(user_id as i64); + let mut user_array = Vec::new(); + for user in users { + let mut container = Vec::new(); + container.push(( + DataType::UserId, + DataValue::SignedNumber(user.user_id as i128), + )); + if let Some(name) = user.user_name { + container.push((DataType::Username, DataValue::Str(name))); + } + if let Some(ts) = user.last_message_at { + container.push((DataType::LastMessageAt, DataValue::SignedNumber(ts as i128))); + } + user_array.push(typed_container(container)); + } + let resp = CommunicationValue::new(CommunicationType::GetChats) + .with_id(cv.get_id()) + .with_receiver(user_id) + .add_typed_default(DataType::UserIds, DataValue::Array(user_array)); + let _ = self.send_message(&resp).await; + } + + async fn handle_add_conversation(self: Arc, cv: &CommunicationValue) { + let user_id = cv.get_sender(); + let other_id = match cv.get_data(DataType::ChatPartnerId).as_number() { + Some(n) => n as i64, + None => cv + .get_data(DataType::ChatPartnerId) + .as_str() + .unwrap_or("0") + .parse() + .unwrap_or(0), + }; + let mut contact = get_user(user_id as i64, other_id).unwrap_or(Contact::new(other_id)); + + if let Some(name) = cv.get_data(DataType::ChatPartnerName).as_str() { + contact.user_name = Some(name.to_string()); + } + + contact.set_last_message_at(now_millis_i64()); + mod_user(user_id as i64, &contact); + let resp = CommunicationValue::new(CommunicationType::AddConversation) + .with_id(cv.get_id()) + .with_receiver(user_id); + let _ = self.send_message(&resp).await; + } + + async fn handle_add_community(self: Arc, cv: &CommunicationValue) { + CommunitiesUtil::add_community( + cv.get_sender() as i64, + cv.get_data(DataType::CommunityAddress) + .as_str() + .unwrap() + .to_string(), + cv.get_data(DataType::CommunityTitle) + .as_str() + .unwrap() + .to_string(), + cv.get_data(DataType::Position) + .as_str() + .unwrap() + .to_string(), + ); + let resp = CommunicationValue::new(CommunicationType::AddCommunity) + .with_id(cv.get_id()) + .with_receiver(cv.get_sender()); + let _ = self.send_message(&resp).await; + } + + async fn handle_get_communities(self: Arc, cv: &CommunicationValue) { + let mut comm_array = Vec::new(); + for c in CommunitiesUtil::get_communities(cv.get_sender() as i64) { + let mut container: Vec<(DataType, DataValue)> = Vec::new(); + container.push(( + DataType::CommunityAddress, + DataValue::Str(c.address.clone()), + )); + container.push((DataType::CommunityTitle, DataValue::Str(c.title.clone()))); + container.push((DataType::Position, DataValue::Str(c.position.clone()))); + comm_array.push(typed_container(container)); + } + + let resp = CommunicationValue::new(CommunicationType::GetCommunities) + .with_id(cv.get_id()) + .with_receiver(cv.get_sender()) + .add_typed_default(DataType::Communities, DataValue::Array(comm_array)); + let _ = self.send_message(&resp).await; + } + + async fn handle_remove_community(self: Arc, cv: &CommunicationValue) { + CommunitiesUtil::remove_community( + cv.get_sender() as i64, + cv.get_data(DataType::CommunityAddress) + .as_str() + .unwrap() + .to_string(), + ); + let resp = CommunicationValue::new(CommunicationType::RemoveCommunity) + .with_id(cv.get_id()) + .with_receiver(cv.get_sender()); + let _ = self.send_message(&resp).await; + } + + async fn handle_global_settings_save(self: Arc, cv: &CommunicationValue) { + let my_id = cv.get_sender(); + let Some(settings_value) = cv.get_data(DataType::Payload).as_str() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) .with_receiver(my_id) - .with_id(cv.get_id()); + .add_typed_default( + DataType::Message, + DataValue::Str("Missing settings payload".to_string()), + ); + let _ = self.send_message(&response).await; + return; + }; + + save_file( + &format!("users/{}", my_id), + "global.settings", + settings_value, + ); + + let mut response = CommunicationValue::new(CommunicationType::GlobalSettingsSave) + .with_receiver(my_id) + .with_id(cv.get_id()); + + if let Some(session_id) = cv.get_data(DataType::SessionId).as_number() { + response = response.add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + } + + let _ = self.send_message(&response).await; + } + + async fn handle_global_settings_load(self: Arc, cv: &CommunicationValue) { + let my_id = cv.get_sender(); + let path = format!("users/{}", my_id); + let name = "global.settings"; + + if !has_file(&path, name) { + let mut response = CommunicationValue::new(CommunicationType::ErrorNotFound) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default(DataType::Path, DataValue::Str(name.to_string())); if let Some(session_id) = cv.get_data(DataType::SessionId).as_number() { response = response.add_typed_default( @@ -1883,126 +1899,182 @@ impl OmikronConnection { ); } - self.send_message(&response).await; + let _ = self.send_message(&response).await; return; } - if cv.is_type(CommunicationType::GlobalSettingsLoad) { - let my_id = cv.get_sender(); - let path = format!("users/{}", my_id); - let name = "global.settings"; + let settings_value_str = load_file(&path, name); + let mut response = CommunicationValue::new(CommunicationType::GlobalSettingsLoad) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default(DataType::Payload, DataValue::Str(settings_value_str)); - if !has_file(&path, name) { - let mut response = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default(DataType::Path, DataValue::Str(name.to_string())); + if let Some(session_id) = cv.get_data(DataType::SessionId).as_number() { + response = response.add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + } - if let Some(session_id) = cv.get_data(DataType::SessionId).as_number() { - response = response.add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - } + let _ = self.send_message(&response).await; + } - self.send_message(&response).await; - return; - } - - let settings_value_str = load_file(&path, name); - let mut response = CommunicationValue::new(CommunicationType::GlobalSettingsLoad) + async fn handle_settings_save(self: Arc, cv: &CommunicationValue) { + let my_id = cv.get_sender(); + let Some(session_id) = cv.get_data(DataType::SessionId).as_number() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) .with_id(cv.get_id()) .with_receiver(my_id) - .add_typed_default(DataType::Payload, DataValue::Str(settings_value_str)); - - if let Some(session_id) = cv.get_data(DataType::SessionId).as_number() { - response = response.add_typed_default( + .add_typed_default( + DataType::Message, + DataValue::Str("Missing session_id".to_string()), + ); + let _ = self.send_message(&response).await; + return; + }; + if session_id == 0 || session_id > 1_000_000 { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Invalid session_id".to_string()), + ) + .add_typed_default( DataType::SessionId, DataValue::SignedNumber(session_id as i128), ); - } + let _ = self.send_message(&response).await; + return; + } + let Some(settings_name) = cv.get_data(DataType::SettingsName).as_str() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Missing settings_name".to_string()), + ) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + let _ = self.send_message(&response).await; + return; + }; + let Some(settings_value) = cv.get_data(DataType::Payload).as_str() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Missing settings payload".to_string()), + ) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + let _ = self.send_message(&response).await; + return; + }; - self.send_message(&response).await; + if !settings_name + .chars() + .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '.') + || settings_name.contains("..") + { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Invalid settings_name".to_string()), + ) + .add_typed_default( + DataType::SettingsName, + DataValue::Str(settings_name.to_string()), + ) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + let _ = self.send_message(&response).await; return; } - if cv.is_type(CommunicationType::SettingsSave) { - let my_id = cv.get_sender(); - let Some(session_id) = cv.get_data(DataType::SessionId).as_number() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing session_id".to_string()), - ); - self.send_message(&response).await; - return; - }; - let Some(settings_name) = cv.get_data(DataType::SettingsName).as_str() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing settings_name".to_string()), - ) - .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - self.send_message(&response).await; - return; - }; - let Some(settings_value) = cv.get_data(DataType::Payload).as_str() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing settings payload".to_string()), - ) - .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - self.send_message(&response).await; - return; - }; + save_file( + &format!("users/{}/settings/{}/", my_id, session_id), + &format!("{}.settings", settings_name), + settings_value, + ); - if !settings_name - .chars() - .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '.') - || settings_name.contains("..") - { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Invalid settings_name".to_string()), - ) - .add_typed_default( - DataType::SettingsName, - DataValue::Str(settings_name.to_string()), - ) - .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - self.send_message(&response).await; - return; - } - - save_file( - &format!("users/{}/settings/{}/", my_id, session_id), - &format!("{}.settings", settings_name), - settings_value, + let response = CommunicationValue::new(CommunicationType::SettingsSave) + .with_receiver(my_id) + .with_id(cv.get_id()) + .add_typed_default( + DataType::SettingsName, + DataValue::Str(settings_name.to_string()), + ) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), ); - let response = CommunicationValue::new(CommunicationType::SettingsSave) - .with_receiver(my_id) + let _ = self.send_message(&response).await; + } + + async fn handle_settings_load(self: Arc, cv: &CommunicationValue) { + let my_id = cv.get_sender(); + let Some(session_id) = cv.get_data(DataType::SessionId).as_number() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Missing session_id".to_string()), + ); + let _ = self.send_message(&response).await; + return; + }; + if session_id == 0 || session_id > 1_000_000 { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Invalid session_id".to_string()), + ); + let _ = self.send_message(&response).await; + return; + } + let Some(settings_name) = cv.get_data(DataType::SettingsName).as_str() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Missing settings_name".to_string()), + ) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + let _ = self.send_message(&response).await; + return; + }; + + if !settings_name + .chars() + .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '.') + || settings_name.contains("..") + { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Invalid settings_name".to_string()), + ) .add_typed_default( DataType::SettingsName, DataValue::Str(settings_name.to_string()), @@ -2011,87 +2083,16 @@ impl OmikronConnection { DataType::SessionId, DataValue::SignedNumber(session_id as i128), ); - - self.send_message(&response).await; + let _ = self.send_message(&response).await; return; } - if cv.is_type(CommunicationType::SettingsLoad) { - let my_id = cv.get_sender(); - let Some(session_id) = cv.get_data(DataType::SessionId).as_number() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing session_id".to_string()), - ); - self.send_message(&response).await; - return; - }; - let Some(settings_name) = cv.get_data(DataType::SettingsName).as_str() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing settings_name".to_string()), - ) - .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - self.send_message(&response).await; - return; - }; - - if !settings_name - .chars() - .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '.') - || settings_name.contains("..") - { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Invalid settings_name".to_string()), - ) - .add_typed_default( - DataType::SettingsName, - DataValue::Str(settings_name.to_string()), - ) - .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - self.send_message(&response).await; - return; - } - - let settings_file = format!("{}.settings", settings_name); - let settings_path = format!("users/{}/settings/{}/", my_id, session_id); - if !has_file(&settings_path, &settings_file) { - let response = CommunicationValue::new(CommunicationType::ErrorNotFound) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::SettingsName, - DataValue::Str(settings_name.to_string()), - ) - .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), - ); - self.send_message(&response).await; - return; - } - - let settings_value_str = load_file(&settings_path, &settings_file); - let response = CommunicationValue::new(CommunicationType::SettingsLoad) + let settings_file = format!("{}.settings", settings_name); + let settings_path = format!("users/{}/settings/{}/", my_id, session_id); + if !has_file(&settings_path, &settings_file) { + let response = CommunicationValue::new(CommunicationType::ErrorNotFound) .with_id(cv.get_id()) .with_receiver(my_id) - .add_typed_default(DataType::Payload, DataValue::Str(settings_value_str)) .add_typed_default( DataType::SettingsName, DataValue::Str(settings_name.to_string()), @@ -2100,59 +2101,78 @@ impl OmikronConnection { DataType::SessionId, DataValue::SignedNumber(session_id as i128), ); - - self.send_message(&response).await; + let _ = self.send_message(&response).await; return; } - if cv.is_type(CommunicationType::SettingsList) { - let my_id = cv.get_sender(); - let Some(session_id) = cv.get_data(DataType::SessionId).as_number() else { - let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) - .with_id(cv.get_id()) - .with_receiver(my_id) - .add_typed_default( - DataType::Message, - DataValue::Str("Missing session_id".to_string()), - ); - self.send_message(&response).await; - return; - }; + let settings_value_str = load_file(&settings_path, &settings_file); + let response = CommunicationValue::new(CommunicationType::SettingsLoad) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default(DataType::Payload, DataValue::Str(settings_value_str)) + .add_typed_default( + DataType::SettingsName, + DataValue::Str(settings_name.to_string()), + ) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); - let settings = get_children(&format!("users/{}/settings/{}/", my_id, session_id)); - let mut settings_json = Vec::new(); - for s in settings { - let s = s.replace(".settings", ""); - if s.is_empty() { - continue; - } - let _ = settings_json.push(DataValue::Str(s)); - } - let response = CommunicationValue::new(CommunicationType::SettingsList) + let _ = self.send_message(&response).await; + } + + async fn handle_settings_list(self: Arc, cv: &CommunicationValue) { + let my_id = cv.get_sender(); + let Some(session_id) = cv.get_data(DataType::SessionId).as_number() else { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) .with_id(cv.get_id()) .with_receiver(my_id) - .add_typed_default(DataType::Settings, DataValue::Array(settings_json)) .add_typed_default( - DataType::SessionId, - DataValue::SignedNumber(session_id as i128), + DataType::Message, + DataValue::Str("Missing session_id".to_string()), ); - - self.send_message(&response).await; + let _ = self.send_message(&response).await; + return; + }; + if session_id == 0 || session_id > 1_000_000 { + let response = CommunicationValue::new(CommunicationType::ErrorInvalidData) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default( + DataType::Message, + DataValue::Str("Invalid session_id".to_string()), + ); + let _ = self.send_message(&response).await; return; } + + let settings = get_children(&format!("users/{}/settings/{}/", my_id, session_id)); + let mut settings_json = Vec::new(); + for s in settings { + let s = s.replace(".settings", ""); + if s.is_empty() { + continue; + } + let _ = settings_json.push(DataValue::Str(s)); + } + let response = CommunicationValue::new(CommunicationType::SettingsList) + .with_id(cv.get_id()) + .with_receiver(my_id) + .add_typed_default(DataType::Settings, DataValue::Array(settings_json)) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id as i128), + ); + + let _ = self.send_message(&response).await; } // ------------------------------------------------------------------------- // Public API // ------------------------------------------------------------------------- - pub async fn send_message(&self, cv: &CommunicationValue) { - if let Err(err) = self.send_message_result(cv).await { - log_t!("send_message_failed", err); - } - } - - async fn send_message_result(&self, cv: &CommunicationValue) -> Result<(), String> { + pub async fn send_message(&self, cv: &CommunicationValue) -> Result<(), String> { let sender_guard = self.sender.read().await; if let Some(sender) = sender_guard.as_ref() { if !sender.is_open() { @@ -2216,24 +2236,21 @@ impl OmikronConnection { cv: &CommunicationValue, timeout_duration: Option, ) -> Result { - let (tx, mut rx) = mpsc::channel(1); + let (tx, rx) = oneshot::channel(); let msg_id = cv.get_id(); WAITING_TASKS.insert( msg_id, WaitingTask { task: Box::new(move |_, response_cv| { - let inner_tx = tx.clone(); - tokio::spawn(async move { - let _ = inner_tx.send(response_cv).await; - }); + let _ = tx.send(response_cv); true }), inserted_at: Instant::now(), }, ); - if let Err(send_err) = self.send_message_result(cv).await { + if let Err(send_err) = self.send_message(cv).await { WAITING_TASKS.remove(&msg_id); return Err(format!( "Request send failed (msg_id={}, reason={})", @@ -2243,8 +2260,8 @@ impl OmikronConnection { let timeout = timeout_duration.unwrap_or(Duration::from_secs(10)); - match tokio::time::timeout(timeout, rx.recv()).await { - Ok(Some(response_cv)) => { + match tokio::time::timeout(timeout, rx).await { + Ok(Ok(response_cv)) => { let is_error = response_cv.is_type(CommunicationType::Error) || response_cv.is_type(CommunicationType::ErrorInternal) || response_cv.is_type(CommunicationType::ErrorNotFound) @@ -2265,7 +2282,7 @@ impl OmikronConnection { Ok(response_cv) } } - Ok(_) => { + Ok(Err(_)) => { WAITING_TASKS.remove(&msg_id); Err("Channel closed while awaiting response".to_string()) } @@ -2284,27 +2301,30 @@ impl OmikronConnection { } pub async fn await_connection(&self, timeout_duration: Option) -> Result<(), String> { - if self.state.read().await.is_connected() { + let mut rx = self.state_watch_tx.subscribe(); + if rx.borrow().is_connected() { return Ok(()); } let timeout = timeout_duration.unwrap_or(CONNECTION_TIMEOUT); - let start = Instant::now(); - loop { - if self.state.read().await.is_connected() { - return Ok(()); + let result: Result<(), String> = tokio::time::timeout(timeout, async { + loop { + rx.changed().await.map_err(|_| "State watch channel closed".to_string())?; + if rx.borrow().is_connected() { + return Ok(()); + } } + }) + .await + .map_err(|_| { + format!( + "Connection not established within {} seconds", + timeout.as_secs() + ) + })?; - if start.elapsed() >= timeout { - return Err(format!( - "Connection not established within {} seconds", - timeout.as_secs() - )); - } - - sleep(Duration::from_millis(100)).await; - } + result } pub async fn has_auth_failure(&self) -> bool { diff --git a/omikron-connector/src/ping_pong_task.rs b/omikron-connector/src/ping_pong_task.rs index 9c81c60..2155b6c 100644 --- a/omikron-connector/src/ping_pong_task.rs +++ b/omikron-connector/src/ping_pong_task.rs @@ -2,6 +2,7 @@ use crate::omikron_connection::OmikronConnection; use dashmap::DashMap; use iota_state::APP_STATE; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; +use std::sync::atomic::Ordering; use std::sync::LazyLock; use std::time::Instant; use tokio::time::Duration; @@ -16,6 +17,8 @@ impl OmikronConnection { PING_TIMES.retain(|_, v| v.elapsed() < Duration::from_secs(30)); + self.missed_pongs.fetch_add(1, Ordering::Relaxed); + let ping_message = CommunicationValue::new(CommunicationType::Ping) .with_id(id) .add_typed_default( @@ -23,10 +26,12 @@ impl OmikronConnection { DataValue::Array(vec![DataValue::SignedNumber(*self.last_ping.lock().await as i128)]), ); - self.send_message(&ping_message).await; + let _ = self.send_message(&ping_message).await; } pub async fn handle_pong(&self, cv: &CommunicationValue) { + self.missed_pongs.store(0, Ordering::Relaxed); + let id = cv.get_id(); if let Some((_, send_time)) = PING_TIMES.remove(&id) {