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

Merge remote-tracking branch 'refs/remotes/origin/main'

Browse source
This commit is contained in:
Alex Emmet 2026-08-28 13:19:25 +02:00
commit dc20a0f261
No known key found for this signature in database
26 changed files with 979 additions and 262 deletions
Download patch file Download diff file Expand all files Collapse all files

48
src/rho/app_connection.rs
View file

@ -1,7 +1,7 @@
use crate::anonymous_clients::anonymous_manager;
use crate::app_state::AppState;
use crate::rho::connection::{
GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat,
GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::rho::relay_router::{self, RelaySource};
use crate::rho::rho_connection::RhoConnection;
@ -89,12 +89,38 @@ impl AppConnection {
/// Handle incoming message from app
pub async fn handle_message(self: Arc<Self>, cv: CommunicationValue) {
tokio::spawn(async move {
let message_id = match cv.require_id() {
Ok(message_id) => message_id,
Err(error) => {
log_err!(
self.user_id as i64,
PrintType::App,
"Rejected malformed message: {}",
error
);
let response =
CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id();
self.send_message(&response).await;
return;
}
};
log_cv_in!(PrintType::App, cv);
if cv.is_type(CommunicationType::Relay) {
let cv = relay_router::ensure_relay_frame_id(cv);
let request_id = cv.get_id();
let next_hop = cv.receiver().unwrap_or_default();
let next_hop = match cv.require_receiver() {
Ok(next_hop) => next_hop,
Err(error) => {
log_err!(
self.user_id as i64,
PrintType::App,
"Rejected malformed relay: {}",
error
);
self.send_error_response(message_id, CommunicationType::ErrorInvalidData)
.await;
return;
}
};
let result = match self.get_rho_connection().await {
Some(rho) => {
relay_router::route_relay(
@ -102,7 +128,7 @@ impl AppConnection {
RelaySource::Client {
iota_id: rho.get_iota_id().await,
},
cv,
relay_router::ensure_relay_frame_id(cv),
)
.await
}
@ -110,7 +136,7 @@ impl AppConnection {
};
let response = match result {
Ok(()) => {
CommunicationValue::new(CommunicationType::Success).with_id(request_id)
CommunicationValue::new(CommunicationType::Success).with_id(message_id)
}
Err(error) => {
log_err!(
@ -121,7 +147,7 @@ impl AppConnection {
error
);
CommunicationValue::new(relay_router::error_response_type(&error))
.with_id(request_id)
.with_id(message_id)
}
};
self.send_message(&response).await;
@ -130,7 +156,7 @@ impl AppConnection {
if cv.is_type(CommunicationType::Success) {
if let Some(rho) = self.get_rho_connection().await {
rho.forward_relay_ack(self.user_id, cv.get_id()).await;
rho.forward_relay_ack(self.user_id, message_id).await;
}
return;
}
@ -139,7 +165,7 @@ impl AppConnection {
relay_router::message_security_class(&cv),
relay_router::MessageSecurityClass::RelayOnly
) {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
self.send_error_response(message_id, CommunicationType::ErrorInvalidData)
.await;
return;
}
@ -155,7 +181,7 @@ impl AppConnection {
}
} {
let response = CommunicationValue::new(CommunicationType::GetUserData)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(
DataType::Username,
DataValue::Str(anonymous.get_user_name().await),
@ -197,7 +223,7 @@ impl AppConnection {
"Rejected unsupported communication type {}",
cv.get_type()
);
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
self.send_error_response(message_id, CommunicationType::ErrorInvalidData)
.await;
});
}

166
src/rho/client_connection.rs
View file

@ -3,7 +3,7 @@ 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, MtpValueCompat, OptionalDataValueCompat,
GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::rho::relay_router::{self, RelaySource};
use crate::rho::rho_connection::RhoConnection;
@ -117,14 +117,41 @@ impl ClientConnection {
};
tokio::spawn(async move {
let _permit = permit;
let message_id = match cv.require_id() {
Ok(message_id) => message_id,
Err(error) => {
log_err!(
self.user_id as i64,
PrintType::Client,
"Rejected malformed message: {}",
error
);
let response =
CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id();
self.send_message(&response).await;
return;
}
};
log_cv_in!(PrintType::Client, cv);
let mut cv = cv;
if cv.is_type(CommunicationType::Relay) {
let next_hop = match cv.require_receiver() {
Ok(next_hop) => next_hop,
Err(error) => {
log_err!(
self.user_id as i64,
PrintType::Client,
"Rejected malformed relay: {}",
error
);
self.send_error_response(message_id, CommunicationType::ErrorInvalidData)
.await;
return;
}
};
cv = relay_router::ensure_relay_frame_id(cv);
let request_id = cv.get_id();
let next_hop = cv.receiver().unwrap_or_default();
let result = match self.get_rho_connection().await {
Some(rho) => {
relay_router::route_relay(
@ -140,7 +167,7 @@ impl ClientConnection {
};
let response = match result {
Ok(()) => {
CommunicationValue::new(CommunicationType::Success).with_id(request_id)
CommunicationValue::new(CommunicationType::Success).with_id(message_id)
}
Err(error) => {
log_err!(
@ -151,7 +178,7 @@ impl ClientConnection {
error
);
CommunicationValue::new(relay_router::error_response_type(&error))
.with_id(request_id)
.with_id(message_id)
}
};
self.send_message(&response).await;
@ -160,7 +187,7 @@ impl ClientConnection {
if cv.is_type(CommunicationType::Success) {
if let Some(rho) = self.get_rho_connection().await {
rho.forward_relay_ack(self.user_id, cv.get_id()).await;
rho.forward_relay_ack(self.user_id, message_id).await;
}
return;
}
@ -169,7 +196,7 @@ impl ClientConnection {
relay_router::message_security_class(&cv),
relay_router::MessageSecurityClass::RelayOnly
) {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
self.send_error_response(message_id, CommunicationType::ErrorInvalidData)
.await;
return;
}
@ -178,11 +205,11 @@ impl ClientConnection {
// 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()
&& cv.get_data(DataType::UserState).is_some()
{
self.handle_set_user_state(
CommunicationValue::new(CommunicationType::ClientChanged)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(
DataType::UserState,
cv.get_data(DataType::UserState)
@ -240,7 +267,7 @@ impl ClientConnection {
}
} {
let response = CommunicationValue::new(CommunicationType::GetUserData)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(
DataType::Username,
DataValue::Str(anonymous.get_user_name().await),
@ -276,19 +303,19 @@ impl ClientConnection {
|| cv.is_type(CommunicationType::DeleteUser)
{
if cv.is_type(CommunicationType::ChangeUserData)
&& cv.get_data_opt(DataType::OnlineStatus).is_some()
&& cv.get_data(DataType::OnlineStatus).is_some()
{
let mut profile_request = cv.clone();
let preference = profile_request
.remove_data(DataType::OnlineStatus)
.unwrap_or(DataValue::Null);
let state_request = CommunicationValue::new(CommunicationType::ClientChanged)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::UserState, preference);
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;
self.send_error_response(message_id, error_type).await;
return;
}
};
@ -311,7 +338,7 @@ impl ClientConnection {
}
Ok(response) => self.send_message(&response).await,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInternal)
self.send_error_response(message_id, CommunicationType::ErrorInternal)
.await;
}
}
@ -336,7 +363,7 @@ impl ClientConnection {
if is_per_device_settings {
let Some(session_id) = session_id else {
let response = CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(self.user_id)
.add_typed_default(
DataType::Message,
@ -352,7 +379,7 @@ impl ClientConnection {
if session_id != expected_session_id {
let response = CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(self.user_id)
.add_typed_default(
DataType::Message,
@ -368,7 +395,7 @@ impl ClientConnection {
} else if let Some(session_id) = session_id {
if session_id != expected_session_id {
let response = CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(self.user_id)
.add_typed_default(
DataType::Message,
@ -396,7 +423,7 @@ impl ClientConnection {
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())
.with_id(message_id)
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(expected_session_id),
@ -417,12 +444,14 @@ impl ClientConnection {
"Rejected unsupported communication type {}",
cv.get_type()
);
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData)
self.send_error_response(message_id, CommunicationType::ErrorInvalidData)
.await;
});
}
async fn handle_omega_forward(self: Arc<Self>, cv: CommunicationValue) {
let request_id = cv.get_id();
let Ok(request_id) = cv.require_id() else {
return;
};
match self.await_omega_response(cv).await {
Ok(response_cv) => self.send_message(&response_cv).await,
Err(_) => {
@ -445,6 +474,9 @@ impl ClientConnection {
&self,
cv: CommunicationValue,
) -> Result<CommunicationValue, CommunicationType> {
let message_id = cv
.require_id()
.map_err(|_| CommunicationType::ErrorInvalidData)?;
if !self.state.omega.is_ready().await {
return Err(CommunicationType::ErrorInternal);
}
@ -459,7 +491,7 @@ impl ClientConnection {
return Err(CommunicationType::ErrorNoIota);
};
let request = CommunicationValue::new(CommunicationType::ClientChanged)
.with_id(cv.get_id())
.with_id(message_id)
.with_sender(self.user_id)
.add_typed_default(
DataType::UserId,
@ -472,22 +504,28 @@ impl ClientConnection {
.await
.map_err(|_| CommunicationType::ErrorInternal)?;
return Ok(CommunicationValue::new(CommunicationType::Success)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::UserState, DataValue::Str(state.to_string())));
}
async fn handle_set_user_state(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
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,
Err(error_type) => self.send_error_response(message_id, error_type).await,
}
}
/// Handle call invite
async fn handle_call_invite(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
let receiver_id: i128 = cv.get_data(DataType::ReceiverId).as_number().unwrap_or(0);
if receiver_id == 0 {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNoUserId)
self.send_error_response(message_id, CommunicationType::ErrorNoUserId)
.await;
return;
}
@ -496,13 +534,13 @@ impl ClientConnection {
Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) {
Ok(id) => id,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
}
},
_ => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNoCallId)
self.send_error_response(message_id, CommunicationType::ErrorNoCallId)
.await;
return;
}
@ -511,7 +549,7 @@ impl ClientConnection {
let secret = match call_invite_secret_from_cv(&cv) {
Some(secret) => secret,
None => {
self.send_error_response(cv.get_id(), CommunicationType::BadRequest)
self.send_error_response(message_id, CommunicationType::BadRequest)
.await;
return;
}
@ -522,7 +560,7 @@ impl ClientConnection {
.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)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
}
@ -532,7 +570,7 @@ impl ClientConnection {
.call_manager
.should_forward_invite(self.user_id, receiver_id as u64)
{
let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
let response = CommunicationValue::new(CommunicationType::Success).with_id(message_id);
self.send_message(&response).await;
return;
}
@ -564,7 +602,7 @@ impl ClientConnection {
});
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(
DataType::ReceiverId,
DataValue::SignedNumber(receiver_id.into()),
@ -594,25 +632,28 @@ impl ClientConnection {
target_rho.message_to_client(forward).await;
let response = CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id());
let response = CommunicationValue::new(CommunicationType::Success).with_id(message_id);
self.send_message(&response).await;
}
/// Handle get call request
async fn handle_get_call(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
let user_id = self.get_user_id().await;
let call_id = match cv.get_data(DataType::CallId) {
Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) {
Ok(id) => id,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
}
},
_ => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNoCallId)
self.send_error_response(message_id, CommunicationType::ErrorNoCallId)
.await;
return;
}
@ -626,7 +667,7 @@ impl ClientConnection {
{
Ok(token) => {
let response = CommunicationValue::new(CommunicationType::CallToken)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(user_id as u64)
.add_typed_default(DataType::CallToken, DataValue::Str(token));
self.send_message(&response).await;
@ -634,26 +675,29 @@ impl ClientConnection {
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())
.with_id(message_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) {
let Ok(message_id) = cv.require_id() else {
return;
};
let user_id = self.get_user_id().await;
let call_id = match cv.get_data(DataType::CallId) {
Some(DataValue::Str(id_str)) => match Uuid::parse_str(id_str) {
Ok(id) => id,
Err(_) => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
}
},
_ => {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNoCallId)
self.send_error_response(message_id, CommunicationType::ErrorNoCallId)
.await;
return;
}
@ -670,20 +714,20 @@ impl ClientConnection {
}
let response = CommunicationValue::new(CommunicationType::CallData)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(user_id as u64)
.add_typed_default(DataType::UserIds, DataValue::Array(user_ids));
self.send_message(&response).await;
} else {
let error_cv = CommunicationValue::new(CommunicationType::ErrorInvalidUserId)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()));
self.send_message(&error_cv).await;
return;
}
} else {
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string()));
self.send_message(&error_cv).await;
return;
@ -691,9 +735,12 @@ impl ClientConnection {
}
async fn handle_call_timeout_user(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or(""))
else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
};
@ -707,13 +754,13 @@ impl ClientConnection {
.unwrap_or(0);
let Some(call) = self.state.call_manager.get_call(call_id).await else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound)
self.send_error_response(message_id, CommunicationType::ErrorNotFound)
.await;
return;
};
let Some(caller) = call.get_caller(self.get_user_id().await).await else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidUserId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId)
.await;
return;
};
@ -729,9 +776,12 @@ impl ClientConnection {
}
}
async fn handle_call_disconnect_user(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or(""))
else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
};
@ -741,12 +791,12 @@ impl ClientConnection {
.unwrap_or(0);
let Some(call) = self.state.call_manager.get_call(call_id).await else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorNotFound)
self.send_error_response(message_id, CommunicationType::ErrorNotFound)
.await;
return;
};
let Some(caller) = call.get_caller(self.get_user_id().await).await else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidUserId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId)
.await;
return;
};
@ -755,9 +805,12 @@ impl ClientConnection {
}
}
async fn handle_call_set_anonymous_joining(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
let Ok(call_id) = Uuid::from_str(cv.get_data(DataType::CallId).as_str().unwrap_or(""))
else {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidCallId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidCallId)
.await;
return;
};
@ -780,7 +833,7 @@ impl ClientConnection {
short_link = call.get_short_link().await;
}
let mut response_cv = CommunicationValue::new(CommunicationType::CallSetAnonymousJoining)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::CallId, DataValue::Str(call_id.to_string()))
.add_typed_default(DataType::Enabled, DataValue::Bool(enable));
if let Some(short_link) = short_link {
@ -790,6 +843,9 @@ impl ClientConnection {
}
async fn handle_load_txt_record(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
if let Some(path) = cv.get_data(DataType::Path).as_str() {
let resolver = match TokioAsyncResolver::tokio_from_system_conf() {
Ok(r) => r,
@ -799,7 +855,7 @@ impl ClientConnection {
.cloned()
.unwrap_or(DataValue::Null);
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
return;
@ -818,7 +874,7 @@ impl ClientConnection {
Ok(text) => text,
Err(_) => {
self.send_error_response(
cv.get_id(),
message_id,
CommunicationType::ErrorInvalidData,
)
.await;
@ -827,8 +883,8 @@ impl ClientConnection {
};
let response = CommunicationValue::new(CommunicationType::LoadTxtRecord)
.with_id(cv.get_id())
.add_typed_default(DataType::Content, DataValue::Str(record_text));
.with_id(message_id)
.add_typed_default(DataType::AppContent, DataValue::Str(record_text));
self.send_message(&response).await;
return;
}
@ -838,7 +894,7 @@ impl ClientConnection {
.cloned()
.unwrap_or(DataValue::Null);
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
}
@ -848,7 +904,7 @@ impl ClientConnection {
.cloned()
.unwrap_or(DataValue::Null);
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
}
@ -862,7 +918,7 @@ impl ClientConnection {
.cloned()
.unwrap_or(DataValue::Null);
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
}

84
src/rho/connection.rs
View file

@ -16,37 +16,38 @@ use crate::{
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::host::AuthState;
use mtp::webserver::{WebMTPConnection, WebMtpReceiver, WebMtpSender};
use thiserror::Error;
pub type MtpSender = WebMtpSender;
pub type MtpReceiver = WebMtpReceiver;
/*
* MTP 0.3 exposes absent frame fields and data entries as Options. These
* adapters keep legacy control handlers explicit while Relay code uses the
* native optional accessors directly.
*/
pub(crate) trait MtpValueCompat {
fn get_id(&self) -> u32;
fn get_sender(&self) -> u64;
fn get_receiver(&self) -> u64;
fn get_data_opt(&self, data_type: DataType) -> Option<&DataValue>;
#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)]
pub enum FrameValidationError {
#[error("message is missing an MTP id")]
MissingId,
#[error("message is missing an MTP sender")]
MissingSender,
#[error("message is missing an MTP receiver")]
MissingReceiver,
}
impl MtpValueCompat for CommunicationValue {
fn get_id(&self) -> u32 {
self.id().unwrap_or_default()
pub trait RequiredMtpFields {
fn require_id(&self) -> Result<u32, FrameValidationError>;
fn require_sender(&self) -> Result<u64, FrameValidationError>;
fn require_receiver(&self) -> Result<u64, FrameValidationError>;
}
impl RequiredMtpFields for CommunicationValue {
fn require_id(&self) -> Result<u32, FrameValidationError> {
self.id().ok_or(FrameValidationError::MissingId)
}
fn get_sender(&self) -> u64 {
self.sender().unwrap_or_default()
fn require_sender(&self) -> Result<u64, FrameValidationError> {
self.sender().ok_or(FrameValidationError::MissingSender)
}
fn get_receiver(&self) -> u64 {
self.receiver().unwrap_or_default()
}
fn get_data_opt(&self, data_type: DataType) -> Option<&DataValue> {
self.get_data(data_type)
fn require_receiver(&self) -> Result<u64, FrameValidationError> {
self.receiver().ok_or(FrameValidationError::MissingReceiver)
}
}
@ -55,6 +56,7 @@ pub(crate) trait OptionalDataValueCompat {
fn as_number(&self) -> Option<i128>;
fn as_signed_number(&self) -> Option<i128>;
fn as_str(&self) -> Option<&str>;
#[allow(dead_code)]
fn as_bytes(&self) -> Option<Vec<u8>>;
}
@ -80,6 +82,42 @@ impl OptionalDataValueCompat for Option<&DataValue> {
}
}
#[cfg(test)]
mod tests {
use super::{FrameValidationError, RequiredMtpFields};
use mtp::codec::{CommunicationType, CommunicationValue};
#[test]
fn required_fields_preserve_missing_field_errors() {
let frame = CommunicationValue::new(CommunicationType::Success)
.without_id()
.without_sender()
.without_receiver();
assert_eq!(frame.require_id(), Err(FrameValidationError::MissingId));
assert_eq!(
frame.require_sender(),
Err(FrameValidationError::MissingSender)
);
assert_eq!(
frame.require_receiver(),
Err(FrameValidationError::MissingReceiver)
);
}
#[test]
fn zero_is_a_present_routing_value() {
let frame = CommunicationValue::new(CommunicationType::Success)
.with_id(0)
.with_sender(0)
.with_receiver(0);
assert_eq!(frame.require_id(), Ok(0));
assert_eq!(frame.require_sender(), Ok(0));
assert_eq!(frame.require_receiver(), Ok(0));
}
}
/*
* How a connection identified itself during the mtp handshake driven by
* `server.rs` ("iota" / "client" authenticated logins, "anonymous"
@ -149,6 +187,10 @@ impl GeneralConnection {
}))
}
pub fn connection_kind(&self) -> ConnectionKind {
self.connection_kind
}
pub async fn handle(self: Arc<Self>) {
log_in!(0, PrintType::General, "General connection handler started");
if self.migrate().await {

112
src/rho/iota_connection.rs
View file

@ -6,7 +6,7 @@ use crate::log_err;
use crate::log_in;
use crate::log_out;
use crate::rho::connection::{
GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat,
GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::rho::relay_router::{self, RelaySource};
use crate::util::data_type_id;
@ -33,8 +33,9 @@ fn contact_snapshot(value: &CommunicationValue) -> Option<(i64, i64, Vec<i64>)>
return None;
}
let user_id = i64::try_from(value.get_receiver())
.ok()
let user_id = value
.receiver()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)?;
let session_id = value
.get_data(DataType::SessionId)
@ -222,6 +223,7 @@ impl IotaConnection {
}
}
#[allow(dead_code)]
pub async fn send_relay(&self, cv: &CommunicationValue) -> Result<(), String> {
self.sender
.send(cv)
@ -235,20 +237,46 @@ impl IotaConnection {
return;
};
let _permit = permit;
let message_id = match cv.require_id() {
Ok(message_id) => message_id,
Err(error) => {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Rejected malformed message: {}",
error
);
let response =
CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id();
self.send_message(&response).await;
return;
}
};
if cv.is_type(CommunicationType::Relay) {
let cv = relay_router::ensure_relay_frame_id(cv);
let request_id = cv.get_id();
let next_hop = cv.receiver().unwrap_or_default();
let next_hop = match cv.require_receiver() {
Ok(next_hop) => next_hop,
Err(error) => {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Rejected malformed relay: {}",
error
);
self.send_error_response(message_id, CommunicationType::ErrorInvalidData, None)
.await;
return;
}
};
let response = match relay_router::route_relay(
&self.state,
RelaySource::Iota {
iota_id: self.iota_id,
},
cv,
relay_router::ensure_relay_frame_id(cv),
)
.await
{
Ok(()) => CommunicationValue::new(CommunicationType::Success).with_id(request_id),
Ok(()) => CommunicationValue::new(CommunicationType::Success).with_id(message_id),
Err(error) => {
log_err!(
self.iota_id as i64,
@ -258,7 +286,7 @@ impl IotaConnection {
error
);
CommunicationValue::new(relay_router::error_response_type(&error))
.with_id(request_id)
.with_id(message_id)
}
};
self.send_message(&response).await;
@ -269,12 +297,12 @@ impl IotaConnection {
crate::rho::relay_router::message_security_class(&cv),
crate::rho::relay_router::MessageSecurityClass::RelayOnly
) {
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData, None)
self.send_error_response(message_id, CommunicationType::ErrorInvalidData, None)
.await;
return;
}
let msg_id = cv.get_id();
let msg_id = message_id;
if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) {
if (task)(self.clone(), cv.clone()) {
return;
@ -309,7 +337,7 @@ impl IotaConnection {
if cv.is_type(CommunicationType::StateSubscribe) {
self.send_error_response(
cv.get_id(),
message_id,
CommunicationType::ErrorInvalidData,
Some("StateSubscribe must come from an authoritative contact snapshot"),
)
@ -335,7 +363,7 @@ impl IotaConnection {
self.iota_id as i64,
PrintType::Omega,
"Forwarding CompleteRegisterUser to Omega (request_id={})",
request.get_id()
message_id
);
let mut response_cv = self
.state
@ -348,7 +376,7 @@ impl IotaConnection {
self.iota_id as i64,
PrintType::Omega,
"CompleteRegisterUser request_id={} failed: {}; retrying once",
request.get_id(),
message_id,
error
);
response_cv = self
@ -363,9 +391,9 @@ impl IotaConnection {
log_in!(
self.iota_id as i64,
PrintType::Omega,
"Omega completed registration (request_id={}, response_id={}, type={})",
request.get_id(),
response_cv.get_id(),
"Omega completed registration (request_id={}, response_id={:?}, type={})",
message_id,
response_cv.id(),
response_cv
.get_comm_type_enum()
.map(|kind| kind.to_string())
@ -408,7 +436,7 @@ impl IotaConnection {
self.add_user_id(user_id as u64).await;
self.send_message(
&CommunicationValue::new(CommunicationType::Success)
.with_id(cv.get_id()),
.with_id(message_id),
)
.await;
return;
@ -416,15 +444,15 @@ impl IotaConnection {
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()
"Registration verification returned an unexpected user (request_id={:?}, response_id={:?})",
verification.id(),
verified.id()
),
Err(verify_error) => log_err!(
self.iota_id as i64,
PrintType::Omega,
"Registration verification failed after request_id={}: {}",
verification.get_id(),
"Registration verification failed after request_id={:?}: {}",
verification.id(),
verify_error
),
}
@ -436,7 +464,7 @@ impl IotaConnection {
error
);
self.send_error_response(
cv.get_id(),
message_id,
CommunicationType::ErrorInternal,
Some(&format!("Omega forwarding failed: {error}")),
)
@ -477,7 +505,7 @@ impl IotaConnection {
"Rejected unsupported communication type {}",
cv.get_type()
);
self.send_error_response(cv.get_id(), CommunicationType::ErrorInvalidData, None)
self.send_error_response(message_id, CommunicationType::ErrorInvalidData, None)
.await;
}
@ -514,6 +542,9 @@ impl IotaConnection {
}
async fn handle_omega_forward_without_sender(self: Arc<Self>, cv: CommunicationValue) {
let Ok(message_id) = cv.require_id() else {
return;
};
let iota_for_closure = self.clone();
let request = cv.clone().add_typed_default(
DataType::IotaId,
@ -530,7 +561,7 @@ impl IotaConnection {
self.iota_id as i64,
PrintType::Omega,
"GetRegister request_id={} failed: {}; retrying once",
request.get_id(),
message_id,
error
);
response_cv = self
@ -550,7 +581,7 @@ impl IotaConnection {
error
);
self.send_error_response(
cv.get_id(),
message_id,
CommunicationType::ErrorInternal,
Some(&format!("Omega forwarding failed: {error}")),
)
@ -560,7 +591,17 @@ impl IotaConnection {
}
/// Handle GET_CHATS message
async fn handle_get_chats(&self, cv: CommunicationValue) {
let user_id = cv.get_sender();
let Ok(message_id) = cv.require_id() else {
return;
};
let Ok(user_id) = cv.require_sender() else {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Rejected get_chats without an MTP sender"
);
return;
};
// Authority check: user must be linked to this Iota
if !self.get_user_ids().await.contains(&user_id) {
@ -582,7 +623,7 @@ impl IotaConnection {
else {
self.forward_to_client(
CommunicationValue::new(CommunicationType::ErrorInvalidData)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(user_id),
)
.await;
@ -763,7 +804,14 @@ impl IotaConnection {
async fn add_call_state(&self, response: CommunicationValue) -> CommunicationValue {
let mut output = response.clone();
let user_id = response.get_receiver();
let Some(user_id) = response.receiver() else {
log_err!(
self.iota_id as i64,
PrintType::Iota,
"Discarded response without an MTP receiver"
);
return output;
};
let typed_data: Vec<_> = response.iter_typed_data().collect();
for (key, value) in typed_data {
@ -807,7 +855,9 @@ impl IotaConnection {
timeout_duration: Option<Duration>,
) -> Result<CommunicationValue, String> {
let (tx, mut rx) = mpsc::channel(1);
let msg_id = cv.get_id();
let msg_id = cv
.require_id()
.map_err(|error| format!("request is missing correlation id: {error}"))?;
let task_tx = tx.clone();
self.waiting_tasks.insert(

3
src/rho/relay_router.rs
View file

@ -26,6 +26,7 @@ pub enum RouteTarget {
}
impl RouteTarget {
#[allow(dead_code)]
pub fn wire_id(self) -> Option<u64> {
let (kind, id) = match self {
Self::User(id) => (USER_TARGET_KIND, id),
@ -251,7 +252,7 @@ fn route_response(response: CommunicationValue) -> Result<(), RelayRouteError> {
}
pub fn ensure_relay_frame_id(frame: CommunicationValue) -> CommunicationValue {
if frame.id().is_some_and(|id| id != 0) {
if frame.id().is_some() {
return frame;
}
let id = NEXT_RELAY_FRAME_ID.fetch_add(1, Ordering::Relaxed).max(1);

16
src/rho/rho_connection.rs
View file

@ -3,10 +3,7 @@ use super::{client_connection::ClientConnection, iota_connection::IotaConnection
use super::relay_router::RouteTarget;
use crate::{
log_err,
rho::{
app_connection::AppConnection,
connection::{MtpValueCompat, OptionalDataValueCompat},
},
rho::{app_connection::AppConnection, connection::OptionalDataValueCompat},
};
use dashmap::DashMap;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
@ -257,7 +254,14 @@ impl RhoConnection {
/// Send message from Iota to specific client
pub async fn message_to_client(&self, cv: CommunicationValue) {
let connections = self.get_client_connections().await;
let receiver_id = cv.get_receiver();
let Some(receiver_id) = cv.receiver() else {
log_err!(
0,
crate::util::logger::PrintType::General,
"Discarded message without an MTP receiver"
);
return;
};
let session_id = cv.get_data(DataType::SessionId).as_number();
for connection in connections.iter() {
@ -273,6 +277,7 @@ impl RhoConnection {
}
}
#[allow(dead_code)]
pub async fn message_to_iota(&self, cv: CommunicationValue) {
self.iota_connection.send_message(&cv).await;
}
@ -308,6 +313,7 @@ impl RhoConnection {
send_error.map_or(Ok(()), Err)
}
#[allow(dead_code)]
pub async fn send_relay_to_iota(&self, cv: &CommunicationValue) -> Result<(), String> {
self.iota_connection.send_relay(cv).await
}

1
src/rho/rho_manager.rs
View file

@ -24,6 +24,7 @@ impl RhoManager {
self.users.get(&user_id).map(|entry| entry.value().clone())
}
#[allow(dead_code)]
pub async fn contains_iota(&self, iota_id: i64) -> bool {
self.connections.contains_key(&iota_id)
}

127
src/rho/server.rs
View file

@ -8,7 +8,7 @@ use crate::{
app_state::AppState,
log, log_err,
omega::omega_connection::OmegaConnection,
rho::connection::{GeneralConnection, OptionalDataValueCompat},
rho::connection::{ConnectionKind, GeneralConnection, OptionalDataValueCompat},
util::{file_util::load_file_vec, logger::PrintType},
};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
@ -16,6 +16,27 @@ use mtp::crypto::PublicKeyBundle;
use mtp::host::{AuthenticationPolicy, HostConfig, Policy, SendMode};
use mtp::webserver::{MTPWebServer, WebServerConfig};
fn web_config(max_connections: usize) -> Result<WebServerConfig, mtp::webserver::RouterError> {
WebServerConfig::new()
.max_connections(max_connections)
.route("/", |_request, response| async move { response.body("OK") })
}
fn rho_policy() -> Policy {
Policy::default()
.with_send_mode(SendMode::SingleStreamPerMessage)
.with_timeouts(
Duration::from_millis(2_000),
Duration::from_millis(2_000),
Duration::from_millis(30_000),
)
.with_keep_alive(Some(Duration::from_secs(6)))
.with_max_idle_timeout(Some(Duration::from_secs(30)))
.with_receiver_queue_capacity(1000)
.with_max_concurrent_stream_tasks(10)
.with_persistent_stream_retries(5, Duration::from_secs(5))
}
/*
* Resolves the PublicKeyBundle mtp needs to verify a login's signed
* challenge response. "iota"/"client" ids are looked up through Omega, the
@ -126,21 +147,7 @@ pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error
cert_pem,
key_pem,
)
.with_policy(
Policy::default()
.with_send_mode(SendMode::SingleStreamPerMessage)
.with_max_message_size(1_000_000_000)
.with_timeouts(
Duration::from_millis(2_000),
Duration::from_millis(2_000),
Duration::from_millis(30_000),
)
.with_keep_alive(Some(Duration::from_secs(6)))
.with_max_idle_timeout(None)
.with_receiver_queue_capacity(1000)
.with_max_concurrent_stream_tasks(10)
.with_persistent_stream_retries(5, Duration::from_secs(5)),
)
.with_policy(rho_policy())
.with_authentication(
state
.keyring_for_host()
@ -160,8 +167,7 @@ pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error
)
.with_authentication_policy(AuthenticationPolicy::AllowAuthentication);
let web_config = WebServerConfig::new()
.route("/", |_request, response| async move { response.body("OK") })?;
let web_config = web_config(state.config.rho_max_connections)?;
let mut host = MTPWebServer::new(host_config, web_config).await?;
log!(
0,
@ -189,9 +195,21 @@ pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error
}
};
let global_permit = match state.rho_connection_limits.all.clone().try_acquire_owned() {
Ok(permit) => permit,
Err(_) => {
log_err!(
0,
PrintType::General,
"Rejected connection: global Rho connection limit reached"
);
continue;
}
};
let peer_ip = conn.remote_addr.map(|address| address.ip());
let state = state.clone();
tokio::spawn(async move {
let Some(conn) = GeneralConnection::new(conn, state) else {
let Some(conn) = GeneralConnection::new(conn, state.clone()) else {
log_err!(
0,
PrintType::General,
@ -199,6 +217,58 @@ pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error
);
return;
};
let anonymous_permit = if conn.connection_kind() == ConnectionKind::AnonymousClient {
match state
.rho_connection_limits
.anonymous
.clone()
.try_acquire_owned()
{
Ok(permit) => Some(permit),
Err(_) => {
log_err!(
0,
PrintType::General,
"Rejected anonymous connection: anonymous limit reached"
);
return;
}
}
} else {
None
};
let anonymous_ip_permit = if conn.connection_kind() == ConnectionKind::AnonymousClient {
let Some(peer_ip) = peer_ip else {
log_err!(
0,
PrintType::General,
"Rejected anonymous connection: peer address unavailable"
);
return;
};
match state
.rho_connection_limits
.try_acquire_anonymous_per_ip(peer_ip)
{
Some(permit) => Some(permit),
None => {
log_err!(
0,
PrintType::General,
"Rejected anonymous connection: per-IP limit reached"
);
return;
}
}
} else {
None
};
let _global_permit = global_permit;
let _anonymous_permit = anonymous_permit;
let _anonymous_ip_permit = anonymous_ip_permit;
conn.handle().await;
});
}
@ -206,3 +276,22 @@ pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error
Ok(())
}
#[cfg(test)]
mod tests {
use super::{rho_policy, web_config};
use std::time::Duration;
#[test]
fn rho_connection_budget_configures_mtp_admission() {
let config = web_config(7).expect("health route is valid");
assert_eq!(config.max_connections, 7);
}
#[test]
fn rho_policy_keeps_idle_peers_alive_and_detects_dead_peers() {
let policy = rho_policy();
assert_eq!(policy.keep_alive_interval, Some(Duration::from_secs(6)));
assert_eq!(policy.max_idle_timeout, Some(Duration::from_secs(30)));
}
}