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

[Fix] Stability

Browse source
This commit is contained in:
Alex 2026-07-27 20:37:30 +02:00
commit 0d5e48ec8f
Signed by: alex
SSH key fingerprint: SHA256:D1+Ub8o0v4K5y1JNivW8IxEOelqLSvPmUzBbDIoZkRQ
12 changed files with 556 additions and 386 deletions
Download patch file Download diff file Expand all files Collapse all files

31
src/rho/client_connection.rs
View file

@ -272,6 +272,28 @@ impl ClientConnection {
}
}
// Every Iota request is bound to the authenticated device. A
// caller may omit the field for compatibility, but can never
// select a different device session.
let expected_session_id = self.session_id as i128;
if let Some(session_id) = cv.get_data(DataType::SessionId).as_signed_number() {
if session_id != expected_session_id {
let response = CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(cv.get_id())
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(expected_session_id),
);
self.send_message(&response).await;
return;
}
} else {
cv = cv.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(expected_session_id),
);
}
// Forward other messages to Iota
self.forward_to_iota(cv).await;
});
@ -330,6 +352,7 @@ impl ClientConnection {
.client_changed(
rho_conn.get_iota_id().await as i64,
user_id as i64,
self.session_id as i64,
user_status,
)
.await;
@ -868,7 +891,7 @@ impl ClientConnection {
user_status
};
let notification = CommunicationValue::new(CommunicationType::ClientChanged)
.add_typed_default(DataType::UserId, DataValue::Str(user_id.to_string()))
.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;
@ -876,12 +899,10 @@ impl ClientConnection {
}
/// Handle connection close
pub async fn handle_close(&self) {
pub async fn handle_close(self: Arc<Self>) {
let user_id = self.get_user_id().await;
if let Some(rho_conn) = self.state.rho.get_for_user(user_id as i64).await {
rho_conn
.close_client_connection(Arc::new(self.clone()))
.await;
rho_conn.close_client_connection(self.clone()).await;
}
}
}

113
src/rho/connection.rs
View file

@ -1,4 +1,5 @@
use std::{sync::Arc, time::Duration};
use strum_macros::Display;
use tokio::sync::RwLock;
use uuid::Uuid;
@ -27,7 +28,7 @@ pub type MtpReceiver = WebMtpReceiver;
* `Identification`/`Register` hello frames; it is kept only so
* `app_connection.rs` still compiles.
*/
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Display)]
#[allow(dead_code)]
pub enum ConnectionKind {
Client,
@ -85,8 +86,20 @@ impl GeneralConnection {
pub async fn handle(self: Arc<Self>) {
log_in!(0, PrintType::General, "General connection handler started");
self.migrate().await;
log_out!(0, PrintType::General, "General connection handler stopped");
if self.migrate().await {
log_out!(
0,
PrintType::General,
"General connection handler stopped, upgraded to {}",
self.connection_kind
);
} else {
log_out!(
0,
PrintType::General,
"General connection handler stopped without upgrade"
);
}
}
async fn migrate(self: &Arc<Self>) -> bool {
@ -102,19 +115,67 @@ impl GeneralConnection {
async fn migrate_client(self: &Arc<Self>) {
let id = self.id;
let user_id = id as i64;
self.notify_user_connected(user_id).await;
let Ok(handshake) = self.receiver.receive().await else {
return;
};
if !handshake.is_type(CommunicationType::ClientConnected) {
let error = CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(handshake.get_id());
let _ = self.sender.send(&error).await;
return;
}
let Some(session_id) = handshake.get_data(DataType::SessionId).as_signed_number() else {
let _ = self
.sender
.send(
&CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(handshake.get_id()),
)
.await;
return;
};
if !(1..=i64::MAX as i128).contains(&session_id) {
let _ = self
.sender
.send(
&CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(handshake.get_id()),
)
.await;
return;
}
let version = handshake
.get_data(DataType::VersionNumber)
.as_signed_number();
if !matches!(version, Some(version) if version >= 0) {
let _ = self
.sender
.send(
&CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(handshake.get_id()),
)
.await;
return;
}
*self.session_id.write().await = session_id as u64;
let client = ClientConnection::from_general(self.clone(), id).await;
let rho = self.find_user_rho(user_id).await;
*self.rho_connection.write().await = rho.clone();
if let Some(rho_conn) = rho {
if let Some(response) = self.request_initial_client_state(&rho_conn, user_id).await {
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;
if let Some(response) = self
.request_initial_client_state(&rho_conn, user_id, handshake)
.await
{
let response = self.add_call_state(response, user_id).await;
log_cv_out!(response);
let _ = self.sender.send(&response).await;
rho_conn.bind_user_id(user_id).await;
rho_conn.add_client_connection(client.clone()).await;
}
} else {
log_err!(
@ -128,9 +189,28 @@ impl GeneralConnection {
client.start();
}
async fn notify_user_connected(&self, user_id: i64) {
async fn notify_user_connected(&self, user_id: i64, iota_id: i64) {
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::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;
}
@ -167,6 +247,7 @@ impl GeneralConnection {
&self,
rho: &Arc<RhoConnection>,
user_id: i64,
handshake: CommunicationValue,
) -> Option<CommunicationValue> {
let session_id = *self.session_id.read().await as i64;
let request = CommunicationValue::new(CommunicationType::ClientConnected)
@ -174,6 +255,18 @@ impl GeneralConnection {
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(session_id.into()),
)
.add_typed_default(
DataType::VersionNumber,
handshake.get_data(DataType::VersionNumber).clone(),
)
.add_typed_default(
DataType::CacheValid,
handshake.get_data(DataType::CacheValid).clone(),
)
.add_typed_default(
DataType::CacheSchemaVersion,
handshake.get_data(DataType::CacheSchemaVersion).clone(),
);
rho.get_iota_connection()
.clone()
@ -187,7 +280,7 @@ impl GeneralConnection {
response: CommunicationValue,
user_id: i64,
) -> CommunicationValue {
let mut output = CommunicationValue::new(CommunicationType::IdentificationResponse);
let mut output = response.clone();
for (key, value) in response.iter_typed_data() {
if key == Some(DataType::Contacts) {

221
src/rho/iota_connection.rs
View file

@ -225,22 +225,124 @@ impl IotaConnection {
}
if cv.is_type(CommunicationType::CompleteRegisterUser) {
let response_cv = self
// Registration carries the authenticated Iota ID separately so
// Omega can bind the allocation lease to this connection rather
// than trusting a client-supplied value.
let request = cv.clone().add_typed_default(
DataType::IotaId,
DataValue::SignedNumber(self.iota_id.into()),
);
log_in!(
self.iota_id as i64,
PrintType::Omega,
"Forwarding CompleteRegisterUser to Omega (request_id={})",
request.get_id()
);
let mut response_cv = self
.state
.omega
.clone()
.await_response(
&cv.clone().with_sender(self.iota_id),
Some(Duration::from_secs(20)),
)
.await_response(&request, Some(Duration::from_secs(8)))
.await;
if let Ok(response_cv) = response_cv {
if response_cv.is_type(CommunicationType::Success) {
if let Some(user_id) = cv.get_data(DataType::UserId).as_number() {
self.add_user_id(user_id as u64).await;
if let Err(error) = &response_cv {
log_err!(
self.iota_id as i64,
PrintType::Omega,
"CompleteRegisterUser request_id={} failed: {}; retrying once",
request.get_id(),
error
);
response_cv = self
.state
.omega
.clone()
.await_response(&request, Some(Duration::from_secs(8)))
.await;
}
match response_cv {
Ok(response_cv) => {
log_in!(
self.iota_id as i64,
PrintType::Omega,
"Omega completed registration (request_id={}, response_id={}, type={})",
request.get_id(),
response_cv.get_id(),
response_cv
.get_comm_type_enum()
.map(|kind| kind.to_string())
.unwrap_or_else(|| response_cv.get_type().to_string())
);
if response_cv.is_type(CommunicationType::Success) {
if let Some(user_id) = cv.get_data(DataType::UserId).as_number() {
self.add_user_id(user_id as u64).await;
}
}
self.send_message(&response_cv).await;
}
Err(error) => {
// Omega may have committed the insert even when its
// Success response was lost in transit. Verify the exact
// generated user ID before reporting failure; GetUserData
// uses the proven request/response path and keeps this
// recovery idempotent.
let user_id = cv.get_data(DataType::UserId).as_number();
if let Some(user_id) = user_id {
let verification = CommunicationValue::new(CommunicationType::GetUserData)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id));
match self
.state
.omega
.clone()
.await_response(&verification, Some(Duration::from_secs(3)))
.await
{
Ok(verified)
if verified.get_data(DataType::UserId).as_number()
== Some(user_id) =>
{
log_in!(
self.iota_id as i64,
PrintType::Omega,
"Registration response was lost; verified user {} remotely",
user_id
);
self.add_user_id(user_id as u64).await;
self.send_message(
&CommunicationValue::new(CommunicationType::Success)
.with_id(cv.get_id()),
)
.await;
return;
}
Ok(verified) => log_err!(
self.iota_id as i64,
PrintType::Omega,
"Registration verification returned an unexpected user (request_id={}, response_id={})",
verification.get_id(),
verified.get_id()
),
Err(verify_error) => log_err!(
self.iota_id as i64,
PrintType::Omega,
"Registration verification failed after request_id={}: {}",
verification.get_id(),
verify_error
),
}
}
log_err!(
self.iota_id as i64,
PrintType::Omega,
"CompleteRegisterUser forwarding failed: {}",
error
);
self.send_error_response(
cv.get_id(),
CommunicationType::ErrorInternal,
Some(&format!("Omega forwarding failed: {error}")),
)
.await;
}
self.send_message(&response_cv).await;
}
return;
}
@ -249,7 +351,6 @@ impl IotaConnection {
|| cv.is_type(CommunicationType::PushNotification)
|| cv.is_type(CommunicationType::GetUserData)
|| cv.is_type(CommunicationType::GetIotaData)
|| cv.is_type(CommunicationType::GetRegister)
|| cv.is_type(CommunicationType::DeleteIota)
{
let sender = self.get_iota_id().await;
@ -258,12 +359,27 @@ impl IotaConnection {
.await;
return;
}
// Register allocation is scoped to this authenticated Iota. Keep the
// source ID in typed data so Omega can create a durable lease.
if cv.is_type(CommunicationType::GetRegister) {
self.handle_omega_forward_without_sender(cv).await;
return;
}
self.forward_to_client(cv).await;
}
#[allow(dead_code)]
async fn send_error_response(&self, message_id: u32, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(message_id);
async fn send_error_response(
&self,
message_id: u32,
error_type: CommunicationType,
detail: Option<&str>,
) {
let mut error = CommunicationValue::new(error_type).with_id(message_id);
if let Some(detail) = detail {
error = error.add_typed_default(DataType::ErrorType, DataValue::Str(detail.into()));
}
self.send_message(&error).await;
}
@ -284,6 +400,52 @@ impl IotaConnection {
iota_for_closure.send_message(&response_cv).await;
}
}
async fn handle_omega_forward_without_sender(self: Arc<Self>, cv: CommunicationValue) {
let iota_for_closure = self.clone();
let request = cv.clone().add_typed_default(
DataType::IotaId,
DataValue::SignedNumber(self.iota_id.into()),
);
let mut response_cv = self
.state
.omega
.clone()
.await_response(&request, Some(Duration::from_secs(8)))
.await;
if let Err(error) = &response_cv {
log_err!(
self.iota_id as i64,
PrintType::Omega,
"GetRegister request_id={} failed: {}; retrying once",
request.get_id(),
error
);
response_cv = self
.state
.omega
.clone()
.await_response(&request, Some(Duration::from_secs(8)))
.await;
}
match response_cv {
Ok(response_cv) => iota_for_closure.send_message(&response_cv).await,
Err(error) => {
log_err!(
self.iota_id as i64,
PrintType::Omega,
"GetRegister forwarding failed: {}",
error
);
self.send_error_response(
cv.get_id(),
CommunicationType::ErrorInternal,
Some(&format!("Omega forwarding failed: {error}")),
)
.await;
}
}
}
/// Handle ping message
async fn handle_ping(&self, cv: CommunicationValue) {
if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) {
@ -397,8 +559,31 @@ impl IotaConnection {
}
let mut interested_ids: Vec<i64> = Vec::new();
let session_id = cv.get_data(DataType::SessionId).as_signed_number();
let tm = TypeMap::latest();
// Presence interest is the complete contact set, independent of
// whether this account currently participates in a call.
if let DataValue::Array(users) = cv.get_data(DataType::UserIds) {
for user in users {
match user {
DataValue::SignedNumber(id) => interested_ids.push(*id as i64),
DataValue::Container(entries) => {
if let Some(DataValue::SignedNumber(id)) =
entries.iter().find_map(|(key, value)| {
(*key == data_type_id(DataType::UserId, &tm)).then_some(value)
})
{
interested_ids.push(*id as i64);
}
}
_ => {}
}
}
interested_ids.sort_unstable();
interested_ids.dedup();
}
// ============================
// Load Calls
// ============================
@ -494,8 +679,6 @@ impl IotaConnection {
if let Some(DataValue::SignedNumber(id)) =
user_map.get(&data_type_id(DataType::UserId, &tm))
{
interested_ids.push(*id as i64);
if let Some(call_list) = invites.get(&(*id as i64))
&& !call_list.is_empty()
{
@ -526,9 +709,11 @@ impl IotaConnection {
// Notify Rho
// ============================
if let Some(rho_conn) = self.get_rho_connection().await {
rho_conn
.set_interested(user_id as i64, interested_ids)
.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;
}
}
// ============================

73
src/rho/rho_connection.rs
View file

@ -2,7 +2,7 @@ use super::{client_connection::ClientConnection, iota_connection::IotaConnection
use crate::{data::user::UserStatus, rho::app_connection::AppConnection};
use dashmap::DashMap;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::codec::{CommunicationValue, DataType};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
@ -128,39 +128,51 @@ impl RhoConnection {
let user_id = connection.user_id as i64;
let session_id = connection.session_id as i64;
let notification = CommunicationValue::new(CommunicationType::ClientConnected)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()));
self.iota_connection.send_message(&notification).await;
let should_notify_online = !self
if let Some((_, old_connection)) = self
.client_connections
.iter()
.any(|entry| entry.key().0 == user_id as u64);
self.client_connections
.remove(&(user_id as u64, session_id as u64));
.remove(&(user_id as u64, session_id as u64))
{
// A reconnect for a persistent device replaces the old transport;
// do this before publishing the new route.
old_connection.close().await;
}
self.client_connections
.insert((user_id as u64, session_id as u64), Arc::clone(&connection));
}
if should_notify_online {
self.iota_connection
.state
.omega
.client_changed(
self.get_iota_id().await as i64,
user_id,
UserStatus::user_online,
)
.await;
pub async fn get_client_connection(
&self,
user_id: i64,
session_id: i64,
) -> Option<Arc<ClientConnection>> {
if user_id < 0 || session_id <= 0 {
return None;
}
self.client_connections
.get(&(user_id as u64, session_id as u64))
.map(|entry| entry.value().clone())
}
/// Remove a client connection
pub async fn close_client_connection(&self, connection: Arc<ClientConnection>) {
let target_user_id = connection.user_id as i64;
let target_session_id = connection.session_id as i64;
self.client_connections
.remove(&(target_user_id as u64, target_session_id as u64));
// Do not let a stale transport tear down the replacement for the same
// persistent session.
let key = (target_user_id as u64, target_session_id as u64);
if self
.client_connections
.get(&key)
.is_some_and(|current| !Arc::ptr_eq(current.value(), &connection))
{
return;
}
self.client_connections.remove(&key);
self.iota_connection
.state
.omega
.client_disconnected(target_user_id, target_session_id)
.await;
let remaining_for_user = self
.client_connections
.iter()
@ -178,6 +190,7 @@ impl RhoConnection {
.client_changed(
self.get_iota_id().await as i64,
target_user_id,
target_session_id,
UserStatus::user_offline,
)
.await;
@ -232,17 +245,9 @@ impl RhoConnection {
}
/// Set interested users for a specific client
pub async fn set_interested(&self, user_id: i64, interested_ids: Vec<i64>) {
let connections = self.get_client_connections().await;
for connection in connections.iter() {
let conn_user_id = connection.get_user_id().await;
if conn_user_id == user_id as u64 {
connection
.clone()
.set_interested_users(interested_ids.clone())
.await;
break;
}
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;
}
}