[Fix] User States

This commit is contained in:
Alex 2026-08-07 23:54:34 +02:00
commit c447e6e863
16 changed files with 2582 additions and 370 deletions

View file

@ -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'
));

@ -1 +1 @@
Subproject commit ece6e2c3b4e925f3cefe46f4a048fbfc8f823093 Subproject commit 486541b9483356ff49ff3ec7016f87d3ecbeaa0e

View file

@ -2,9 +2,11 @@ use crate::{
db::pool, db::pool,
error::{OmegaError, Result}, error::{OmegaError, Result},
models::{IotaId, User, UserId}, models::{IotaId, User, UserId},
sql::connection_status::UserStatus,
}; };
use mtp::crypto::PublicKeyBundle; 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; pub const MAX_PROTOCOL_ID: i64 = (1_i64 << 48) - 1;
const ID_ALLOCATION_ATTEMPTS: usize = 16; 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_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, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE id = ?"; 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 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_COLUMNS: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users";
#[derive(FromRow)] #[derive(FromRow)]
struct UserRow { struct UserRow {
@ -103,6 +105,7 @@ struct UserRow {
username: Vec<u8>, username: Vec<u8>,
display: Option<Vec<u8>>, display: Option<Vec<u8>>,
status: Option<Vec<u8>>, status: Option<Vec<u8>>,
presence_preference: Vec<u8>,
about: Option<Vec<u8>>, about: Option<Vec<u8>>,
avatar: Option<Vec<u8>>, avatar: Option<Vec<u8>>,
sub_level: i32, sub_level: i32,
@ -125,6 +128,7 @@ impl TryFrom<UserRow> for User {
username: decode(row.username)?, username: decode(row.username)?,
display: row.display.map(decode).transpose()?, display: row.display.map(decode).transpose()?,
status: row.status.map(decode).transpose()?, status: row.status.map(decode).transpose()?,
presence_preference: decode(row.presence_preference)?,
about: row.about.map(decode).transpose()?, about: row.about.map(decode).transpose()?,
avatar: row.avatar, avatar: row.avatar,
sub_level: row.sub_level, sub_level: row.sub_level,
@ -154,15 +158,129 @@ pub async fn get_by_user_id(id: UserId) -> Result<User> {
} }
pub async fn get_users_by_iota_id(id: IotaId) -> Result<Vec<User>> { pub async fn get_users_by_iota_id(id: IotaId) -> Result<Vec<User>> {
let rows = sqlx::query_as::<_, UserRow>(USERS_BY_IOTA_ID_QUERY) get_users_by_iota_ids(&[id.0]).await
.bind(id.0) }
fn normalized_ids(ids: &[i64]) -> Vec<i64> {
let mut ids = ids
.iter()
.copied()
.filter(|id| valid_protocol_id(*id))
.collect::<Vec<_>>();
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<Vec<User>> {
query
.build_query_as::<UserRow>()
.fetch_all(&pool().await?) .fetch_all(&pool().await?)
.await?; .await?
rows.into_iter() .into_iter()
.map(|row| row.try_into().map_err(OmegaError::from)) .map(|row| row.try_into().map_err(OmegaError::from))
.collect() .collect()
} }
pub async fn get_users_by_ids(ids: &[i64]) -> Result<Vec<User>> {
let ids = normalized_ids(ids);
if ids.is_empty() {
return Ok(Vec::new());
}
let mut query = QueryBuilder::<MySql>::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<Vec<User>> {
let ids = normalized_ids(ids);
if ids.is_empty() {
return Ok(Vec::new());
}
let mut query = QueryBuilder::<MySql>::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<Vec<User>> {
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::<MySql>::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<u8>,
}
pub async fn get_presence_preferences(ids: &[i64]) -> Result<HashMap<i64, UserStatus>> {
let ids = normalized_ids(ids);
if ids.is_empty() {
return Ok(HashMap::new());
}
let mut query =
QueryBuilder::<MySql>::new("SELECT id, presence_preference FROM users WHERE id IN ");
append_in_clause(&mut query, &ids);
let rows = query
.build_query_as::<PresencePreferenceRow>()
.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( async fn update(
id: UserId, id: UserId,
query: &'static str, query: &'static str,
@ -217,6 +335,15 @@ pub async fn change_status(id: UserId, value: String) -> Result<()> {
.await .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<()> { pub async fn change_iota_id(id: UserId, value: IotaId) -> Result<()> {
sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?") sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?")
.bind(value.0) .bind(value.0)

View file

@ -5,12 +5,14 @@ pub mod error;
mod models; mod models;
mod server; mod server;
mod sql; mod sql;
mod state;
mod transport; mod transport;
mod util; mod util;
pub use error::{OmegaError, Result}; pub use error::{OmegaError, Result};
use crate::db::initialize; use crate::db::initialize;
use crate::state::OmegaState;
use crate::transport::omikron_connection; use crate::transport::omikron_connection;
use crate::util::file_util::get_directory; use crate::util::file_util::get_directory;
use crate::util::logger::PrintType; use crate::util::logger::PrintType;
@ -94,7 +96,7 @@ async fn main() {
.unwrap_or(443); .unwrap_or(443);
tokio::select! { tokio::select! {
result = omikron_connection::start(port) => { result = omikron_connection::start(port, OmegaState::new()) => {
if let Err(e) = result { if let Err(e) = result {
log_err!(0, PrintType::General, "Server error: {:?}", e); log_err!(0, PrintType::General, "Server error: {:?}", e);
} }

View file

@ -8,6 +8,7 @@ pub struct User {
pub username: String, pub username: String,
pub display: Option<String>, pub display: Option<String>,
pub status: Option<String>, pub status: Option<String>,
pub presence_preference: String,
pub about: Option<String>, pub about: Option<String>,
pub avatar: Option<Vec<u8>>, pub avatar: Option<Vec<u8>>,
pub sub_level: i32, pub sub_level: i32,

View file

@ -14,8 +14,10 @@ use crate::server::{
middleware, middleware,
validation::{parse_positive_id, validate_non_empty}, 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::{
use crate::transport::omikron_manager::{get_connected_omikron, get_random_omikron}; get_all_connections, get_connected_omikron, get_iota_primary_omikron_connection,
get_random_omikron,
};
use crate::util::file_util::get_directory; use crate::util::file_util::get_directory;
use base64::Engine as _; use base64::Engine as _;
use bytes::Bytes; use bytes::Bytes;

View file

@ -1,7 +1,4 @@
use strum::IntoEnumIterator; #[derive(Debug, Clone, PartialEq, Eq)]
use strum_macros::EnumIter;
#[derive(Debug, Clone, PartialEq, EnumIter, Eq)]
#[allow(unused, non_camel_case_types)] #[allow(unused, non_camel_case_types)]
pub enum UserStatus { pub enum UserStatus {
user_offline, user_offline,
@ -20,12 +17,70 @@ impl UserStatus {
pub fn to_string(&self) -> String { pub fn to_string(&self) -> String {
format!("{:?}", self) format!("{:?}", self)
} }
pub fn from_str(s: &str) -> Option<UserStatus> { pub fn from_client_preference(s: &str) -> Option<Self> {
for sel in UserStatus::iter() { match s {
if &sel.to_string() == s { "user_online" => Some(Self::user_online),
return Some(sel); "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> {
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);
} }
} }

File diff suppressed because it is too large Load diff

72
src/state.rs Normal file
View file

@ -0,0 +1,72 @@
use crate::sql::user_online_tracker::PresenceTracker;
use std::sync::Arc;
pub struct OmegaState {
pub presence: Arc<PresenceTracker>,
}
impl Default for OmegaState {
fn default() -> Self {
Self {
presence: Arc::new(PresenceTracker::default()),
}
}
}
impl OmegaState {
pub fn new() -> Arc<Self> {
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
);
}
}

View file

@ -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<Self, ()> {
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<PeerCapabilities, ()> {
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());
}
}

View file

@ -1,69 +1,207 @@
use super::super::omikron_connection::{OmikronConnection, OmikronResult}; use super::super::omikron_connection::{OmikronConnection, OmikronResult};
use crate::{ use crate::{
db::user_repo, db::user_repo, log_in, models::IotaId, sql::connection_status::UserStatus, state::OmegaState,
log_in,
models::IotaId,
sql::{connection_status::UserStatus, user_online_tracker},
}; };
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use std::sync::Arc; use std::{
collections::{BTreeMap, HashMap, HashSet},
sync::Arc,
};
pub async fn user_connected( fn parse_subscription(value: &CommunicationValue) -> Result<(i64, i64, Vec<i64>), &'static str> {
_connection: Arc<OmikronConnection>, let user_id = i64::try_from(value.get_sender())
value: CommunicationValue, .ok()
omikron_id: i64, .filter(|id| *id > 0)
) -> OmikronResult<()> { .ok_or("user_id")?;
log_in!(crate::util::logger::PrintType::Omega, "User connected"); let session_id = value
if let Some(user_id) = value.get_data(DataType::UserId).as_number() { .get_data(DataType::SessionId)
let status = value .as_number()
.get_data(DataType::UserState) .and_then(|id| i64::try_from(id).ok())
.as_str() .filter(|id| *id > 0)
.and_then(UserStatus::from_str) .ok_or("session_id")?;
.unwrap_or(UserStatus::user_online); let DataValue::Array(values) = value.get_data(DataType::UserIds) else {
if let Ok(user_id) = i64::try_from(user_id) { return Err("user_ids");
if let Some(session_id) = value };
.get_data(DataType::SessionId)
.as_number() let mut user_ids = Vec::with_capacity(values.len());
.and_then(|id| i64::try_from(id).ok()) for value in values {
.filter(|id| *id > 0) let DataValue::SignedNumber(user_id) = value else {
{ return Err("user_ids");
user_online_tracker::track_user_session_status( };
user_id, session_id, status, omikron_id, let Ok(user_id) = i64::try_from(*user_id) else {
); return Err("user_ids");
} else { };
user_online_tracker::track_user_status(user_id, status, omikron_id); 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( fn apply_preferences(state: &OmegaState, preferences: HashMap<i64, UserStatus>) {
_: Arc<OmikronConnection>, state.presence.set_preferences(preferences);
value: CommunicationValue, }
omikron_id: i64,
) -> OmikronResult<()> { fn states_for_users(state: &OmegaState, users: &[crate::models::User]) -> HashMap<i64, UserStatus> {
log_in!(crate::util::logger::PrintType::Omega, "User disconnected"); users
if let Some(user_id) = value.get_data(DataType::UserId).as_number() { .iter()
if let Some(session_id) = value .map(|user| {
.get_data(DataType::SessionId) (
.as_number() user.id.0,
.and_then(|id| i64::try_from(id).ok()) state
.filter(|id| *id > 0) .presence
.resolve_public_state(user.id.0, user.iota_id.0),
)
})
.collect()
}
fn changed_states(
state: &OmegaState,
before: &HashMap<i64, UserStatus>,
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::<Vec<_>>();
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::<i64, Vec<CommunicationValue>>::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( log_in!(
user_id as i64, crate::util::logger::PrintType::General,
session_id, "Failed to deliver presence state batch to Omikron {}: {}",
omikron_id, 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<i64, UserStatus>,
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::<i64, Vec<CommunicationValue>>::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<OmegaState>,
connection: Arc<OmikronConnection>,
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<OmegaState>,
_: Arc<OmikronConnection>, _: Arc<OmikronConnection>,
value: CommunicationValue, value: CommunicationValue,
_: i64, _: i64,
@ -72,29 +210,266 @@ pub async fn client_changed(
.get_data(DataType::UserId) .get_data(DataType::UserId)
.as_number() .as_number()
.and_then(|id| i64::try_from(id).ok()) .and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
else { else {
return Ok(()); return Ok(());
}; };
let Some(status) = value let Some(status) = value
.get_data(DataType::UserState) .get_data(DataType::UserState)
.as_str() .as_str()
.and_then(UserStatus::from_str) .and_then(UserStatus::from_client_preference)
else { else {
return Ok(()); return Ok(());
}; };
// Connectivity is derived from routes. Clients may choose only public state.presence.set_preference(user_id, status);
// 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);
Ok(()) Ok(())
} }
pub async fn user_connected(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
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<OmegaState>,
connection: Arc<OmikronConnection>,
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<OmegaState>,
connection: Arc<OmikronConnection>,
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( pub async fn iota_connected(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>, connection: Arc<OmikronConnection>,
value: CommunicationValue, value: CommunicationValue,
omikron_id: i64, omikron_id: i64,
@ -105,36 +480,32 @@ pub async fn iota_connected(
.as_number() .as_number()
.map(|id| id as i64) .map(|id| id as i64)
else { 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 users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?;
let mut user_ids = Vec::new(); let ids = users.iter().map(|user| user.id.0).collect::<Vec<_>>();
match user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await { apply_preferences(&state, user_repo::get_presence_preferences(&ids).await?);
Ok(users) => { let before = states_for_users(&state, &users);
for user in users { state.presence.connect_iota(iota_id, omikron_id);
user_ids.push(DataValue::SignedNumber(user.id.0.into())); let user_ids = users
user_online_tracker::track_user_status( .iter()
user.id.0, .map(|user| DataValue::SignedNumber(user.id.0.into()))
UserStatus::user_offline, .collect();
omikron_id,
);
}
}
Err(_) => log_in!(
crate::util::logger::PrintType::General,
"SQL error loading users for IOTA"
),
}
let response = CommunicationValue::new(CommunicationType::IotaUserData) 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::IotaId, DataValue::SignedNumber(iota_id.into()))
.add_typed_default(DataType::UserIds, DataValue::Array(user_ids)); .add_typed_default(DataType::UserIds, DataValue::Array(user_ids));
let _ = connection.send(&response).await; connection.clone().send(&response).await?;
Ok(()) publish_changed_states(&state, &before, &users).await;
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.await
} }
pub async fn iota_disconnected( pub async fn iota_disconnected(
_: Arc<OmikronConnection>, state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
value: CommunicationValue, value: CommunicationValue,
omikron_id: i64, omikron_id: i64,
) -> OmikronResult<()> { ) -> OmikronResult<()> {
@ -144,40 +515,317 @@ pub async fn iota_disconnected(
.as_number() .as_number()
.map(|id| id as i64) .map(|id| id as i64)
else { 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) { let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?;
if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await { let ids = users.iter().map(|user| user.id.0).collect::<Vec<_>>();
user_online_tracker::untrack_many_users( apply_preferences(&state, user_repo::get_presence_preferences(&ids).await?);
&users.iter().map(|user| user.id.0).collect::<Vec<_>>(), let before = states_for_users(&state, &users);
); state.presence.untrack_iota_connection(iota_id, omikron_id);
} publish_changed_states(&state, &before, &users).await;
} connection
Ok(()) .send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.await
} }
pub async fn sync_status( pub async fn sync_status(
_: Arc<OmikronConnection>, state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
value: CommunicationValue, value: CommunicationValue,
omikron_id: i64, omikron_id: i64,
) -> OmikronResult<()> { ) -> OmikronResult<()> {
if let DataValue::Array(ids) = value.get_data(DataType::UserIds) { let request_id = value.get_id();
for id in ids { let DataValue::Array(iota_values) = value.get_data(DataType::IotaIds) else {
if let DataValue::SignedNumber(id) = id { return connection
user_online_tracker::track_user_status( .send_error_response(request_id, CommunicationType::ErrorInvalidData)
*id as i64, .await;
UserStatus::user_offline, };
omikron_id, 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 !connection.peer_capabilities().session_snapshot_v1 {
if let DataValue::SignedNumber(id) = id { let DataValue::Array(user_values) = value.get_data(DataType::UserIds) else {
user_online_tracker::track_iota_connection(*id as i64, omikron_id, true); 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::<HashSet<_>>();
let users = user_repo::get_users_by_ids(&user_ids).await?;
let returned_user_ids = users.iter().map(|user| user.id.0).collect::<Vec<_>>();
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::<Vec<_>>(),
)
.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::<HashSet<_>>();
let new_session_user_ids = sessions
.iter()
.map(|(user_id, _, _)| *user_id)
.collect::<HashSet<_>>();
let affected_iota_ids = previous_iota_ids
.iter()
.chain(iota_ids.iter())
.copied()
.collect::<HashSet<_>>();
let users = user_repo::get_users_by_ids_and_iota_ids(
&previous_session_user_ids
.iter()
.chain(new_session_user_ids.iter())
.copied()
.collect::<Vec<_>>(),
&affected_iota_ids.iter().copied().collect::<Vec<_>>(),
)
.await?;
let user_ids = users.iter().map(|user| user.id.0).collect::<Vec<_>>();
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<OmegaState>, 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::<Vec<_>>();
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::<Vec<_>>();
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(())
} }

View file

@ -1,46 +1,146 @@
use super::super::omikron_connection::{OmikronConnection, OmikronResult}; use super::super::omikron_connection::{OmikronConnection, OmikronResult};
use crate::sql::{connection_status::UserStatus, user_online_tracker}; use crate::db::user_repo;
use mtp::{ use mtp::{
codec::{CommunicationType, CommunicationValue, DataType, DataValue}, codec::{CommunicationType, CommunicationValue, DataType, DataValue},
type_map::TypeMap, type_map::TypeMap,
}; };
use std::sync::Arc; use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
async fn send_error(
connection: Arc<OmikronConnection>,
request_id: u32,
error_type: CommunicationType,
session_id: Option<i128>,
) -> 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( pub async fn get(
connection: Arc<OmikronConnection>, connection: Arc<OmikronConnection>,
value: CommunicationValue, value: CommunicationValue,
) -> OmikronResult<()> { ) -> 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 { 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 tm = TypeMap::latest();
let states = ids let mut requested_user_ids = Vec::new();
.iter() let mut requested_set = HashSet::new();
.filter_map(|id| { for id in ids {
let DataValue::SignedNumber(id) = id else { let DataValue::SignedNumber(id) = id else {
return None; return send_error(
}; connection,
let status = user_online_tracker::get_user_status(*id as i64) value.get_id(),
.map(|status| { CommunicationType::ErrorInvalidData,
if status.connection_type == UserStatus::user_invisible { session_id,
UserStatus::user_offline.to_string() )
} else { .await;
status.connection_type.to_string() };
} let Ok(user_id) = i64::try_from(*id) else {
}) return send_error(
.unwrap_or_else(|| UserStatus::iota_offline.to_string()); connection,
let mut map = Vec::new(); value.get_id(),
if let Some(kind) = DataType::UserId.try_to_id(&tm) { CommunicationType::ErrorInvalidData,
map.push((kind, DataValue::SignedNumber((*id as i64).into()))); session_id,
} )
if let Some(kind) = DataType::UserState.try_to_id(&tm) { .await;
map.push((kind, DataValue::Str(status))); };
} if user_id <= 0 {
Some(DataValue::Container(map)) return send_error(
}) connection,
.collect(); 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) let response = CommunicationValue::new(CommunicationType::GetStates)
.with_id(value.get_id()) .with_id(value.get_id())
.add_typed_default(DataType::UserStates, DataValue::Array(states)); .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 connection.send(&response).await
} }

View file

@ -2,7 +2,6 @@ use super::super::omikron_connection::{OmikronConnection, OmikronResult};
use crate::{ use crate::{
db::{iota_repo, user_repo}, db::{iota_repo, user_repo},
models::{IotaId, UserId}, models::{IotaId, UserId},
sql::{connection_status::UserStatus, user_online_tracker},
}; };
use base64::{Engine as _, engine::general_purpose::STANDARD}; use base64::{Engine as _, engine::general_purpose::STANDARD};
use mtp::{ use mtp::{
@ -11,9 +10,12 @@ use mtp::{
}; };
use std::sync::Arc; use std::sync::Arc;
fn connections(iota_id: i64) -> DataValue { fn connections(connection: &OmikronConnection, iota_id: i64) -> DataValue {
DataValue::Array( DataValue::Array(
user_online_tracker::get_iota_omikron_connections(iota_id) connection
.state()
.presence
.iota_connections(iota_id)
.unwrap_or_default() .unwrap_or_default()
.into_iter() .into_iter()
.map(|id| DataValue::SignedNumber(id.into())) .map(|id| DataValue::SignedNumber(id.into()))
@ -25,6 +27,7 @@ pub async fn get_user(
connection: Arc<OmikronConnection>, connection: Arc<OmikronConnection>,
value: CommunicationValue, value: CommunicationValue,
) -> OmikronResult<()> { ) -> OmikronResult<()> {
let state = connection.state();
let user = if let Some(id) = value.get_data(DataType::UserId).as_number() { let user = if let Some(id) = value.get_data(DataType::UserId).as_number() {
user_repo::get_by_user_id(UserId::from(id as i64)) user_repo::get_by_user_id(UserId::from(id as i64))
.await .await
@ -74,28 +77,36 @@ pub async fn get_user(
response = response =
response.add_typed_default(DataType::Avatar, DataValue::Str(STANDARD.encode(avatar))); 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 response = response
.add_typed_default( .add_typed_default(
DataType::OnlineStatus, DataType::OnlineStatus,
DataValue::Str( DataValue::Str(resolved_status.to_string()),
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()),
),
) )
.add_typed_default(DataType::OmikronConnections, connections(iota_id)); .add_typed_default(
if let Some(status) = online { DataType::OmikronConnections,
connections(&connection, iota_id),
);
if let Some(route) = route {
response = response.add_typed_default( response = response.add_typed_default(
DataType::OmikronId, DataType::OmikronId,
DataValue::SignedNumber(status.omikron_id.into()), DataValue::SignedNumber(route.omikron_id.into()),
); );
} }
connection.send(&response).await connection.send(&response).await
@ -147,7 +158,7 @@ pub async fn get_iota(
.with_id(value.get_id()) .with_id(value.get_id())
.add_typed_default(DataType::PublicKey, DataValue::Str(key.to_base64())) .add_typed_default(DataType::PublicKey, DataValue::Str(key.to_base64()))
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into())) .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 { if let Some(user_id) = user_id {
response = response =
response.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())); response.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()));

View file

@ -1,3 +1,4 @@
pub mod capabilities;
pub mod connection; pub mod connection;
pub mod handlers; pub mod handlers;
pub mod omikron_connection; pub mod omikron_connection;

View file

@ -1,6 +1,8 @@
use super::capabilities::{OmegaCapabilities, PeerCapabilities};
use crate::models::OmikronId; use crate::models::OmikronId;
use crate::{ use crate::{
load_keyring, log, log_cv_in, log_cv_out, log_err, log_in, server, load_keyring, log, log_cv_in, log_cv_out, log_err, log_in, server,
state::OmegaState,
transport::omikron_manager, transport::omikron_manager,
util::{file_util::load_file_vec, logger::PrintType}, util::{file_util::load_file_vec, logger::PrintType},
}; };
@ -48,9 +50,11 @@ pub struct WaitingTask {
pub struct OmikronConnection { pub struct OmikronConnection {
id: u64, id: u64,
state: Arc<OmegaState>,
sender: Mutex<Option<WebMtpSender>>, sender: Mutex<Option<WebMtpSender>>,
waiting_tasks: DashMap<u32, WaitingTask>, waiting_tasks: DashMap<u32, WaitingTask>,
cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>, cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
peer_capabilities: PeerCapabilities,
} }
impl Drop for OmikronConnection { impl Drop for OmikronConnection {
fn drop(&mut self) { fn drop(&mut self) {
@ -61,13 +65,26 @@ impl Drop for OmikronConnection {
} }
impl OmikronConnection { impl OmikronConnection {
pub fn new(sender: WebMtpSender, id: u64) -> Arc<Self> { pub fn new(
Arc::new(Self { sender: WebMtpSender,
id: u64,
description: Option<&str>,
state: Arc<OmegaState>,
) -> Option<Arc<Self>> {
let peer_capabilities =
PeerCapabilities::from_identification_description(description).ok()?;
Some(Arc::new(Self {
id, id,
state,
sender: Mutex::new(Some(sender)), sender: Mutex::new(Some(sender)),
waiting_tasks: DashMap::new(), waiting_tasks: DashMap::new(),
cleanup_handle: std::sync::Mutex::new(None), cleanup_handle: std::sync::Mutex::new(None),
}) peer_capabilities,
}))
}
pub fn peer_capabilities(&self) -> &PeerCapabilities {
&self.peer_capabilities
} }
pub async fn handle(self: Arc<Self>, receiver: &mut WebMtpReceiver) { pub async fn handle(self: Arc<Self>, receiver: &mut WebMtpReceiver) {
@ -76,6 +93,23 @@ impl OmikronConnection {
PrintType::Omega, PrintType::Omega,
"Omikron connection started" "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(); let cleanup_conn = self.clone();
*self.cleanup_handle.lock().unwrap() = Some(tokio::spawn(async move { *self.cleanup_handle.lock().unwrap() = Some(tokio::spawn(async move {
let mut ticker = interval(CLEANUP_INTERVAL); let mut ticker = interval(CLEANUP_INTERVAL);
@ -132,27 +166,30 @@ impl OmikronConnection {
async fn dispatch(self: Arc<Self>, value: CommunicationValue) -> OmikronResult<()> { async fn dispatch(self: Arc<Self>, value: CommunicationValue) -> OmikronResult<()> {
let id = self.id as i64; let id = self.id as i64;
let state = self.state.clone();
match value.get_comm_type_enum() { match value.get_comm_type_enum() {
Some(CommunicationType::ShortenLink) => { Some(CommunicationType::ShortenLink) => {
crate::transport::handlers::links::shorten(self, value).await crate::transport::handlers::links::shorten(self, value).await
} }
Some(CommunicationType::UserConnected) => { 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) => { 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) => { Some(CommunicationType::SetUserState) => {
crate::transport::handlers::presence::client_changed(self, value, id).await crate::transport::handlers::presence::set_user_state(state, self, value, id).await
} }
Some(CommunicationType::IotaConnected) => { 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) => { 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) => { 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) => { Some(CommunicationType::GetUserData) => {
crate::transport::handlers::user_data::get_user(self, value).await crate::transport::handlers::user_data::get_user(self, value).await
@ -193,6 +230,13 @@ impl OmikronConnection {
Some(CommunicationType::GetStates) => { Some(CommunicationType::GetStates) => {
crate::transport::handlers::states::get(self, value).await 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!( log_err!(
0, 0,
@ -216,6 +260,24 @@ impl OmikronConnection {
.await .await
.map_err(|error| crate::error::OmegaError::SendError(error.to_string())) .map_err(|error| crate::error::OmegaError::SendError(error.to_string()))
} }
pub(crate) async fn send_messages(
self: Arc<Self>,
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( pub(crate) async fn send_error_response(
self: Arc<Self>, self: Arc<Self>,
message_id: u32, message_id: u32,
@ -224,6 +286,22 @@ impl OmikronConnection {
self.send(&CommunicationValue::new(error_type).with_id(message_id)) self.send(&CommunicationValue::new(error_type).with_id(message_id))
.await .await
} }
pub(crate) async fn send_error_response_with_detail(
self: Arc<Self>,
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<Self>) { pub async fn close(self: Arc<Self>) {
log_in!( log_in!(
self.id as i64, self.id as i64,
@ -238,7 +316,11 @@ impl OmikronConnection {
if self.id != 0 { if self.id != 0 {
log_in!(self.id as i64, PrintType::Omega, "Omikron disconnected"); log_in!(self.id as i64, PrintType::Omega, "Omikron disconnected");
if omikron_manager::remove_omikron(self.id as i64, &self).await { 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() { if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
@ -248,12 +330,19 @@ impl OmikronConnection {
pub async fn get_omikron_id(self: Arc<Self>) -> Option<i64> { pub async fn get_omikron_id(self: Arc<Self>) -> Option<i64> {
Some(self.id as i64) Some(self.id as i64)
} }
pub fn state(&self) -> Arc<OmegaState> {
self.state.clone()
}
pub async fn send_message(self: Arc<Self>, value: &CommunicationValue) -> OmikronResult<()> { pub async fn send_message(self: Arc<Self>, value: &CommunicationValue) -> OmikronResult<()> {
self.send(value).await self.send(value).await
} }
} }
pub async fn get_by_omikron_id(omikron_id: u64, _: Option<String>) -> Option<PublicKeyBundle> { pub async fn get_by_omikron_id(
omikron_id: u64,
description: Option<String>,
) -> Option<PublicKeyBundle> {
PeerCapabilities::from_identification_description(description.as_deref()).ok()?;
crate::db::omikron_repo::get_omikron_by_id(OmikronId::from(omikron_id as i64)) crate::db::omikron_repo::get_omikron_by_id(OmikronId::from(omikron_id as i64))
.await .await
.ok() .ok()
@ -263,7 +352,7 @@ pub async fn complete_register(_: PublicKeyBundle, _: Option<String>) -> u64 {
0 0
} }
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> { pub async fn start(port: u16, state: Arc<OmegaState>) -> Result<(), Box<dyn std::error::Error>> {
let cert_pem = load_file_vec("certs", "cert.pem")?; let cert_pem = load_file_vec("certs", "cert.pem")?;
let key_pem = load_file_vec("certs", "key.pem")?; let key_pem = load_file_vec("certs", "key.pem")?;
let web_config = server::server::build_web_config()? let web_config = server::server::build_web_config()?
@ -331,7 +420,19 @@ pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
); );
continue; 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 { tokio::spawn(async move {
let _guard = ConnectionLimitGuard(peer_ip); let _guard = ConnectionLimitGuard(peer_ip);
omikron_manager::add_omikron(connection.clone()).await; omikron_manager::add_omikron(connection.clone()).await;

View file

@ -1,4 +1,7 @@
use crate::db::user_repo;
use crate::state::OmegaState;
use crate::transport::connection::OmikronConnection; use crate::transport::connection::OmikronConnection;
use crate::transport::omikron_connection::OmikronResult;
use dashmap::DashMap; use dashmap::DashMap;
use mtp::codec::CommunicationValue; use mtp::codec::CommunicationValue;
use once_cell::sync::Lazy; use once_cell::sync::Lazy;
@ -34,6 +37,61 @@ pub fn get_connected_omikron(omikron_id: i64) -> Option<Arc<OmikronConnection>>
.map(|connection| connection.clone()) .map(|connection| connection.clone())
} }
pub fn get_state() -> Option<Arc<OmegaState>> {
OMIKRON_CONNECTIONS
.iter()
.next()
.map(|connection| connection.value().state())
}
pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
get_state().and_then(|state| state.presence.primary_iota_route(iota_id))
}
pub async fn get_all_connections()
-> Result<std::collections::HashMap<i64, std::collections::HashMap<i64, Vec<i64>>>, ()> {
match get_state() {
Some(state) => {
let mut result = state.presence.connection_routes();
let iota_ids = state
.presence
.all_iota_routes()
.keys()
.copied()
.collect::<Vec<_>>();
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<CommunicationValue>,
) -> OmikronResult<()> {
let connection =
get_connected_omikron(omikron_id).ok_or(crate::error::OmegaError::NotConnected)?;
connection.send_messages(&notifications).await
}
pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> { pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect(); let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
@ -47,9 +105,11 @@ pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
} }
pub async fn send_to_user(user_id: i64, cv: &CommunicationValue) { 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(state) = get_state() {
if let Some(omikron_conn) = OMIKRON_CONNECTIONS.get(&user_conn.omikron_id) { for user_route in state.presence.routes_for_user(user_id) {
let _ = omikron_conn.value().clone().send_message(cv).await; if let Some(omikron_conn) = OMIKRON_CONNECTIONS.get(&user_route.omikron_id) {
let _ = omikron_conn.value().clone().send_message(cv).await;
}
} }
} }
} }