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

mtp update

Browse source
This commit is contained in:
Alex Emmet 2026-07-20 15:27:52 +02:00
commit a642afce5a
12 changed files with 632 additions and 354 deletions
Download patch file Download diff file Expand all files Collapse all files

7
src/anonymous_clients/anonymous_client_connection.rs
View file

@ -1,5 +1,4 @@
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap};
use mtp::host::{Receiver, Sender};
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
@ -10,7 +9,7 @@ use crate::anonymous_clients::anonymous_manager::{self, generate_username};
use crate::calls::{call_group::call_invite_secret_from_cv, call_manager};
use crate::data::user::UserStatus;
use crate::omega::omega_connection::{OmegaConnection, get_omega_connection};
use crate::rho::connection::GeneralConnection;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::rho::rho_manager;
use crate::util::logger::PrintType;
use crate::{log_cv_in, log_cv_out, log_out};
@ -18,8 +17,8 @@ use crate::{log_cv_in, log_cv_out, log_out};
pub struct AnonymousClientConnection {
user_id: u64,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub sender: Arc<MtpSender>,
pub receiver: Arc<MtpReceiver>,
pub ping: Arc<RwLock<i64>>,
pub interested_users: Arc<RwLock<Vec<i64>>>,
is_open: Arc<RwLock<bool>>,

8
src/main.rs
View file

@ -16,7 +16,7 @@ pub static WORKING_DIR: Lazy<PathBuf> =
use rustls::crypto::aws_lc_rs::default_provider;
use mtp::crypto::Keyring;
use mtp::files::{load_keyring as load_keyring_file, save_keyring, save_public_key_bundle};
use mtp::files::{load_keyring_raw, save_keyring_raw, save_public_key_bundle};
use crate::{
calls::call_util::garbage_collect_calls, omega::omega_connection::get_omega_connection,
@ -27,9 +27,9 @@ const KEYRING_PATH: &str = "./omikron.mk";
const PUBLIC_KEY_PATH: &str = "./omikron.mpkb";
static KEYRING: Lazy<Keyring> = Lazy::new(|| {
load_keyring_file(KEYRING_PATH).unwrap_or_else(|_| {
load_keyring_raw(KEYRING_PATH).unwrap_or_else(|_| {
let kr = Keyring::generate();
save_keyring(&kr, KEYRING_PATH).expect("Failed to save generated keyring");
save_keyring_raw(&kr, KEYRING_PATH).expect("Failed to save generated keyring");
save_public_key_bundle(&kr.public_key_bundle(), PUBLIC_KEY_PATH)
.expect("Failed to save generated public key bundle");
eprintln!("Generated new keyring at {}", KEYRING_PATH);
@ -56,7 +56,7 @@ async fn main() {
let rho_port = env::var("RHO_PORT")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(959);
.unwrap_or(443);
get_omega_connection();
tokio::spawn(async move {

31
src/omega/omega_connection.rs
View file

@ -263,22 +263,21 @@ impl OmegaConnection {
.parse::<u64>()
.unwrap_or(0),
)
.with_policy(Policy {
send_mode: SendMode::SingleStreamPerMessage,
max_message_size: 1_000_000_000,
close_frame_len: u32::MAX,
application_close_code: 0,
open_stream_timeout: Duration::from_millis(2_000),
write_timeout: Duration::from_millis(2_000),
accept_stream_timeout: Duration::from_millis(10_000),
read_timeout: Duration::from_millis(30_000),
keep_alive_interval: Some(Duration::from_secs(6)),
max_idle_timeout: Some(Duration::from_secs(30)),
force_close_delay: Duration::from_millis(300),
max_transient_recv_errors: 20,
transient_recv_backoff: Duration::from_millis(100),
receiver_queue_capacity: 1000,
});
.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(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)),
);
let host_public_key = load_public_key_bundle("./omega.mpkb")
.map_err(|e| format!("Failed to load omega.mpkb: {}", e))?;

7
src/rho/app_connection.rs
View file

@ -1,11 +1,10 @@
use crate::anonymous_clients::anonymous_manager;
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::rho::{rho_connection::RhoConnection, rho_manager};
use crate::util::logger::PrintType;
use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::host::{Receiver, Sender};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
@ -17,8 +16,8 @@ pub struct AppConnection {
pub app_session: Uuid,
pub client_version: String,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub sender: Arc<MtpSender>,
pub receiver: Arc<MtpReceiver>,
pub ping: Arc<RwLock<i64>>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,

67
src/rho/client_connection.rs
View file

@ -1,17 +1,17 @@
use crate::anonymous_clients::anonymous_manager;
use crate::calls::{call_group::call_invite_secret_from_cv, call_manager, call_util};
use crate::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::rho::{rho_connection::RhoConnection, rho_manager};
use crate::util::logger::PrintType;
use crate::{data::user::UserStatus, omega::omega_connection::OmegaConnection};
use crate::{log_cv_in, log_cv_out, log_err, log_in, log_out};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::host::{Receiver, Sender};
use std::str::FromStr;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use trust_dns_resolver::TokioAsyncResolver;
use uuid::Uuid;
pub struct ClientConnection {
@ -19,8 +19,8 @@ pub struct ClientConnection {
pub session_id: u64,
pub client_version: String,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub sender: Arc<MtpSender>,
pub receiver: Arc<MtpReceiver>,
pub ping: Arc<RwLock<i64>>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,
pub rho_connection: Arc<RwLock<Option<Arc<RhoConnection>>>>,
@ -611,27 +611,52 @@ impl ClientConnection {
async fn handle_load_txt_record(self: Arc<Self>, cv: CommunicationValue) {
if let Some(path) = cv.get_data(DataType::Path).as_str() {
if let Ok(builder) = hickory_resolver::Resolver::builder_tokio() {
let resolver = builder.build();
if let Ok(lookup) = resolver.txt_lookup(path).await {
for record in lookup.iter() {
for txt_data in record.txt_data() {
if let Ok(s) = std::str::from_utf8(txt_data) {
let response =
CommunicationValue::new(CommunicationType::LoadTxtRecord)
.with_id(cv.get_id())
.add_typed_default(
DataType::Content,
DataValue::Str(s.to_string()),
);
self.send_message(&response).await;
return;
}
}
let resolver = match TokioAsyncResolver::tokio_from_system_conf() {
Ok(r) => r,
Err(_) => {
let path_data = cv.get_data(DataType::Path).clone();
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
return;
}
};
match resolver.txt_lookup(path).await {
Ok(txt_lookup) => {
if let Some(txt_record) = txt_lookup.iter().next() {
let record_text: String = txt_record
.txt_data()
.iter()
.map(|b| String::from_utf8_lossy(b))
.collect();
let response = CommunicationValue::new(CommunicationType::LoadTxtRecord)
.with_id(cv.get_id())
.add_typed_default(DataType::Content, DataValue::Str(record_text));
self.send_message(&response).await;
return;
}
let path_data = cv.get_data(DataType::Path).clone();
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
}
Err(_) => {
let path_data = cv.get_data(DataType::Path).clone();
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())
.add_typed_default(DataType::Path, path_data);
self.send_message(&error_cv).await;
}
}
return;
}
let path_data = cv.get_data(DataType::Path).clone();
let error_cv = CommunicationValue::new(CommunicationType::ErrorNotFound)
.with_id(cv.get_id())

14
src/rho/connection.rs
View file

@ -14,7 +14,11 @@ use crate::{
util::logger::PrintType,
};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue};
use mtp::host::{AuthState, Connection as MTPHostConnection, Receiver, Sender};
use mtp::host::AuthState;
use mtp::webserver::{WebMTPConnection, WebMtpReceiver, WebMtpSender};
pub type MtpSender = WebMtpSender;
pub type MtpReceiver = WebMtpReceiver;
/*
* How a connection identified itself during the mtp handshake driven by
@ -34,8 +38,8 @@ pub enum ConnectionKind {
}
pub struct GeneralConnection {
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub sender: Arc<MtpSender>,
pub receiver: Arc<MtpReceiver>,
connection_kind: ConnectionKind,
id: u64,
@ -50,14 +54,14 @@ pub struct GeneralConnection {
impl GeneralConnection {
/*
* `conn` has already been authenticated (or deliberately left
* unauthenticated) by `mtp::host::Host::accept`, via the
* unauthenticated) by `mtp::webserver::MTPWebServer::accept`, via the
* `get_by_connector_id`/`complete_register` callbacks in `server.rs`
* keyed off `conn.description`. There is no separate application-level
* challenge step anymore; a connection whose description doesn't resolve
* to a known, appropriately-authenticated kind is rejected here instead
* of being handed off to a connection handler.
*/
pub fn new(conn: MTPHostConnection) -> Option<Arc<Self>> {
pub fn new(conn: WebMTPConnection) -> Option<Arc<Self>> {
let kind = match (conn.description.as_deref(), &conn.auth_state) {
(Some("iota"), AuthState::Authenticated) => ConnectionKind::Iota,
(Some("client"), AuthState::Authenticated) => ConnectionKind::Client,

8
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::omega::omega_connection::get_omega_connection;
use crate::rho::connection::GeneralConnection;
use crate::rho::connection::{GeneralConnection, MtpReceiver, MtpSender};
use crate::util::logger::PrintType;
use dashmap::DashMap;
use mtp::codec::CommunicationType;
@ -16,8 +16,6 @@ use mtp::codec::DataTypeId;
use mtp::codec::DataValue;
use mtp::codec::TypeMap;
use mtp::crypto::KemPublicKey;
use mtp::host::Receiver;
use mtp::host::Sender;
use std::collections::BTreeMap;
use std::{collections::HashMap, sync::Arc, sync::LazyLock, time::Duration};
use tokio::sync::RwLock;
@ -33,8 +31,8 @@ static PENDING_CHAT_SECRETS: LazyLock<DashMap<u64, Vec<CommunicationValue>>> =
pub struct IotaConnection {
pub iota_id: u64,
pub client_version: String,
pub sender: Arc<Sender>,
pub receiver: Arc<Receiver>,
pub sender: Arc<MtpSender>,
pub receiver: Arc<MtpReceiver>,
pub user_ids: Arc<RwLock<Vec<u64>>>,
pub ping: Arc<RwLock<i64>>,
pub_key: Arc<RwLock<Option<Vec<u8>>>>,

38
src/rho/server.rs
View file

@ -13,7 +13,8 @@ use crate::{
};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::crypto::PublicKeyBundle;
use mtp::host::{AuthenticationPolicy, Host, HostConfig, Policy, SendMode};
use mtp::host::{AuthenticationPolicy, HostConfig, Policy, SendMode};
use mtp::webserver::{MTPWebServer, WebServerConfig};
/*
* Resolves the PublicKeyBundle mtp needs to verify a login's signed
@ -107,22 +108,21 @@ pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
cert_pem,
key_pem,
)
.with_policy(Policy {
send_mode: SendMode::SingleStreamPerMessage,
max_message_size: 1_000_000_000,
close_frame_len: u32::MAX,
application_close_code: 0,
open_stream_timeout: Duration::from_millis(2_000),
write_timeout: Duration::from_millis(2_000),
accept_stream_timeout: Duration::from_millis(10_000),
read_timeout: Duration::from_millis(30_000),
keep_alive_interval: Some(Duration::from_secs(6)),
max_idle_timeout: Some(Duration::from_secs(30)),
force_close_delay: Duration::from_millis(300),
max_transient_recv_errors: 20,
transient_recv_backoff: Duration::from_millis(100),
receiver_queue_capacity: 1000,
})
.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(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)),
)
.with_authentication(
load_keyring(),
Box::new(|user_id, description| Box::pin(get_by_connector_id(user_id, description))),
@ -130,7 +130,9 @@ pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
)
.with_authentication_policy(AuthenticationPolicy::AllowAuthentication);
let mut host: Host = Host::new(host_config).await?;
let web_config = WebServerConfig::new()
.route("/", |_request, response| async move { response.body("OK") })?;
let mut host = MTPWebServer::new(host_config, web_config).await?;
log!(0, PrintType::General, "Server listening on port {}.", port);
loop {