From ed6f948df083130428bde7885cec53039116022a Mon Sep 17 00:00:00 2001 From: Rasensprenger Date: Wed, 5 Aug 2026 21:00:52 +0200 Subject: [PATCH 1/5] Update Rust crate base64 to 0.23.0 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 3b4f6a3..bb9bca8 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ mtp = { git = "https://git.methanium.net/methanium/mtp", features = [ ] } ansi_term = "0.12.1" -base64 = "0.22.1" +base64 = "0.23.0" bytes = "1" dashmap = "6.2.1" dotenv = "0.15.0" From 0de811544e4766a3faa9b72091b435103fe3845e Mon Sep 17 00:00:00 2001 From: Rasensprenger Date: Wed, 5 Aug 2026 22:01:11 +0200 Subject: [PATCH 2/5] Update Rust crate http to v1.5.0 --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 34e4f21..626eeaa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1077,9 +1077,9 @@ checksum = "1a9fcbcc408c5526c3ab80d534e5c86e7967c1fb7aa0a8c76abd1edc27deb877" [[package]] name = "http" -version = "1.4.2" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" dependencies = [ "bytes", "itoa", From b65b22dfbc322487226f8a3c0480759b6f3ed7b9 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 7 Aug 2026 23:54:34 +0200 Subject: [PATCH 3/5] [Fix] User States --- migrations/004_presence_preference.sql | 22 + mtp-type-maps | 2 +- src/db/user_repo.rs | 143 ++- src/main.rs | 4 +- src/models/user.rs | 1 + src/server/api.rs | 6 +- src/sql/connection_status.rs | 75 +- src/sql/user_online_tracker.rs | 1277 ++++++++++++++++++++---- src/state.rs | 72 ++ src/transport/capabilities.rs | 147 +++ src/transport/handlers/presence.rs | 862 ++++++++++++++-- src/transport/handlers/states.rs | 156 ++- src/transport/handlers/user_data.rs | 51 +- src/transport/mod.rs | 1 + src/transport/omikron_connection.rs | 129 ++- src/transport/omikron_manager.rs | 66 +- 16 files changed, 2613 insertions(+), 401 deletions(-) create mode 100644 migrations/004_presence_preference.sql create mode 100644 src/state.rs create mode 100644 src/transport/capabilities.rs diff --git a/migrations/004_presence_preference.sql b/migrations/004_presence_preference.sql new file mode 100644 index 0000000..163a623 --- /dev/null +++ b/migrations/004_presence_preference.sql @@ -0,0 +1,22 @@ +ALTER TABLE users + ADD COLUMN presence_preference VARBINARY(32) NOT NULL DEFAULT 'user_online'; + +UPDATE users +SET presence_preference = 'user_online' +WHERE presence_preference NOT IN ( + 'user_online', + 'user_idle', + 'user_dnd', + 'user_wc', + 'user_invisible' +); + +ALTER TABLE users + ADD CONSTRAINT chk_users_presence_preference + CHECK (presence_preference IN ( + 'user_online', + 'user_idle', + 'user_dnd', + 'user_wc', + 'user_invisible' + )); diff --git a/mtp-type-maps b/mtp-type-maps index ece6e2c..486541b 160000 --- a/mtp-type-maps +++ b/mtp-type-maps @@ -1 +1 @@ -Subproject commit ece6e2c3b4e925f3cefe46f4a048fbfc8f823093 +Subproject commit 486541b9483356ff49ff3ec7016f87d3ecbeaa0e diff --git a/src/db/user_repo.rs b/src/db/user_repo.rs index 8554566..01b0e6e 100644 --- a/src/db/user_repo.rs +++ b/src/db/user_repo.rs @@ -2,9 +2,11 @@ use crate::{ db::pool, error::{OmegaError, Result}, models::{IotaId, User, UserId}, + sql::connection_status::UserStatus, }; use mtp::crypto::PublicKeyBundle; -use sqlx::{FromRow, Row}; +use sqlx::{FromRow, MySql, QueryBuilder, Row}; +use std::collections::HashMap; pub const MAX_PROTOCOL_ID: i64 = (1_i64 << 48) - 1; const ID_ALLOCATION_ATTEMPTS: usize = 16; @@ -92,9 +94,9 @@ pub(crate) fn is_duplicate_key(error: &sqlx::Error) -> bool { }) } -const USER_BY_USERNAME_QUERY: &str = "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE username = ?"; -const USER_BY_ID_QUERY: &str = "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE id = ?"; -const USERS_BY_IOTA_ID_QUERY: &str = "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE iota_id = ?"; +const USER_BY_USERNAME_QUERY: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE username = ?"; +const USER_BY_ID_QUERY: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE id = ?"; +const USER_COLUMNS: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users"; #[derive(FromRow)] struct UserRow { @@ -103,6 +105,7 @@ struct UserRow { username: Vec, display: Option>, status: Option>, + presence_preference: Vec, about: Option>, avatar: Option>, sub_level: i32, @@ -125,6 +128,7 @@ impl TryFrom for User { username: decode(row.username)?, display: row.display.map(decode).transpose()?, status: row.status.map(decode).transpose()?, + presence_preference: decode(row.presence_preference)?, about: row.about.map(decode).transpose()?, avatar: row.avatar, sub_level: row.sub_level, @@ -154,15 +158,129 @@ pub async fn get_by_user_id(id: UserId) -> Result { } pub async fn get_users_by_iota_id(id: IotaId) -> Result> { - let rows = sqlx::query_as::<_, UserRow>(USERS_BY_IOTA_ID_QUERY) - .bind(id.0) + get_users_by_iota_ids(&[id.0]).await +} + +fn normalized_ids(ids: &[i64]) -> Vec { + let mut ids = ids + .iter() + .copied() + .filter(|id| valid_protocol_id(*id)) + .collect::>(); + ids.sort_unstable(); + ids.dedup(); + ids +} + +fn append_in_clause(query: &mut QueryBuilder<'_, MySql>, ids: &[i64]) { + query.push("("); + for (index, id) in ids.iter().enumerate() { + if index > 0 { + query.push(", "); + } + query.push_bind(*id); + } + query.push(")"); +} + +async fn fetch_users(mut query: QueryBuilder<'_, MySql>) -> Result> { + query + .build_query_as::() .fetch_all(&pool().await?) - .await?; - rows.into_iter() + .await? + .into_iter() .map(|row| row.try_into().map_err(OmegaError::from)) .collect() } +pub async fn get_users_by_ids(ids: &[i64]) -> Result> { + let ids = normalized_ids(ids); + if ids.is_empty() { + return Ok(Vec::new()); + } + let mut query = QueryBuilder::::new(USER_COLUMNS); + query.push(" WHERE id IN "); + append_in_clause(&mut query, &ids); + fetch_users(query).await +} + +pub async fn get_users_by_iota_ids(ids: &[i64]) -> Result> { + let ids = normalized_ids(ids); + if ids.is_empty() { + return Ok(Vec::new()); + } + let mut query = QueryBuilder::::new(USER_COLUMNS); + query.push(" WHERE iota_id IN "); + append_in_clause(&mut query, &ids); + fetch_users(query).await +} + +pub async fn get_users_by_ids_and_iota_ids( + user_ids: &[i64], + iota_ids: &[i64], +) -> Result> { + let user_ids = normalized_ids(user_ids); + let iota_ids = normalized_ids(iota_ids); + if user_ids.is_empty() && iota_ids.is_empty() { + return Ok(Vec::new()); + } + + let mut query = QueryBuilder::::new(USER_COLUMNS); + query.push(" WHERE "); + if !user_ids.is_empty() { + query.push("id IN "); + append_in_clause(&mut query, &user_ids); + } + if !iota_ids.is_empty() { + if !user_ids.is_empty() { + query.push(" OR "); + } + query.push("iota_id IN "); + append_in_clause(&mut query, &iota_ids); + } + fetch_users(query).await +} + +#[derive(FromRow)] +struct PresencePreferenceRow { + id: i64, + presence_preference: Vec, +} + +pub async fn get_presence_preferences(ids: &[i64]) -> Result> { + let ids = normalized_ids(ids); + if ids.is_empty() { + return Ok(HashMap::new()); + } + + let mut query = + QueryBuilder::::new("SELECT id, presence_preference FROM users WHERE id IN "); + append_in_clause(&mut query, &ids); + let rows = query + .build_query_as::() + .fetch_all(&pool().await?) + .await?; + let mut preferences = HashMap::with_capacity(rows.len()); + for row in rows { + let status = String::from_utf8(row.presence_preference) + .ok() + .and_then(|value| UserStatus::from_client_preference(&value)); + let status = match status { + Some(status) => status, + None => { + crate::log_in!( + crate::util::logger::PrintType::General, + "Invalid persisted presence preference for user {}, using user_online", + row.id + ); + UserStatus::user_online + } + }; + preferences.insert(row.id, status); + } + Ok(preferences) +} + async fn update( id: UserId, query: &'static str, @@ -217,6 +335,15 @@ pub async fn change_status(id: UserId, value: String) -> Result<()> { .await } +pub async fn change_presence_preference(id: UserId, value: String) -> Result<()> { + update( + id, + "UPDATE users SET presence_preference = ? WHERE id = ?", + value.into_bytes(), + ) + .await +} + pub async fn change_iota_id(id: UserId, value: IotaId) -> Result<()> { sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?") .bind(value.0) diff --git a/src/main.rs b/src/main.rs index 251713e..16b850e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,12 +5,14 @@ pub mod error; mod models; mod server; mod sql; +mod state; mod transport; mod util; pub use error::{OmegaError, Result}; use crate::db::initialize; +use crate::state::OmegaState; use crate::transport::omikron_connection; use crate::util::file_util::get_directory; use crate::util::logger::PrintType; @@ -94,7 +96,7 @@ async fn main() { .unwrap_or(443); tokio::select! { - result = omikron_connection::start(port) => { + result = omikron_connection::start(port, OmegaState::new()) => { if let Err(e) = result { log_err!(0, PrintType::General, "Server error: {:?}", e); } diff --git a/src/models/user.rs b/src/models/user.rs index 0dda84c..2843b6d 100644 --- a/src/models/user.rs +++ b/src/models/user.rs @@ -8,6 +8,7 @@ pub struct User { pub username: String, pub display: Option, pub status: Option, + pub presence_preference: String, pub about: Option, pub avatar: Option>, pub sub_level: i32, diff --git a/src/server/api.rs b/src/server/api.rs index ed5bb80..498685a 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -14,8 +14,10 @@ use crate::server::{ middleware, validation::{parse_positive_id, validate_non_empty}, }; -use crate::sql::user_online_tracker::{get_all_connections, get_iota_primary_omikron_connection}; -use crate::transport::omikron_manager::{get_connected_omikron, get_random_omikron}; +use crate::transport::omikron_manager::{ + get_all_connections, get_connected_omikron, get_iota_primary_omikron_connection, + get_random_omikron, +}; use crate::util::file_util::get_directory; use base64::Engine as _; use bytes::Bytes; diff --git a/src/sql/connection_status.rs b/src/sql/connection_status.rs index 8587ffe..104e316 100644 --- a/src/sql/connection_status.rs +++ b/src/sql/connection_status.rs @@ -1,7 +1,4 @@ -use strum::IntoEnumIterator; -use strum_macros::EnumIter; - -#[derive(Debug, Clone, PartialEq, EnumIter, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] #[allow(unused, non_camel_case_types)] pub enum UserStatus { user_offline, @@ -20,12 +17,70 @@ impl UserStatus { pub fn to_string(&self) -> String { format!("{:?}", self) } - pub fn from_str(s: &str) -> Option { - for sel in UserStatus::iter() { - if &sel.to_string() == s { - return Some(sel); - } + pub fn from_client_preference(s: &str) -> Option { + match s { + "user_online" => Some(Self::user_online), + "user_idle" => Some(Self::user_idle), + "user_dnd" => Some(Self::user_dnd), + "user_wc" => Some(Self::user_wc), + "user_invisible" => Some(Self::user_invisible), + _ => None, } - None + } + + pub fn public_value(&self) -> Self { + match self { + Self::user_invisible => Self::user_offline, + value => value.clone(), + } + } + + /// Parse a value received from a client or persisted as an account + /// preference. Derived connectivity and diagnostic states are never valid + /// preferences. + pub fn from_str(s: &str) -> Option { + Self::from_client_preference(s) + } +} + +#[cfg(test)] +mod tests { + use super::UserStatus; + + #[test] + fn accepts_only_client_preferences() { + for value in [ + "user_online", + "user_idle", + "user_dnd", + "user_wc", + "user_invisible", + ] { + assert!( + UserStatus::from_client_preference(value).is_some(), + "{value}" + ); + } + + for value in [ + "user_offline", + "iota_offline", + "iota_online", + "user_borked", + "iota_borked", + "unknown", + ] { + assert_eq!(UserStatus::from_client_preference(value), None, "{value}"); + assert_eq!(UserStatus::from_str(value), None, "{value}"); + } + } + + #[test] + fn invisible_is_publicly_offline() { + assert_eq!( + UserStatus::user_invisible.public_value(), + UserStatus::user_offline + ); + assert_eq!(UserStatus::user_dnd.public_value(), UserStatus::user_dnd); } } diff --git a/src/sql/user_online_tracker.rs b/src/sql/user_online_tracker.rs index b0e5f91..8017c3d 100644 --- a/src/sql/user_online_tracker.rs +++ b/src/sql/user_online_tracker.rs @@ -1,231 +1,1094 @@ -use crate::db::user_repo; -use crate::models::IotaId; use crate::sql::connection_status::UserStatus; -use dashmap::DashMap; -use once_cell::sync::Lazy; +use std::collections::{HashMap, HashSet}; +use std::sync::RwLock; -#[derive(Debug, Clone)] -pub struct UserConnection { - pub connection_type: UserStatus, +fn select_primary_route(previous_primary: Option, connected: &HashSet) -> Option { + previous_primary + .filter(|id| connected.contains(id)) + .or_else(|| connected.iter().min().copied()) +} + +fn sorted_vec(values: HashSet) -> Vec { + let mut values = values.into_iter().collect::>(); + values.sort_unstable(); + values +} + +fn sorted_pairs(values: HashSet<(i64, i64)>) -> Vec<(i64, i64)> { + let mut values = values.into_iter().collect::>(); + values.sort_unstable(); + values +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ClientSessionRoute { + pub omikron_id: i64, + pub iota_id: i64, +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub struct PresenceSubscriber { + pub user_id: i64, + pub session_id: i64, pub omikron_id: i64, } -// IotaID -> Primary OmikronID -static IOTA_PRIMARY_OMIKRON_CONNECTION: Lazy> = Lazy::new(DashMap::new); - -// IotaID -> Vec -static IOTA_OMIKRON_CONNECTIONS: Lazy>> = Lazy::new(DashMap::new); - -// UserID -> UserStatus -static USER_STATUS_MAP: Lazy> = Lazy::new(DashMap::new); - -// The legacy account map is kept for Iota-only/offline compatibility. Client -// transports are tracked independently: two devices must never overwrite each -// other's route merely because they authenticate as the same account. -static USER_SESSION_STATUS_MAP: Lazy> = Lazy::new(DashMap::new); - -pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) { - let mut entry = IOTA_OMIKRON_CONNECTIONS - .entry(iota_id) - .or_insert_with(Vec::new); - - if !entry.contains(&omikron_id) { - entry.push(omikron_id); - } - - if primary { - IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, omikron_id); - } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct PresenceIndices { + pub sessions_by_user: HashMap>, + pub sessions_by_omikron: HashMap>, + pub targets_by_subscriber: HashMap<(i64, i64), HashSet>, + pub subscribers_by_omikron: HashMap>, + pub iotas_by_omikron: HashMap>, } -pub fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool { - let mut replacement = None; - let mut connections_empty = false; - if let Some(mut entry) = IOTA_OMIKRON_CONNECTIONS.get_mut(&iota_id) { - entry.retain(|&id| id != omikron_id); - connections_empty = entry.is_empty(); - replacement = entry.first().copied(); - } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct IotaRoute { + pub primary_omikron_id: Option, + pub connected_omikron_ids: HashSet, +} - if connections_empty { - IOTA_OMIKRON_CONNECTIONS.remove_if(&iota_id, |_, connections| connections.is_empty()); - } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct PresenceRoutes { + pub iotas: HashMap, + pub sessions: HashMap<(i64, i64), ClientSessionRoute>, + pub subscriptions: HashMap>, + pub indices: PresenceIndices, +} - if IOTA_PRIMARY_OMIKRON_CONNECTION - .get(&iota_id) - .is_some_and(|primary| *primary == omikron_id) - { - match replacement { - Some(omikron_id) => { - IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, omikron_id); - } - None => { - IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id); - } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RemovedOmikronState { + pub iota_ids: Vec, + pub session_user_ids: Vec, + pub subscriber_sessions: Vec<(i64, i64)>, +} + +/// All live presence state is kept together so routing cannot be overwritten +/// by a preference update. Preferences are account-scoped, routes are +/// session-scoped, and effective state is always derived. +pub struct PresenceTracker { + pub(crate) routes: RwLock, + preferences: RwLock>, +} + +impl Default for PresenceTracker { + fn default() -> Self { + Self { + routes: RwLock::new(PresenceRoutes::default()), + preferences: RwLock::new(HashMap::new()), } } - - connections_empty } -pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option { - IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id).map(|v| *v) -} - -pub fn get_iota_omikron_connections(iota_id: i64) -> Option> { - IOTA_OMIKRON_CONNECTIONS.get(&iota_id).map(|v| v.clone()) -} - -pub async fn get_all_connections() --> Result>>, ()> { - let mut result: std::collections::HashMap>> = - std::collections::HashMap::new(); - - for entry in IOTA_OMIKRON_CONNECTIONS.iter() { - let iota_id = *entry.key(); - for omikron_id in entry.value().iter() { - result - .entry(*omikron_id) - .or_insert_with(std::collections::HashMap::new) - .entry(iota_id) - .or_insert_with(Vec::new); - } - } - - for entry in IOTA_OMIKRON_CONNECTIONS.iter() { - let iota_id = *entry.key(); - let omikron_ids = entry.value().clone(); - - if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await { - for user in users { - let user_id = user.id.0; - if let Some(conn) = USER_STATUS_MAP.get(&user_id) { - let user_omikron_id = conn.omikron_id; - if omikron_ids.contains(&user_omikron_id) { - if let Some(iota_map) = result.get_mut(&user_omikron_id) { - if let Some(user_vec) = iota_map.get_mut(&iota_id) { - user_vec.push(user_id); - } - } +impl PresenceTracker { + fn remove_subscription_locked(routes: &mut PresenceRoutes, key: (i64, i64)) { + let Some(targets) = routes.indices.targets_by_subscriber.remove(&key) else { + return; + }; + let mut omikron_ids = HashSet::new(); + for target_user_id in targets { + if let Some(subscribers) = routes.subscriptions.get_mut(&target_user_id) { + subscribers.retain(|subscriber| { + if (subscriber.user_id, subscriber.session_id) == key { + omikron_ids.insert(subscriber.omikron_id); + false + } else { + true } + }); + if subscribers.is_empty() { + routes.subscriptions.remove(&target_user_id); + } + } + } + for omikron_id in omikron_ids { + if let Some(sessions) = routes.indices.subscribers_by_omikron.get_mut(&omikron_id) { + sessions.remove(&key); + if sessions.is_empty() { + routes.indices.subscribers_by_omikron.remove(&omikron_id); } } } } - Ok(result) -} - -pub fn track_user_status(user_id: i64, status: UserStatus, omikron_id: i64) { - USER_STATUS_MAP.insert( - user_id, - UserConnection { - connection_type: status, - omikron_id, - }, - ); -} - -pub fn track_user_session_status( - user_id: i64, - session_id: i64, - status: UserStatus, - omikron_id: i64, -) { - USER_SESSION_STATUS_MAP.insert( - (user_id, session_id), - UserConnection { - connection_type: status, - omikron_id, - }, - ); -} - -pub fn untrack_user_status(user_id: i64, omikron_id: i64) { - USER_STATUS_MAP.remove_if(&user_id, |_, connection| { - connection.omikron_id == omikron_id - }); -} - -pub fn untrack_user_session_status(user_id: i64, session_id: i64, omikron_id: i64) { - USER_SESSION_STATUS_MAP.remove_if(&(user_id, session_id), |_, connection| { - connection.omikron_id == omikron_id - }); -} - -pub fn update_user_session_status(user_id: i64, status: UserStatus) { - for mut entry in USER_SESSION_STATUS_MAP.iter_mut() { - if entry.key().0 == user_id { - entry.connection_type = status.clone(); + fn remove_session_locked( + routes: &mut PresenceRoutes, + key: (i64, i64), + expected_omikron_id: Option, + ) -> bool { + let Some(route) = routes.sessions.get(&key) else { + return false; + }; + if expected_omikron_id.is_some_and(|id| route.omikron_id != id) { + return false; } - } - // Preserve the account preference for legacy routes as well. - if let Some(mut entry) = USER_STATUS_MAP.get_mut(&user_id) { - entry.connection_type = status; - } -} - -pub fn get_user_status(user_id: i64) -> Option { - // A connected visible session is preferred. Invisible sessions remain - // routable but are intentionally presented as offline when they are the - // only active routes. - let sessions: Vec = USER_SESSION_STATUS_MAP - .iter() - .filter(|entry| entry.key().0 == user_id) - .map(|entry| entry.value().clone()) - .collect(); - if let Some(status) = sessions - .iter() - .find(|status| status.connection_type != UserStatus::user_invisible) - { - return Some(status.clone()); - } - if let Some(status) = sessions.first() { - return Some(status.clone()); - } - USER_STATUS_MAP.get(&user_id).map(|v| v.clone()) -} - -pub fn untrack_many_users(user_ids: &[i64]) { - for user_id in user_ids { - USER_STATUS_MAP.remove(user_id); - } -} - -pub async fn untrack_omikron(omikron_id: i64) { - let mut offline_iotas = Vec::new(); - let mut primary_replacements = Vec::new(); - - for mut entry in IOTA_OMIKRON_CONNECTIONS.iter_mut() { - let iota_id = *entry.key(); - entry.retain(|&id| id != omikron_id); - if entry.is_empty() { - offline_iotas.push(iota_id); - } else if IOTA_PRIMARY_OMIKRON_CONNECTION - .get(&iota_id) - .is_some_and(|primary| *primary == omikron_id) - { - primary_replacements.push((iota_id, entry[0])); - } - } - - for iota_id in &offline_iotas { - IOTA_OMIKRON_CONNECTIONS.remove_if(iota_id, |_, connections| connections.is_empty()); - IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id); - } - - for (iota_id, replacement) in primary_replacements { - IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, replacement); - } - - USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id); - USER_SESSION_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id); - - for iota_id in offline_iotas { - if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await { - for user in users { - USER_STATUS_MAP.remove(&user.id.0); + let route = routes.sessions.remove(&key).unwrap(); + if let Some(session_ids) = routes.indices.sessions_by_user.get_mut(&key.0) { + session_ids.remove(&key.1); + if session_ids.is_empty() { + routes.indices.sessions_by_user.remove(&key.0); } } + if let Some(session_keys) = routes + .indices + .sessions_by_omikron + .get_mut(&route.omikron_id) + { + session_keys.remove(&key); + if session_keys.is_empty() { + routes.indices.sessions_by_omikron.remove(&route.omikron_id); + } + } + Self::remove_subscription_locked(routes, key); + true + } + + fn insert_session_locked( + routes: &mut PresenceRoutes, + key: (i64, i64), + route: ClientSessionRoute, + ) { + Self::remove_session_locked(routes, key, None); + routes.sessions.insert(key, route.clone()); + routes + .indices + .sessions_by_user + .entry(key.0) + .or_default() + .insert(key.1); + routes + .indices + .sessions_by_omikron + .entry(route.omikron_id) + .or_default() + .insert(key); + } + + fn remove_iota_owner_locked( + routes: &mut PresenceRoutes, + iota_id: i64, + omikron_id: i64, + ) -> bool { + let Some(route) = routes.iotas.get_mut(&iota_id) else { + return false; + }; + route.connected_omikron_ids.remove(&omikron_id); + route.primary_omikron_id = + select_primary_route(route.primary_omikron_id, &route.connected_omikron_ids); + let empty = route.connected_omikron_ids.is_empty(); + if let Some(iotas) = routes.indices.iotas_by_omikron.get_mut(&omikron_id) { + iotas.remove(&iota_id); + if iotas.is_empty() { + routes.indices.iotas_by_omikron.remove(&omikron_id); + } + } + if empty { + routes.iotas.remove(&iota_id); + } + empty + } + + #[cfg(debug_assertions)] + fn debug_assert_indices(routes: &PresenceRoutes) { + debug_assert!(Self::check_indices(routes).is_ok()); + } + + #[cfg(not(debug_assertions))] + fn debug_assert_indices(_: &PresenceRoutes) {} + + fn check_indices(routes: &PresenceRoutes) -> Result<(), String> { + for (key, route) in &routes.sessions { + if !routes + .indices + .sessions_by_user + .get(&key.0) + .is_some_and(|sessions| sessions.contains(&key.1)) + { + return Err(format!("session {key:?} missing user index")); + } + if !routes + .indices + .sessions_by_omikron + .get(&route.omikron_id) + .is_some_and(|sessions| sessions.contains(key)) + { + return Err(format!("session {key:?} missing Omikron index")); + } + } + for (user_id, session_ids) in &routes.indices.sessions_by_user { + for session_id in session_ids { + if !routes.sessions.contains_key(&(*user_id, *session_id)) { + return Err(format!("stale user session index {user_id}:{session_id}")); + } + } + } + for (omikron_id, session_keys) in &routes.indices.sessions_by_omikron { + for key in session_keys { + if routes.sessions.get(key).map(|route| route.omikron_id) != Some(*omikron_id) { + return Err(format!("stale Omikron session index {omikron_id}:{key:?}")); + } + } + } + for (target_user_id, subscribers) in &routes.subscriptions { + for subscriber in subscribers { + let key = (subscriber.user_id, subscriber.session_id); + if !routes + .indices + .targets_by_subscriber + .get(&key) + .is_some_and(|targets| targets.contains(target_user_id)) + { + return Err(format!("subscription {key:?} missing target index")); + } + if !routes + .indices + .subscribers_by_omikron + .get(&subscriber.omikron_id) + .is_some_and(|sessions| sessions.contains(&key)) + { + return Err(format!("subscription {key:?} missing Omikron index")); + } + } + } + for (key, targets) in &routes.indices.targets_by_subscriber { + for target in targets { + if !routes.subscriptions.get(target).is_some_and(|subscribers| { + subscribers + .iter() + .any(|subscriber| (subscriber.user_id, subscriber.session_id) == *key) + }) { + return Err(format!("stale target index {key:?}:{target}")); + } + } + } + for (omikron_id, iota_ids) in &routes.indices.iotas_by_omikron { + for iota_id in iota_ids { + if !routes + .iotas + .get(iota_id) + .is_some_and(|route| route.connected_omikron_ids.contains(omikron_id)) + { + return Err(format!("stale Iota index {omikron_id}:{iota_id}")); + } + } + } + Ok(()) + } + + pub fn check_index_consistency(&self) -> Result<(), String> { + Self::check_indices(&self.routes.read().unwrap()) + } + + pub fn set_preference(&self, user_id: i64, status: UserStatus) { + self.preferences.write().unwrap().insert(user_id, status); + } + + pub fn set_preferences(&self, preferences: HashMap) { + self.preferences.write().unwrap().extend(preferences); + } + + pub fn load_preference(&self, user_id: i64, raw: &str) -> bool { + let Some(status) = UserStatus::from_client_preference(raw) else { + self.set_preference(user_id, UserStatus::user_online); + return false; + }; + self.set_preference(user_id, status); + true + } + + pub fn preference(&self, user_id: i64) -> UserStatus { + self.preferences + .read() + .unwrap() + .get(&user_id) + .cloned() + .unwrap_or(UserStatus::user_online) + } + + pub fn has_iota_route(&self, iota_id: i64) -> bool { + self.routes + .read() + .unwrap() + .iotas + .get(&iota_id) + .is_some_and(|route| !route.connected_omikron_ids.is_empty()) + } + + pub fn has_active_session_for_iota(&self, user_id: i64, iota_id: i64) -> bool { + let routes = self.routes.read().unwrap(); + routes + .indices + .sessions_by_user + .get(&user_id) + .into_iter() + .flat_map(|session_ids| session_ids.iter()) + .any(|session_id| { + routes + .sessions + .get(&(user_id, *session_id)) + .is_some_and(|route| { + route.iota_id == iota_id + && routes.iotas.get(&route.iota_id).is_some_and(|iota| { + iota.connected_omikron_ids.contains(&route.omikron_id) + }) + }) + }) + } + + pub fn track_iota(&self, iota_id: i64, omikron_id: i64) { + self.connect_iota(iota_id, omikron_id); + } + + pub fn untrack_iota(&self, iota_id: i64, omikron_id: i64) -> bool { + self.disconnect_iota(iota_id, omikron_id) + } + + pub fn track_session(&self, user_id: i64, session_id: i64, omikron_id: i64, iota_id: i64) { + let mut routes = self.routes.write().unwrap(); + Self::insert_session_locked( + &mut routes, + (user_id, session_id), + ClientSessionRoute { + omikron_id, + iota_id, + }, + ); + Self::debug_assert_indices(&routes); + } + + pub fn remove_session(&self, user_id: i64, session_id: i64, omikron_id: i64) -> bool { + let mut routes = self.routes.write().unwrap(); + let removed = + Self::remove_session_locked(&mut routes, (user_id, session_id), Some(omikron_id)); + Self::debug_assert_indices(&routes); + removed + } + + pub fn owns_session(&self, user_id: i64, session_id: i64, omikron_id: i64) -> bool { + self.routes + .read() + .unwrap() + .sessions + .get(&(user_id, session_id)) + .is_some_and(|route| route.omikron_id == omikron_id) + } + + pub fn session_route(&self, user_id: i64, session_id: i64) -> Option { + self.routes + .read() + .unwrap() + .sessions + .get(&(user_id, session_id)) + .map(|route| route.clone()) + } + + pub fn resolve_public_state(&self, user_id: i64, iota_id: i64) -> UserStatus { + if !self.has_iota_route(iota_id) { + return UserStatus::iota_offline; + } + if !self.has_active_session_for_iota(user_id, iota_id) { + return UserStatus::user_offline; + } + self.preference(user_id).public_value() + } + + pub fn resolve_private_state(&self, user_id: i64) -> UserStatus { + self.preference(user_id) + } + + pub fn remove_omikron(&self, omikron_id: i64) -> RemovedOmikronState { + let mut routes = self.routes.write().unwrap(); + let iota_ids = routes + .indices + .iotas_by_omikron + .remove(&omikron_id) + .unwrap_or_default(); + for iota_id in &iota_ids { + Self::remove_iota_owner_locked(&mut routes, *iota_id, omikron_id); + } + + let subscriber_sessions = routes + .indices + .subscribers_by_omikron + .remove(&omikron_id) + .unwrap_or_default(); + + let session_keys = routes + .indices + .sessions_by_omikron + .remove(&omikron_id) + .unwrap_or_default(); + let session_user_ids = session_keys + .iter() + .map(|(user_id, _)| *user_id) + .collect::>(); + for key in &session_keys { + Self::remove_session_locked(&mut routes, *key, Some(omikron_id)); + } + for key in &subscriber_sessions { + Self::remove_subscription_locked(&mut routes, *key); + } + Self::debug_assert_indices(&routes); + + RemovedOmikronState { + iota_ids: sorted_vec(iota_ids), + session_user_ids: sorted_vec(session_user_ids), + subscriber_sessions: sorted_pairs(subscriber_sessions), + } + } + + pub fn remove_sessions_for_iota(&self, iota_id: i64, omikron_id: i64) { + let mut routes = self.routes.write().unwrap(); + let keys = routes + .indices + .sessions_by_omikron + .get(&omikron_id) + .into_iter() + .flat_map(|keys| keys.iter()) + .filter(|key| { + routes + .sessions + .get(key) + .is_some_and(|route| route.iota_id == iota_id) + }) + .copied() + .collect::>(); + for key in &keys { + Self::remove_session_locked(&mut routes, *key, Some(omikron_id)); + } + Self::debug_assert_indices(&routes); + } + + pub fn replace_subscription( + &self, + user_id: i64, + session_id: i64, + omikron_id: i64, + user_ids: Vec, + ) { + let mut routes = self.routes.write().unwrap(); + let key = (user_id, session_id); + Self::remove_subscription_locked(&mut routes, key); + for target_user_id in user_ids { + let subscriber = PresenceSubscriber { + user_id, + session_id, + omikron_id, + }; + routes + .subscriptions + .entry(target_user_id) + .or_default() + .insert(subscriber.clone()); + routes + .indices + .targets_by_subscriber + .entry(key) + .or_default() + .insert(target_user_id); + routes + .indices + .subscribers_by_omikron + .entry(omikron_id) + .or_default() + .insert(key); + } + Self::debug_assert_indices(&routes); + } + + pub fn remove_subscription(&self, user_id: i64, session_id: i64) { + let mut routes = self.routes.write().unwrap(); + Self::remove_subscription_locked(&mut routes, (user_id, session_id)); + Self::debug_assert_indices(&routes); + } + + pub fn subscribers(&self, target_user_id: i64) -> Vec { + let mut subscribers = self + .routes + .read() + .unwrap() + .subscriptions + .get(&target_user_id) + .cloned() + .unwrap_or_default() + .into_iter() + .collect::>(); + subscribers.sort_by_key(|subscriber| { + ( + subscriber.omikron_id, + subscriber.user_id, + subscriber.session_id, + ) + }); + subscribers + } + + pub fn sessions_owned_by(&self, omikron_id: i64) -> Vec<(i64, i64, ClientSessionRoute)> { + let routes = self.routes.read().unwrap(); + let mut sessions = routes + .indices + .sessions_by_omikron + .get(&omikron_id) + .into_iter() + .flat_map(|keys| keys.iter()) + .filter_map(|key| { + routes + .sessions + .get(key) + .map(|route| (key.0, key.1, route.clone())) + }) + .collect::>(); + sessions.sort_by_key(|(user_id, session_id, _)| (*user_id, *session_id)); + sessions + } + + pub fn replace_routes_owned_by( + &self, + omikron_id: i64, + iota_ids: &[i64], + sessions: &[(i64, i64, i64)], + ) { + self.replace_omikron_snapshot(omikron_id, iota_ids, sessions); + } + + pub fn replace_omikron_snapshot( + &self, + omikron_id: i64, + iota_ids: &[i64], + sessions: &[(i64, i64, i64)], + ) { + let mut routes = self.routes.write().unwrap(); + let old_sessions = routes + .indices + .sessions_by_omikron + .get(&omikron_id) + .cloned() + .unwrap_or_default(); + let old_subscriber_sessions = routes + .indices + .subscribers_by_omikron + .remove(&omikron_id) + .unwrap_or_default(); + for key in &old_subscriber_sessions { + Self::remove_subscription_locked(&mut routes, *key); + } + for key in &old_sessions { + Self::remove_session_locked(&mut routes, *key, Some(omikron_id)); + } + let old_iota_ids = routes + .indices + .iotas_by_omikron + .get(&omikron_id) + .cloned() + .unwrap_or_default(); + for iota_id in old_iota_ids { + Self::remove_iota_owner_locked(&mut routes, iota_id, omikron_id); + } + for iota_id in iota_ids { + let route = routes.iotas.entry(*iota_id).or_default(); + route.connected_omikron_ids.insert(omikron_id); + route.primary_omikron_id = + select_primary_route(route.primary_omikron_id, &route.connected_omikron_ids); + routes + .indices + .iotas_by_omikron + .entry(omikron_id) + .or_default() + .insert(*iota_id); + } + for (user_id, session_id, iota_id) in sessions { + Self::insert_session_locked( + &mut routes, + (*user_id, *session_id), + ClientSessionRoute { + omikron_id, + iota_id: *iota_id, + }, + ); + } + Self::debug_assert_indices(&routes); + } + + pub fn iota_ids_owned_by(&self, omikron_id: i64) -> Vec { + let routes = self.routes.read().unwrap(); + sorted_vec( + routes + .indices + .iotas_by_omikron + .get(&omikron_id) + .cloned() + .unwrap_or_default(), + ) + } + + pub fn track_iota_connection(&self, iota_id: i64, omikron_id: i64, primary: bool) { + self.connect_iota_with_priority(iota_id, omikron_id, primary); + } + + pub fn untrack_iota_connection(&self, iota_id: i64, omikron_id: i64) -> bool { + let connections_empty = self.disconnect_iota(iota_id, omikron_id); + self.remove_sessions_for_iota(iota_id, omikron_id); + connections_empty + } + + pub fn primary_iota_route(&self, iota_id: i64) -> Option { + self.routes + .read() + .unwrap() + .iotas + .get(&iota_id) + .and_then(|route| route.primary_omikron_id) + } + + pub fn iota_connections(&self, iota_id: i64) -> Option> { + self.all_iota_routes().get(&iota_id).map(|route| { + let mut ids = route + .connected_omikron_ids + .iter() + .copied() + .collect::>(); + ids.sort_unstable(); + ids + }) + } + + pub fn routes_for_user(&self, user_id: i64) -> Vec { + let routes = self.routes.read().unwrap(); + let mut user_routes = routes + .indices + .sessions_by_user + .get(&user_id) + .into_iter() + .flat_map(|session_ids| session_ids.iter()) + .filter_map(|session_id| { + routes + .sessions + .get(&(user_id, *session_id)) + .map(|route| (*session_id, route.clone())) + }) + .collect::>(); + user_routes.sort_by_key(|(session_id, route)| (*session_id, route.omikron_id)); + user_routes.into_iter().map(|(_, route)| route).collect() + } + + pub fn sessions_for_user(&self, user_id: i64) -> Vec<(i64, ClientSessionRoute)> { + let routes = self.routes.read().unwrap(); + let mut sessions = routes + .indices + .sessions_by_user + .get(&user_id) + .into_iter() + .flat_map(|session_ids| session_ids.iter()) + .filter_map(|session_id| { + routes + .sessions + .get(&(user_id, *session_id)) + .map(|route| (*session_id, route.clone())) + }) + .collect::>(); + sessions.sort_by_key(|(session_id, route)| (*session_id, route.omikron_id)); + sessions + } + + pub fn route_for_session(&self, user_id: i64, session_id: i64) -> Option { + self.session_route(user_id, session_id) + } + + pub fn user_route(&self, user_id: i64) -> Option { + self.routes_for_user(user_id).into_iter().next() + } + + pub fn connect_iota(&self, iota_id: i64, omikron_id: i64) { + self.connect_iota_with_priority(iota_id, omikron_id, false); + } + + fn connect_iota_with_priority(&self, iota_id: i64, omikron_id: i64, primary: bool) { + let mut routes = self.routes.write().unwrap(); + let route = routes.iotas.entry(iota_id).or_default(); + route.connected_omikron_ids.insert(omikron_id); + if primary || route.primary_omikron_id.is_none() { + route.primary_omikron_id = Some(omikron_id); + } + routes + .indices + .iotas_by_omikron + .entry(omikron_id) + .or_default() + .insert(iota_id); + Self::debug_assert_indices(&routes); + } + + pub fn disconnect_iota(&self, iota_id: i64, omikron_id: i64) -> bool { + let mut routes = self.routes.write().unwrap(); + let empty = Self::remove_iota_owner_locked(&mut routes, iota_id, omikron_id); + Self::debug_assert_indices(&routes); + empty + } + + pub fn all_iota_routes(&self) -> HashMap { + self.routes.read().unwrap().iotas.clone() + } + + pub fn connection_routes(&self) -> HashMap>> { + let mut result: std::collections::HashMap>> = + std::collections::HashMap::new(); + + let iotas = self.all_iota_routes(); + for (iota_id, route) in &iotas { + for omikron_id in &route.connected_omikron_ids { + result + .entry(*omikron_id) + .or_default() + .entry(*iota_id) + .or_default(); + } + } + result + } + + pub fn remove_omikron_and_offline_users(&self, omikron_id: i64) -> RemovedOmikronState { + self.remove_omikron(omikron_id) + } +} + +#[cfg(test)] +mod tests { + use super::{ClientSessionRoute, PresenceTracker}; + use crate::sql::connection_status::UserStatus; + + #[test] + fn effective_state_is_derived_from_routes_and_preference() { + let tracker = PresenceTracker::default(); + tracker.set_preference(7, UserStatus::user_online); + assert_eq!( + tracker.resolve_public_state(7, 11), + UserStatus::iota_offline + ); + + tracker.track_iota(11, 42); + assert_eq!( + tracker.resolve_public_state(7, 11), + UserStatus::user_offline + ); + + tracker.track_session(7, 3, 42, 11); + assert_eq!(tracker.resolve_public_state(7, 11), UserStatus::user_online); + tracker.set_preference(7, UserStatus::user_invisible); + assert_eq!( + tracker.resolve_public_state(7, 11), + UserStatus::user_offline + ); + } + + #[test] + fn session_must_match_assigned_iota_and_live_route() { + let tracker = PresenceTracker::default(); + tracker.track_iota(11, 42); + tracker.track_session(7, 3, 42, 12); + assert!(!tracker.has_active_session_for_iota(7, 11)); + assert_eq!( + tracker.resolve_public_state(7, 11), + UserStatus::user_offline + ); + + tracker.track_session(7, 3, 43, 11); + assert!(!tracker.has_active_session_for_iota(7, 11)); + + tracker.track_session(7, 3, 42, 11); + assert!(tracker.has_active_session_for_iota(7, 11)); + tracker.disconnect_iota(11, 42); + assert!(!tracker.has_active_session_for_iota(7, 11)); + assert_eq!( + tracker.resolve_public_state(7, 11), + UserStatus::iota_offline + ); + } + + #[test] + fn resolver_covers_all_public_preference_states() { + let tracker = PresenceTracker::default(); + tracker.track_iota(11, 42); + tracker.track_session(7, 3, 42, 11); + + for (preference, expected) in [ + (UserStatus::user_online, UserStatus::user_online), + (UserStatus::user_idle, UserStatus::user_idle), + (UserStatus::user_dnd, UserStatus::user_dnd), + (UserStatus::user_wc, UserStatus::user_wc), + (UserStatus::user_invisible, UserStatus::user_offline), + ] { + tracker.set_preference(7, preference); + assert_eq!(tracker.resolve_public_state(7, 11), expected); + } + } + + #[test] + fn session_routes_replace_by_session_without_replacing_preferences() { + let tracker = PresenceTracker::default(); + tracker.set_preference(7, UserStatus::user_dnd); + tracker.track_iota(11, 42); + tracker.track_iota(12, 43); + tracker.track_session(7, 3, 42, 11); + tracker.track_session(7, 4, 43, 12); + + assert_eq!(tracker.preference(7), UserStatus::user_dnd); + assert_eq!( + tracker.session_route(7, 3), + Some(ClientSessionRoute { + omikron_id: 42, + iota_id: 11, + }) + ); + + tracker.remove_session(7, 3, 42); + assert!(!tracker.owns_session(7, 3, 42)); + assert!(tracker.owns_session(7, 4, 43)); + assert_eq!(tracker.preference(7), UserStatus::user_dnd); + } + + #[test] + fn subscriptions_are_authoritative_and_cleanup_is_scoped() { + let tracker = PresenceTracker::default(); + tracker.replace_subscription(7, 3, 42, vec![20, 21, 20]); + tracker.replace_subscription(7, 3, 42, vec![21, 22]); + tracker.replace_subscription(8, 4, 43, vec![21]); + + assert_eq!(tracker.subscribers(20).len(), 0); + assert_eq!(tracker.subscribers(21).len(), 2); + assert_eq!(tracker.subscribers(22).len(), 1); + + tracker.remove_omikron(42); + assert!(tracker.subscribers(22).is_empty()); + assert_eq!(tracker.subscribers(21).len(), 1); + } + + #[test] + fn losing_one_iota_route_keeps_iota_online() { + let tracker = PresenceTracker::default(); + tracker.track_iota(11, 42); + tracker.track_iota(11, 43); + assert!(tracker.has_iota_route(11)); + + assert!(!tracker.untrack_iota(11, 42)); + assert!(tracker.has_iota_route(11)); + assert!(tracker.untrack_iota(11, 43)); + assert!(!tracker.has_iota_route(11)); + } + + #[test] + fn iota_connection_routes_track_primary_and_replacement() { + let tracker = PresenceTracker::default(); + tracker.track_iota_connection(11, 42, true); + tracker.track_iota_connection(11, 43, false); + + assert_eq!(tracker.iota_connections(11), Some(vec![42, 43])); + assert_eq!(tracker.primary_iota_route(11), Some(42)); + + tracker.untrack_iota_connection(11, 42); + + assert_eq!(tracker.iota_connections(11), Some(vec![43])); + assert_eq!(tracker.primary_iota_route(11), Some(43)); + } + + #[test] + fn snapshot_preserves_primary_and_uses_deterministic_fallback() { + let tracker = PresenceTracker::default(); + tracker.connect_iota(11, 20); + tracker.connect_iota(11, 30); + assert_eq!(tracker.primary_iota_route(11), Some(20)); + + tracker.replace_omikron_snapshot(30, &[11], &[]); + assert_eq!(tracker.primary_iota_route(11), Some(20)); + + tracker.replace_omikron_snapshot(20, &[11], &[]); + assert_eq!(tracker.primary_iota_route(11), Some(30)); + tracker.replace_omikron_snapshot(30, &[11], &[]); + assert_eq!(tracker.primary_iota_route(11), Some(20)); + } + + #[test] + fn user_routes_are_sorted_by_session_then_omikron() { + let tracker = PresenceTracker::default(); + tracker.track_session(7, 20, 30, 11); + tracker.track_session(7, 10, 40, 11); + assert_eq!( + tracker + .routes_for_user(7) + .iter() + .map(|route| route.omikron_id) + .collect::>(), + vec![40, 30] + ); + + tracker.remove_session(7, 10, 40); + assert_eq!(tracker.routes_for_user(7).len(), 1); + } + + #[test] + fn stale_disconnect_cannot_remove_current_session_subscription() { + let tracker = PresenceTracker::default(); + tracker.track_session(7, 3, 42, 11); + tracker.replace_subscription(7, 3, 42, vec![20]); + + assert!(!tracker.remove_session(7, 3, 99)); + assert!(tracker.owns_session(7, 3, 42)); + assert_eq!(tracker.subscribers(20).len(), 1); + + assert!(tracker.remove_session(7, 3, 42)); + assert!(tracker.subscribers(20).is_empty()); + } + + #[test] + fn private_state_is_not_affected_by_connectivity() { + let tracker = PresenceTracker::default(); + tracker.set_preference(7, UserStatus::user_invisible); + assert_eq!(tracker.resolve_private_state(7), UserStatus::user_invisible); + } + + #[test] + fn invalid_persisted_preference_falls_back_to_online() { + let tracker = PresenceTracker::default(); + assert!(!tracker.load_preference(7, "user_offline")); + assert_eq!(tracker.preference(7), UserStatus::user_online); + assert!(!tracker.load_preference(7, "iota_offline")); + assert_eq!(tracker.preference(7), UserStatus::user_online); + assert!(tracker.load_preference(7, "user_invisible")); + assert_eq!(tracker.preference(7), UserStatus::user_invisible); + assert!(tracker.load_preference(7, "user_dnd")); + assert_eq!(tracker.preference(7), UserStatus::user_dnd); + } + + #[test] + fn active_session_can_use_one_of_several_live_iota_routes() { + let tracker = PresenceTracker::default(); + tracker.track_iota(11, 42); + tracker.track_iota(11, 43); + tracker.track_session(7, 3, 43, 11); + + assert!(tracker.has_active_session_for_iota(7, 11)); + tracker.disconnect_iota(11, 43); + assert!(!tracker.has_active_session_for_iota(7, 11)); + } + + #[test] + fn route_snapshot_replaces_only_one_omikron() { + let tracker = PresenceTracker::default(); + tracker.track_iota(11, 42); + tracker.track_iota(12, 43); + tracker.track_session(7, 3, 42, 11); + tracker.track_session(8, 4, 43, 12); + tracker.replace_subscription(7, 3, 42, vec![20]); + + tracker.replace_routes_owned_by(42, &[12], &[(9, 5, 12)]); + + assert!(!tracker.owns_session(7, 3, 42)); + assert!(tracker.owns_session(8, 4, 43)); + assert!(tracker.owns_session(9, 5, 42)); + assert!(tracker.subscribers(20).is_empty()); + assert!(tracker.has_iota_route(12)); + } + + #[test] + fn omikron_cleanup_preserves_other_routes_and_preferences() { + let tracker = PresenceTracker::default(); + tracker.set_preference(7, UserStatus::user_dnd); + tracker.track_iota(11, 42); + tracker.track_iota(11, 43); + tracker.track_session(7, 3, 42, 11); + tracker.track_session(7, 4, 43, 11); + tracker.replace_subscription(7, 3, 42, vec![20]); + tracker.replace_subscription(7, 4, 43, vec![21]); + + tracker.remove_omikron(42); + + assert!(tracker.has_iota_route(11)); + assert!(!tracker.owns_session(7, 3, 42)); + assert!(tracker.owns_session(7, 4, 43)); + assert!(tracker.subscribers(20).is_empty()); + assert_eq!(tracker.subscribers(21).len(), 1); + assert_eq!(tracker.preference(7), UserStatus::user_dnd); + } + + #[test] + fn secondary_indices_follow_session_and_subscription_mutations() { + let tracker = PresenceTracker::default(); + tracker.track_session(7, 3, 42, 11); + tracker.replace_subscription(7, 3, 42, vec![20, 21, 20]); + assert_eq!( + tracker + .routes + .read() + .unwrap() + .indices + .sessions_by_user + .get(&7) + .cloned(), + Some([3].into_iter().collect()) + ); + assert_eq!( + tracker + .routes + .read() + .unwrap() + .indices + .targets_by_subscriber + .get(&(7, 3)) + .cloned(), + Some([20, 21].into_iter().collect()) + ); + assert!(tracker.check_index_consistency().is_ok()); + + tracker.replace_subscription(7, 3, 42, vec![22]); + assert!(tracker.subscribers(20).is_empty()); + assert!(tracker.subscribers(21).is_empty()); + assert_eq!(tracker.subscribers(22).len(), 1); + tracker.remove_session(7, 3, 42); + assert!(tracker.check_index_consistency().is_ok()); + assert!(tracker + .routes + .read() + .unwrap() + .indices + .sessions_by_user + .is_empty()); + assert!(tracker + .routes + .read() + .unwrap() + .indices + .targets_by_subscriber + .is_empty()); + } + + #[test] + fn omikron_removal_returns_all_affected_identifiers() { + let tracker = PresenceTracker::default(); + tracker.connect_iota(11, 42); + tracker.track_session(7, 3, 42, 11); + tracker.replace_subscription(7, 3, 42, vec![20]); + + let removed = tracker.remove_omikron(42); + + assert_eq!(removed.iota_ids, vec![11]); + assert_eq!(removed.session_user_ids, vec![7]); + assert_eq!(removed.subscriber_sessions, vec![(7, 3)]); + assert!(tracker.check_index_consistency().is_ok()); + } + + #[test] + fn consistency_check_reports_corrupted_reverse_index() { + let tracker = PresenceTracker::default(); + tracker.track_session(7, 3, 42, 11); + tracker + .routes + .write() + .unwrap() + .indices + .sessions_by_user + .get_mut(&7) + .unwrap() + .clear(); + + assert!(tracker.check_index_consistency().is_err()); } } diff --git a/src/state.rs b/src/state.rs new file mode 100644 index 0000000..dc91021 --- /dev/null +++ b/src/state.rs @@ -0,0 +1,72 @@ +use crate::sql::user_online_tracker::PresenceTracker; +use std::sync::Arc; + +pub struct OmegaState { + pub presence: Arc, +} + +impl Default for OmegaState { + fn default() -> Self { + Self { + presence: Arc::new(PresenceTracker::default()), + } + } +} + +impl OmegaState { + pub fn new() -> Arc { + Arc::new(Self::default()) + } +} + +#[cfg(test)] +mod tests { + use super::OmegaState; + use crate::sql::connection_status::UserStatus; + + #[test] + fn state_instances_have_independent_presence_trackers() { + let first = OmegaState::new(); + let second = OmegaState::new(); + + first.presence.track_iota_connection(11, 42, true); + + assert!(first.presence.has_iota_route(11)); + assert!(!second.presence.has_iota_route(11)); + assert_eq!(first.presence.primary_iota_route(11), Some(42)); + assert_eq!(second.presence.primary_iota_route(11), None); + } + + #[test] + fn two_session_private_and_public_presence_flow_is_authoritative() { + let state = OmegaState::new(); + state.presence.set_preference(7, UserStatus::user_online); + state.presence.set_preference(8, UserStatus::user_online); + state.presence.track_iota_connection(11, 42, true); + state.presence.track_session(7, 100, 42, 11); + state.presence.track_session(7, 101, 42, 11); + state.presence.replace_subscription(7, 100, 42, vec![8]); + + assert_eq!( + state.presence.resolve_public_state(8, 11), + UserStatus::user_offline + ); + state.presence.set_preference(8, UserStatus::user_invisible); + assert_eq!( + state.presence.resolve_public_state(8, 11), + UserStatus::user_offline + ); + assert_eq!( + state.presence.resolve_private_state(8), + UserStatus::user_invisible + ); + + state.presence.remove_session(7, 100, 42); + assert!(state.presence.owns_session(7, 101, 42)); + state.presence.remove_session(7, 101, 42); + assert_eq!( + state.presence.resolve_public_state(7, 11), + UserStatus::user_offline + ); + } +} diff --git a/src/transport/capabilities.rs b/src/transport/capabilities.rs new file mode 100644 index 0000000..0a3c685 --- /dev/null +++ b/src/transport/capabilities.rs @@ -0,0 +1,147 @@ +use std::collections::BTreeSet; + +const OMIKRON_PREFIX: &str = "omikron;caps="; +const OMEGA_PREFIX: &str = "omega;caps="; +const SET_USER_STATE: &str = "set_user_state_v1"; +const STATE_SUBSCRIBE: &str = "state_subscribe_v1"; +const SESSION_SNAPSHOT: &str = "session_snapshot_v1"; +const CLIENT_STATE_PUSH: &str = "client_state_push_v1"; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct PeerCapabilities { + pub set_user_state_v1: bool, + pub state_subscribe_v1: bool, + pub session_snapshot_v1: bool, + pub client_state_push_v1: bool, +} + +impl PeerCapabilities { + /// A missing descriptor is the legacy protocol: tuple route snapshots, + /// GetStates-only subscription refreshes, and ClientChanged pushes. + pub fn from_identification_description(description: Option<&str>) -> Result { + parse_capabilities(description, OMIKRON_PREFIX) + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct OmegaCapabilities { + pub set_user_state_v1: bool, + pub state_subscribe_v1: bool, + pub session_snapshot_v1: bool, + pub client_state_push_v1: bool, +} + +impl OmegaCapabilities { + pub fn current() -> Self { + Self { + set_user_state_v1: true, + state_subscribe_v1: true, + session_snapshot_v1: true, + client_state_push_v1: true, + } + } + + pub fn identification_description(&self) -> String { + let mut names = Vec::new(); + if self.set_user_state_v1 { + names.push(SET_USER_STATE); + } + if self.state_subscribe_v1 { + names.push(STATE_SUBSCRIBE); + } + if self.session_snapshot_v1 { + names.push(SESSION_SNAPSHOT); + } + if self.client_state_push_v1 { + names.push(CLIENT_STATE_PUSH); + } + format!("{OMEGA_PREFIX}{}", names.join(",")) + } +} + +fn parse_capabilities(description: Option<&str>, prefix: &str) -> Result { + let Some(description) = description else { + return Ok(PeerCapabilities::default()); + }; + if description == "omikron" { + return Ok(PeerCapabilities::default()); + } + let Some(capabilities) = description.strip_prefix(prefix) else { + return Err(()); + }; + let mut seen = BTreeSet::new(); + for capability in capabilities.split(',') { + if capability.is_empty() || !seen.insert(capability) { + return Err(()); + } + } + if seen.iter().any(|capability| { + !matches!( + *capability, + SET_USER_STATE | STATE_SUBSCRIBE | SESSION_SNAPSHOT | CLIENT_STATE_PUSH + ) + }) { + return Err(()); + } + Ok(PeerCapabilities { + set_user_state_v1: seen.contains(SET_USER_STATE), + state_subscribe_v1: seen.contains(STATE_SUBSCRIBE), + session_snapshot_v1: seen.contains(SESSION_SNAPSHOT), + client_state_push_v1: seen.contains(CLIENT_STATE_PUSH), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn advertised_capabilities_are_parsed() { + let capabilities = PeerCapabilities::from_identification_description(Some( + "omikron;caps=set_user_state_v1,state_subscribe_v1,session_snapshot_v1,client_state_push_v1", + )).unwrap(); + assert!(capabilities.set_user_state_v1); + assert!(capabilities.state_subscribe_v1); + assert!(capabilities.session_snapshot_v1); + assert!(capabilities.client_state_push_v1); + } + + #[test] + fn invalid_capability_values_fail_identification() { + assert!( + PeerCapabilities::from_identification_description(Some("omikron;caps=unsupported")) + .is_err() + ); + } + + #[test] + fn legacy_peer_has_no_version_specific_features() { + let capabilities = PeerCapabilities::from_identification_description(None).unwrap(); + assert!(!capabilities.set_user_state_v1); + assert!(!capabilities.state_subscribe_v1); + assert!(!capabilities.session_snapshot_v1); + assert!(!capabilities.client_state_push_v1); + assert_eq!( + PeerCapabilities::from_identification_description(Some("omikron")), + Ok(PeerCapabilities::default()) + ); + } + + #[test] + fn reconnecting_with_the_same_identification_is_stable() { + let description = Some( + "omikron;caps=set_user_state_v1,state_subscribe_v1,session_snapshot_v1,client_state_push_v1", + ); + assert_eq!( + PeerCapabilities::from_identification_description(description), + PeerCapabilities::from_identification_description(description) + ); + } + + #[test] + fn omega_capability_description_is_distinct_from_omikron_capabilities() { + let description = OmegaCapabilities::current().identification_description(); + assert!(description.starts_with(OMEGA_PREFIX)); + assert!(PeerCapabilities::from_identification_description(Some(&description)).is_err()); + } +} diff --git a/src/transport/handlers/presence.rs b/src/transport/handlers/presence.rs index adaa1a2..0c31904 100644 --- a/src/transport/handlers/presence.rs +++ b/src/transport/handlers/presence.rs @@ -1,69 +1,207 @@ use super::super::omikron_connection::{OmikronConnection, OmikronResult}; use crate::{ - db::user_repo, - log_in, - models::IotaId, - sql::{connection_status::UserStatus, user_online_tracker}, + db::user_repo, log_in, models::IotaId, sql::connection_status::UserStatus, state::OmegaState, }; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; -use std::sync::Arc; +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + sync::Arc, +}; -pub async fn user_connected( - _connection: Arc, - value: CommunicationValue, - omikron_id: i64, -) -> OmikronResult<()> { - log_in!(crate::util::logger::PrintType::Omega, "User connected"); - if let Some(user_id) = value.get_data(DataType::UserId).as_number() { - let status = value - .get_data(DataType::UserState) - .as_str() - .and_then(UserStatus::from_str) - .unwrap_or(UserStatus::user_online); - if let Ok(user_id) = i64::try_from(user_id) { - if let Some(session_id) = value - .get_data(DataType::SessionId) - .as_number() - .and_then(|id| i64::try_from(id).ok()) - .filter(|id| *id > 0) - { - user_online_tracker::track_user_session_status( - user_id, session_id, status, omikron_id, - ); - } else { - user_online_tracker::track_user_status(user_id, status, omikron_id); - } +fn parse_subscription(value: &CommunicationValue) -> Result<(i64, i64, Vec), &'static str> { + let user_id = i64::try_from(value.get_sender()) + .ok() + .filter(|id| *id > 0) + .ok_or("user_id")?; + let session_id = value + .get_data(DataType::SessionId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + .ok_or("session_id")?; + let DataValue::Array(values) = value.get_data(DataType::UserIds) else { + return Err("user_ids"); + }; + + let mut user_ids = Vec::with_capacity(values.len()); + for value in values { + let DataValue::SignedNumber(user_id) = value else { + return Err("user_ids"); + }; + let Ok(user_id) = i64::try_from(*user_id) else { + return Err("user_ids"); + }; + if user_id <= 0 { + return Err("user_ids"); + } + if !user_ids.contains(&user_id) { + user_ids.push(user_id); } } - Ok(()) + Ok((user_id, session_id, user_ids)) } -pub async fn user_disconnected( - _: Arc, - value: CommunicationValue, - omikron_id: i64, -) -> OmikronResult<()> { - log_in!(crate::util::logger::PrintType::Omega, "User disconnected"); - if let Some(user_id) = value.get_data(DataType::UserId).as_number() { - if let Some(session_id) = value - .get_data(DataType::SessionId) - .as_number() - .and_then(|id| i64::try_from(id).ok()) - .filter(|id| *id > 0) +fn apply_preferences(state: &OmegaState, preferences: HashMap) { + state.presence.set_preferences(preferences); +} + +fn states_for_users(state: &OmegaState, users: &[crate::models::User]) -> HashMap { + users + .iter() + .map(|user| { + ( + user.id.0, + state + .presence + .resolve_public_state(user.id.0, user.iota_id.0), + ) + }) + .collect() +} + +fn changed_states( + state: &OmegaState, + before: &HashMap, + users: &[crate::models::User], +) -> Vec<(i64, UserStatus)> { + let mut changes = users + .iter() + .filter_map(|user| { + let after = state + .presence + .resolve_public_state(user.id.0, user.iota_id.0); + (before.get(&user.id.0) != Some(&after)).then_some((user.id.0, after)) + }) + .collect::>(); + changes.sort_by_key(|(user_id, _)| *user_id); + changes.dedup_by_key(|(user_id, _)| *user_id); + changes +} + +fn state_notification( + subscriber: &crate::sql::user_online_tracker::PresenceSubscriber, + user_id: i64, + user_state: &UserStatus, +) -> CommunicationValue { + CommunicationValue::new(CommunicationType::ClientChanged) + .with_receiver(subscriber.user_id as u64) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(subscriber.session_id.into()), + ) + .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())) + .add_typed_default(DataType::UserState, DataValue::Str(user_state.to_string())) +} + +fn private_state_notification( + user_id: i64, + session_id: i64, + user_state: &UserStatus, +) -> CommunicationValue { + CommunicationValue::new(CommunicationType::ClientChanged) + .with_receiver(user_id as u64) + .add_typed_default( + DataType::SessionId, + DataValue::SignedNumber(session_id.into()), + ) + .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())) + .add_typed_default(DataType::UserState, DataValue::Str(user_state.to_string())) +} + +async fn publish_state_changes(state: &OmegaState, changes: &[(i64, UserStatus)]) { + let mut grouped = BTreeMap::>::new(); + for (user_id, user_state) in changes { + for subscriber in state.presence.subscribers(*user_id) { + grouped + .entry(subscriber.omikron_id) + .or_default() + .push(state_notification(&subscriber, *user_id, user_state)); + } + } + for (omikron_id, notifications) in grouped { + if let Err(error) = + crate::transport::omikron_manager::send_state_batch(omikron_id, notifications).await { - user_online_tracker::untrack_user_session_status( - user_id as i64, - session_id, + log_in!( + crate::util::logger::PrintType::General, + "Failed to deliver presence state batch to Omikron {}: {}", omikron_id, + error ); - } else { - user_online_tracker::untrack_user_status(user_id as i64, omikron_id); } } - Ok(()) } -pub async fn client_changed( +async fn publish_changed_states( + state: &OmegaState, + before: &HashMap, + users: &[crate::models::User], +) { + publish_state_changes(state, &changed_states(state, before, users)).await; +} + +async fn publish_private_state(state: &OmegaState, user_id: i64, user_state: &UserStatus) { + let mut grouped = BTreeMap::>::new(); + for (session_id, route) in state.presence.sessions_for_user(user_id) { + grouped + .entry(route.omikron_id) + .or_default() + .push(private_state_notification(user_id, session_id, user_state)); + } + for (omikron_id, notifications) in grouped { + if let Err(error) = + crate::transport::omikron_manager::send_state_batch(omikron_id, notifications).await + { + log_in!( + crate::util::logger::PrintType::General, + "Failed to deliver private presence state batch to Omikron {}: {}", + omikron_id, + error + ); + } + } +} + +pub async fn state_subscribe( + state: Arc, + connection: Arc, + value: CommunicationValue, + omikron_id: i64, +) -> OmikronResult<()> { + let (user_id, session_id, user_ids) = match parse_subscription(&value) { + Ok(subscription) => subscription, + Err("user_id") => { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoUserId) + .await; + } + Err(detail) => { + return connection + .send_error_response_with_detail( + value.get_id(), + CommunicationType::ErrorInvalidData, + detail, + ) + .await; + } + }; + if !state.presence.owns_session(user_id, session_id, omikron_id) { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoIota) + .await; + } + state + .presence + .replace_subscription(user_id, session_id, omikron_id, user_ids); + connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())) + .await +} + +/// Legacy state-change contract used by Omegas that predate SetUserState. +/// The payload is ClientChanged with UserId and UserState only. +pub async fn client_changed_legacy( + state: Arc, _: Arc, value: CommunicationValue, _: i64, @@ -72,29 +210,266 @@ pub async fn client_changed( .get_data(DataType::UserId) .as_number() .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) else { return Ok(()); }; let Some(status) = value .get_data(DataType::UserState) .as_str() - .and_then(UserStatus::from_str) + .and_then(UserStatus::from_client_preference) else { return Ok(()); }; - // Connectivity is derived from routes. Clients may choose only public - // presence preferences, never server/offline states. - if matches!( - status, - UserStatus::user_offline | UserStatus::iota_offline | UserStatus::iota_online - ) { - return Ok(()); - } - user_online_tracker::update_user_session_status(user_id, status); + state.presence.set_preference(user_id, status); Ok(()) } +pub async fn user_connected( + state: Arc, + connection: Arc, + value: CommunicationValue, + omikron_id: i64, +) -> OmikronResult<()> { + log_in!(crate::util::logger::PrintType::Omega, "User connected"); + let Some(user_id) = value + .get_data(DataType::UserId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; + }; + let Some(session_id) = value + .get_data(DataType::SessionId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; + }; + let Some(iota_id) = value + .get_data(DataType::IotaId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; + }; + let user = match user_repo::get_by_user_id(user_id.into()).await { + Ok(user) => user, + Err(crate::error::OmegaError::NotFound) => { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNotFound) + .await; + } + Err(error) => return Err(error.into()), + }; + let preferences = match user_repo::get_presence_preferences(&[user_id]).await { + Ok(preferences) => preferences, + Err(error) => return Err(error.into()), + }; + if user.iota_id.0 != iota_id || !state.presence.has_iota_route(iota_id) { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoIota) + .await; + } + apply_preferences(&state, preferences); + let users = [user]; + let before = states_for_users(&state, &users); + state + .presence + .track_session(user_id, session_id, omikron_id, iota_id); + publish_changed_states(&state, &before, &users).await; + connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())) + .await +} + +pub async fn user_disconnected( + state: Arc, + connection: Arc, + value: CommunicationValue, + omikron_id: i64, +) -> OmikronResult<()> { + log_in!(crate::util::logger::PrintType::Omega, "User disconnected"); + let Some(user_id) = value + .get_data(DataType::UserId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; + }; + let Some(session_id) = value + .get_data(DataType::SessionId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; + }; + if let Ok(user) = user_repo::get_by_user_id(user_id.into()).await { + let preferences = user_repo::get_presence_preferences(&[user_id]).await?; + apply_preferences(&state, preferences); + let users = [user]; + let before = states_for_users(&state, &users); + state + .presence + .remove_session(user_id, session_id, omikron_id); + publish_changed_states(&state, &before, &users).await; + } else { + state + .presence + .remove_session(user_id, session_id, omikron_id); + } + connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())) + .await +} + +pub async fn set_user_state( + state: Arc, + connection: Arc, + value: CommunicationValue, + omikron_id: i64, +) -> OmikronResult<()> { + let Some(user_id) = i64::try_from(value.get_sender()).ok().filter(|id| *id > 0) else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoUserId) + .await; + }; + if let Some(requested_user) = value.get_data_opt(DataType::UserId) { + let Some(requested_user_id) = requested_user + .as_number() + .and_then(|id| i64::try_from(id).ok()) + else { + return connection + .send_error_response_with_detail( + value.get_id(), + CommunicationType::ErrorInvalidData, + "user_id", + ) + .await; + }; + if requested_user_id != user_id { + return connection + .send_error_response_with_detail( + value.get_id(), + CommunicationType::ErrorInvalidData, + "user_id", + ) + .await; + } + } + let Some(iota_id) = value + .get_data(DataType::IotaId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + else { + return connection + .send_error_response_with_detail( + value.get_id(), + CommunicationType::ErrorInvalidData, + "iota_id", + ) + .await; + }; + let Some(requested_state) = value + .get_data(DataType::UserState) + .as_str() + .and_then(UserStatus::from_client_preference) + else { + return connection + .send_error_response_with_detail( + value.get_id(), + CommunicationType::ErrorInvalidData, + "user_state", + ) + .await; + }; + if !state.presence.has_iota_route(iota_id) { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoIota) + .await; + } + let Some(session_id) = value + .get_data(DataType::SessionId) + .as_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response_with_detail( + value.get_id(), + CommunicationType::ErrorInvalidData, + "session_id", + ) + .await; + }; + let Some(route) = state.presence.session_route(user_id, session_id) else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoIota) + .await; + }; + if route.omikron_id != omikron_id || route.iota_id != iota_id { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; + } + if !state.presence.has_active_session_for_iota(user_id, iota_id) { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNoIota) + .await; + } + let previous_preference = state.presence.preference(user_id); + let previous_state = state.presence.resolve_public_state(user_id, iota_id); + if let Err(error) = + user_repo::change_presence_preference(user_id.into(), requested_state.to_string()).await + { + log_in!( + crate::util::logger::PrintType::General, + "Failed to persist presence preference: {}", + error + ); + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInternal) + .await; + } + state + .presence + .set_preference(user_id, requested_state.clone()); + let new_state = state.presence.resolve_public_state(user_id, iota_id); + if requested_state != previous_preference { + publish_private_state(&state, user_id, &requested_state).await; + } + if requested_state != previous_preference && new_state != previous_state { + publish_state_changes(&state, &[(user_id, new_state)]).await; + } + connection + .send( + &CommunicationValue::new(CommunicationType::Success) + .with_id(value.get_id()) + .add_typed_default( + DataType::UserState, + DataValue::Str(requested_state.to_string()), + ), + ) + .await +} + pub async fn iota_connected( + state: Arc, connection: Arc, value: CommunicationValue, omikron_id: i64, @@ -105,36 +480,32 @@ pub async fn iota_connected( .as_number() .map(|id| id as i64) else { - return Ok(()); + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; }; - user_online_tracker::track_iota_connection(iota_id, omikron_id, true); - let mut user_ids = Vec::new(); - match user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await { - Ok(users) => { - for user in users { - user_ids.push(DataValue::SignedNumber(user.id.0.into())); - user_online_tracker::track_user_status( - user.id.0, - UserStatus::user_offline, - omikron_id, - ); - } - } - Err(_) => log_in!( - crate::util::logger::PrintType::General, - "SQL error loading users for IOTA" - ), - } + let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?; + let ids = users.iter().map(|user| user.id.0).collect::>(); + apply_preferences(&state, user_repo::get_presence_preferences(&ids).await?); + let before = states_for_users(&state, &users); + state.presence.connect_iota(iota_id, omikron_id); + let user_ids = users + .iter() + .map(|user| DataValue::SignedNumber(user.id.0.into())) + .collect(); let response = CommunicationValue::new(CommunicationType::IotaUserData) - .with_id(value.get_id()) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())) .add_typed_default(DataType::UserIds, DataValue::Array(user_ids)); - let _ = connection.send(&response).await; - Ok(()) + connection.clone().send(&response).await?; + publish_changed_states(&state, &before, &users).await; + connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())) + .await } pub async fn iota_disconnected( - _: Arc, + state: Arc, + connection: Arc, value: CommunicationValue, omikron_id: i64, ) -> OmikronResult<()> { @@ -144,40 +515,317 @@ pub async fn iota_disconnected( .as_number() .map(|id| id as i64) else { - return Ok(()); + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData) + .await; }; - if user_online_tracker::untrack_iota_connection(iota_id, omikron_id) { - if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await { - user_online_tracker::untrack_many_users( - &users.iter().map(|user| user.id.0).collect::>(), - ); - } - } - Ok(()) + let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?; + let ids = users.iter().map(|user| user.id.0).collect::>(); + apply_preferences(&state, user_repo::get_presence_preferences(&ids).await?); + let before = states_for_users(&state, &users); + state.presence.untrack_iota_connection(iota_id, omikron_id); + publish_changed_states(&state, &before, &users).await; + connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())) + .await } pub async fn sync_status( - _: Arc, + state: Arc, + connection: Arc, value: CommunicationValue, omikron_id: i64, ) -> OmikronResult<()> { - if let DataValue::Array(ids) = value.get_data(DataType::UserIds) { - for id in ids { - if let DataValue::SignedNumber(id) = id { - user_online_tracker::track_user_status( - *id as i64, - UserStatus::user_offline, - omikron_id, - ); - } + let request_id = value.get_id(); + let DataValue::Array(iota_values) = value.get_data(DataType::IotaIds) else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + let DataValue::Array(session_values) = value.get_data(DataType::UserStates) else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + + let mut iota_ids = Vec::with_capacity(iota_values.len()); + for item in iota_values { + let DataValue::SignedNumber(id) = item else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + let Ok(id) = i64::try_from(*id) else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + if id <= 0 { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + } + if !iota_ids.contains(&id) { + iota_ids.push(id); } } - if let DataValue::Array(ids) = value.get_data(DataType::IotaIds) { - for id in ids { - if let DataValue::SignedNumber(id) = id { - user_online_tracker::track_iota_connection(*id as i64, omikron_id, true); + + if !connection.peer_capabilities().session_snapshot_v1 { + let DataValue::Array(user_values) = value.get_data(DataType::UserIds) else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + let mut user_ids = Vec::with_capacity(user_values.len()); + for item in user_values { + let DataValue::SignedNumber(user_id) = item else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + let Ok(user_id) = i64::try_from(*user_id) else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + if user_id <= 0 { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + } + if !user_ids.contains(&user_id) { + user_ids.push(user_id); } } + + let previous_iota_ids = state.presence.iota_ids_owned_by(omikron_id); + let affected_iota_ids = previous_iota_ids + .iter() + .chain(iota_ids.iter()) + .copied() + .collect::>(); + let users = user_repo::get_users_by_ids(&user_ids).await?; + let returned_user_ids = users.iter().map(|user| user.id.0).collect::>(); + apply_preferences( + &state, + user_repo::get_presence_preferences(&returned_user_ids).await?, + ); + let before = states_for_users(&state, &users); + state + .presence + .replace_omikron_snapshot(omikron_id, &iota_ids, &[]); + let affected_users = user_repo::get_users_by_ids_and_iota_ids( + &returned_user_ids, + &affected_iota_ids.iter().copied().collect::>(), + ) + .await?; + publish_changed_states(&state, &before, &affected_users).await; + return connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(request_id)) + .await; + } + + let tm = mtp::type_map::TypeMap::latest(); + let mut sessions = Vec::with_capacity(session_values.len()); + for item in session_values { + let (user_id, session_id, iota_id) = if connection.peer_capabilities().session_snapshot_v1 { + let DataValue::Container(entries) = item else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + let find = |kind| { + entries.iter().find_map(|(key, value)| { + (Some(*key) == DataType::try_to_id(kind, &tm)).then_some(value) + }) + }; + let ( + Some(DataValue::SignedNumber(user_id)), + Some(DataValue::SignedNumber(session_id)), + Some(DataValue::SignedNumber(iota_id)), + ) = ( + find(DataType::UserId), + find(DataType::SessionId), + find(DataType::IotaId), + ) + else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + (*user_id, *session_id, *iota_id) + } else { + let DataValue::Array(values) = item else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + let [ + DataValue::SignedNumber(user_id), + DataValue::SignedNumber(session_id), + DataValue::SignedNumber(iota_id), + ] = values.as_slice() + else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + (*user_id, *session_id, *iota_id) + }; + let (Ok(user_id), Ok(session_id), Ok(iota_id)) = ( + i64::try_from(user_id), + i64::try_from(session_id), + i64::try_from(iota_id), + ) else { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + }; + if user_id <= 0 || session_id <= 0 || iota_id <= 0 || !iota_ids.contains(&iota_id) { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + } + if sessions.iter().any(|(existing_user, existing_session, _)| { + *existing_user == user_id && *existing_session == session_id + }) { + return connection + .send_error_response(request_id, CommunicationType::ErrorInvalidData) + .await; + } + sessions.push((user_id, session_id, iota_id)); + } + + let previous_iota_ids = state.presence.iota_ids_owned_by(omikron_id); + let previous_session_user_ids = state + .presence + .sessions_owned_by(omikron_id) + .into_iter() + .map(|(user_id, _, _)| user_id) + .collect::>(); + let new_session_user_ids = sessions + .iter() + .map(|(user_id, _, _)| *user_id) + .collect::>(); + let affected_iota_ids = previous_iota_ids + .iter() + .chain(iota_ids.iter()) + .copied() + .collect::>(); + let users = user_repo::get_users_by_ids_and_iota_ids( + &previous_session_user_ids + .iter() + .chain(new_session_user_ids.iter()) + .copied() + .collect::>(), + &affected_iota_ids.iter().copied().collect::>(), + ) + .await?; + let user_ids = users.iter().map(|user| user.id.0).collect::>(); + apply_preferences( + &state, + user_repo::get_presence_preferences(&user_ids).await?, + ); + let before = states_for_users(&state, &users); + state + .presence + .replace_omikron_snapshot(omikron_id, &iota_ids, &sessions); + publish_changed_states(&state, &before, &users).await; + connection + .send(&CommunicationValue::new(CommunicationType::Success).with_id(request_id)) + .await +} + +pub async fn omikron_disconnected(state: Arc, omikron_id: i64) { + let iota_ids = state.presence.iota_ids_owned_by(omikron_id); + let session_user_ids = state + .presence + .sessions_owned_by(omikron_id) + .into_iter() + .map(|(user_id, _, _)| user_id) + .collect::>(); + let users = match user_repo::get_users_by_ids_and_iota_ids(&session_user_ids, &iota_ids).await { + Ok(users) => users, + Err(error) => { + log_in!( + crate::util::logger::PrintType::General, + "Failed to load users before Omikron {} cleanup: {}", + omikron_id, + error + ); + return; + } + }; + let user_ids = users.iter().map(|user| user.id.0).collect::>(); + if let Err(error) = user_repo::get_presence_preferences(&user_ids) + .await + .map(|preferences| apply_preferences(&state, preferences)) + { + log_in!( + crate::util::logger::PrintType::General, + "Failed to load preferences before Omikron {} cleanup: {}", + omikron_id, + error + ); + return; + } + let before = states_for_users(&state, &users); + let removed = state.presence.remove_omikron(omikron_id); + debug_assert_eq!(removed.iota_ids, { + let mut ids = iota_ids.clone(); + ids.sort_unstable(); + ids.dedup(); + ids + }); + publish_changed_states(&state, &before, &users).await; +} + +#[cfg(test)] +mod tests { + use super::parse_subscription; + use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; + + fn request(user_ids: DataValue) -> CommunicationValue { + CommunicationValue::new(CommunicationType::StateSubscribe) + .with_sender(7) + .add_typed_default(DataType::SessionId, DataValue::SignedNumber(11)) + .add_typed_default(DataType::UserIds, user_ids) + } + + #[test] + fn subscription_parser_deduplicates_valid_targets() { + let parsed = parse_subscription(&request(DataValue::Array(vec![ + DataValue::SignedNumber(20), + DataValue::SignedNumber(21), + DataValue::SignedNumber(20), + ]))) + .unwrap(); + assert_eq!(parsed, (7, 11, vec![20, 21])); + } + + #[test] + fn subscription_parser_rejects_missing_or_malformed_fields() { + let missing_users = CommunicationValue::new(CommunicationType::StateSubscribe) + .with_sender(7) + .add_typed_default(DataType::SessionId, DataValue::SignedNumber(11)); + assert_eq!(parse_subscription(&missing_users), Err("user_ids")); + + let malformed_users = request(DataValue::Array(vec![DataValue::Str("bad".into())])); + assert_eq!(parse_subscription(&malformed_users), Err("user_ids")); + + let invalid_session = CommunicationValue::new(CommunicationType::StateSubscribe) + .with_sender(7) + .add_typed_default(DataType::SessionId, DataValue::SignedNumber(0)) + .add_typed_default( + DataType::UserIds, + DataValue::Array(vec![DataValue::SignedNumber(20)]), + ); + assert_eq!(parse_subscription(&invalid_session), Err("session_id")); + } + + #[test] + fn empty_subscription_is_valid_and_authoritative() { + let parsed = parse_subscription(&request(DataValue::Array(Vec::new()))).unwrap(); + assert_eq!(parsed, (7, 11, Vec::new())); } - Ok(()) } diff --git a/src/transport/handlers/states.rs b/src/transport/handlers/states.rs index be116b1..690a887 100644 --- a/src/transport/handlers/states.rs +++ b/src/transport/handlers/states.rs @@ -1,46 +1,146 @@ use super::super::omikron_connection::{OmikronConnection, OmikronResult}; -use crate::sql::{connection_status::UserStatus, user_online_tracker}; +use crate::db::user_repo; use mtp::{ codec::{CommunicationType, CommunicationValue, DataType, DataValue}, type_map::TypeMap, }; -use std::sync::Arc; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; + +async fn send_error( + connection: Arc, + request_id: u32, + error_type: CommunicationType, + session_id: Option, +) -> OmikronResult<()> { + let mut response = CommunicationValue::new(error_type).with_id(request_id); + if let Some(session_id) = session_id { + response = + response.add_typed_default(DataType::SessionId, DataValue::SignedNumber(session_id)); + } + connection.send(&response).await +} pub async fn get( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { + let state = connection.state(); + let legacy_peer = !connection.peer_capabilities().client_state_push_v1; let DataValue::Array(ids) = value.get_data(DataType::UserIds) else { - return Ok(()); + return send_error( + connection, + value.get_id(), + CommunicationType::ErrorInvalidData, + None, + ) + .await; }; + let session_id = value + .get_data(DataType::SessionId) + .as_number() + .filter(|id| *id > 0); + if session_id.is_none() && !legacy_peer { + return send_error( + connection, + value.get_id(), + CommunicationType::ErrorInvalidData, + None, + ) + .await; + } let tm = TypeMap::latest(); - let states = ids - .iter() - .filter_map(|id| { - let DataValue::SignedNumber(id) = id else { - return None; - }; - let status = user_online_tracker::get_user_status(*id as i64) - .map(|status| { - if status.connection_type == UserStatus::user_invisible { - UserStatus::user_offline.to_string() - } else { - status.connection_type.to_string() - } - }) - .unwrap_or_else(|| UserStatus::iota_offline.to_string()); - let mut map = Vec::new(); - if let Some(kind) = DataType::UserId.try_to_id(&tm) { - map.push((kind, DataValue::SignedNumber((*id as i64).into()))); - } - if let Some(kind) = DataType::UserState.try_to_id(&tm) { - map.push((kind, DataValue::Str(status))); - } - Some(DataValue::Container(map)) - }) - .collect(); + let mut requested_user_ids = Vec::new(); + let mut requested_set = HashSet::new(); + for id in ids { + let DataValue::SignedNumber(id) = id else { + return send_error( + connection, + value.get_id(), + CommunicationType::ErrorInvalidData, + session_id, + ) + .await; + }; + let Ok(user_id) = i64::try_from(*id) else { + return send_error( + connection, + value.get_id(), + CommunicationType::ErrorInvalidData, + session_id, + ) + .await; + }; + if user_id <= 0 { + return send_error( + connection, + value.get_id(), + CommunicationType::ErrorInvalidData, + session_id, + ) + .await; + } + if !requested_set.insert(user_id) { + continue; + } + requested_user_ids.push(user_id); + } + + let users = match user_repo::get_users_by_ids(&requested_user_ids).await { + Ok(users) => users, + Err(_) => { + return send_error( + connection, + value.get_id(), + CommunicationType::ErrorInternal, + session_id, + ) + .await; + } + }; + let users_by_id: HashMap<_, _> = users.into_iter().map(|user| (user.id.0, user)).collect(); + let mut states = Vec::new(); + let mut missing_user_ids = Vec::new(); + for user_id in requested_user_ids { + let Some(user) = users_by_id.get(&user_id) else { + missing_user_ids.push(user_id); + continue; + }; + let status = state + .presence + .resolve_public_state(user_id, user.iota_id.0) + .to_string(); + let mut map = Vec::new(); + if let Some(kind) = DataType::UserId.try_to_id(&tm) { + map.push((kind, DataValue::SignedNumber(user_id.into()))); + } + if let Some(kind) = DataType::UserState.try_to_id(&tm) { + map.push((kind, DataValue::Str(status))); + } + states.push(DataValue::Container(map)); + } let response = CommunicationValue::new(CommunicationType::GetStates) .with_id(value.get_id()) .add_typed_default(DataType::UserStates, DataValue::Array(states)); + let response = if let Some(session_id) = session_id { + response.add_typed_default(DataType::SessionId, DataValue::SignedNumber(session_id)) + } else { + response + }; + let response = if legacy_peer { + response + } else { + response.add_typed_default( + DataType::MissingUserIds, + DataValue::Array( + missing_user_ids + .into_iter() + .map(|id| DataValue::SignedNumber(id.into())) + .collect(), + ), + ) + }; connection.send(&response).await } diff --git a/src/transport/handlers/user_data.rs b/src/transport/handlers/user_data.rs index 27c9dbb..7e79bee 100644 --- a/src/transport/handlers/user_data.rs +++ b/src/transport/handlers/user_data.rs @@ -2,7 +2,6 @@ use super::super::omikron_connection::{OmikronConnection, OmikronResult}; use crate::{ db::{iota_repo, user_repo}, models::{IotaId, UserId}, - sql::{connection_status::UserStatus, user_online_tracker}, }; use base64::{Engine as _, engine::general_purpose::STANDARD}; use mtp::{ @@ -11,9 +10,12 @@ use mtp::{ }; use std::sync::Arc; -fn connections(iota_id: i64) -> DataValue { +fn connections(connection: &OmikronConnection, iota_id: i64) -> DataValue { DataValue::Array( - user_online_tracker::get_iota_omikron_connections(iota_id) + connection + .state() + .presence + .iota_connections(iota_id) .unwrap_or_default() .into_iter() .map(|id| DataValue::SignedNumber(id.into())) @@ -25,6 +27,7 @@ pub async fn get_user( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { + let state = connection.state(); let user = if let Some(id) = value.get_data(DataType::UserId).as_number() { user_repo::get_by_user_id(UserId::from(id as i64)) .await @@ -74,28 +77,36 @@ pub async fn get_user( response = response.add_typed_default(DataType::Avatar, DataValue::Str(STANDARD.encode(avatar))); } - let online = user_online_tracker::get_user_status(id); + let route = state.presence.user_route(id); + let private_request = value.get_sender() as i64 == id; + let resolved_status = if private_request { + if !state + .presence + .load_preference(id, &user.presence_preference) + { + crate::log_in!( + crate::util::logger::PrintType::General, + "Invalid persisted presence preference for user {}, using user_online", + id + ); + } + state.presence.resolve_private_state(id) + } else { + state.presence.resolve_public_state(id, iota_id) + }; response = response .add_typed_default( DataType::OnlineStatus, - DataValue::Str( - online - .as_ref() - .map(|status| { - if status.connection_type == UserStatus::user_invisible { - UserStatus::user_offline.to_string() - } else { - status.connection_type.to_string() - } - }) - .unwrap_or_else(|| UserStatus::iota_offline.to_string()), - ), + DataValue::Str(resolved_status.to_string()), ) - .add_typed_default(DataType::OmikronConnections, connections(iota_id)); - if let Some(status) = online { + .add_typed_default( + DataType::OmikronConnections, + connections(&connection, iota_id), + ); + if let Some(route) = route { response = response.add_typed_default( DataType::OmikronId, - DataValue::SignedNumber(status.omikron_id.into()), + DataValue::SignedNumber(route.omikron_id.into()), ); } connection.send(&response).await @@ -147,7 +158,7 @@ pub async fn get_iota( .with_id(value.get_id()) .add_typed_default(DataType::PublicKey, DataValue::Str(key.to_base64())) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into())) - .add_typed_default(DataType::OmikronConnections, connections(id)); + .add_typed_default(DataType::OmikronConnections, connections(&connection, id)); if let Some(user_id) = user_id { response = response.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())); diff --git a/src/transport/mod.rs b/src/transport/mod.rs index 8171bf2..52c5fe8 100644 --- a/src/transport/mod.rs +++ b/src/transport/mod.rs @@ -1,3 +1,4 @@ +pub mod capabilities; pub mod connection; pub mod handlers; pub mod omikron_connection; diff --git a/src/transport/omikron_connection.rs b/src/transport/omikron_connection.rs index e86f4a2..5001ff3 100644 --- a/src/transport/omikron_connection.rs +++ b/src/transport/omikron_connection.rs @@ -1,6 +1,8 @@ +use super::capabilities::{OmegaCapabilities, PeerCapabilities}; use crate::models::OmikronId; use crate::{ load_keyring, log, log_cv_in, log_cv_out, log_err, log_in, server, + state::OmegaState, transport::omikron_manager, util::{file_util::load_file_vec, logger::PrintType}, }; @@ -48,9 +50,11 @@ pub struct WaitingTask { pub struct OmikronConnection { id: u64, + state: Arc, sender: Mutex>, waiting_tasks: DashMap, cleanup_handle: std::sync::Mutex>>, + peer_capabilities: PeerCapabilities, } impl Drop for OmikronConnection { fn drop(&mut self) { @@ -61,13 +65,26 @@ impl Drop for OmikronConnection { } impl OmikronConnection { - pub fn new(sender: WebMtpSender, id: u64) -> Arc { - Arc::new(Self { + pub fn new( + sender: WebMtpSender, + id: u64, + description: Option<&str>, + state: Arc, + ) -> Option> { + let peer_capabilities = + PeerCapabilities::from_identification_description(description).ok()?; + Some(Arc::new(Self { id, + state, sender: Mutex::new(Some(sender)), waiting_tasks: DashMap::new(), cleanup_handle: std::sync::Mutex::new(None), - }) + peer_capabilities, + })) + } + + pub fn peer_capabilities(&self) -> &PeerCapabilities { + &self.peer_capabilities } pub async fn handle(self: Arc, receiver: &mut WebMtpReceiver) { @@ -76,6 +93,23 @@ impl OmikronConnection { PrintType::Omega, "Omikron connection started" ); + let capabilities = CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default( + mtp::codec::DataType::Description, + mtp::codec::DataValue::Str( + OmegaCapabilities::current().identification_description(), + ), + ); + if let Err(error) = self.clone().send(&capabilities).await { + log_err!( + self.id as i64, + PrintType::Omega, + "Failed to send Omega capabilities: {}", + error + ); + self.clone().cleanup().await; + return; + } let cleanup_conn = self.clone(); *self.cleanup_handle.lock().unwrap() = Some(tokio::spawn(async move { let mut ticker = interval(CLEANUP_INTERVAL); @@ -132,27 +166,30 @@ impl OmikronConnection { async fn dispatch(self: Arc, value: CommunicationValue) -> OmikronResult<()> { let id = self.id as i64; + let state = self.state.clone(); match value.get_comm_type_enum() { Some(CommunicationType::ShortenLink) => { crate::transport::handlers::links::shorten(self, value).await } Some(CommunicationType::UserConnected) => { - crate::transport::handlers::presence::user_connected(self, value, id).await + crate::transport::handlers::presence::user_connected(state, self, value, id).await } Some(CommunicationType::UserDisconnected) => { - crate::transport::handlers::presence::user_disconnected(self, value, id).await + crate::transport::handlers::presence::user_disconnected(state, self, value, id) + .await } - Some(CommunicationType::ClientChanged) => { - crate::transport::handlers::presence::client_changed(self, value, id).await + Some(CommunicationType::SetUserState) => { + crate::transport::handlers::presence::set_user_state(state, self, value, id).await } Some(CommunicationType::IotaConnected) => { - crate::transport::handlers::presence::iota_connected(self, value, id).await + crate::transport::handlers::presence::iota_connected(state, self, value, id).await } Some(CommunicationType::IotaDisconnected) => { - crate::transport::handlers::presence::iota_disconnected(self, value, id).await + crate::transport::handlers::presence::iota_disconnected(state, self, value, id) + .await } Some(CommunicationType::SyncClientIotaStatus) => { - crate::transport::handlers::presence::sync_status(self, value, id).await + crate::transport::handlers::presence::sync_status(state, self, value, id).await } Some(CommunicationType::GetUserData) => { crate::transport::handlers::user_data::get_user(self, value).await @@ -193,6 +230,13 @@ impl OmikronConnection { Some(CommunicationType::GetStates) => { crate::transport::handlers::states::get(self, value).await } + Some(CommunicationType::StateSubscribe) => { + crate::transport::handlers::presence::state_subscribe(state, self, value, id).await + } + Some(CommunicationType::ClientChanged) => { + crate::transport::handlers::presence::client_changed_legacy(state, self, value, id) + .await + } _ => { log_err!( 0, @@ -216,6 +260,24 @@ impl OmikronConnection { .await .map_err(|error| crate::error::OmegaError::SendError(error.to_string())) } + + pub(crate) async fn send_messages( + self: Arc, + values: &[CommunicationValue], + ) -> OmikronResult<()> { + let guard = self.sender.lock().await; + let sender = guard + .as_ref() + .ok_or(crate::error::OmegaError::NotConnected)?; + for value in values { + log_cv_out!(PrintType::Omikron, value); + sender + .send(value) + .await + .map_err(|error| crate::error::OmegaError::SendError(error.to_string()))?; + } + Ok(()) + } pub(crate) async fn send_error_response( self: Arc, message_id: u32, @@ -224,6 +286,22 @@ impl OmikronConnection { self.send(&CommunicationValue::new(error_type).with_id(message_id)) .await } + pub(crate) async fn send_error_response_with_detail( + self: Arc, + message_id: u32, + error_type: CommunicationType, + detail: &'static str, + ) -> OmikronResult<()> { + self.send( + &CommunicationValue::new(error_type) + .with_id(message_id) + .add_typed_default( + mtp::codec::DataType::ErrorType, + mtp::codec::DataValue::Str(detail.to_string()), + ), + ) + .await + } pub async fn close(self: Arc) { log_in!( self.id as i64, @@ -238,7 +316,11 @@ impl OmikronConnection { if self.id != 0 { log_in!(self.id as i64, PrintType::Omega, "Omikron disconnected"); if omikron_manager::remove_omikron(self.id as i64, &self).await { - crate::sql::user_online_tracker::untrack_omikron(self.id as i64).await; + crate::transport::handlers::presence::omikron_disconnected( + self.state.clone(), + self.id as i64, + ) + .await; } } if let Some(handle) = self.cleanup_handle.lock().unwrap().take() { @@ -248,12 +330,19 @@ impl OmikronConnection { pub async fn get_omikron_id(self: Arc) -> Option { Some(self.id as i64) } + pub fn state(&self) -> Arc { + self.state.clone() + } pub async fn send_message(self: Arc, value: &CommunicationValue) -> OmikronResult<()> { self.send(value).await } } -pub async fn get_by_omikron_id(omikron_id: u64, _: Option) -> Option { +pub async fn get_by_omikron_id( + omikron_id: u64, + description: Option, +) -> Option { + PeerCapabilities::from_identification_description(description.as_deref()).ok()?; crate::db::omikron_repo::get_omikron_by_id(OmikronId::from(omikron_id as i64)) .await .ok() @@ -263,7 +352,7 @@ pub async fn complete_register(_: PublicKeyBundle, _: Option) -> u64 { 0 } -pub async fn start(port: u16) -> Result<(), Box> { +pub async fn start(port: u16, state: Arc) -> Result<(), Box> { let cert_pem = load_file_vec("certs", "cert.pem")?; let key_pem = load_file_vec("certs", "key.pem")?; let web_config = server::server::build_web_config()? @@ -331,7 +420,19 @@ pub async fn start(port: u16) -> Result<(), Box> { ); continue; } - let connection = OmikronConnection::new(conn.sender, conn.client_id); + let Some(connection) = OmikronConnection::new( + conn.sender, + conn.client_id, + conn.description.as_deref(), + state.clone(), + ) else { + log_err!( + 0, + PrintType::Omega, + "Rejected Omikron connection with invalid capabilities" + ); + continue; + }; tokio::spawn(async move { let _guard = ConnectionLimitGuard(peer_ip); omikron_manager::add_omikron(connection.clone()).await; diff --git a/src/transport/omikron_manager.rs b/src/transport/omikron_manager.rs index 0e1575d..254ced8 100644 --- a/src/transport/omikron_manager.rs +++ b/src/transport/omikron_manager.rs @@ -1,4 +1,7 @@ +use crate::db::user_repo; +use crate::state::OmegaState; use crate::transport::connection::OmikronConnection; +use crate::transport::omikron_connection::OmikronResult; use dashmap::DashMap; use mtp::codec::CommunicationValue; use once_cell::sync::Lazy; @@ -34,6 +37,61 @@ pub fn get_connected_omikron(omikron_id: i64) -> Option> .map(|connection| connection.clone()) } +pub fn get_state() -> Option> { + OMIKRON_CONNECTIONS + .iter() + .next() + .map(|connection| connection.value().state()) +} + +pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option { + get_state().and_then(|state| state.presence.primary_iota_route(iota_id)) +} + +pub async fn get_all_connections() +-> Result>>, ()> { + match get_state() { + Some(state) => { + let mut result = state.presence.connection_routes(); + let iota_ids = state + .presence + .all_iota_routes() + .keys() + .copied() + .collect::>(); + let users = user_repo::get_users_by_iota_ids(&iota_ids) + .await + .map_err(|_| ())?; + for user in users { + for route in state.presence.routes_for_user(user.id.0) { + if let Some(iotas) = result.get_mut(&route.omikron_id) { + if let Some(users) = iotas.get_mut(&user.iota_id.0) { + users.push(user.id.0); + } + } + } + } + for iotas in result.values_mut() { + for users in iotas.values_mut() { + users.sort_unstable(); + users.dedup(); + } + } + Ok(result) + } + None => Ok(std::collections::HashMap::new()), + } +} + +pub async fn send_state_batch( + omikron_id: i64, + notifications: Vec, +) -> OmikronResult<()> { + let connection = + get_connected_omikron(omikron_id).ok_or(crate::error::OmegaError::NotConnected)?; + connection.send_messages(¬ifications).await +} + pub async fn get_random_omikron() -> Result, ()> { let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect(); @@ -47,9 +105,11 @@ pub async fn get_random_omikron() -> Result, ()> { } pub async fn send_to_user(user_id: i64, cv: &CommunicationValue) { - if let Some(user_conn) = crate::sql::user_online_tracker::get_user_status(user_id) { - if let Some(omikron_conn) = OMIKRON_CONNECTIONS.get(&user_conn.omikron_id) { - let _ = omikron_conn.value().clone().send_message(cv).await; + if let Some(state) = get_state() { + for user_route in state.presence.routes_for_user(user_id) { + if let Some(omikron_conn) = OMIKRON_CONNECTIONS.get(&user_route.omikron_id) { + let _ = omikron_conn.value().clone().send_message(cv).await; + } } } } From d3bcf56b59915d81234da62bc4b9a5a9f47cbfb5 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 9 Aug 2026 02:50:54 +0200 Subject: [PATCH 4/5] [Fix] User deletion & migration --- Cargo.lock | 228 ++++++++++--------- migrations/005_nullable_user_iota.sql | 11 + migrations/006_pending_iota_user_erasure.sql | 8 + src/api/response.rs | 4 +- src/db/user_repo.rs | 46 +++- src/models/user.rs | 2 +- src/server/api.rs | 14 +- src/state.rs | 29 +++ src/transport/handlers/account.rs | 165 +++++++++++++- src/transport/handlers/presence.rs | 7 +- src/transport/handlers/states.rs | 2 +- src/transport/handlers/user_data.rs | 41 ++-- src/transport/omikron_connection.rs | 19 ++ src/transport/omikron_manager.rs | 30 ++- 14 files changed, 457 insertions(+), 149 deletions(-) create mode 100644 migrations/005_nullable_user_iota.sql create mode 100644 migrations/006_pending_iota_user_erasure.sql diff --git a/Cargo.lock b/Cargo.lock index 34e4f21..1e03155 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -56,7 +56,7 @@ dependencies = [ "nom", "num-traits", "rusticata-macros", - "thiserror 2.0.19", + "thiserror", "time", ] @@ -117,9 +117,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-lc-rs" -version = "1.17.3" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00bdb5da18dac48ca2cc7cd4a98e533e8635a58e2361d13a1a4ee3888e0d72f1" +checksum = "ce2b2dcc879c3bae0d371e77c99f2238400ef24ec001394befa67b6e543add9e" dependencies = [ "aws-lc-sys", "untrusted 0.7.1", @@ -128,9 +128,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.43.0" +version = "0.44.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43103168cc76fe62678a375e722fc9cb3a0146159ac5828bc4f0dfd755c2224c" +checksum = "f09fae7be8bb3174e05c6afdb34199e6dc0c7c04ba9fa237b1967adfbde27483" dependencies = [ "cc", "cmake", @@ -145,6 +145,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" @@ -217,9 +223,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.4.0" +version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" +checksum = "9066c49992464636f92905fa096ec58baaa4d57ec19a5c096c68d3e25ef3d136" dependencies = [ "find-msvc-tools", "jobserver", @@ -467,8 +473,23 @@ dependencies = [ "cfg-if", "cpufeatures 0.2.17", "curve25519-dalek-derive", - "digest 0.10.7", - "fiat-crypto", + "fiat-crypto 0.2.9", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5eed333089e2e1c1ac8c6c0398e5e2497b4c9926ca6d0365ed1e099afa5bc23" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "curve25519-dalek-derive", + "digest 0.11.3", + "fiat-crypto 0.3.0", "rustc_version", "subtle", "zeroize", @@ -501,9 +522,9 @@ dependencies = [ [[package]] name = "data-encoding" -version = "2.11.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" [[package]] name = "deflate64" @@ -518,7 +539,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ "const-oid 0.9.6", - "pem-rfc7468", + "pem-rfc7468 0.7.0", "zeroize", ] @@ -529,6 +550,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69dedd701da44b0536442edf09c81a64b0ab97a7a4a5e3d1971f00027cbc63d" dependencies = [ "const-oid 0.10.2", + "pem-rfc7468 1.0.0", "zeroize", ] @@ -608,24 +630,25 @@ checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" [[package]] name = "ed25519" -version = "2.2.3" +version = "3.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +checksum = "29fcf32e6c73d1079f83ab4d782de2d81620346a5f38c6237a86a22f8368980a" dependencies = [ - "pkcs8 0.10.2", - "signature 2.2.0", + "pkcs8 0.11.0", + "signature 3.0.0", ] [[package]] name = "ed25519-dalek" -version = "2.2.0" +version = "3.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +checksum = "6ebaa1a2bf1290ab3bfe5a7b771d050ebffab2711c19a81691c683a5144a25de" dependencies = [ - "curve25519-dalek", + "curve25519-dalek 5.0.0", "ed25519", "serde", - "sha2 0.10.9", + "sha2 0.11.0", + "signature 3.0.0", "subtle", "zeroize", ] @@ -710,10 +733,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" [[package]] -name = "find-msvc-tools" -version = "0.1.9" +name = "fiat-crypto" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +checksum = "64cd1e32ddd350061ae6edb1b082d7c54915b5c672c389143b9a63403a109f24" + +[[package]] +name = "find-msvc-tools" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b73573e6edcd2af0cdf47bd6cb58f0b3839491263c314eaad1ccf24430e1de" [[package]] name = "flate2" @@ -1077,9 +1106,9 @@ checksum = "1a9fcbcc408c5526c3ab80d534e5c86e7967c1fb7aa0a8c76abd1edc27deb877" [[package]] name = "http" -version = "1.4.2" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" dependencies = [ "bytes", "itoa", @@ -1122,9 +1151,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hybrid-array" -version = "0.4.13" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +checksum = "707114b52a152fa7bdb290cd7cd5912d9467273b6d74e21b8d81aca1f8533f6b" dependencies = [ "ctutils", "typenum", @@ -1173,7 +1202,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1325,9 +1354,9 @@ dependencies = [ [[package]] name = "ipnet" -version = "2.12.0" +version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +checksum = "6a756c3fac73139e83f14c2d742155dd2b78d3ee56597b419a0579b7bdd6dd78" [[package]] name = "itoa" @@ -1347,7 +1376,7 @@ dependencies = [ "jni-sys", "log", "simd_cesu8", - "thiserror 2.0.19", + "thiserror", "walkdir", "windows-link", ] @@ -1416,9 +1445,9 @@ dependencies = [ [[package]] name = "keccak" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e24a010dd405bd7ed803e5253182815b41bf2e6a80cc3bfc066658e03a198aa" +checksum = "ffd9697dc4a9a62e2da93389f34400b77a28f0287711263cabb203b3ccb9c0e4" dependencies = [ "cfg-if", "cpufeatures 0.3.0", @@ -1453,9 +1482,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.18" +version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +checksum = "2026a5056764a10b2bf5d56488cba40da507f5493a6a429340e2004d9ed085fa" dependencies = [ "bitflags", "libc", @@ -1614,7 +1643,7 @@ dependencies = [ [[package]] name = "mtp" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "mtp-client", "mtp-codec", @@ -1630,7 +1659,7 @@ dependencies = [ [[package]] name = "mtp-client" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "mtp-codec", "mtp-common", @@ -1643,9 +1672,9 @@ dependencies = [ [[package]] name = "mtp-codec" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ - "base64", + "base64 0.23.1", "byteorder", "mtp-common", "mtp-crypto", @@ -1656,20 +1685,20 @@ dependencies = [ [[package]] name = "mtp-common" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "quinn", "rustls", - "thiserror 2.0.19", + "thiserror", "wtransport", ] [[package]] name = "mtp-crypto" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ - "base64", + "base64 0.23.1", "chacha20poly1305", "ed25519-dalek", "getrandom 0.4.3", @@ -1681,7 +1710,7 @@ dependencies = [ "rustls", "serde", "sha2 0.11.0", - "thiserror 1.0.69", + "thiserror", "tokio", "zeroize", ] @@ -1689,24 +1718,24 @@ dependencies = [ [[package]] name = "mtp-files" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "mtp-crypto", "rand 0.10.2", - "thiserror 1.0.69", + "thiserror", "zeroize", ] [[package]] name = "mtp-host" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "mtp-codec", "mtp-common", "mtp-crypto", "mtp-transport", - "rand 0.8.7", + "rand 0.10.2", "tokio", "tracing", "wtransport", @@ -1715,7 +1744,7 @@ dependencies = [ [[package]] name = "mtp-transport" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "async-trait", "mtp-codec", @@ -1733,7 +1762,7 @@ dependencies = [ [[package]] name = "mtp-type-map" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "serde", "serde_yaml", @@ -1742,7 +1771,7 @@ dependencies = [ [[package]] name = "mtp-webserver" version = "0.2.0" -source = "git+https://git.methanium.net/methanium/mtp#a692bed326dbfc8eac1a05825f4a287cbab6fd3e" +source = "git+https://git.methanium.net/methanium/mtp#b067614a684eb1856bc5db7b3fd82148c036ce6b" dependencies = [ "async-trait", "bytes", @@ -1761,7 +1790,7 @@ dependencies = [ "quinn", "rand 0.10.2", "rustls", - "thiserror 2.0.19", + "thiserror", "tokio", "tokio-rustls", "tokio-stream", @@ -1859,7 +1888,7 @@ name = "omega" version = "0.1.0" dependencies = [ "ansi_term", - "base64", + "base64 0.22.1", "bytes", "dashmap", "dotenv", @@ -1874,7 +1903,7 @@ dependencies = [ "sqlx", "strum", "strum_macros", - "thiserror 2.0.19", + "thiserror", "tokio", "uuid", "zip", @@ -1943,7 +1972,7 @@ version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ - "base64", + "base64 0.22.1", "serde_core", ] @@ -1956,6 +1985,15 @@ dependencies = [ "base64ct", ] +[[package]] +name = "pem-rfc7468" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9" +dependencies = [ + "base64ct", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -2082,7 +2120,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror 2.0.19", + "thiserror", "tokio", "tracing", "web-time", @@ -2107,7 +2145,7 @@ dependencies = [ "rustls-pki-types", "rustls-platform-verifier", "slab", - "thiserror 2.0.19", + "thiserror", "tinyvec", "tracing", "web-time", @@ -2237,7 +2275,7 @@ version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -2329,9 +2367,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.42" +version = "0.23.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" dependencies = [ "aws-lc-rs", "log", @@ -2598,7 +2636,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09057cb2149ad4cbd2da1e26b351f9a4c354219421229c69c3063e6f61947c4a" dependencies = [ "digest 0.11.3", - "keccak 0.2.0", + "keccak 0.2.1", "sponge-cursor", ] @@ -2745,7 +2783,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "crc", "crossbeam-queue", @@ -2766,7 +2804,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "smallvec", - "thiserror 2.0.19", + "thiserror", "tokio", "tokio-stream", "tracing", @@ -2818,7 +2856,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa003f0038df784eb8fecbbac13affe3da23b45194bd57dba231c8f48199c526" dependencies = [ "atoi", - "base64", + "base64 0.22.1", "bitflags", "byteorder", "bytes", @@ -2848,7 +2886,7 @@ dependencies = [ "smallvec", "sqlx-core", "stringprep", - "thiserror 2.0.19", + "thiserror", "tracing", "whoami", ] @@ -2860,7 +2898,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db58fcd5a53cf07c184b154801ff91347e4c30d17a3562a635ff028ad5deda46" dependencies = [ "atoi", - "base64", + "base64 0.22.1", "bitflags", "byteorder", "crc", @@ -2885,7 +2923,7 @@ dependencies = [ "smallvec", "sqlx-core", "stringprep", - "thiserror 2.0.19", + "thiserror", "tracing", "whoami", ] @@ -2909,7 +2947,7 @@ dependencies = [ "serde", "serde_urlencoded", "sqlx-core", - "thiserror 2.0.19", + "thiserror", "tracing", "url", ] @@ -3018,33 +3056,13 @@ dependencies = [ "libc", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl 2.0.19", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] @@ -3060,9 +3078,9 @@ dependencies = [ [[package]] name = "time" -version = "0.3.54" +version = "0.3.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e1d5e639ff6bab73cb6885cc7e7b1de96c3f32c68ec55f3952614bec1092244" +checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134" dependencies = [ "deranged", "js-sys", @@ -3133,13 +3151,13 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.7.1" +version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -3515,7 +3533,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -3728,7 +3746,7 @@ dependencies = [ "rustls-pki-types", "sha2 0.11.0", "socket2", - "thiserror 2.0.19", + "thiserror", "time", "tokio", "tracing", @@ -3745,7 +3763,7 @@ checksum = "d5867c629e4252f7439d82315923daaf27f4fa442410d51b78ab93ef4c432a11" dependencies = [ "httlib-huffman", "octets", - "thiserror 2.0.19", + "thiserror", "url", ] @@ -3755,7 +3773,7 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" dependencies = [ - "curve25519-dalek", + "curve25519-dalek 4.1.3", "rand_core 0.6.4", "serde", "zeroize", @@ -3776,7 +3794,7 @@ dependencies = [ "oid-registry", "ring", "rusticata-macros", - "thiserror 2.0.19", + "thiserror", "time", ] @@ -3815,18 +3833,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.55" +version = "0.8.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.55" +version = "0.8.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" dependencies = [ "proc-macro2", "quote", @@ -3936,9 +3954,9 @@ dependencies = [ [[package]] name = "zlib-rs" -version = "0.6.6" +version = "0.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" [[package]] name = "zmij" diff --git a/migrations/005_nullable_user_iota.sql b/migrations/005_nullable_user_iota.sql new file mode 100644 index 0000000..66e0edb --- /dev/null +++ b/migrations/005_nullable_user_iota.sql @@ -0,0 +1,11 @@ +ALTER TABLE users + DROP FOREIGN KEY fk_users_iota; + +ALTER TABLE users + MODIFY iota_id BIGINT NULL; + +ALTER TABLE users + ADD CONSTRAINT fk_users_iota + FOREIGN KEY (iota_id) + REFERENCES iotas (id) + ON DELETE SET NULL; diff --git a/migrations/006_pending_iota_user_erasure.sql b/migrations/006_pending_iota_user_erasure.sql new file mode 100644 index 0000000..eb72c7b --- /dev/null +++ b/migrations/006_pending_iota_user_erasure.sql @@ -0,0 +1,8 @@ +CREATE TABLE pending_iota_user_erasure ( + user_id BIGINT NOT NULL, + iota_id BIGINT NOT NULL, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (user_id, iota_id), + CONSTRAINT fk_pending_iota_user_erasure_iota + FOREIGN KEY (iota_id) REFERENCES iotas(id) ON DELETE CASCADE +); diff --git a/src/api/response.rs b/src/api/response.rs index e8489b5..965d85e 100644 --- a/src/api/response.rs +++ b/src/api/response.rs @@ -28,7 +28,7 @@ pub struct UserResponse { pub username: String, pub public_key: String, pub user_id: i64, - pub iota_id: i64, + pub iota_id: Option, pub sub_level: i32, pub sub_end: i64, #[serde(skip_serializing_if = "Option::is_none")] @@ -47,7 +47,7 @@ pub struct UsernameResponse { pub username: String, pub public_key: String, pub user_id: i64, - pub iota_id: i64, + pub iota_id: Option, pub sub_level: i32, pub sub_end: i64, } diff --git a/src/db/user_repo.rs b/src/db/user_repo.rs index 01b0e6e..df99a18 100644 --- a/src/db/user_repo.rs +++ b/src/db/user_repo.rs @@ -101,7 +101,7 @@ const USER_COLUMNS: &str = "SELECT id, iota_id, username, display, status, prese #[derive(FromRow)] struct UserRow { id: i64, - iota_id: i64, + iota_id: Option, username: Vec, display: Option>, status: Option>, @@ -124,7 +124,7 @@ impl TryFrom for User { |value| String::from_utf8(value).map_err(|error| sqlx::Error::Decode(Box::new(error))); Ok(User { id: row.id.into(), - iota_id: row.iota_id.into(), + iota_id: row.iota_id.map(IotaId::from), username: decode(row.username)?, display: row.display.map(decode).transpose()?, status: row.status.map(decode).transpose()?, @@ -344,9 +344,9 @@ pub async fn change_presence_preference(id: UserId, value: String) -> Result<()> .await } -pub async fn change_iota_id(id: UserId, value: IotaId) -> Result<()> { +pub async fn change_iota_id(id: UserId, value: Option) -> Result<()> { sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?") - .bind(value.0) + .bind(value.map(|id| id.0)) .bind(id.0) .execute(&pool().await?) .await?; @@ -368,6 +368,42 @@ pub async fn delete_user(id: UserId) -> Result<()> { Ok(()) } +/// Delete the central identity while retaining a durable instruction for the +/// last hosting Iota. The pending row is intentionally independent of users: +/// it must outlive the account row. +pub async fn delete_user_with_pending_erasure(id: UserId) -> Result> { + let mut tx = pool().await?.begin().await?; + let row = sqlx::query("SELECT iota_id FROM users WHERE id = ? FOR UPDATE") + .bind(id.0) + .fetch_optional(&mut *tx) + .await? + .ok_or(OmegaError::NotFound)?; + let iota_id: Option = row.get("iota_id"); + if let Some(iota_id) = iota_id { + sqlx::query("INSERT IGNORE INTO pending_iota_user_erasure (user_id, iota_id) VALUES (?, ?)") + .bind(id.0) + .bind(iota_id) + .execute(&mut *tx) + .await?; + } + sqlx::query("DELETE FROM registration_leases WHERE user_id = ?").bind(id.0).execute(&mut *tx).await?; + sqlx::query("DELETE FROM users WHERE id = ?").bind(id.0).execute(&mut *tx).await?; + tx.commit().await?; + Ok(iota_id.map(IotaId::from)) +} + +pub async fn pending_erasures_for_iota(iota_id: IotaId) -> Result> { + let rows = sqlx::query("SELECT user_id FROM pending_iota_user_erasure WHERE iota_id = ?") + .bind(iota_id.0).fetch_all(&pool().await?).await?; + Ok(rows.into_iter().map(|row| UserId::from(row.get::("user_id"))).collect()) +} + +pub async fn acknowledge_pending_erasure(user_id: UserId, iota_id: IotaId) -> Result { + let result = sqlx::query("DELETE FROM pending_iota_user_erasure WHERE user_id = ? AND iota_id = ?") + .bind(user_id.0).bind(iota_id.0).execute(&pool().await?).await?; + Ok(result.rows_affected() == 1) +} + pub async fn change_keys(id: UserId, public_key: PublicKeyBundle) -> Result<()> { sqlx::query("UPDATE users SET public_key = ? WHERE id = ?") .bind(public_key.as_bytes()) @@ -438,7 +474,7 @@ pub async fn register_complete_user( .map_err(OmegaError::from)? { Some(existing) - if existing.iota_id == iota_id + if existing.iota_id == Some(iota_id) && existing.username == username && existing.public_key.as_bytes() == public_key.as_bytes() && existing.token == token => diff --git a/src/models/user.rs b/src/models/user.rs index 2843b6d..a9595d9 100644 --- a/src/models/user.rs +++ b/src/models/user.rs @@ -4,7 +4,7 @@ use mtp::crypto::PublicKeyBundle; #[derive(Clone, Debug, serde::Serialize)] pub struct User { pub id: UserId, - pub iota_id: IotaId, + pub iota_id: Option, pub username: String, pub display: Option, pub status: Option, diff --git a/src/server/api.rs b/src/server/api.rs index 498685a..0319607 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -41,7 +41,7 @@ fn user_response(user: crate::models::User) -> UserResponse { username: user.username, public_key: user.public_key.to_base64(), user_id: user.id.0, - iota_id: user.iota_id.0, + iota_id: user.iota_id.map(|id| id.0), sub_level: user.sub_level, sub_end: user.sub_end, display: user.display, @@ -83,7 +83,15 @@ async fn route(path_parts: &[&str]) -> Result<(StatusCode, String)> { omikron_id } else { let user = get_by_user_id(UserId::from(id)).await?; - get_iota_primary_omikron_connection(user.iota_id.0).ok_or(OmegaError::NotFound)? + match user.iota_id { + Some(iota_id) => get_iota_primary_omikron_connection(iota_id.0), + None => get_random_omikron() + .await + .map_err(|_| OmegaError::NotFound)? + .get_omikron_id() + .await, + } + .ok_or(OmegaError::NotFound)? }; // Database rows describe registered Omikrons. The public discovery @@ -150,7 +158,7 @@ async fn route(path_parts: &[&str]) -> Result<(StatusCode, String)> { username: user.username, public_key: user.public_key.to_base64(), user_id: user.id.0, - iota_id: user.iota_id.0, + iota_id: user.iota_id.map(|id| id.0), sub_level: user.sub_level, sub_end: user.sub_end, }), diff --git a/src/state.rs b/src/state.rs index dc91021..075d268 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,18 +1,47 @@ use crate::sql::user_online_tracker::PresenceTracker; use std::sync::Arc; +use dashmap::DashMap; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum AccountChallengeOperation { Attach, Delete } + +#[derive(Clone, Debug)] +pub struct AccountChallenge { + pub operation: AccountChallengeOperation, + pub user_id: i64, + pub requester_iota_id: i64, + pub nonce: u64, + pub created_at: Instant, +} pub struct OmegaState { pub presence: Arc, + challenges: DashMap<(AccountChallengeOperation, i64, i64), AccountChallenge>, } impl Default for OmegaState { fn default() -> Self { Self { presence: Arc::new(PresenceTracker::default()), + challenges: DashMap::new(), } } } +impl OmegaState { + pub fn issue_challenge(&self, operation: AccountChallengeOperation, user_id: i64, requester_iota_id: i64) -> u64 { + let nonce = rand::random::(); + self.challenges.insert((operation, user_id, requester_iota_id), AccountChallenge { operation, user_id, requester_iota_id, nonce, created_at: Instant::now() }); + nonce + } + + pub fn consume_challenge(&self, operation: AccountChallengeOperation, user_id: i64, requester_iota_id: i64, nonce: u64) -> bool { + self.challenges.remove(&(operation, user_id, requester_iota_id)).is_some_and(|(_, value)| + value.nonce == nonce && value.created_at.elapsed() <= Duration::from_secs(120)) + } +} + impl OmegaState { pub fn new() -> Arc { Arc::new(Self::default()) diff --git a/src/transport/handlers/account.rs b/src/transport/handlers/account.rs index 1927b14..360076d 100644 --- a/src/transport/handlers/account.rs +++ b/src/transport/handlers/account.rs @@ -2,8 +2,9 @@ use super::super::omikron_connection::{OmikronConnection, OmikronResult}; use crate::{ db::{iota_repo, user_repo}, models::{IotaId, UserId}, + state::AccountChallengeOperation, }; -use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; +use mtp::{codec::{CommunicationType, CommunicationValue, DataType, DataValue}, crypto::{verify_ed25519, verify_ml_dsa}}; use std::sync::Arc; async fn delete( @@ -23,12 +24,8 @@ pub async fn user( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { - delete( - connection, - value.clone(), - user_repo::delete_user(UserId::from(value.get_sender() as i64)), - ) - .await + let user_id = UserId::from(value.get_sender() as i64); + complete_delete(connection, value, user_id).await } pub async fn iota( connection: Arc, @@ -41,3 +38,157 @@ pub async fn iota( ) .await } + +pub async fn release_from_iota( + connection: Arc, + value: CommunicationValue, +) -> OmikronResult<()> { + let Some(user_id) = value + .get_data(DataType::UserId) + .as_signed_number() + .and_then(|id| i64::try_from(id).ok()) + .filter(|id| *id > 0) + else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId) + .await; + }; + let requester = IotaId::from(value.get_sender() as i64); + let Ok(user) = user_repo::get_by_user_id(UserId::from(user_id)).await else { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNotFound) + .await; + }; + if user.iota_id != Some(requester) { + return connection + .send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated) + .await; + } + let previous_iota = user.iota_id; + match user_repo::change_iota_id(user.id, None).await { + Ok(()) => { + if let Some(iota) = previous_iota { crate::transport::omikron_manager::publish_iota_user_snapshot(iota.0).await; } + connection.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())).await + }, + Err(error) => connection + .send(&CommunicationValue::new(CommunicationType::ErrorInternal) + .with_id(value.get_id()) + .add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string()))) + .await, + } +} + +fn lifecycle_payload(domain: &[u8], user_id: i64, iota_id: i64, nonce: u64) -> Vec { + let mut payload = Vec::with_capacity(domain.len() + 32); + payload.extend_from_slice(domain); + payload.extend_from_slice(&user_id.to_be_bytes()); + payload.extend_from_slice(&iota_id.to_be_bytes()); + payload.extend_from_slice(&nonce.to_be_bytes()); + payload +} + +pub async fn attach_begin(connection: Arc, value: CommunicationValue) -> OmikronResult<()> { + let Some(user_id) = value.get_data(DataType::UserId).as_signed_number().and_then(|v| i64::try_from(v).ok()).filter(|v| *v > 0) else { + return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId).await; + }; + if user_repo::get_by_user_id(UserId::from(user_id)).await.is_err() { + return connection.send_error_response(value.get_id(), CommunicationType::ErrorNotFound).await; + } + let requester = value.get_sender() as i64; + let nonce = connection.state().issue_challenge(AccountChallengeOperation::Attach, user_id, requester); + connection.send(&CommunicationValue::new(CommunicationType::AttachUserChallenge).with_id(value.get_id()) + .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())) + .add_typed_default(DataType::ServerNonce, DataValue::SignedNumber(nonce.into()))).await +} + +pub async fn attach_complete(connection: Arc, value: CommunicationValue) -> OmikronResult<()> { + let Some(user_id) = value.get_data(DataType::UserId).as_signed_number().and_then(|v| i64::try_from(v).ok()).filter(|v| *v > 0) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId).await; }; + let requester = value.get_sender() as i64; + let Some(nonce) = value.get_data(DataType::ServerNonce).as_signed_number().and_then(|v| u64::try_from(v).ok()) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge).await; }; + let signature = value.get_data(DataType::Signature).as_bytes(); + let pq_signature = value.get_data(DataType::PqSignature).as_bytes(); + let (Some(signature), Some(pq_signature)) = (signature, pq_signature) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge).await; }; + if !connection.state().consume_challenge(AccountChallengeOperation::Attach, user_id, requester, nonce) { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge).await; } + let Ok(user) = user_repo::get_by_user_id(UserId::from(user_id)).await else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorNotFound).await; }; + let payload = lifecycle_payload(b"tensamin:user-attach:v1\0", user_id, requester, nonce); + if verify_ed25519(&user.public_key.sig_cl_public_key, &payload, &signature).is_err() || verify_ml_dsa(&user.public_key.sig_pq_public_key, &payload, &pq_signature).is_err() { return connection.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated).await; } + let previous_iota = user.iota_id; + match user_repo::change_iota_id(user.id, Some(IotaId::from(requester))).await { + Ok(()) => { + if let Some(iota) = previous_iota.filter(|id| id.0 != requester) { crate::transport::omikron_manager::publish_iota_user_snapshot(iota.0).await; } + crate::transport::omikron_manager::publish_iota_user_snapshot(requester).await; + connection.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())).await + }, + Err(_) => connection.send_error_response(value.get_id(), CommunicationType::ErrorInternal).await, + } +} + +async fn complete_delete(connection: Arc, value: CommunicationValue, user_id: UserId) -> OmikronResult<()> { + match user_repo::delete_user_with_pending_erasure(user_id).await { + Ok(iota_id) => { + let cleanup_pending = iota_id.is_some(); + if let Some(iota_id) = iota_id { + crate::transport::omikron_manager::publish_iota_user_snapshot(iota_id.0).await; + crate::transport::omikron_manager::deliver_pending_erasures(iota_id.0).await; + } + connection.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()) + .add_typed_default(DataType::CleanupPending, DataValue::Bool(cleanup_pending))).await + } + Err(crate::error::OmegaError::NotFound) => connection.send_error_response(value.get_id(), CommunicationType::ErrorNotFound).await, + Err(error) => connection.send(&CommunicationValue::new(CommunicationType::ErrorInternal).with_id(value.get_id()).add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string()))).await, + } +} + +pub async fn delete_credential_begin(connection: Arc, value: CommunicationValue) -> OmikronResult<()> { + let Some(user_id) = value.get_data(DataType::UserId).as_signed_number().and_then(|v| i64::try_from(v).ok()).filter(|v| *v > 0) else { + return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId).await; + }; + if user_repo::get_by_user_id(UserId::from(user_id)).await.is_err() { return connection.send_error_response(value.get_id(), CommunicationType::ErrorNotFound).await; } + let requester = value.get_sender() as i64; + let nonce = connection.state().issue_challenge(AccountChallengeOperation::Delete, user_id, requester); + connection.send(&CommunicationValue::new(CommunicationType::DeleteUserCredentialChallenge).with_id(value.get_id()) + .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())) + .add_typed_default(DataType::ServerNonce, DataValue::SignedNumber(nonce.into()))).await +} + +pub async fn delete_credential_complete(connection: Arc, value: CommunicationValue) -> OmikronResult<()> { + let Some(user_id) = value.get_data(DataType::UserId).as_signed_number().and_then(|v| i64::try_from(v).ok()).filter(|v| *v > 0) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId).await; }; + let requester = value.get_sender() as i64; + let Some(nonce) = value.get_data(DataType::ServerNonce).as_signed_number().and_then(|v| u64::try_from(v).ok()) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge).await; }; + let (Some(signature), Some(pq_signature)) = (value.get_data(DataType::Signature).as_bytes(), value.get_data(DataType::PqSignature).as_bytes()) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge).await; }; + if !connection.state().consume_challenge(AccountChallengeOperation::Delete, user_id, requester, nonce) { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge).await; } + let Ok(user) = user_repo::get_by_user_id(UserId::from(user_id)).await else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorNotFound).await; }; + let payload = lifecycle_payload(b"tensamin:user-delete:v1\0", user_id, requester, nonce); + if verify_ed25519(&user.public_key.sig_cl_public_key, &payload, &signature).is_err() || verify_ml_dsa(&user.public_key.sig_pq_public_key, &payload, &pq_signature).is_err() { return connection.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated).await; } + complete_delete(connection, value, user.id).await +} + +pub async fn erase_hosted_user_data_ack(connection: Arc, value: CommunicationValue) -> OmikronResult<()> { + let Some(user_id) = value.get_data(DataType::UserId).as_signed_number().and_then(|v| i64::try_from(v).ok()).filter(|v| *v > 0) else { return connection.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId).await; }; + let iota_id = IotaId::from(value.get_sender() as i64); + match user_repo::acknowledge_pending_erasure(UserId::from(user_id), iota_id).await { + Ok(true) => connection.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())).await, + Ok(false) => connection.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated).await, + Err(_) => connection.send_error_response(value.get_id(), CommunicationType::ErrorInternal).await, + } +} + +/// New lifecycle operation names are intentionally fail-closed until their +/// proof and durable-erasure handlers are enabled. This explicit dispatch +/// prevents either a bare Iota request or the legacy DeleteUser path from +/// acquiring account-deletion authority during a staged rollout. +pub async fn lifecycle_unavailable( + connection: Arc, + value: CommunicationValue, +) -> OmikronResult<()> { + connection + .send( + &CommunicationValue::new(CommunicationType::ErrorNotAuthenticated) + .with_id(value.get_id()) + .add_typed_default( + DataType::ErrorType, + DataValue::Str("user lifecycle proof handler is not enabled".into()), + ), + ) + .await +} diff --git a/src/transport/handlers/presence.rs b/src/transport/handlers/presence.rs index 0c31904..de8a275 100644 --- a/src/transport/handlers/presence.rs +++ b/src/transport/handlers/presence.rs @@ -53,7 +53,7 @@ fn states_for_users(state: &OmegaState, users: &[crate::models::User]) -> HashMa user.id.0, state .presence - .resolve_public_state(user.id.0, user.iota_id.0), + .resolve_public_state(user.id.0, user.iota_id.map(|id| id.0).unwrap_or_default()), ) }) .collect() @@ -69,7 +69,7 @@ fn changed_states( .filter_map(|user| { let after = state .presence - .resolve_public_state(user.id.0, user.iota_id.0); + .resolve_public_state(user.id.0, user.iota_id.map(|id| id.0).unwrap_or_default()); (before.get(&user.id.0) != Some(&after)).then_some((user.id.0, after)) }) .collect::>(); @@ -275,7 +275,7 @@ pub async fn user_connected( Ok(preferences) => preferences, Err(error) => return Err(error.into()), }; - if user.iota_id.0 != iota_id || !state.presence.has_iota_route(iota_id) { + if user.iota_id.map(|id| id.0) != Some(iota_id) || !state.presence.has_iota_route(iota_id) { return connection .send_error_response(value.get_id(), CommunicationType::ErrorNoIota) .await; @@ -497,6 +497,7 @@ pub async fn iota_connected( .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())) .add_typed_default(DataType::UserIds, DataValue::Array(user_ids)); connection.clone().send(&response).await?; + crate::transport::omikron_manager::deliver_pending_erasures(iota_id).await; publish_changed_states(&state, &before, &users).await; connection .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id())) diff --git a/src/transport/handlers/states.rs b/src/transport/handlers/states.rs index 690a887..d147f27 100644 --- a/src/transport/handlers/states.rs +++ b/src/transport/handlers/states.rs @@ -110,7 +110,7 @@ pub async fn get( }; let status = state .presence - .resolve_public_state(user_id, user.iota_id.0) + .resolve_public_state(user_id, user.iota_id.map(|id| id.0).unwrap_or_default()) .to_string(); let mut map = Vec::new(); if let Some(kind) = DataType::UserId.try_to_id(&tm) { diff --git a/src/transport/handlers/user_data.rs b/src/transport/handlers/user_data.rs index 7e79bee..e331125 100644 --- a/src/transport/handlers/user_data.rs +++ b/src/transport/handlers/user_data.rs @@ -43,7 +43,7 @@ pub async fn get_user( .await; }; let id = user.id.0; - let iota_id = user.iota_id.0; + let iota_id = user.iota_id.map(|id| id.0); let username = user.username.clone(); let display = user .display @@ -57,7 +57,6 @@ pub async fn get_user( DataValue::Str(user.public_key.to_base64()), ) .add_typed_default(DataType::UserId, DataValue::SignedNumber(id.into())) - .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())) .add_typed_default(DataType::Display, DataValue::Str(display)) .add_typed_default( DataType::SubLevel, @@ -92,7 +91,7 @@ pub async fn get_user( } state.presence.resolve_private_state(id) } else { - state.presence.resolve_public_state(id, iota_id) + iota_id.map(|iota_id| state.presence.resolve_public_state(id, iota_id)).unwrap_or(crate::sql::connection_status::UserStatus::user_offline) }; response = response .add_typed_default( @@ -101,8 +100,11 @@ pub async fn get_user( ) .add_typed_default( DataType::OmikronConnections, - connections(&connection, iota_id), + iota_id.map(|iota_id| connections(&connection, iota_id)).unwrap_or_else(|| DataValue::Array(Vec::new())), ); + if let Some(iota_id) = iota_id { + response = response.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())); + } if let Some(route) = route { response = response.add_typed_default( DataType::OmikronId, @@ -123,26 +125,25 @@ pub async fn get_iota( .map(|iota| (iota.id.0, iota.public_key, None, None)) } else if let Some(id) = value.get_data(DataType::UserId).as_number() { if let Ok(user) = user_repo::get_by_user_id(UserId::from(id as i64)).await { - iota_repo::get_iota_by_id(user.iota_id) - .await - .ok() - .map(|iota| (iota.id.0, iota.public_key, Some(user.id.0), None)) + match user.iota_id { + Some(iota_id) => iota_repo::get_iota_by_id(iota_id) + .await + .ok() + .map(|iota| (iota.id.0, iota.public_key, Some(user.id.0), None)), + None => None, + } } else { None } } else if let Some(name) = value.get_data(DataType::Username).as_str() { if let Ok(user) = user_repo::get_by_username(name).await { - iota_repo::get_iota_by_id(user.iota_id) - .await - .ok() - .map(|iota| { - ( - iota.id.0, - iota.public_key, - Some(user.id.0), - Some(name.to_owned()), - ) - }) + match user.iota_id { + Some(iota_id) => iota_repo::get_iota_by_id(iota_id) + .await + .ok() + .map(|iota| (iota.id.0, iota.public_key, Some(user.id.0), Some(name.to_owned()))), + None => None, + } } else { None } @@ -269,7 +270,7 @@ pub async fn change_iota( .await; } let result = - match user_repo::change_iota_id(user_id, IotaId::from(value.get_sender() as i64)).await { + match user_repo::change_iota_id(user_id, Some(IotaId::from(value.get_sender() as i64))).await { Ok(()) => user_repo::change_token(user_id, new_token.to_owned()).await, Err(error) => Err(error), }; diff --git a/src/transport/omikron_connection.rs b/src/transport/omikron_connection.rs index 5001ff3..7ad0136 100644 --- a/src/transport/omikron_connection.rs +++ b/src/transport/omikron_connection.rs @@ -87,6 +87,7 @@ impl OmikronConnection { &self.peer_capabilities } + pub async fn handle(self: Arc, receiver: &mut WebMtpReceiver) { log_in!( self.id as i64, @@ -215,6 +216,24 @@ impl OmikronConnection { Some(CommunicationType::DeleteUser) => { crate::transport::handlers::account::user(self, value).await } + Some(CommunicationType::AttachUserBegin) => { + crate::transport::handlers::account::attach_begin(self, value).await + } + Some(CommunicationType::AttachUserComplete) => { + crate::transport::handlers::account::attach_complete(self, value).await + } + Some(CommunicationType::DeleteUserCredentialBegin) => { + crate::transport::handlers::account::delete_credential_begin(self, value).await + } + Some(CommunicationType::DeleteUserCredentialComplete) => { + crate::transport::handlers::account::delete_credential_complete(self, value).await + } + Some(CommunicationType::EraseHostedUserDataAck) => { + crate::transport::handlers::account::erase_hosted_user_data_ack(self, value).await + } + Some(CommunicationType::ReleaseUserFromIota) => { + crate::transport::handlers::account::release_from_iota(self, value).await + } Some(CommunicationType::DeleteIota) => { crate::transport::handlers::account::iota(self, value).await } diff --git a/src/transport/omikron_manager.rs b/src/transport/omikron_manager.rs index 254ced8..9020556 100644 --- a/src/transport/omikron_manager.rs +++ b/src/transport/omikron_manager.rs @@ -3,7 +3,7 @@ use crate::state::OmegaState; use crate::transport::connection::OmikronConnection; use crate::transport::omikron_connection::OmikronResult; use dashmap::DashMap; -use mtp::codec::CommunicationValue; +use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use once_cell::sync::Lazy; use rand::prelude::IteratorRandom; use std::sync::Arc; @@ -65,7 +65,8 @@ pub async fn get_all_connections() for user in users { for route in state.presence.routes_for_user(user.id.0) { if let Some(iotas) = result.get_mut(&route.omikron_id) { - if let Some(users) = iotas.get_mut(&user.iota_id.0) { + if let Some(iota_id) = user.iota_id + && let Some(users) = iotas.get_mut(&iota_id.0) { users.push(user.id.0); } } @@ -113,3 +114,28 @@ pub async fn send_to_user(user_id: i64, cv: &CommunicationValue) { } } } + +/// Publish the authoritative membership list after an attach, migration, or +/// release. Omikron replaces its full local index from this snapshot. +pub async fn publish_iota_user_snapshot(iota_id: i64) { + let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) else { return; }; + let Some(connection) = get_connected_omikron(omikron_id) else { return; }; + let Ok(users) = user_repo::get_users_by_iota_id(crate::models::IotaId::from(iota_id)).await else { return; }; + let user_ids = users.into_iter().map(|user| DataValue::SignedNumber(user.id.0.into())).collect(); + let snapshot = CommunicationValue::new(CommunicationType::IotaUserData) + .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())) + .add_typed_default(DataType::UserIds, DataValue::Array(user_ids)); + let _ = connection.send(&snapshot).await; +} + +pub async fn deliver_pending_erasures(iota_id: i64) { + let Ok(users) = user_repo::pending_erasures_for_iota(crate::models::IotaId::from(iota_id)).await else { return; }; + let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) else { return; }; + let Some(connection) = get_connected_omikron(omikron_id) else { return; }; + for user_id in users { + let request = CommunicationValue::new(CommunicationType::EraseHostedUserData) + .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.0.into())) + .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())); + let _ = connection.clone().send(&request).await; + } +} From a3cb1729f9b177d2db96ef636a995a1de545a533 Mon Sep 17 00:00:00 2001 From: Rasensprenger Date: Sun, 9 Aug 2026 21:01:32 +0200 Subject: [PATCH 5/5] Update Rust crate base64 to 0.23.0 --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1e03155..f367654 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1888,7 +1888,7 @@ name = "omega" version = "0.1.0" dependencies = [ "ansi_term", - "base64 0.22.1", + "base64 0.23.1", "bytes", "dashmap", "dotenv", diff --git a/Cargo.toml b/Cargo.toml index 3b4f6a3..bb9bca8 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ mtp = { git = "https://git.methanium.net/methanium/mtp", features = [ ] } ansi_term = "0.12.1" -base64 = "0.22.1" +base64 = "0.23.0" bytes = "1" dashmap = "6.2.1" dotenv = "0.15.0"