From b65b22dfbc322487226f8a3c0480759b6f3ed7b9 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 7 Aug 2026 23:54:34 +0200 Subject: [PATCH] [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; + } } } }