tensamin/omikron
Watch
6
0
Fork
You've already forked omikron
 0
Code Issues Pull requests 3 Projects Releases Packages Wiki Activity Actions

[Fix] User States

Browse source
This commit is contained in:
Alex 2026-08-07 23:54:28 +02:00
commit da5a5a5dff
Signed by: alex
SSH key fingerprint: SHA256:D1+Ub8o0v4K5y1JNivW8IxEOelqLSvPmUzBbDIoZkRQ
12 changed files with 1156 additions and 232 deletions
Download patch file Download diff file Expand all files Collapse all files

202
src/rho/client_connection.rs
View file

@ -24,7 +24,6 @@ pub struct ClientConnection {
pub receiver: Arc<MtpReceiver>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
pub interested_users: Arc<RwLock<Vec<i64>>>,
is_open: Arc<RwLock<bool>>,
message_slots: Arc<tokio::sync::Semaphore>,
}
@ -35,7 +34,6 @@ impl ClientConnection {
state: general.state.clone(),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: general.rho_connection.clone(),
interested_users: Arc::new(RwLock::new(Vec::new())),
is_open: Arc::new(RwLock::new(true)),
sender: general.sender.clone(),
receiver: general.receiver.clone(),
@ -104,9 +102,26 @@ impl ClientConnection {
let mut cv = cv;
// Handle client status changes
if cv.is_type(CommunicationType::ClientChanged) {
self.handle_client_changed(cv).await;
if cv.is_type(CommunicationType::SetUserState) {
self.handle_set_user_state(cv).await;
return;
}
// Compatibility for clients predating SetUserState. The target
// user fields, if present, are deliberately ignored: an
// authenticated connection may only change its own state.
if cv.is_type(CommunicationType::ClientChanged)
&& cv.get_data_opt(DataType::UserState).is_some()
{
self.handle_set_user_state(
CommunicationValue::new(CommunicationType::SetUserState)
.with_id(cv.get_id())
.add_typed_default(
DataType::UserState,
cv.get_data(DataType::UserState).clone(),
),
)
.await;
return;
}
@ -191,6 +206,46 @@ impl ClientConnection {
|| cv.is_type(CommunicationType::GetIotaData)
|| cv.is_type(CommunicationType::DeleteUser)
{
if cv.is_type(CommunicationType::ChangeUserData)
&& cv.get_data_opt(DataType::OnlineStatus).is_some()
{
let mut profile_request = cv.clone();
let preference = profile_request.remove_data(DataType::OnlineStatus);
let state_request = CommunicationValue::new(CommunicationType::SetUserState)
.with_id(cv.get_id())
.add_typed_default(DataType::UserState, preference.unwrap());
let state_response = match self.request_set_user_state(state_request).await {
Ok(response) => response,
Err(error_type) => {
self.send_error_response(cv.get_id(), error_type).await;
return;
}
};
if !state_response.is_type(CommunicationType::Success) {
self.send_message(&state_response).await;
return;
}
if profile_request.data_len() == 0 {
self.send_message(&state_response).await;
return;
}
match self
.await_omega_response(profile_request.with_sender(self.user_id))
.await
{
Ok(response) if response.is_type(CommunicationType::Success) => {
self.send_message(&response).await;
}
Ok(response) => self.send_message(&response).await,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInternal)
.await;
}
}
return;
}
let sender = self.get_user_id().await;
self.handle_omega_forward(cv.with_sender(sender as u64))
.await;
@ -290,34 +345,91 @@ impl ClientConnection {
});
}
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let client_for_closure = self.clone();
let response_cv = self
.state
.omega
.clone()
.await_response(&cv.with_sender(self.user_id), Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
client_for_closure.send_message(&response_cv).await;
let request_id = cv.get_id();
match self.await_omega_response(cv).await {
Ok(response_cv) => self.send_message(&response_cv).await,
Err(_) => {
self.send_error_response(request_id, CommunicationType::ErrorInternal)
.await;
}
}
}
/// Handle client status change
async fn handle_client_changed(self: Arc<Self>, cv: CommunicationValue) {
let user_id = self.get_user_id().await;
if let DataValue::Str(status_str) = cv.get_data(DataType::UserState) {
let user_status = UserStatus::from_str(&status_str).unwrap_or(UserStatus::user_online);
if let Some(rho_conn) = self.get_rho_connection().await {
self.state
.omega
.client_changed(
rho_conn.get_iota_id().await as i64,
user_id as i64,
self.session_id as i64,
user_status,
)
.await;
}
async fn await_omega_response(&self, cv: CommunicationValue) -> Result<CommunicationValue, ()> {
self.state
.omega
.clone()
.await_response(&cv.with_sender(self.user_id), Some(Duration::from_secs(20)))
.await
.map_err(|_| ())
}
async fn request_set_user_state(
&self,
cv: CommunicationValue,
) -> Result<CommunicationValue, CommunicationType> {
if !self.state.omega.is_ready().await {
return Err(CommunicationType::ErrorInternal);
}
let Some(state) = cv
.get_data(DataType::UserState)
.as_str()
.and_then(UserStatus::from_client_preference)
else {
return Err(CommunicationType::ErrorInvalidData);
};
let Some(rho) = self.get_rho_connection().await else {
return Err(CommunicationType::ErrorNoIota);
};
let request = if self.state.omega.supports_set_user_state().await {
CommunicationValue::new(CommunicationType::SetUserState)
.with_id(cv.get_id())
.with_sender(self.user_id)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(self.user_id as i128),
)
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(self.session_id as i128),
)
.add_typed_default(
DataType::IotaId,
DataValue::SignedNumber(rho.get_iota_id().await as i128),
)
.add_typed_default(DataType::UserState, DataValue::Str(state.to_string()))
} else {
// Legacy Omega accepts the original ClientChanged request. It is
// intentionally limited to the authenticated user's ID and a
// writable preference, with no derived connectivity state.
CommunicationValue::new(CommunicationType::ClientChanged)
.with_id(cv.get_id())
.with_sender(self.user_id)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(self.user_id as i128),
)
.add_typed_default(DataType::UserState, DataValue::Str(state.to_string()))
};
if !self.state.omega.supports_set_user_state().await {
self.state
.omega
.send_message_result(&request)
.await
.map_err(|_| CommunicationType::ErrorInternal)?;
return Ok(CommunicationValue::new(CommunicationType::Success)
.with_id(cv.get_id())
.add_typed_default(DataType::UserState, DataValue::Str(state.to_string())));
}
self.await_omega_response(request)
.await
.map_err(|_| CommunicationType::ErrorInternal)
}
async fn handle_set_user_state(self: Arc<Self>, cv: CommunicationValue) {
match self.request_set_user_state(cv.clone()).await {
Ok(response) => self.send_message(&response).await,
Err(error_type) => self.send_error_response(cv.get_id(), error_type).await,
}
}
@ -830,35 +942,6 @@ impl ClientConnection {
let _ = self.sender.close();
}
/// Set interested users list
pub async fn set_interested_users(self: Arc<Self>, interested_ids: Vec<i64>) {
let mut interested_guard = self.interested_users.write().await;
*interested_guard = interested_ids;
}
#[allow(dead_code)]
pub async fn get_interested_users(self: Arc<Self>) -> Vec<i64> {
let interested_guard = self.interested_users.read().await;
interested_guard.clone()
}
/// Check if interested in a user and send notification
#[allow(dead_code)]
pub async fn are_you_interested(self: Arc<Self>, user_id: i64, user_status: &str) {
let interested_guard = self.clone().get_interested_users().await;
if interested_guard.contains(&user_id) {
let status = if user_status == "user_invisible" {
"user_offline"
} else {
user_status
};
let notification = CommunicationValue::new(CommunicationType::ClientChanged)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
.add_typed_default(DataType::UserState, DataValue::Str(status.to_string()));
self.send_message(&notification).await;
}
}
/// Handle connection close
pub async fn handle_close(self: Arc<Self>) {
let user_id = self.get_user_id().await;
@ -880,7 +963,6 @@ impl Clone for ClientConnection {
client_version: self.client_version.clone(),
pub_key: Arc::clone(&self.pub_key),
rho_connection: Arc::clone(&self.rho_connection),
interested_users: Arc::clone(&self.interested_users),
is_open: Arc::clone(&self.is_open),
message_slots: Arc::clone(&self.message_slots),
}

60
src/rho/connection.rs
View file

@ -136,10 +136,23 @@ impl GeneralConnection {
if let Some(rho_conn) = rho {
rho_conn.bind_user_id(user_id).await;
rho_conn.add_client_connection(client.clone()).await;
self.notify_user_connected(user_id, rho_conn.get_iota_id().await as i64)
.await;
self.send_initial_client_state_request(&rho_conn, user_id)
.await;
match self
.notify_user_connected(user_id, rho_conn.get_iota_id().await as i64)
.await
{
Ok(()) => {
self.send_initial_client_state_request(&rho_conn, user_id)
.await;
}
Err(error) => {
log_err!(
user_id,
PrintType::Client,
"UserConnected acknowledgement failed: {}",
error
);
}
}
} else {
log_err!(
user_id,
@ -155,29 +168,13 @@ impl GeneralConnection {
true
}
async fn notify_user_connected(&self, user_id: i64, iota_id: i64) {
async fn notify_user_connected(&self, user_id: i64, iota_id: i64) -> Result<(), String> {
let session_id = *self.session_id.read().await as i64;
let notify = CommunicationValue::new(CommunicationType::UserConnected)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(session_id.into()),
)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()))
.add_typed_default(
DataType::UserState,
DataValue::Str("user_online".to_string()),
)
.add_typed_default(
DataType::UpdatedAt,
DataValue::SignedNumber(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as i128,
),
);
self.state.omega.clone().send_message(&notify).await;
self.state
.omega
.clone()
.user_connected(user_id, session_id, iota_id)
.await
}
async fn find_user_rho(&self, user_id: i64) -> Option<Arc<RhoConnection>> {
@ -227,9 +224,6 @@ impl GeneralConnection {
async fn migrate_iota(self: &Arc<Self>) {
let id = self.id;
log_in!(id as i64, PrintType::Iota, "Iota {} connected", id);
let notify = CommunicationValue::new(CommunicationType::IotaConnected)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into()));
self.state.omega.clone().send_message(&notify).await;
let iota = IotaConnection::from_general(self.clone(), id).await;
let previous = self.state.rho.get_by_iota(id as i64).await;
@ -246,6 +240,14 @@ impl GeneralConnection {
}
}
self.state.rho.add(rho).await;
if let Err(error) = self.state.omega.clone().iota_connected(id as i64).await {
log_err!(
id as i64,
PrintType::Iota,
"IotaConnected acknowledgement failed: {}",
error
);
}
self.load_iota_users(id).await;
iota.start();
}

178
src/rho/iota_connection.rs
View file

@ -23,6 +23,47 @@ use tokio::sync::mpsc;
use super::rho_connection::RhoConnection;
fn contact_snapshot(value: &CommunicationValue) -> Option<(i64, i64, Vec<i64>)> {
if !value.is_type(CommunicationType::ClientStateSync)
&& !value.is_type(CommunicationType::AddConversation)
{
return None;
}
let user_id = i64::try_from(value.get_receiver())
.ok()
.filter(|id| *id > 0)?;
let session_id = value
.get_data(DataType::SessionId)
.as_signed_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)?;
let DataValue::Array(values) = value.get_data(DataType::UserIds) else {
return None;
};
let user_id_type = data_type_id(DataType::UserId, &TypeMap::latest());
let mut contact_ids = Vec::with_capacity(values.len());
for value in values {
let contact_id = match value {
DataValue::SignedNumber(id) => i64::try_from(*id).ok(),
DataValue::Container(entries) => entries.iter().find_map(|(key, value)| {
(*key == user_id_type)
.then(|| value.as_signed_number())
.flatten()
.and_then(|id| i64::try_from(id).ok())
}),
_ => None,
}
.filter(|id| *id > 0)?;
contact_ids.push(contact_id);
}
contact_ids.sort_unstable();
contact_ids.dedup();
Some((user_id, session_id, contact_ids))
}
#[allow(dead_code)]
pub struct IotaConnection {
pub state: Arc<AppState>,
@ -172,6 +213,29 @@ impl IotaConnection {
}
}
async fn replace_state_subscription(
&self,
user_id: i64,
session_id: i64,
contact_ids: Vec<i64>,
) -> Result<(), String> {
let Some(rho_connection) = self.get_rho_connection().await else {
return Err("Rho connection is unavailable".to_string());
};
if rho_connection
.get_client_connection(user_id, session_id)
.await
.is_none()
{
return Err("target client session is no longer connected".to_string());
}
self.state
.omega
.replace_state_subscription(user_id, session_id, contact_ids)
.await
}
/// Send a CommunicationValue to the Iota
pub async fn send_message(&self, cv: &CommunicationValue) {
log_cv_out!(PrintType::Iota, cv);
@ -206,6 +270,34 @@ impl IotaConnection {
cv
};
// Iota's UserIds field is an authoritative replacement snapshot for
// the requesting client session, including contact mutations.
if let Some((user_id, session_id, contact_ids)) = contact_snapshot(&cv) {
if let Err(error) = self
.replace_state_subscription(user_id, session_id, contact_ids)
.await
{
log_err!(
self.iota_id as i64,
PrintType::Omega,
"Failed to replace presence subscription for user {} session {}: {}",
user_id,
session_id,
error
);
}
}
if cv.is_type(CommunicationType::StateSubscribe) {
self.send_error_response(
cv.get_id(),
CommunicationType::ErrorInvalidData,
Some("StateSubscribe must come from an authoritative contact snapshot"),
)
.await;
return;
}
// Handle GET_CHATS
if cv.is_type(CommunicationType::GetChats) {
self.handle_get_chats(cv).await;
@ -517,7 +609,20 @@ impl IotaConnection {
}
let mut interested_ids: Vec<i64> = Vec::new();
let session_id = cv.get_data(DataType::SessionId).as_signed_number();
let Some(session_id) = cv
.get_data(DataType::SessionId)
.as_signed_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
else {
self.forward_to_client(
CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(cv.get_id())
.with_receiver(user_id),
)
.await;
return;
};
let tm = TypeMap::latest();
// Presence interest is the complete contact set, independent of
@ -658,20 +763,18 @@ impl IotaConnection {
// ============================
// Notify Omega
// ============================
self.state
.omega
.user_states(user_id as i64, interested_ids.clone())
.await;
// ============================
// Notify Rho
// ============================
if let Some(rho_conn) = self.get_rho_connection().await {
if let Some(session_id) = session_id.and_then(|id| i64::try_from(id).ok()) {
rho_conn
.set_interested(user_id as i64, session_id, interested_ids)
.await;
}
if let Err(error) = self
.replace_state_subscription(user_id as i64, session_id, interested_ids.clone())
.await
{
log_err!(
self.iota_id as i64,
PrintType::Omega,
"Failed to replace presence subscription for user {} session {}: {}",
user_id,
session_id,
error
);
}
// ============================
@ -779,3 +882,48 @@ impl std::fmt::Debug for IotaConnection {
.finish()
}
}
#[cfg(test)]
mod tests {
use super::contact_snapshot;
use crate::util::data_type_id;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap};
fn snapshot(user_ids: DataValue) -> CommunicationValue {
CommunicationValue::new(CommunicationType::ClientStateSync)
.with_receiver(7)
.add_typed_default(DataType::SessionId, DataValue::SignedNumber(11))
.add_typed_default(DataType::UserIds, user_ids)
}
#[test]
fn contact_snapshot_deduplicates_flat_and_typed_ids() {
let user_id_type = data_type_id(DataType::UserId, &TypeMap::latest());
let parsed = contact_snapshot(&snapshot(DataValue::Array(vec![
DataValue::SignedNumber(20),
DataValue::Container(vec![(user_id_type, DataValue::SignedNumber(21))]),
DataValue::SignedNumber(20),
])))
.unwrap();
assert_eq!(parsed, (7, 11, vec![20, 21]));
}
#[test]
fn contact_snapshot_accepts_empty_contact_sets() {
assert_eq!(
contact_snapshot(&snapshot(DataValue::Array(Vec::new()))),
Some((7, 11, Vec::new()))
);
}
#[test]
fn contact_snapshot_rejects_malformed_contact_sets() {
assert_eq!(
contact_snapshot(&snapshot(DataValue::Array(vec![DataValue::Str(
"bad".into(),
)]))),
None
);
}
}

54
src/rho/rho_connection.rs
View file

@ -1,6 +1,6 @@
use super::{client_connection::ClientConnection, iota_connection::IotaConnection};
use crate::{data::user::UserStatus, rho::app_connection::AppConnection};
use crate::{log_err, rho::app_connection::AppConnection};
use dashmap::DashMap;
use mtp::codec::{CommunicationValue, DataType};
use std::sync::Arc;
@ -178,11 +178,21 @@ impl RhoConnection {
return;
}
self.client_connections.remove(&key);
self.iota_connection
if let Err(error) = self
.iota_connection
.state
.omega
.client_disconnected(target_user_id, target_session_id)
.await;
.await
{
log_err!(
target_user_id,
crate::util::logger::PrintType::Client,
"UserDisconnected acknowledgement failed for session {}: {}",
target_session_id,
error
);
}
let remaining_for_user = self
.client_connections
.iter()
@ -194,16 +204,6 @@ impl RhoConnection {
.rho
.remove_user_binding(target_user_id, self.get_iota_id().await as i64)
.await;
self.iota_connection
.state
.omega
.client_changed(
self.get_iota_id().await as i64,
target_user_id,
target_session_id,
UserStatus::user_offline,
)
.await;
}
}
@ -228,7 +228,14 @@ impl RhoConnection {
}
// Notify OmegaConnection
self.iota_connection.state.omega.close_iota(iota_id).await;
if let Err(error) = self.iota_connection.state.omega.close_iota(iota_id).await {
log_err!(
iota_id,
crate::util::logger::PrintType::Iota,
"IotaDisconnected acknowledgement failed: {}",
error
);
}
}
/// Send message from Iota to specific client
@ -255,25 +262,6 @@ impl RhoConnection {
self.iota_connection.send_message(&cv).await;
}
/// Set interested users for a specific client
pub async fn set_interested(&self, user_id: i64, session_id: i64, interested_ids: Vec<i64>) {
if let Some(connection) = self.get_client_connection(user_id, session_id).await {
connection.set_interested_users(interested_ids).await;
}
}
/// Check if clients are interested in a user
#[allow(dead_code)]
pub async fn are_they_interested(&self, user_id: i64, user_status: &str) {
let connections = self.get_client_connections().await;
for connection in connections.iter() {
connection
.clone()
.are_you_interested(user_id, user_status)
.await;
}
}
/// Check if this RhoConnection contains a specific user ID
#[allow(dead_code)]
pub async fn contains_user(&self, user_id: &i64) -> bool {

13
src/rho/rho_manager.rs
View file

@ -135,4 +135,17 @@ impl RhoManager {
.map(|entry| entry.value().clone())
.collect()
}
pub async fn get_client_connection(
&self,
user_id: i64,
session_id: i64,
) -> Option<Arc<super::client_connection::ClientConnection>> {
for rho in self.connections().await {
if let Some(client) = rho.get_client_connection(user_id, session_id).await {
return Some(client);
}
}
None
}
}