omega/src/transport/handlers/presence.rs
2026-08-20 17:05:37 +02:00

672 lines
24 KiB
Rust

use super::super::connection::{
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::{
db::user_repo, log_in, models::IotaId, sql::connection_status::UserStatus, state::OmegaState,
};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use std::{
collections::{BTreeMap, HashMap, HashSet},
sync::Arc,
};
fn parse_subscription(value: &CommunicationValue) -> Result<(i64, i64, Vec<i64>), &'static str> {
let Some(sender) = value.sender() else {
return Err("user_id");
};
let user_id = i64::try_from(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 Some(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((user_id, session_id, user_ids))
}
fn apply_preferences(state: &OmegaState, preferences: HashMap<i64, UserStatus>) {
state.presence.set_preferences(preferences);
}
fn states_for_users(state: &OmegaState, users: &[crate::models::User]) -> HashMap<i64, UserStatus> {
users
.iter()
.map(|user| {
let status = user
.iota_id
.map(|iota_id| state.presence.resolve_public_state(user.id.0, iota_id.0))
.unwrap_or(UserStatus::user_offline);
(user.id.0, status)
})
.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 = user
.iota_id
.map(|iota_id| state.presence.resolve_public_state(user.id.0, iota_id.0))
.unwrap_or(UserStatus::user_offline);
(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()))
}
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
{
log_in!(
crate::util::logger::PrintType::General,
"Failed to deliver presence state batch to Omikron {}: {}",
omikron_id,
error
);
}
}
}
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;
}
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.require_id()?, CommunicationType::ErrorNoUserId)
.await;
}
Err(detail) => {
return connection
.send_error_response_with_detail(
value.require_id()?,
CommunicationType::ErrorInvalidData,
detail,
)
.await;
}
};
if !state.presence.owns_session(user_id, session_id, omikron_id) {
return connection
.send_error_response(value.require_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.require_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>,
value: CommunicationValue,
_: i64,
) -> OmikronResult<()> {
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 Ok(());
};
let Some(status) = value
.get_data(DataType::UserState)
.as_str()
.and_then(UserStatus::from_client_preference)
else {
return Ok(());
};
state.presence.set_preference(user_id, status);
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.require_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.require_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.require_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.require_id()?, CommunicationType::ErrorNotFound)
.await;
}
Err(error) => return Err(error),
};
let preferences = match user_repo::get_presence_preferences(&[user_id]).await {
Ok(preferences) => preferences,
Err(error) => return Err(error),
};
if user.iota_id.map(|id| id.0) != Some(iota_id) || !state.presence.has_iota_route(iota_id) {
return connection
.send_error_response(value.require_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.require_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.require_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.require_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.require_id()?))
.await
}
pub async fn iota_connected(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
value: CommunicationValue,
omikron_id: i64,
) -> OmikronResult<()> {
log_in!(crate::util::logger::PrintType::Omega, "IOTA connected");
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.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?;
let ids = users.iter().map(|user| user.id.0).collect::<Vec<_>>();
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)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()))
.add_typed_default(DataType::UserIds, DataValue::Array(user_ids));
connection.clone().send(&response).await?;
crate::transport::omikron_manager::deliver_pending_erasures(iota_id).await;
publish_changed_states(&state, &before, &users).await;
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.require_id()?))
.await
}
pub async fn iota_disconnected(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
value: CommunicationValue,
omikron_id: i64,
) -> OmikronResult<()> {
log_in!(crate::util::logger::PrintType::Omega, "IOTA disconnected");
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.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?;
let ids = users.iter().map(|user| user.id.0).collect::<Vec<_>>();
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.require_id()?))
.await
}
pub async fn sync_status(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
value: CommunicationValue,
omikron_id: i64,
) -> OmikronResult<()> {
let request_id = value.require_id()?;
let Some(DataValue::Array(iota_values)) = value.get_data(DataType::IotaIds) else {
return connection
.send_error_response(request_id, CommunicationType::ErrorInvalidData)
.await;
};
let Some(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 !connection.peer_capabilities().session_snapshot_v1 {
let Some(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::<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()));
}
}