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

[Add] Structure

Browse source
This commit is contained in:
Alex Emmet 2026-07-20 22:22:12 +02:00
commit c363ea48d0
27 changed files with 1730 additions and 1400 deletions
Download patch file Download diff file Expand all files Collapse all files

21
src/rho/app_connection.rs
View file

@ -1,7 +1,7 @@
use crate::anonymous_clients::anonymous_manager;
use crate::omega::omega_connection::get_omega_connection;
use crate::app_state::AppState;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::rho::{rho_connection::RhoConnection, rho_manager};
use crate::rho::rho_connection::RhoConnection;
use crate::util::logger::PrintType;
use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
@ -11,6 +11,7 @@ use tokio::sync::RwLock;
use uuid::Uuid;
pub struct AppConnection {
pub state: Arc<AppState>,
pub user_id: u64,
pub app_identifier: String,
pub app_session: Uuid,
@ -27,6 +28,7 @@ pub struct AppConnection {
impl AppConnection {
pub async fn from_general(general: Arc<GeneralConnection>, user_id: u64) -> Arc<Self> {
Arc::new(Self {
state: general.state.clone(),
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: general.rho_connection.clone(),
@ -154,7 +156,10 @@ impl AppConnection {
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let app_for_closure = self.clone();
tokio::spawn(async move {
let response_cv = get_omega_connection()
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 {
@ -169,7 +174,7 @@ impl AppConnection {
if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) {
let current = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.unwrap_or_default()
.as_millis();
let mut ping_guard = self.ping.write().await;
*ping_guard = (current as i128 - *last_ping) as i64;
@ -227,7 +232,10 @@ impl AppConnection {
return;
}
let load_uuid_response = get_omega_connection()
let load_uuid_response = self
.state
.omega
.clone()
.await_response(
&CommunicationValue::new(CommunicationType::GetUserData)
.with_id(cv.clone().get_id())
@ -328,7 +336,7 @@ impl AppConnection {
/// Handle connection close
pub async fn handle_close(&self) {
let user_id = self.get_user_id().await;
if let Some(rho_conn) = rho_manager::get_rho_con_for_user(user_id as i64).await {
if let Some(rho_conn) = self.state.rho.get_for_user(user_id as i64).await {
rho_conn.close_app_connection(Arc::new(self.clone())).await;
}
}
@ -338,6 +346,7 @@ impl AppConnection {
impl Clone for AppConnection {
fn clone(&self) -> Self {
Self {
state: self.state.clone(),
sender: Arc::clone(&self.sender),
receiver: Arc::clone(&self.receiver),
user_id: self.user_id,

144
src/rho/client_connection.rs
View file

@ -1,10 +1,10 @@
use crate::anonymous_clients::anonymous_manager;
use crate::calls::{call_group::call_invite_secret_from_cv, call_manager, call_util};
use crate::omega::omega_connection::get_omega_connection;
use crate::app_state::AppState;
use crate::calls::call_group::call_invite_secret_from_cv;
use crate::data::user::UserStatus;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::rho::{rho_connection::RhoConnection, rho_manager};
use crate::rho::rho_connection::RhoConnection;
use crate::util::logger::PrintType;
use crate::{data::user::UserStatus, omega::omega_connection::OmegaConnection};
use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use std::str::FromStr;
@ -15,6 +15,7 @@ use trust_dns_resolver::TokioAsyncResolver;
use uuid::Uuid;
pub struct ClientConnection {
pub state: Arc<AppState>,
pub user_id: u64,
pub session_id: u64,
pub client_version: String,
@ -26,11 +27,13 @@ pub struct ClientConnection {
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>,
}
impl ClientConnection {
pub async fn from_general(general: Arc<GeneralConnection>, user_id: u64) -> Arc<Self> {
Arc::new(Self {
state: general.state.clone(),
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: general.rho_connection.clone(),
@ -41,6 +44,7 @@ impl ClientConnection {
user_id: user_id,
session_id: general.session_id.read().await.clone(),
client_version: general.client_version.read().await.clone(),
message_slots: Arc::new(tokio::sync::Semaphore::new(32)),
})
}
pub fn start(self: Arc<Self>) {
@ -96,7 +100,11 @@ impl ClientConnection {
/// Handle incoming message from client
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
let Ok(permit) = self.message_slots.clone().acquire_owned().await else {
return;
};
tokio::spawn(async move {
let _permit = permit;
if cv.is_type(CommunicationType::Ping) {
self.handle_ping(cv).await;
return;
@ -270,14 +278,15 @@ impl ClientConnection {
}
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let client_for_closure = self.clone();
tokio::spawn(async move {
let response_cv = get_omega_connection()
.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 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;
}
}
/// Handle ping message
@ -286,7 +295,7 @@ impl ClientConnection {
if let DataValue::SignedNumber(last_ping) = cv.get_data(DataType::LastPing) {
let current = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.unwrap_or_default()
.as_millis();
let mut ping_guard = self.ping.write().await;
*ping_guard = (current as i128 - *last_ping) as i64;
@ -316,12 +325,14 @@ impl ClientConnection {
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 {
OmegaConnection::client_changed(
rho_conn.get_iota_id().await as i64,
user_id as i64,
user_status,
)
.await;
self.state
.omega
.client_changed(
rho_conn.get_iota_id().await as i64,
user_id as i64,
user_status,
)
.await;
}
}
}
@ -359,23 +370,29 @@ impl ClientConnection {
return;
}
};
let invited =
call_manager::add_invite(call_id, self.user_id, receiver_id as u64, secret.clone())
.await;
let invited = self
.state
.call_manager
.add_invite(call_id, self.user_id, receiver_id as u64, secret.clone())
.await;
if !invited {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
.await;
return;
}
if !call_manager::should_forward_invite(self.user_id, receiver_id as u64) {
if !self
.state
.call_manager
.should_forward_invite(self.user_id, receiver_id as u64)
{
let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
self.send_message(&response).await;
return;
}
// Find target RhoConnection
let target_rho = match rho_manager::get_rho_con_for_user(receiver_id as i64).await {
let target_rho = match self.state.rho.get_for_user(receiver_id as i64).await {
Some(rho) => rho,
_ => {
// Get sender user ID
@ -394,7 +411,7 @@ impl ClientConnection {
DataValue::Str("call_invite".to_string()),
);
let omega_conn = get_omega_connection();
let omega_conn = self.state.omega.clone();
// Send fire-and-forget, don't await to avoid blocking
tokio::spawn(async move {
let _ = omega_conn.send_message(&push_cv).await;
@ -455,18 +472,26 @@ impl ClientConnection {
}
};
if let Some(token) = call_manager::get_call_token(user_id, call_id).await {
let response = CommunicationValue::new(CommunicationType::CallToken)
.with_id(cv.get_id())
.with_receiver(user_id as u64)
.add_typed_default(DataType::CallToken, DataValue::Str(token));
self.send_message(&response).await;
} else {
let error_cv = CommunicationValue::new(CommunicationType::ErrorNoCallId)
.with_id(cv.get_id())
.add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string()));
self.send_message(&error_cv).await;
return;
match self
.state
.call_manager
.get_call_token(user_id, call_id)
.await
{
Ok(token) => {
let response = CommunicationValue::new(CommunicationType::CallToken)
.with_id(cv.get_id())
.with_receiver(user_id as u64)
.add_typed_default(DataType::CallToken, DataValue::Str(token));
self.send_message(&response).await;
}
Err(error) => {
log::warn!("Unable to create call token for {}: {}", call_id, error);
let error_cv = CommunicationValue::new(CommunicationType::ErrorNoCallId)
.with_id(cv.get_id())
.add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string()));
self.send_message(&error_cv).await;
}
}
}
async fn handle_get_call_data(self: Arc<Self>, cv: CommunicationValue) {
@ -488,7 +513,7 @@ impl ClientConnection {
}
};
if let Some(call) = call_manager::get_call(call_id).await {
if let Some(call) = self.state.call_manager.get_call(call_id).await {
if let Some(_) = call.get_caller(user_id).await {
let mut user_ids: Vec<DataValue> = Vec::new();
let members = call.members.read().await.clone();
@ -535,7 +560,7 @@ impl ClientConnection {
.as_signed_number()
.unwrap_or(0);
let Some(call) = call_manager::get_call(call_id).await else {
let Some(call) = self.state.call_manager.get_call(call_id).await else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound)
.await;
return;
@ -547,7 +572,11 @@ impl ClientConnection {
return;
};
if caller.has_admin() {
let _ = call_util::remove_participant(call_id, user_id as u64).await;
let _ = self
.state
.livekit
.remove_participant(call_id, user_id as u64)
.await;
if let Some(target) = call.get_caller(user_id as u64).await {
target.set_timeout(untill as i64).await;
}
@ -565,7 +594,7 @@ impl ClientConnection {
.as_signed_number()
.unwrap_or(0);
let Some(call) = call_manager::get_call(call_id).await else {
let Some(call) = self.state.call_manager.get_call(call_id).await else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound)
.await;
return;
@ -588,13 +617,18 @@ impl ClientConnection {
};
let enable = cv.get_data(DataType::Enabled).as_bool().unwrap_or(true);
let call = call_manager::get_call(call_id).await;
let call = self.state.call_manager.get_call(call_id).await;
let mut short_link = None;
if let Some(call) = call {
if let Some(caller) = call.get_caller(self.get_user_id().await).await {
if caller.has_admin() {
call.set_anonymous_joining(enable).await;
call.set_anonymous_joining(
enable,
&self.state.omega,
self.state.config.omikron_id,
)
.await;
}
}
short_link = call.get_short_link().await;
@ -626,11 +660,22 @@ impl ClientConnection {
match resolver.txt_lookup(path).await {
Ok(txt_lookup) => {
if let Some(txt_record) = txt_lookup.iter().next() {
let record_text: String = txt_record
let record_bytes: Vec<u8> = txt_record
.txt_data()
.iter()
.map(|b| String::from_utf8_lossy(b))
.flat_map(|chunk| chunk.iter().copied())
.collect();
let record_text = match String::from_utf8(record_bytes) {
Ok(text) => text,
Err(_) => {
self.send_error_response(
cv.get_id(),
CommunicationType::ErrorInvalidData,
)
.await;
return;
}
};
let response = CommunicationValue::new(CommunicationType::LoadTxtRecord)
.with_id(cv.get_id())
@ -701,7 +746,10 @@ impl ClientConnection {
return;
}
let load_uuid_response = get_omega_connection()
let load_uuid_response = self
.state
.omega
.clone()
.await_response(
&CommunicationValue::new(CommunicationType::GetUserData)
.with_id(cv.clone().get_id())
@ -830,7 +878,7 @@ impl ClientConnection {
/// Handle connection close
pub async fn handle_close(&self) {
let user_id = self.get_user_id().await;
if let Some(rho_conn) = rho_manager::get_rho_con_for_user(user_id as i64).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;
@ -842,6 +890,7 @@ impl ClientConnection {
impl Clone for ClientConnection {
fn clone(&self) -> Self {
Self {
state: self.state.clone(),
sender: Arc::clone(&self.sender),
receiver: Arc::clone(&self.receiver),
user_id: self.user_id,
@ -852,6 +901,7 @@ impl Clone for ClientConnection {
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),
}
}
}

471
src/rho/connection.rs
View file

@ -1,19 +1,18 @@
use std::{collections::BTreeMap, collections::HashMap, sync::Arc, time::Duration};
use std::{sync::Arc, time::Duration};
use tokio::sync::RwLock;
use uuid::Uuid;
use crate::{
anonymous_clients::anonymous_client_connection::AnonymousClientConnection,
calls::call_manager,
app_state::AppState,
log_cv_out, log_err, log_in, log_out,
omega::omega_connection::get_omega_connection,
rho::{
app_connection::AppConnection, client_connection::ClientConnection,
iota_connection::IotaConnection, rho_connection::RhoConnection, rho_manager,
iota_connection::IotaConnection, rho_connection::RhoConnection,
},
util::logger::PrintType,
};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::host::AuthState;
use mtp::webserver::{WebMTPConnection, WebMtpReceiver, WebMtpSender};
@ -38,6 +37,7 @@ pub enum ConnectionKind {
}
pub struct GeneralConnection {
pub state: Arc<AppState>,
pub sender: Arc<MtpSender>,
pub receiver: Arc<MtpReceiver>,
@ -61,7 +61,7 @@ impl GeneralConnection {
* to a known, appropriately-authenticated kind is rejected here instead
* of being handed off to a connection handler.
*/
pub fn new(conn: WebMTPConnection) -> Option<Arc<Self>> {
pub fn new(conn: WebMTPConnection, state: Arc<AppState>) -> Option<Arc<Self>> {
let kind = match (conn.description.as_deref(), &conn.auth_state) {
(Some("iota"), AuthState::Authenticated) => ConnectionKind::Iota,
(Some("client"), AuthState::Authenticated) => ConnectionKind::Client,
@ -70,6 +70,7 @@ impl GeneralConnection {
};
Some(Arc::new(Self {
state,
sender: Arc::new(conn.sender),
receiver: Arc::new(conn.receiver),
connection_kind: kind,
@ -89,293 +90,181 @@ impl GeneralConnection {
}
async fn migrate(self: &Arc<Self>) -> bool {
let id = self.id;
match self.connection_kind {
ConnectionKind::Client => {
let notify = CommunicationValue::new(CommunicationType::UserConnected)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber((id as i64).into()),
);
get_omega_connection().send_message(&notify).await;
let user_id = id as i64;
let client = ClientConnection::from_general(self.clone(), id).await;
let mut rho = rho_manager::get_rho_con_for_user(user_id).await;
if rho.is_none() {
let get_user_msg = CommunicationValue::new(CommunicationType::GetUserData)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(user_id.into()),
);
if let Ok(user_data_cv) = get_omega_connection()
.await_response(&get_user_msg, Some(Duration::from_secs(20)))
.await
{
if let DataValue::SignedNumber(iota_id) =
user_data_cv.get_data(DataType::IotaId)
{
if let Some(bound_rho) =
rho_manager::bind_user_to_iota(user_id, *iota_id as i64).await
{
bound_rho.bind_user_id(user_id).await;
rho = Some(bound_rho);
}
}
}
}
*self.rho_connection.write().await = rho.clone();
if let Some(rho_conn) = rho {
let session_id = *self.session_id.read().await as i64;
let iota_msg = CommunicationValue::new(CommunicationType::ClientConnected)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(user_id.into()),
)
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(session_id.into()),
);
if let Ok(resp) = rho_conn
.get_iota_connection()
.clone()
.await_response(&iota_msg, Some(Duration::from_secs(20)))
.await
{
/*
* mtp's handshake already confirmed this login; this
* message is an unsolicited push carrying the initial
* contacts/calls state, not a reply to a client-side
* request id (the old raw hello frame this used to
* reply to no longer reaches application code).
*/
let mut ident_resp =
CommunicationValue::new(CommunicationType::IdentificationResponse);
let tm = mtp::codec::TypeMap::latest();
for (k, v) in resp.iter_typed_data() {
let value_to_add = if k == Some(DataType::Contacts) {
if let Some(contacts) = v.as_array() {
let call_groups =
call_manager::get_call_groups(user_id as u64).await;
let mut invites: HashMap<i64, Vec<DataValue>> = HashMap::new();
let mut global_calls: Vec<DataValue> = Vec::new();
for call in call_groups {
let members = call.members.read().await;
let call_self =
members.iter().find(|m| m.user_id == user_id as u64);
if let Some(call_self) = call_self {
let timeout = *call_self.timeout.read().await;
let admin = call_self.has_admin();
// List of all members in the call
let member_ids: Vec<DataValue> = members
.iter()
.map(|m| DataValue::SignedNumber(m.user_id.into()))
.collect();
let mut base_call_map: BTreeMap<DataTypeId, DataValue> =
BTreeMap::new();
base_call_map.insert(
DataType::CallId.to_id(&tm),
DataValue::Str(call.call_id.to_string()),
);
base_call_map.insert(
DataType::CallMembers.to_id(&tm),
DataValue::Array(member_ids),
);
if timeout > 0 {
base_call_map.insert(
DataType::Timeout.to_id(&tm),
DataValue::SignedNumber(
(timeout as i64).into(),
),
);
}
if admin {
base_call_map.insert(
DataType::HasAdmin.to_id(&tm),
DataValue::Bool(true),
);
}
if let Some(secret) =
call.get_secret_for_user(user_id as u64).await
{
base_call_map.insert(
DataType::CallSecret.to_id(&tm),
secret.to_data_value(),
);
}
// Add to global calls with only this user's recipient-specific secret.
global_calls.push(DataValue::container_from_map(
&base_call_map,
));
for member in members.iter() {
let member_id = member.user_id;
if member_id == user_id as u64 {
continue;
}
let contact_call_map = base_call_map.clone();
invites
.entry(member_id as i64)
.or_insert_with(Vec::new)
.push(DataValue::container_from_map(
&contact_call_map,
));
}
}
}
let mut new_contacts: Vec<DataValue> = Vec::new();
for contact in contacts {
if let Some(mut contact_map) = contact.as_map() {
if let Some(DataValue::SignedNumber(id)) =
contact_map.get(&DataType::UserId.to_id(&tm))
{
if let Some(call_list) = invites.get(&(*id as i64))
{
contact_map.insert(
DataType::Calls.to_id(&tm),
DataValue::Array(call_list.clone()),
);
}
}
new_contacts
.push(DataValue::container_from_map(&contact_map));
} else {
new_contacts.push(contact.clone());
}
}
ident_resp = ident_resp.add_typed_default(
DataType::Calls,
DataValue::Array(global_calls),
);
DataValue::Array(new_contacts)
} else {
v.clone()
}
} else {
v.clone()
};
if let Some(dt) = k {
ident_resp = ident_resp.add_typed_default(dt, value_to_add);
}
}
log_cv_out!(ident_resp);
let _ = self.sender.send(&ident_resp).await;
rho_conn.bind_user_id(user_id).await;
rho_conn.add_client_connection(client.clone()).await;
}
} else {
log_err!(
user_id,
PrintType::Client,
"No RhoConnection found for user {}, client not attached to iota",
id
);
}
client.start();
}
ConnectionKind::Iota => {
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()));
get_omega_connection().send_message(&notify).await;
let iota = IotaConnection::from_general(self.clone(), id).await;
let rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
iota.set_rho_connection(rho.clone()).await;
rho_manager::add_rho(rho).await;
let get_iota_msg = CommunicationValue::new(CommunicationType::GetIotaData)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into()));
if let Ok(iota_data_cv) = get_omega_connection()
.await_response(&get_iota_msg, Some(Duration::from_secs(20)))
.await
{
if let DataValue::Array(users) = iota_data_cv.get_data(DataType::UserIds) {
let mut user_ids: Vec<u64> = Vec::new();
for value in users {
if let DataValue::SignedNumber(user_id) = value {
user_ids.push(*user_id as u64);
}
}
iota.set_user_ids(user_ids).await;
}
}
// mtp's handshake already confirmed this login/registration to the Iota.
iota.clone().start();
}
ConnectionKind::AnonymousClient => {
let client = AnonymousClientConnection::from_general(self.clone(), id).await;
client.start();
}
ConnectionKind::Phi => {
let user_id = id as i64;
let mut rho = rho_manager::get_rho_con_for_user(user_id).await;
if rho.is_none() {
let get_user_msg = CommunicationValue::new(CommunicationType::GetUserData)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(user_id.into()),
);
if let Ok(user_data_cv) = get_omega_connection()
.await_response(&get_user_msg, Some(Duration::from_secs(20)))
.await
{
if let DataValue::SignedNumber(iota_id) =
user_data_cv.get_data(DataType::IotaId)
{
if let Some(bound_rho) =
rho_manager::bind_user_to_iota(user_id, *iota_id as i64).await
{
bound_rho.bind_user_id(user_id).await;
rho = Some(bound_rho);
}
}
}
}
*self.rho_connection.write().await = rho.clone();
let app_conn = AppConnection::from_general(self.clone(), id).await;
if let Some(rho_conn) = rho {
rho_conn.bind_user_id(user_id).await;
rho_conn.add_app_connection(app_conn.clone()).await;
}
app_conn.start();
}
ConnectionKind::Client => self.migrate_client().await,
ConnectionKind::Iota => self.migrate_iota().await,
ConnectionKind::AnonymousClient => self.migrate_anonymous_client().await,
ConnectionKind::Phi => self.migrate_phi().await,
}
true
}
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 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 {
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!(
user_id,
PrintType::Client,
"No RhoConnection found for user {}, client not attached to iota",
id
);
}
client.start();
}
async fn notify_user_connected(&self, user_id: i64) {
let notify = CommunicationValue::new(CommunicationType::UserConnected)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()));
self.state.omega.clone().send_message(&notify).await;
}
async fn find_user_rho(&self, user_id: i64) -> Option<Arc<RhoConnection>> {
if let Some(rho) = self.state.rho.get_for_user(user_id).await {
return Some(rho);
}
let get_user_msg = CommunicationValue::new(CommunicationType::GetUserData)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()));
let Ok(user_data_cv) = self
.state
.omega
.clone()
.await_response(&get_user_msg, Some(Duration::from_secs(20)))
.await
else {
return None;
};
let DataValue::SignedNumber(iota_id) = user_data_cv.get_data(DataType::IotaId) else {
return None;
};
let rho = self
.state
.rho
.bind_user_to_iota(user_id, *iota_id as i64)
.await?;
rho.bind_user_id(user_id).await;
Some(rho)
}
async fn request_initial_client_state(
&self,
rho: &Arc<RhoConnection>,
user_id: i64,
) -> Option<CommunicationValue> {
let session_id = *self.session_id.read().await as i64;
let request = CommunicationValue::new(CommunicationType::ClientConnected)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(session_id.into()),
);
rho.get_iota_connection()
.clone()
.await_response(&request, Some(Duration::from_secs(20)))
.await
.ok()
}
async fn add_call_state(
&self,
response: CommunicationValue,
user_id: i64,
) -> CommunicationValue {
let mut output = CommunicationValue::new(CommunicationType::IdentificationResponse);
for (key, value) in response.iter_typed_data() {
if key == Some(DataType::Contacts) {
if let Some(contacts) = value.as_array() {
let (contacts, global_calls) = self
.state
.call_state_aggregator
.augment_contacts(user_id as u64, contacts.clone())
.await;
output =
output.add_typed_default(DataType::Contacts, DataValue::Array(contacts));
output =
output.add_typed_default(DataType::Calls, DataValue::Array(global_calls));
continue;
}
}
if let Some(data_type) = key {
output = output.add_typed_default(data_type, value.clone());
}
}
output
}
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 rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
iota.set_rho_connection(rho.clone()).await;
self.state.rho.add(rho).await;
self.load_iota_users(&iota, id).await;
iota.start();
}
async fn load_iota_users(&self, iota: &Arc<IotaConnection>, iota_id: u64) {
let request = CommunicationValue::new(CommunicationType::GetIotaData)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()));
let Ok(response) = self
.state
.omega
.clone()
.await_response(&request, Some(Duration::from_secs(20)))
.await
else {
return;
};
let DataValue::Array(users) = response.get_data(DataType::UserIds) else {
return;
};
let user_ids = users
.iter()
.filter_map(|value| match value {
DataValue::SignedNumber(id) => Some(*id as u64),
_ => None,
})
.collect();
iota.set_user_ids(user_ids).await;
}
async fn migrate_anonymous_client(self: &Arc<Self>) {
let client = AnonymousClientConnection::from_general(self.clone(), self.id).await;
client.start();
}
async fn migrate_phi(self: &Arc<Self>) {
let user_id = self.id as i64;
let rho = self.find_user_rho(user_id).await;
*self.rho_connection.write().await = rho.clone();
let app_conn = AppConnection::from_general(self.clone(), self.id).await;
if let Some(rho_conn) = rho {
rho_conn.bind_user_id(user_id).await;
rho_conn.add_app_connection(app_conn.clone()).await;
}
app_conn.start();
}
}

74
src/rho/iota_connection.rs
View file

@ -1,12 +1,12 @@
use crate::app_state::AppState;
use crate::calls::call_group::CallGroup;
use crate::calls::call_manager;
use crate::log_cv_in;
use crate::log_cv_out;
use crate::log_err;
use crate::log_in;
use crate::log_out;
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::util::data_type_id;
use crate::util::logger::PrintType;
use dashmap::DashMap;
use mtp::codec::CommunicationType;
@ -17,18 +17,15 @@ use mtp::codec::DataValue;
use mtp::codec::TypeMap;
use mtp::crypto::KemPublicKey;
use std::collections::BTreeMap;
use std::{collections::HashMap, sync::Arc, sync::LazyLock, time::Duration};
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
use super::{rho_connection::RhoConnection, rho_manager};
use crate::omega::omega_connection::OmegaConnection;
static PENDING_CHAT_SECRETS: LazyLock<DashMap<u64, Vec<CommunicationValue>>> =
LazyLock::new(DashMap::new);
use super::rho_connection::RhoConnection;
#[allow(dead_code)]
pub struct IotaConnection {
pub state: Arc<AppState>,
pub iota_id: u64,
pub client_version: String,
pub sender: Arc<MtpSender>,
@ -39,11 +36,14 @@ pub struct IotaConnection {
pub waiting_tasks:
DashMap<u32, Box<dyn Fn(Arc<IotaConnection>, CommunicationValue) -> bool + Send + Sync>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
pending_chat_secrets: DashMap<u64, Vec<CommunicationValue>>,
message_slots: Arc<tokio::sync::Semaphore>,
}
impl IotaConnection {
pub async fn from_general(general: Arc<GeneralConnection>, iota_id: u64) -> Arc<Self> {
Arc::new(Self {
state: general.state.clone(),
ping: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
rho_connection: general.rho_connection.clone(),
@ -53,6 +53,8 @@ impl IotaConnection {
iota_id: iota_id,
client_version: general.client_version.read().await.clone(),
waiting_tasks: DashMap::new(),
pending_chat_secrets: DashMap::new(),
message_slots: Arc::new(tokio::sync::Semaphore::new(32)),
})
}
pub fn start(self: Arc<Self>) {
@ -128,7 +130,7 @@ impl IotaConnection {
}
async fn flush_pending_chat_secrets(&self, user_id: u64) {
let Some((_, messages)) = PENDING_CHAT_SECRETS.remove(&user_id) else {
let Some((_, messages)) = self.pending_chat_secrets.remove(&user_id) else {
return;
};
@ -137,13 +139,13 @@ impl IotaConnection {
}
}
fn store_pending_chat_secret(cv: CommunicationValue) {
fn store_pending_chat_secret(&self, cv: CommunicationValue) {
let receiver_id = cv.get_receiver();
if receiver_id == 0 || !cv.is_type(CommunicationType::SetChatSecret) {
return;
}
PENDING_CHAT_SECRETS
self.pending_chat_secrets
.entry(receiver_id)
.or_default()
.push(cv);
@ -187,6 +189,10 @@ impl IotaConnection {
/// Handle incoming message from Iota
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
let Ok(permit) = self.message_slots.clone().acquire_owned().await else {
return;
};
let _permit = permit;
let msg_id = cv.get_id();
if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) {
if (task)(self.clone(), cv.clone()) {
@ -219,7 +225,10 @@ impl IotaConnection {
}
if cv.is_type(CommunicationType::CompleteRegisterUser) {
let response_cv = get_omega_connection()
let response_cv = self
.state
.omega
.clone()
.await_response(
&cv.clone().with_sender(self.iota_id),
Some(Duration::from_secs(20)),
@ -265,7 +274,10 @@ impl IotaConnection {
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let iota_for_closure = self.clone();
let response_cv = get_omega_connection()
let response_cv = self
.state
.omega
.clone()
.await_response(&cv.with_sender(self.iota_id), Some(Duration::from_secs(20)))
.await;
if let Ok(response_cv) = response_cv {
@ -293,10 +305,13 @@ impl IotaConnection {
.map(|(k, v)| {
let mut map = BTreeMap::new();
if let Ok(uid) = k.parse::<i128>() {
map.insert(DataType::UserId.to_id(&tm), DataValue::SignedNumber(uid));
map.insert(
data_type_id(DataType::UserId, &tm),
DataValue::SignedNumber(uid),
);
}
map.insert(
DataType::LastPing.to_id(&tm),
data_type_id(DataType::LastPing, &tm),
DataValue::SignedNumber(v.into()),
);
DataValue::container_from_map(&map)
@ -327,11 +342,11 @@ impl IotaConnection {
);
if my_user_ids.contains(&(sender_id as u64)) {
if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id as i64).await {
if let Some(target_rho) = self.state.rho.get_for_user(receiver_id as i64).await {
target_rho.message_to_iota(cv).await;
} else {
if cv.is_type(CommunicationType::SetChatSecret) {
Self::store_pending_chat_secret(cv.clone());
self.store_pending_chat_secret(cv.clone());
let success = CommunicationValue::new(CommunicationType::Success)
.with_id(cv.get_id())
.with_sender(cv.get_sender())
@ -387,7 +402,7 @@ impl IotaConnection {
// ============================
// Load Calls
// ============================
let calls: Vec<Arc<CallGroup>> = call_manager::get_call_groups(user_id).await;
let calls: Vec<Arc<CallGroup>> = self.state.call_manager.get_call_groups(user_id).await;
let mut invites: HashMap<i64, Vec<DataValue>> = HashMap::new();
let mut global_calls: Vec<DataValue> = Vec::new();
@ -412,27 +427,31 @@ impl IotaConnection {
// Build base call container
let mut base_call_map: BTreeMap<DataTypeId, DataValue> = BTreeMap::new();
base_call_map.insert(
DataType::CallId.to_id(&tm),
data_type_id(DataType::CallId, &tm),
DataValue::Str(call.call_id.to_string()),
);
base_call_map.insert(
DataType::CallMembers.to_id(&tm),
data_type_id(DataType::CallMembers, &tm),
DataValue::Array(member_ids),
);
if timeout > 0 {
base_call_map.insert(
DataType::Timeout.to_id(&tm),
data_type_id(DataType::Timeout, &tm),
DataValue::SignedNumber(timeout.into()),
);
}
if admin {
base_call_map.insert(DataType::HasAdmin.to_id(&tm), DataValue::Bool(true));
base_call_map
.insert(data_type_id(DataType::HasAdmin, &tm), DataValue::Bool(true));
}
if let Some(secret) = call.get_secret_for_user(user_id).await {
base_call_map.insert(DataType::CallSecret.to_id(&tm), secret.to_data_value());
base_call_map.insert(
data_type_id(DataType::CallSecret, &tm),
secret.to_data_value(),
);
}
// Add to global calls with only this user's recipient-specific secret.
@ -473,7 +492,7 @@ impl IotaConnection {
entries.iter().cloned().collect();
if let Some(DataValue::SignedNumber(id)) =
user_map.get(&DataType::UserId.to_id(&tm))
user_map.get(&data_type_id(DataType::UserId, &tm))
{
interested_ids.push(*id as i64);
@ -481,7 +500,7 @@ impl IotaConnection {
&& !call_list.is_empty()
{
user_map.insert(
DataType::Calls.to_id(&tm),
data_type_id(DataType::Calls, &tm),
DataValue::Array(call_list.clone()),
);
}
@ -498,7 +517,10 @@ impl IotaConnection {
// ============================
// Notify Omega
// ============================
OmegaConnection::user_states(user_id as i64, interested_ids.clone()).await;
self.state
.omega
.user_states(user_id as i64, interested_ids.clone())
.await;
// ============================
// Notify Rho

164
src/rho/rho_connection.rs
View file

@ -1,7 +1,7 @@
use super::{client_connection::ClientConnection, iota_connection::IotaConnection, rho_manager};
use super::{client_connection::ClientConnection, iota_connection::IotaConnection};
use crate::omega::omega_connection::OmegaConnection;
use crate::{data::user::UserStatus, rho::app_connection::AppConnection};
use dashmap::DashMap;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use std::collections::HashMap;
use std::sync::Arc;
@ -11,8 +11,8 @@ use uuid::Uuid;
pub struct RhoConnection {
iota_connection: Arc<IotaConnection>,
user_ids: Arc<RwLock<Vec<i64>>>,
client_connections: Arc<RwLock<Vec<Arc<ClientConnection>>>>,
app_connections: Arc<RwLock<Vec<Arc<AppConnection>>>>,
client_connections: DashMap<(u64, u64), Arc<ClientConnection>>,
app_connections: DashMap<(u64, String, Uuid), Arc<AppConnection>>,
}
impl RhoConnection {
@ -21,8 +21,8 @@ impl RhoConnection {
let rho_connection = Self {
iota_connection,
user_ids: Arc::new(RwLock::new(user_ids.clone())),
client_connections: Arc::new(RwLock::new(Vec::new())),
app_connections: Arc::new(RwLock::new(Vec::new())),
client_connections: DashMap::new(),
app_connections: DashMap::new(),
};
rho_connection
@ -58,8 +58,10 @@ impl RhoConnection {
}
pub async fn get_client_connections(&self) -> Vec<Arc<ClientConnection>> {
let connections = self.client_connections.read().await;
connections.clone()
self.client_connections
.iter()
.map(|entry| entry.value().clone())
.collect()
}
/// Get client connections for a specific user
@ -67,14 +69,11 @@ impl RhoConnection {
&self,
user_id: i64,
) -> Vec<Arc<ClientConnection>> {
let connections = self.client_connections.read().await;
let mut collections = Vec::new();
for con in connections.iter() {
if con.get_user_id().await == user_id as u64 {
collections.push(con.clone());
}
}
collections
self.client_connections
.iter()
.filter(|entry| entry.key().0 == user_id as u64)
.map(|entry| entry.value().clone())
.collect()
}
#[allow(unused)]
@ -84,10 +83,10 @@ impl RhoConnection {
app_identifier: Option<String>,
app_session: Option<Uuid>,
) -> Vec<Arc<AppConnection>> {
let connections = self.app_connections.read().await;
connections
self.app_connections
.iter()
.filter(|conn| {
.filter(|entry| {
let conn = entry.value();
if let Some(uid) = userid {
if conn.user_id != uid as u64 {
return false;
@ -105,18 +104,22 @@ impl RhoConnection {
}
true
})
.cloned()
.map(|entry| entry.value().clone())
.collect()
}
pub async fn add_app_connection(&self, connection: Arc<AppConnection>) {
let mut connections = self.app_connections.write().await;
connections.push(connection);
let key = (
connection.user_id,
connection.app_identifier.clone(),
connection.app_session,
);
self.app_connections.insert(key, connection);
}
pub async fn close_app_connection(&self, connection: Arc<AppConnection>) {
let mut connections = self.app_connections.write().await;
connections.retain(|c| c.app_session != connection.app_session);
self.app_connections
.retain(|_, c| c.app_session != connection.app_session);
}
/// Add a client connection
@ -130,37 +133,25 @@ impl RhoConnection {
self.iota_connection.send_message(&notification).await;
let mut should_notify_online = false;
{
let mut connections = self.client_connections.write().await;
let mut keep = Vec::new();
let mut had_user = false;
for con in connections.drain(..) {
if con.user_id as i64 == user_id {
had_user = true;
if con.session_id as i64 == session_id {
continue;
}
}
keep.push(con);
}
if !had_user {
should_notify_online = true;
}
keep.push(Arc::clone(&connection));
*connections = keep;
}
let should_notify_online = !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));
self.client_connections
.insert((user_id as u64, session_id as u64), Arc::clone(&connection));
if should_notify_online {
OmegaConnection::client_changed(
self.get_iota_id().await as i64,
user_id,
UserStatus::user_online,
)
.await;
self.iota_connection
.state
.omega
.client_changed(
self.get_iota_id().await as i64,
user_id,
UserStatus::user_online,
)
.await;
}
}
@ -168,32 +159,28 @@ impl RhoConnection {
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;
let mut remaining_for_user = false;
{
let mut connections = self.client_connections.write().await;
let mut keep = Vec::new();
for con in connections.drain(..) {
if con.user_id as i64 == target_user_id
&& con.session_id as i64 == target_session_id
{
continue;
}
if con.user_id as i64 == target_user_id {
remaining_for_user = true;
}
keep.push(con);
}
*connections = keep;
}
self.client_connections
.remove(&(target_user_id as u64, target_session_id as u64));
let remaining_for_user = self
.client_connections
.iter()
.any(|entry| entry.key().0 == target_user_id as u64);
if !remaining_for_user {
OmegaConnection::client_changed(
self.get_iota_id().await as i64,
target_user_id,
UserStatus::user_offline,
)
.await;
self.iota_connection
.state
.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,
UserStatus::user_offline,
)
.await;
}
}
@ -206,15 +193,23 @@ impl RhoConnection {
}
// Remove from manager
rho_manager::remove_rho(self.get_iota_id().await as i64).await;
self.iota_connection
.state
.rho
.remove(self.get_iota_id().await as i64)
.await;
// Notify OmegaConnection
OmegaConnection::close_iota(self.get_iota_id().await as i64).await;
self.iota_connection
.state
.omega
.close_iota(self.get_iota_id().await as i64)
.await;
}
/// Send message from Iota to specific client
pub async fn message_to_client(&self, cv: CommunicationValue) {
let connections = self.client_connections.read().await;
let connections = self.get_client_connections().await;
let receiver_id = cv.get_receiver();
let session_id = cv.get_data(DataType::SessionId).as_number();
@ -238,7 +233,7 @@ 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.client_connections.read().await;
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 {
@ -254,7 +249,7 @@ impl RhoConnection {
/// 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.client_connections.read().await;
let connections = self.get_client_connections().await;
for connection in connections.iter() {
connection
.clone()
@ -265,7 +260,7 @@ impl RhoConnection {
/// Get ping information for all clients
pub async fn get_client_pings(&self) -> HashMap<String, i64> {
let connections = self.client_connections.read().await;
let connections = self.get_client_connections().await;
let mut pings = HashMap::new();
for connection in connections.iter() {
@ -285,7 +280,6 @@ impl RhoConnection {
/// Get count of active client connections
#[allow(dead_code)]
pub async fn client_count(&self) -> usize {
let connections = self.client_connections.read().await;
connections.len()
self.client_connections.len()
}
}

188
src/rho/rho_manager.rs
View file

@ -1,83 +1,129 @@
use std::sync::Arc;
use dashmap::DashMap;
use super::rho_connection::RhoConnection;
use crate::log_in;
use crate::util::logger::PrintType;
use std::{
collections::HashMap,
sync::{Arc, LazyLock},
};
use tokio::sync::RwLock;
pub static RHO_CONNECTIONS: LazyLock<Arc<RwLock<HashMap<i64, Arc<RhoConnection>>>>> =
LazyLock::new(|| Arc::new(RwLock::new(HashMap::new())));
/*
* Owns both indices for active Iota connections. Keeping the user index next
* to the Iota index makes binding updates atomic from the manager's caller's
* perspective and lets separate Omikron instances keep separate routing state.
*/
#[derive(Default)]
pub struct RhoManager {
connections: DashMap<i64, Arc<RhoConnection>>,
users: DashMap<i64, Arc<RhoConnection>>,
}
pub async fn get_rho_con_for_user(user_id: i64) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
for rho_connection in connections.values() {
let rho_user_ids = rho_connection.get_user_ids().await;
log_in!(
user_id,
PrintType::Client,
"Comparing user IDs: {:?}",
rho_user_ids
);
if rho_user_ids.contains(&user_id) {
return Some(Arc::clone(rho_connection));
impl RhoManager {
pub fn new() -> Self {
Self::default()
}
pub async fn get_for_user(&self, user_id: i64) -> Option<Arc<RhoConnection>> {
self.users.get(&user_id).map(|entry| entry.value().clone())
}
pub async fn contains_iota(&self, iota_id: i64) -> bool {
self.connections.contains_key(&iota_id)
}
pub async fn bind_user_to_iota(
&self,
user_id: i64,
iota_id: i64,
) -> Option<Arc<RhoConnection>> {
let rho = self
.connections
.get(&iota_id)
.map(|entry| entry.value().clone());
if let Some(rho) = rho {
rho.add_user_id(user_id).await;
self.users.insert(user_id, rho.clone());
Some(rho)
} else {
None
}
}
None
}
#[allow(dead_code)]
pub async fn contains_iota(iota_id: i64) -> bool {
let connections = RHO_CONNECTIONS.read().await;
connections.contains_key(&iota_id)
}
pub async fn replace_users_for_iota(&self, iota_id: i64, user_ids: Vec<i64>) {
let rho = self
.connections
.get(&iota_id)
.map(|entry| entry.value().clone());
let Some(rho) = rho else {
return;
};
/// Bind a user ID to an already tracked iota/rho connection.
pub async fn bind_user_to_iota(user_id: i64, iota_id: i64) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
if let Some(rho_connection) = connections.get(&iota_id) {
let rho = Arc::clone(rho_connection);
drop(connections);
let previous_users = rho.get_user_ids().await;
rho.set_user_ids(user_ids.clone()).await;
rho.get_iota_connection()
.set_user_ids(
user_ids
.iter()
.filter_map(|user_id| u64::try_from(*user_id).ok())
.collect(),
)
.await;
rho.add_user_id(user_id).await;
for user_id in previous_users {
if self
.users
.get(&user_id)
.is_some_and(|entry| Arc::ptr_eq(entry.value(), &rho))
{
self.users.remove(&user_id);
}
}
for user_id in user_ids {
self.users.insert(user_id, rho.clone());
}
}
log_in!(
user_id,
PrintType::Client,
"Bound user {} to iota {}",
user_id,
iota_id
);
pub async fn remove_user_binding(&self, user_id: i64, iota_id: i64) {
let rho = self.users.get(&user_id).map(|entry| entry.value().clone());
if let Some(rho) = rho {
if rho.get_iota_id().await as i64 == iota_id {
self.users.remove(&user_id);
}
}
}
Some(rho)
} else {
None
pub async fn remove(&self, iota_id: i64) -> Option<Arc<RhoConnection>> {
let rho = self.connections.remove(&iota_id).map(|(_, rho)| rho);
if let Some(rho) = rho.as_ref() {
self.users.retain(|_, mapped| !Arc::ptr_eq(mapped, rho));
}
rho
}
pub async fn add(&self, rho_connection: Arc<RhoConnection>) {
let iota_id = rho_connection.get_iota_id().await as i64;
let user_ids = rho_connection.get_user_ids().await;
self.connections.insert(iota_id, rho_connection.clone());
for user_id in user_ids {
self.users.insert(user_id, rho_connection.clone());
}
}
pub async fn get_by_iota(&self, iota_id: i64) -> Option<Arc<RhoConnection>> {
self.connections
.get(&iota_id)
.map(|entry| entry.value().clone())
}
pub async fn connection_count(&self) -> usize {
self.connections.len()
}
pub async fn iota_ids(&self) -> Vec<i64> {
self.connections.iter().map(|entry| *entry.key()).collect()
}
pub async fn connections(&self) -> Vec<Arc<RhoConnection>> {
self.connections
.iter()
.map(|entry| entry.value().clone())
.collect()
}
}
/// Remove a RhoConnection by Iota ID
pub async fn remove_rho(iota_id: i64) -> Option<Arc<RhoConnection>> {
let mut connections = RHO_CONNECTIONS.write().await;
connections.remove(&iota_id)
}
/// Add a RhoConnection to the manager
pub async fn add_rho(rho_connection: Arc<RhoConnection>) {
let mut connections = RHO_CONNECTIONS.write().await;
let iota_id = rho_connection.get_iota_id().await;
connections.insert(iota_id as i64, rho_connection);
}
/// Get a RhoConnection by Iota ID directly
#[allow(dead_code)]
pub async fn get_rho_by_iota(iota_id: i64) -> Option<Arc<RhoConnection>> {
let connections = RHO_CONNECTIONS.read().await;
connections.get(&iota_id).map(Arc::clone)
}
/// Get the count of active connections
pub async fn connection_count() -> usize {
let connections = RHO_CONNECTIONS.read().await;
connections.len()
}

48
src/rho/server.rs
View file

@ -1,13 +1,14 @@
use std::net::{IpAddr, Ipv4Addr};
use std::sync::Arc;
use std::time::Duration;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64_STD;
use crate::load_keyring;
use crate::{
app_state::AppState,
log, log_err,
omega::omega_connection::get_omega_connection,
omega::omega_connection::OmegaConnection,
rho::connection::GeneralConnection,
util::{file_util::load_file_vec, logger::PrintType},
};
@ -24,6 +25,7 @@ use mtp::webserver::{MTPWebServer, WebServerConfig};
* through mtp's unauthenticated fallback instead of rejecting it outright.
*/
pub async fn get_by_connector_id(
omega: Arc<OmegaConnection>,
client_id: u64,
description: Option<String>,
) -> Option<PublicKeyBundle> {
@ -38,7 +40,7 @@ pub async fn get_by_connector_id(
_ => return None,
};
let response = match get_omega_connection()
let response = match omega
.await_response(&request, Some(Duration::from_secs(20)))
.await
{
@ -62,7 +64,11 @@ pub async fn get_by_connector_id(
}
/* Only Iota registration goes through mtp's Register flow; users are registered out of band. */
pub async fn complete_register(pub_key: PublicKeyBundle, description: Option<String>) -> u64 {
pub async fn complete_register(
omega: Arc<OmegaConnection>,
pub_key: PublicKeyBundle,
description: Option<String>,
) -> u64 {
println!("Iota register start");
if description.as_deref() != Some("iota") {
return 0;
@ -75,7 +81,7 @@ pub async fn complete_register(pub_key: PublicKeyBundle, description: Option<Str
DataValue::Str(BASE64_STD.encode(pub_key.as_bytes())),
);
let response = match get_omega_connection()
let response = match omega
.await_response(&request, Some(Duration::from_secs(20)))
.await
{
@ -97,14 +103,14 @@ pub async fn complete_register(pub_key: PublicKeyBundle, description: Option<Str
}
}
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error>> {
let cert_pem = load_file_vec("certs", "cert.pem").expect("Error loading Pemfile");
let key_pem = load_file_vec("certs", "key.pem").expect("Error loading Keyfile");
let host_config = HostConfig::new(
IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)),
port,
state.config.rho_port,
cert_pem,
key_pem,
)
@ -124,16 +130,33 @@ pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
.with_persistent_stream_retries(5, Duration::from_secs(5)),
)
.with_authentication(
load_keyring(),
Box::new(|user_id, description| Box::pin(get_by_connector_id(user_id, description))),
Box::new(|pub_key, description| Box::pin(complete_register(pub_key, description))),
state
.keyring_for_host()
.map_err(|error| format!("Unable to copy keyring for host: {error}"))?,
Box::new({
let omega = state.omega.clone();
move |user_id, description| {
Box::pin(get_by_connector_id(omega.clone(), user_id, description))
}
}),
Box::new({
let omega = state.omega.clone();
move |pub_key, description| {
Box::pin(complete_register(omega.clone(), pub_key, description))
}
}),
)
.with_authentication_policy(AuthenticationPolicy::AllowAuthentication);
let web_config = WebServerConfig::new()
.route("/", |_request, response| async move { response.body("OK") })?;
let mut host = MTPWebServer::new(host_config, web_config).await?;
log!(0, PrintType::General, "Server listening on port {}.", port);
log!(
0,
PrintType::General,
"Server listening on port {}.",
state.config.rho_port
);
loop {
let conn = match host.accept().await {
@ -153,8 +176,9 @@ pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
}
};
let state = state.clone();
tokio::spawn(async move {
let Some(conn) = GeneralConnection::new(conn) else {
let Some(conn) = GeneralConnection::new(conn, state) else {
log_err!(
0,
PrintType::General,