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

86
src/anonymous_clients/anonymous_client_connection.rs
View file

@ -9,7 +9,7 @@ use crate::anonymous_clients::anonymous_manager::{self, generate_username};
use crate::app_state::AppState;
use crate::calls::call_group::call_invite_secret_from_cv;
use crate::rho::connection::{
GeneralConnection, MtpReceiver, MtpSender, MtpValueCompat, OptionalDataValueCompat,
GeneralConnection, MtpReceiver, MtpSender, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::util::data_type_id;
use crate::util::logger::PrintType;
@ -109,10 +109,25 @@ impl AnonymousClientConnection {
};
tokio::spawn(async move {
let _permit = permit;
let message_id = match cv.require_id() {
Ok(message_id) => message_id,
Err(error) => {
let response =
CommunicationValue::new(CommunicationType::ErrorInvalidData).without_id();
log_out!(
self.user_id as i64,
PrintType::Client,
"Rejected malformed message: {}",
error
);
self.send_message(&response).await;
return;
}
};
log_cv_in!(PrintType::Client, &cv);
if cv.is_type(CommunicationType::Relay) {
self.send_error_response(&cv.get_id(), CommunicationType::ErrorNotAuthenticated)
self.send_error_response(message_id, CommunicationType::ErrorNotAuthenticated)
.await;
return;
}
@ -126,18 +141,15 @@ impl AnonymousClientConnection {
call
} else {
self.send_error_response(
&cv.get_id(),
message_id,
CommunicationType::ErrorNotAuthenticated,
)
.await;
return;
}
} else {
self.send_error_response(
&cv.get_id(),
CommunicationType::ErrorNotAuthenticated,
)
.await;
self.send_error_response(message_id, CommunicationType::ErrorNotAuthenticated)
.await;
return;
};
@ -210,7 +222,7 @@ impl AnonymousClientConnection {
self.clone()
.send_message(
&&CommunicationValue::new(CommunicationType::IdentificationResponse)
.with_id(cv.get_id())
.with_id(message_id)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(self.user_id.into()),
@ -245,7 +257,7 @@ impl AnonymousClientConnection {
// Presence is account-scoped and anonymous sessions have no
// persisted account preference to change.
if cv.is_type(CommunicationType::ClientChanged) {
self.send_error_response(&cv.get_id(), CommunicationType::ErrorNoUserId)
self.send_error_response(message_id, CommunicationType::ErrorNoUserId)
.await;
return;
}
@ -291,7 +303,7 @@ impl AnonymousClientConnection {
}
} {
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),
@ -344,9 +356,12 @@ impl AnonymousClientConnection {
/// 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: i64 = cv.get_data(DataType::ReceiverId).as_number().unwrap_or(0) as i64;
if receiver_id == 0 {
self.send_error_response(&cv.get_id(), CommunicationType::ErrorNoUserId)
self.send_error_response(message_id, CommunicationType::ErrorNoUserId)
.await;
return;
}
@ -355,13 +370,13 @@ impl AnonymousClientConnection {
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;
}
@ -370,7 +385,7 @@ impl AnonymousClientConnection {
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;
}
@ -381,7 +396,7 @@ impl AnonymousClientConnection {
.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;
}
@ -391,7 +406,7 @@ impl AnonymousClientConnection {
.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;
}
@ -423,7 +438,7 @@ impl AnonymousClientConnection {
});
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()),
@ -453,25 +468,28 @@ impl AnonymousClientConnection {
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();
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;
}
@ -485,7 +503,7 @@ impl AnonymousClientConnection {
{
Ok(token) => {
let response = CommunicationValue::new(CommunicationType::CallToken)
.with_id(cv.get_id())
.with_id(message_id)
.with_receiver(user_id)
.add_typed_default(DataType::CallToken, DataValue::Str(token));
self.send_message(&response).await;
@ -497,16 +515,19 @@ impl AnonymousClientConnection {
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_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;
};
@ -520,12 +541,12 @@ impl AnonymousClientConnection {
.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 else {
self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidUserId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId)
.await;
return;
};
@ -536,9 +557,12 @@ impl AnonymousClientConnection {
}
}
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;
};
@ -548,12 +572,12 @@ impl AnonymousClientConnection {
.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 else {
self.send_error_response(&cv.get_id(), CommunicationType::ErrorInvalidUserId)
self.send_error_response(message_id, CommunicationType::ErrorInvalidUserId)
.await;
return;
};
@ -563,8 +587,8 @@ impl AnonymousClientConnection {
}
/// Send error response
async fn send_error_response(self: Arc<Self>, message_id: &u32, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(*message_id);
async fn send_error_response(self: Arc<Self>, message_id: u32, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(message_id);
self.send_message(&error).await;
}
/// Close the connection

63
src/app_state.rs
View file

@ -1,6 +1,8 @@
use std::sync::Arc;
use std::{net::IpAddr, sync::Arc};
use dashmap::DashMap;
use mtp::crypto::Keyring;
use tokio::sync::Semaphore;
use crate::{
calls::{call_manager::CallManager, call_util::LiveKitService},
@ -10,6 +12,58 @@ use crate::{
services::call_state::CallStateAggregator,
};
pub struct RhoConnectionLimits {
pub all: Arc<Semaphore>,
pub anonymous: Arc<Semaphore>,
anonymous_by_ip: Arc<DashMap<IpAddr, usize>>,
max_anonymous_per_ip: usize,
}
impl RhoConnectionLimits {
pub fn new(max_all: usize, max_anonymous: usize, max_anonymous_per_ip: usize) -> Self {
Self {
all: Arc::new(Semaphore::new(max_all)),
anonymous: Arc::new(Semaphore::new(max_anonymous)),
anonymous_by_ip: Arc::new(DashMap::new()),
max_anonymous_per_ip,
}
}
pub fn try_acquire_anonymous_per_ip(&self, ip: IpAddr) -> Option<RhoIpConnectionPermit> {
let mut count = self.anonymous_by_ip.entry(ip).or_insert(0);
if *count >= self.max_anonymous_per_ip {
return None;
}
*count += 1;
drop(count);
Some(RhoIpConnectionPermit {
ip,
active_by_ip: self.anonymous_by_ip.clone(),
})
}
#[cfg(test)]
fn anonymous_count_for_ip(&self, ip: IpAddr) -> usize {
self.anonymous_by_ip.get(&ip).map_or(0, |count| *count)
}
}
pub struct RhoIpConnectionPermit {
ip: IpAddr,
active_by_ip: Arc<DashMap<IpAddr, usize>>,
}
impl Drop for RhoIpConnectionPermit {
fn drop(&mut self) {
if let Some(mut count) = self.active_by_ip.get_mut(&self.ip) {
*count = count.saturating_sub(1);
}
self.active_by_ip
.remove_if(&self.ip, |_, count| *count == 0);
}
}
/*
* Holds startup dependencies so listener and connection constructors can take
* one explicit handle while the remaining manager migrations are completed.
@ -22,6 +76,7 @@ pub struct AppState {
pub call_manager: Arc<CallManager>,
pub call_state_aggregator: Arc<CallStateAggregator>,
pub livekit: Arc<LiveKitService>,
pub rho_connection_limits: Arc<RhoConnectionLimits>,
}
impl AppState {
@ -33,6 +88,11 @@ impl AppState {
call_manager: Arc<CallManager>,
livekit: Arc<LiveKitService>,
) -> Arc<Self> {
let rho_connection_limits = Arc::new(RhoConnectionLimits::new(
config.rho_max_connections,
config.rho_max_anonymous_connections,
config.rho_max_anonymous_connections_per_ip,
));
Arc::new(Self {
config,
keyring,
@ -41,6 +101,7 @@ impl AppState {
call_state_aggregator: Arc::new(CallStateAggregator::new(call_manager.clone())),
call_manager,
livekit,
rho_connection_limits,
})
}

23
src/bin/migrate_raw_keyring.rs Normal file
View file

@ -0,0 +1,23 @@
#[path = "../identity.rs"]
mod identity;
use identity::{
KEYRING_PATH, PUBLIC_KEY_PATH, identity_secret_from_environment, migrate_raw_keyring,
};
fn main() {
let secret = match identity_secret_from_environment() {
Ok(secret) => secret,
Err(error) => {
eprintln!("Unable to load Omikron identity secret: {error}");
std::process::exit(1);
}
};
if let Err(error) = migrate_raw_keyring(secret.as_slice(), KEYRING_PATH, PUBLIC_KEY_PATH) {
eprintln!("Unable to migrate Omikron identity: {error}");
std::process::exit(1);
}
println!("Migrated {KEYRING_PATH} to protected keyring storage");
}

2
src/calls/call_group.rs
View file

@ -15,6 +15,7 @@ use crate::{
pub struct CallGroup {
pub call_id: Uuid,
pub members: RwLock<Vec<Arc<Caller>>>,
#[allow(dead_code)]
pub show: RwLock<bool>,
pub anonymous_joining: RwLock<bool>,
pub short_link: RwLock<Option<String>>,
@ -100,6 +101,7 @@ pub fn call_invite_secret_from_cv(cv: &CommunicationValue) -> Option<CallSecretE
}
impl CallGroup {
#[allow(dead_code)]
pub fn new(call_id: Uuid, user: Arc<Caller>) -> Self {
Self::new_with_service(call_id, user, Arc::new(LiveKitService::new(None)))
}

2
src/calls/call_util.rs
View file

@ -134,6 +134,7 @@ impl LiveKitService {
Ok(())
}
#[allow(dead_code)]
pub fn garbage_collect_calls(self: Arc<Self>, manager: Arc<CallManager>) {
tokio::spawn(async move {
loop {
@ -148,6 +149,7 @@ impl LiveKitService {
}
}
#[allow(dead_code)]
pub async fn clean_calls(manager: &CallManager, room_service: RoomClient) {
let rooms =
match tokio::time::timeout(LIVEKIT_REQUEST_TIMEOUT, room_service.list_rooms(Vec::new()))

1
src/calls/caller.rs
View file

@ -39,6 +39,7 @@ impl Caller {
pub async fn set_timeout(&self, timeout: i64) {
*self.timeout.write().await = timeout;
}
#[allow(dead_code)]
pub fn create_token(&self, livekit: &LiveKitService) -> Result<String, CallError> {
livekit.create_token(self.user_id, self.call_id, self.has_admin())
}

22
src/config.rs
View file

@ -8,6 +8,9 @@ const DEFAULT_OMEGA_HOST: &str = "tensamin.net";
const DEFAULT_OMEGA_PORT: u16 = 9187;
const DEFAULT_OMEGA_SYNC_TIMEOUT_SECONDS: u64 = 20;
const DEFAULT_OMEGA_SYNC_RETRIES: u32 = 3;
const DEFAULT_RHO_MAX_CONNECTIONS: usize = 256;
const DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS: usize = 128;
const DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP: usize = 16;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct LiveKitConfig {
@ -28,6 +31,12 @@ pub struct Config {
/// Number of synchronization requests before the transport is closed and
/// the normal reconnect loop starts.
pub omega_sync_retries: u32,
/// Maximum number of application sessions accepted by the Rho listener.
pub rho_max_connections: usize,
/// Maximum number of anonymous application sessions.
pub rho_max_anonymous_connections: usize,
/// Maximum number of anonymous sessions from one peer IP address.
pub rho_max_anonymous_connections_per_ip: usize,
pub livekit: Option<LiveKitConfig>,
}
@ -57,6 +66,16 @@ impl Config {
)?);
let omega_sync_retries =
parse_or_default("OMEGA_SYNC_RETRIES", DEFAULT_OMEGA_SYNC_RETRIES)?.max(1);
let rho_max_connections =
parse_positive_or_default("RHO_MAX_CONNECTIONS", DEFAULT_RHO_MAX_CONNECTIONS)?;
let rho_max_anonymous_connections = parse_positive_or_default(
"RHO_MAX_ANONYMOUS_CONNECTIONS",
DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS,
)?;
let rho_max_anonymous_connections_per_ip = parse_positive_or_default(
"RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP",
DEFAULT_RHO_MAX_ANONYMOUS_CONNECTIONS_PER_IP,
)?;
let omega_host = env::var("OMEGA_HOST")
.unwrap_or_else(|_| DEFAULT_OMEGA_HOST.to_string())
.trim()
@ -77,6 +96,9 @@ impl Config {
omikron_id,
omega_sync_timeout,
omega_sync_retries,
rho_max_connections,
rho_max_anonymous_connections,
rho_max_anonymous_connections_per_ip,
livekit: livekit_from_environment()?,
})
}

181
src/identity.rs Normal file
View file

@ -0,0 +1,181 @@
use std::{env, io::ErrorKind, path::Path};
use mtp::crypto::Keyring;
use mtp::files::{FileError, load_keyring, save_keyring, save_public_key_bundle};
use zeroize::Zeroizing;
pub const KEYRING_PATH: &str = "./omikron.mk";
pub const PUBLIC_KEY_PATH: &str = "./omikron.mpkb";
const IDENTITY_SECRET_ENV: &str = "OMIKRON_IDENTITY_SECRET";
pub fn identity_secret_from_environment() -> Result<Zeroizing<Vec<u8>>, String> {
let secret = env::var(IDENTITY_SECRET_ENV).map_err(|_| {
format!("{IDENTITY_SECRET_ENV} must be set before Omikron networking starts")
})?;
if secret.trim().is_empty() {
return Err(format!("{IDENTITY_SECRET_ENV} must not be empty"));
}
Ok(Zeroizing::new(secret.into_bytes()))
}
#[allow(dead_code)]
pub fn load_or_create_keyring(
passphrase: &[u8],
keyring_path: impl AsRef<Path>,
public_key_path: impl AsRef<Path>,
) -> Result<Keyring, String> {
let keyring_path = keyring_path.as_ref();
let public_key_path = public_key_path.as_ref();
match load_keyring(keyring_path, passphrase) {
Ok(keyring) => Ok(keyring),
Err(FileError::Io(error)) if error.kind() == ErrorKind::NotFound => {
let keyring = Keyring::generate();
save_keyring(&keyring, keyring_path, passphrase).map_err(|error| {
format!("unable to persist Omikron keyring at {keyring_path:?}: {error}")
})?;
save_public_key_bundle(&keyring.public_key_bundle(), public_key_path).map_err(
|error| {
format!(
"unable to persist Omikron public key bundle at {public_key_path:?}: {error}"
)
},
)?;
eprintln!(
"Generated new protected keyring at {}",
keyring_path.display()
);
Ok(keyring)
}
Err(error) => Err(format!(
"unable to load existing Omikron identity from {}: {error}",
keyring_path.display()
)),
}
}
#[cfg(feature = "raw-migration")]
#[allow(dead_code)]
pub fn migrate_raw_keyring(
passphrase: &[u8],
keyring_path: impl AsRef<Path>,
public_key_path: impl AsRef<Path>,
) -> Result<Keyring, String> {
use mtp::files::load_keyring_raw;
let keyring = load_keyring_raw(keyring_path.as_ref())
.map_err(|error| format!("raw keyring migration failed: {error}"))?;
save_keyring(&keyring, keyring_path.as_ref(), passphrase)
.map_err(|error| format!("protected keyring write failed: {error}"))?;
save_public_key_bundle(&keyring.public_key_bundle(), public_key_path.as_ref())
.map_err(|error| format!("public key bundle write failed: {error}"))?;
Ok(keyring)
}
#[cfg(test)]
mod tests {
use super::*;
use std::{
fs,
time::{SystemTime, UNIX_EPOCH},
};
struct TestPaths {
directory: std::path::PathBuf,
keyring: std::path::PathBuf,
public_key: std::path::PathBuf,
}
impl TestPaths {
fn new() -> Self {
let suffix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let directory = env::temp_dir().join(format!("omikron-identity-{suffix}"));
fs::create_dir_all(&directory).unwrap();
Self {
keyring: directory.join("omikron.mk"),
public_key: directory.join("omikron.mpkb"),
directory,
}
}
}
impl Drop for TestPaths {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.directory);
}
}
#[test]
fn missing_identity_is_created_and_can_be_loaded_again() {
let paths = TestPaths::new();
let passphrase = b"test identity secret";
let first = load_or_create_keyring(passphrase, &paths.keyring, &paths.public_key).unwrap();
let first_public = first.public_key_bundle().try_as_bytes().unwrap();
let second = load_or_create_keyring(passphrase, &paths.keyring, &paths.public_key).unwrap();
assert_eq!(
second.public_key_bundle().try_as_bytes().unwrap(),
first_public
);
}
#[test]
fn corrupt_identity_fails_closed() {
let paths = TestPaths::new();
fs::write(&paths.keyring, b"not a keyring").unwrap();
let error =
load_or_create_keyring(b"test identity secret", &paths.keyring, &paths.public_key)
.unwrap_err();
assert!(error.contains("unable to load existing Omikron identity"));
}
#[test]
fn unreadable_identity_fails_closed() {
let paths = TestPaths::new();
fs::create_dir(&paths.keyring).unwrap();
let error =
load_or_create_keyring(b"test identity secret", &paths.keyring, &paths.public_key)
.unwrap_err();
assert!(error.contains("unable to load existing Omikron identity"));
}
#[test]
fn public_key_persistence_failure_fails_closed() {
let paths = TestPaths::new();
fs::create_dir(&paths.public_key).unwrap();
let error =
load_or_create_keyring(b"test identity secret", &paths.keyring, &paths.public_key)
.unwrap_err();
assert!(error.contains("unable to persist Omikron public key bundle"));
}
#[cfg(feature = "raw-migration")]
#[test]
fn raw_identity_migration_rewrites_protected_storage() {
use mtp::files::{load_keyring, save_keyring_raw};
let paths = TestPaths::new();
let original = Keyring::generate();
save_keyring_raw(&original, &paths.keyring).unwrap();
let migrated =
migrate_raw_keyring(b"test identity secret", &paths.keyring, &paths.public_key)
.unwrap();
let loaded = load_keyring(&paths.keyring, b"test identity secret").unwrap();
assert_eq!(
migrated.public_key_bundle().try_as_bytes().unwrap(),
original.public_key_bundle().try_as_bytes().unwrap()
);
assert_eq!(
loaded.public_key_bundle().try_as_bytes().unwrap(),
original.public_key_bundle().try_as_bytes().unwrap()
);
}
}

28
src/main.rs
View file

@ -3,6 +3,7 @@ mod app_state;
mod calls;
mod config;
mod data;
mod identity;
mod omega;
mod rho;
mod services;
@ -18,35 +19,20 @@ pub static WORKING_DIR: Lazy<PathBuf> =
use rustls::crypto::aws_lc_rs::default_provider;
use mtp::crypto::Keyring;
use mtp::files::{load_keyring_raw, save_keyring_raw, save_public_key_bundle};
use crate::{
app_state::AppState,
calls::{call_manager::CallManager, call_util::LiveKitService},
config::Config,
identity::{
KEYRING_PATH, PUBLIC_KEY_PATH, identity_secret_from_environment, load_or_create_keyring,
},
omega::omega_connection::{OmegaConnection, start_task_cleanup_loop},
rho::rho_manager::RhoManager,
rho::server::start,
util::logger::{PrintType, startup},
};
const KEYRING_PATH: &str = "./omikron.mk";
const PUBLIC_KEY_PATH: &str = "./omikron.mpkb";
fn load_keyring() -> Result<Keyring, String> {
match load_keyring_raw(KEYRING_PATH) {
Ok(keyring) => Ok(keyring),
Err(_) => {
let kr = Keyring::generate();
save_keyring_raw(&kr, KEYRING_PATH).map_err(|error| error.to_string())?;
save_public_key_bundle(&kr.public_key_bundle(), PUBLIC_KEY_PATH)
.map_err(|error| error.to_string())?;
eprintln!("Generated new keyring at {}", KEYRING_PATH);
Ok(kr)
}
}
}
#[tokio::main]
async fn main() {
if let Err(_) = default_provider().install_default() {
@ -64,10 +50,10 @@ async fn main() {
}
};
let keyring = match load_keyring() {
Ok(keyring) => keyring,
let identity_secret = match identity_secret_from_environment() {
Ok(secret) => secret,
Err(error) => {
eprintln!("Unable to load keyring: {error}");
eprintln!("Unable to load Omikron identity secret: {error}");
return;
}
};

1
src/omega/capabilities.rs
View file

@ -42,6 +42,7 @@ impl PeerCapabilities {
format!("{OMIKRON_PREFIX}{}", names.join(","))
}
#[allow(dead_code)]
pub fn from_identification_description(description: Option<&str>) -> Result<Self, ()> {
parse_capabilities(description, OMIKRON_PREFIX)
}

219
src/omega/omega_connection.rs
View file

@ -2,7 +2,7 @@ use super::capabilities::PeerCapabilities;
use crate::{
config::Config,
log_cv_in, log_cv_out, log_err, log_in, log_out,
rho::connection::{MtpValueCompat, OptionalDataValueCompat},
rho::connection::{OptionalDataValueCompat, RequiredMtpFields},
rho::relay_router,
rho::rho_manager::RhoManager,
util::{data_type_id, logger::PrintType},
@ -16,6 +16,7 @@ use mtp::{
host::{Policy, SendMode},
};
use mtp_transport::ConnectionHandle;
use rand::RngExt;
use std::{sync::Arc, time::Duration};
use tokio::{
sync::{Mutex, RwLock, mpsc, oneshot, watch},
@ -27,6 +28,7 @@ use uuid::Uuid;
const RECONNECT_DELAY: Duration = Duration::from_secs(5);
const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(300);
const RECONNECT_JITTER: Duration = Duration::from_secs(1);
const CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
const CAPABILITY_NEGOTIATION_TIMEOUT: Duration = Duration::from_secs(1);
const PING_INTERVAL: Duration = Duration::from_secs(5);
@ -36,13 +38,29 @@ const MAX_CONCURRENT_REQUESTS: usize = 128;
const CIRCUIT_BREAKER_FAILURE_THRESHOLD: u32 = 3;
const CIRCUIT_BREAKER_COOLDOWN: Duration = Duration::from_secs(30);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ConnectionOutcome {
HealthySessionEnded,
FailedBeforeHealthy,
}
fn reconnect_base_after_outcome(current: Duration, outcome: ConnectionOutcome) -> Duration {
match outcome {
ConnectionOutcome::HealthySessionEnded => RECONNECT_DELAY,
ConnectionOutcome::FailedBeforeHealthy => current,
}
}
fn reconnect_delay_with_jitter(delay: Duration) -> Duration {
let max_jitter_ms = std::cmp::min(delay, RECONNECT_JITTER).as_millis() as u64;
delay + Duration::from_millis(rand::rng().random_range(0..=max_jitter_ms))
}
fn client_changed_target(value: &CommunicationValue) -> Option<(i64, i64)> {
if !value.is_type(CommunicationType::ClientChanged) {
return None;
}
let receiver = i64::try_from(value.get_receiver())
.ok()
.filter(|id| *id > 0)?;
let receiver = i64::try_from(value.receiver()?).ok().filter(|id| *id > 0)?;
let session_id = value
.get_data(DataType::SessionId)
.as_signed_number()
@ -284,33 +302,45 @@ impl OmegaConnection {
}
match self.clone().connect_once().await {
Ok(()) => {
// Connection closed gracefully, check if we should reconnect
if *self.reconnect_on_close.read().await {
log_err!(
0,
PrintType::Omega,
"Connection lost, reconnecting in {:?}...",
reconnect_delay
);
} else {
log_in!(0, PrintType::Omega, "Connection closed, not reconnecting");
break;
Ok(outcome) => {
reconnect_delay = reconnect_base_after_outcome(reconnect_delay, outcome);
match outcome {
ConnectionOutcome::HealthySessionEnded => {
log_err!(
0,
PrintType::Omega,
"Healthy Omega connection ended; reconnecting"
);
}
ConnectionOutcome::FailedBeforeHealthy => {
log_err!(
0,
PrintType::Omega,
"Omega connection failed before synchronization"
);
}
}
}
Err(e) => {
log_err!(
0,
PrintType::Omega,
"Connection failed: {}, retrying in {:?}...",
e,
reconnect_delay
);
Err(error) => {
log_err!(0, PrintType::Omega, "Connection failed: {}", error);
}
}
if !*self.reconnect_on_close.read().await {
log_in!(0, PrintType::Omega, "Connection closed, not reconnecting");
break;
}
let retry_delay = reconnect_delay_with_jitter(reconnect_delay);
log_err!(
0,
PrintType::Omega,
"Retrying Omega connection in {:?}...",
retry_delay
);
tokio::select! {
_ = sleep(reconnect_delay) => {}
_ = sleep(retry_delay) => {}
_ = shutdown_rx.changed() => {
if *shutdown_rx.borrow() {
break;
@ -322,7 +352,7 @@ impl OmegaConnection {
}
}
async fn connect_once(self: Arc<Self>) -> Result<(), String> {
async fn connect_once(self: Arc<Self>) -> Result<ConnectionOutcome, String> {
*self.state.write().await = ConnectionState::Connecting;
let client_config = ClientConfig::new(format!("https://{}:{}", self.host, self.port))
@ -331,7 +361,6 @@ impl OmegaConnection {
.with_policy(
Policy::default()
.with_send_mode(SendMode::SingleStreamPerMessage)
.with_max_message_size(1_000_000_000)
.with_timeouts(
Duration::from_millis(5_000),
Duration::from_millis(5_000),
@ -343,8 +372,7 @@ impl OmegaConnection {
.with_max_concurrent_stream_tasks(64)
.with_persistent_stream_retries(5, Duration::from_secs(5)),
)
.with_ping_interval(PING_INTERVAL)
.with_max_missed_pings(0);
.with_ping_interval(PING_INTERVAL);
let host_public_key = load_public_key_bundle("./omega.mpkb")
.map_err(|e| format!("Failed to load omega.mpkb: {}", e))?;
@ -389,7 +417,7 @@ impl OmegaConnection {
self.fail_synchronization("capability negotiation", error)
.await;
let _ = read_handle.await;
return Err("Omega capability negotiation failed".to_string());
return Ok(ConnectionOutcome::FailedBeforeHealthy);
}
Ok(Err(_)) | Err(_) => {
// An Omega from before capability negotiation sends no second
@ -403,7 +431,7 @@ impl OmegaConnection {
// Tell omega our current state now that we're actually connected -
// doing this after teardown (as before) sent into a sender that had
// already been cleared, silently dropping the sync every time.
self.clone().sync_client_iota_status().await;
let reached_healthy_state = self.clone().sync_client_iota_status().await;
// Wait for read loop to complete (connection closed)
let result = read_handle.await;
@ -428,19 +456,13 @@ impl OmegaConnection {
*self.state.write().await = ConnectionState::Disconnected;
match result {
Ok(()) => {
// Check if we should reconnect
if *self.reconnect_on_close.read().await {
Err("Connection closed, will reconnect".to_string())
} else {
Ok(())
}
}
Ok(()) if reached_healthy_state => Ok(ConnectionOutcome::HealthySessionEnded),
Ok(()) => Ok(ConnectionOutcome::FailedBeforeHealthy),
Err(e) => Err(format!("Read loop error: {}", e)),
}
}
async fn sync_client_iota_status(&self) {
async fn sync_client_iota_status(&self) -> bool {
*self.state.write().await = ConnectionState::SynchronizingRoutes;
let mut connected_iota_ids: Vec<DataValue> = Vec::new();
let mut connected_sessions: Vec<DataValue> = Vec::new();
@ -497,16 +519,16 @@ impl OmegaConnection {
if let Err(error) = self.send_message_result(&sync_msg).await {
self.fail_synchronization("legacy route synchronization", error)
.await;
return;
return false;
}
*self.state.write().await = ConnectionState::SynchronizingSubscriptions;
if let Err(error) = self.restore_presence_subscriptions().await {
self.fail_synchronization("subscription restoration", error)
.await;
return;
return false;
}
*self.state.write().await = ConnectionState::Ready;
return;
return true;
}
let mut response = Err("route synchronization did not start".to_string());
@ -535,7 +557,7 @@ impl OmegaConnection {
);
self.fail_synchronization("subscription restoration", error)
.await;
return;
return false;
}
*self.state.write().await = ConnectionState::Ready;
}
@ -546,14 +568,15 @@ impl OmegaConnection {
);
self.fail_synchronization("route synchronization", error)
.await;
return;
return false;
}
Err(error) => {
self.fail_synchronization("route synchronization", error)
.await;
return;
return false;
}
}
true
}
async fn fail_synchronization(&self, phase: &str, error: String) {
@ -600,8 +623,39 @@ impl OmegaConnection {
}
if cv.is_type(CommunicationType::Relay) {
let destination_iota = cv.receiver().unwrap_or_default();
let request_id = cv.get_id();
let request_id = match cv.require_id() {
Ok(request_id) => request_id,
Err(error) => {
log_err!(
self.omikron_id as i64,
PrintType::Omega,
"Rejected malformed relay: {}",
error
);
let response =
CommunicationValue::new(CommunicationType::ErrorInvalidData)
.without_id();
let _ = self.send_message_result(&response).await;
continue;
}
};
let destination_iota = match cv.require_receiver() {
Ok(destination_iota) => destination_iota,
Err(error) => {
log_err!(
self.omikron_id as i64,
PrintType::Omega,
"Rejected malformed relay: {}",
error
);
let response = CommunicationValue::new(
CommunicationType::ErrorInvalidData,
)
.with_id(request_id);
let _ = self.send_message_result(&response).await;
continue;
}
};
let response = match relay_router::route_from_omega(&self.rho, cv).await {
Ok(()) => CommunicationValue::new(CommunicationType::Success)
.with_id(request_id),
@ -630,7 +684,18 @@ impl OmegaConnection {
continue;
}
let msg_id = cv.get_id();
let msg_id = match cv.require_id() {
Ok(msg_id) => msg_id,
Err(error) => {
log_err!(
self.omikron_id as i64,
PrintType::Omega,
"Discarded response without an MTP id: {}",
error
);
continue;
}
};
if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) {
if (task.task)(self.clone(), cv.clone()) {
continue;
@ -749,6 +814,7 @@ impl OmegaConnection {
self.try_send_message(cv).await
}
#[allow(dead_code)]
pub async fn supports_set_user_state(&self) -> bool {
self.peer_capabilities.read().await.set_user_state_v1
}
@ -807,10 +873,12 @@ impl OmegaConnection {
cv: &CommunicationValue,
timeout_duration: Option<Duration>,
) -> Result<CommunicationValue, String> {
let msg_id = cv
.require_id()
.map_err(|error| format!("request is missing correlation id: {error}"))?;
self.await_connection(timeout_duration).await?;
let (tx, mut rx) = mpsc::channel(1);
let msg_id = cv.get_id();
self.waiting_tasks.insert(
msg_id,
@ -819,9 +887,9 @@ impl OmegaConnection {
log_in!(
0,
PrintType::Omega,
"Matched Omega response (request_id={}, response_id={}, type={})",
"Matched Omega response (request_id={}, response_id={:?}, type={})",
msg_id,
response_cv.get_id(),
response_cv.id(),
response_cv
.get_comm_type_enum()
.map(|kind| kind.to_string())
@ -880,7 +948,7 @@ impl OmegaConnection {
pub async fn close_iota(&self, iota_id: i64) -> Result<(), String> {
let cv = CommunicationValue::new(CommunicationType::IotaDisconnected)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()));
let request_id = cv.get_id();
let request_id = cv.require_id().map_err(|error| error.to_string())?;
let result = self.lifecycle_request(&cv).await;
if let Err(error) = &result {
self.log_lifecycle_failure("IotaDisconnected", iota_id, None, request_id, error);
@ -897,7 +965,7 @@ impl OmegaConnection {
DataType::SessionId,
DataValue::SignedNumber(session_id.into()),
);
let request_id = cv.get_id();
let request_id = cv.require_id().map_err(|error| error.to_string())?;
let result = self.lifecycle_request(&cv).await;
if let Err(error) = &result {
self.log_lifecycle_failure(
@ -938,7 +1006,7 @@ impl OmegaConnection {
.as_millis() as i128,
),
);
let request_id = request.get_id();
let request_id = request.require_id().map_err(|error| error.to_string())?;
let result = self.lifecycle_request(&request).await;
if let Err(error) = &result {
self.log_lifecycle_failure(
@ -956,7 +1024,7 @@ impl OmegaConnection {
pub async fn iota_connected(&self, iota_id: i64) -> Result<(), String> {
let request = CommunicationValue::new(CommunicationType::IotaConnected)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()));
let request_id = request.get_id();
let request_id = request.require_id().map_err(|error| error.to_string())?;
let result = self.lifecycle_request(&request).await;
if let Err(error) = &result {
self.log_lifecycle_failure("IotaConnected", iota_id, None, request_id, error);
@ -965,6 +1033,7 @@ impl OmegaConnection {
result
}
#[allow(dead_code)]
pub async fn reconcile_routes(&self) {
self.sync_client_iota_status().await;
}
@ -1174,7 +1243,11 @@ impl OmegaConnection {
#[cfg(test)]
mod tests {
use super::{PeerCapabilities, client_changed_target, parse_omega_capability_response};
use super::{
ConnectionOutcome, MAX_RECONNECT_DELAY, PeerCapabilities, RECONNECT_DELAY,
client_changed_target, parse_omega_capability_response, reconnect_base_after_outcome,
reconnect_delay_with_jitter,
};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
fn notification() -> CommunicationValue {
@ -1196,6 +1269,9 @@ mod tests {
let mut missing_state = notification();
missing_state.remove_data(DataType::UserState);
assert_eq!(client_changed_target(&missing_state), None);
let missing_receiver = notification().without_receiver();
assert_eq!(client_changed_target(&missing_receiver), None);
}
#[test]
@ -1208,6 +1284,35 @@ mod tests {
assert_eq!(client_changed_target(&invalid), None);
}
#[test]
fn reconnect_jitter_is_bounded_without_changing_the_base_delay() {
for _ in 0..32 {
let delay = reconnect_delay_with_jitter(RECONNECT_DELAY);
assert!(delay >= RECONNECT_DELAY);
assert!(delay <= RECONNECT_DELAY + std::time::Duration::from_secs(1));
}
assert!(MAX_RECONNECT_DELAY > RECONNECT_DELAY);
}
#[test]
fn healthy_session_resets_accumulated_reconnect_backoff() {
let mut delay = RECONNECT_DELAY;
for _ in 0..4 {
delay = std::cmp::min(delay * 2, MAX_RECONNECT_DELAY);
}
assert!(delay > RECONNECT_DELAY);
assert_eq!(
reconnect_base_after_outcome(delay, ConnectionOutcome::FailedBeforeHealthy),
delay
);
assert_eq!(
reconnect_base_after_outcome(delay, ConnectionOutcome::HealthySessionEnded),
RECONNECT_DELAY
);
}
#[test]
fn capability_response_negotiates_new_omega() {
let response = CommunicationValue::new(CommunicationType::IdentificationResponse)

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)));
}
}

2
src/services/call_service.rs
View file

@ -7,10 +7,12 @@ use crate::{
calls::{call_group::CallGroup, error::CallError},
};
#[allow(dead_code)]
pub struct CallService {
state: Arc<AppState>,
}
#[allow(dead_code)]
impl CallService {
pub fn new(state: Arc<AppState>) -> Self {
Self { state }

2
src/services/routing_service.rs
View file

@ -4,10 +4,12 @@ use mtp::codec::CommunicationValue;
use crate::{app_state::AppState, rho::rho_connection::RhoConnection};
#[allow(dead_code)]
pub struct RoutingService {
state: Arc<AppState>,
}
#[allow(dead_code)]
impl RoutingService {
pub fn new(state: Arc<AppState>) -> Self {
Self { state }

2
src/services/user_service.rs
View file

@ -2,10 +2,12 @@ use std::sync::Arc;
use crate::{app_state::AppState, rho::rho_connection::RhoConnection};
#[allow(dead_code)]
pub struct UserService {
state: Arc<AppState>,
}
#[allow(dead_code)]
impl UserService {
pub fn new(state: Arc<AppState>) -> Self {
Self { state }

18
src/util/logger.rs
View file

@ -9,8 +9,6 @@ use std::{
use ansi_term::Color;
use mtp::codec::{CommunicationType, CommunicationValue, DataTypeId, DataValue, Version};
use crate::rho::connection::MtpValueCompat;
static LOGGER: OnceLock<mpsc::Sender<LogMessage>> = OnceLock::new();
#[allow(dead_code)]
@ -153,7 +151,9 @@ pub fn log_cv_internal(
let formatted = format_cv(cv);
log_internal(
cv.get_sender() as i64,
cv.sender()
.and_then(|sender| i64::try_from(sender).ok())
.unwrap_or(0),
print_type.unwrap_or(PrintType::General),
prefix,
false,
@ -164,14 +164,14 @@ pub fn log_cv_internal(
pub fn format_cv(cv: &CommunicationValue) -> String {
let mut parts = Vec::new();
let sender = cv.get_sender();
let receiver = cv.get_receiver();
let sender = cv.sender();
let receiver = cv.receiver();
if sender > 0 && receiver > 0 {
if let (Some(sender), Some(receiver)) = (sender, receiver) {
parts.push(format!("{} > {}", sender, receiver));
} else if sender > 0 {
} else if let Some(sender) = sender {
parts.push(format!("{}", sender));
} else if receiver > 0 {
} else if let Some(receiver) = receiver {
parts.push(format!("> {}", receiver));
}
@ -179,7 +179,7 @@ pub fn format_cv(cv: &CommunicationValue) -> String {
.get_comm_type_enum()
.map(|kind| kind.to_string())
.unwrap_or_else(|| cv.get_type().to_string());
parts.push(format!("{} (id={})", comm_type, cv.get_id()));
parts.push(format!("{} (id={:?})", comm_type, cv.id()));
if cv.is_type(CommunicationType::Relay) {
parts.push("<opaque relay payload>".to_string());