From 18817bdad479fcdd7e0f79a3ea573610f56c9a44 Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:57:31 +0200 Subject: [PATCH] MTP migration --- Cargo.lock | 17 +- src/main.rs | 10 - src/notifications/mod.rs | 1 - src/notifications/tauri.rs | 157 --------------- src/server/api.rs | 12 +- src/sql/sql.rs | 63 +++--- src/transport/omikron_connection.rs | 288 +++++++--------------------- 7 files changed, 123 insertions(+), 425 deletions(-) delete mode 100644 src/notifications/mod.rs delete mode 100644 src/notifications/tauri.rs diff --git a/Cargo.lock b/Cargo.lock index 36da87b..592912f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2280,7 +2280,7 @@ dependencies = [ [[package]] name = "mtp" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "mtp-client", "mtp-codec", @@ -2294,7 +2294,7 @@ dependencies = [ [[package]] name = "mtp-client" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "mtp-codec", "mtp-common", @@ -2307,7 +2307,7 @@ dependencies = [ [[package]] name = "mtp-codec" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "base64", "byteorder", @@ -2320,7 +2320,7 @@ dependencies = [ [[package]] name = "mtp-common" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "quinn", "rustls", @@ -2331,8 +2331,9 @@ dependencies = [ [[package]] name = "mtp-crypto" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ + "base64", "chacha20poly1305", "ed25519-dalek", "getrandom 0.4.3", @@ -2349,7 +2350,7 @@ dependencies = [ [[package]] name = "mtp-host" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "mtp-codec", "mtp-common", @@ -2362,7 +2363,7 @@ dependencies = [ [[package]] name = "mtp-transport" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "log", "mtp-codec", @@ -2376,7 +2377,7 @@ dependencies = [ [[package]] name = "mtp-type-map" version = "0.1.0" -source = "git+https://git.methanium.net/methanium/mtp#8dc8b54b26544554f98d7af32200d6bf9a01d300" +source = "git+https://git.methanium.net/methanium/mtp#44ff1d8781b4645dbff3dce2406155e4d9d43a4c" dependencies = [ "serde", "serde_yaml", diff --git a/src/main.rs b/src/main.rs index c54fba3..8672167 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,8 @@ -mod notifications; mod server; mod sql; mod transport; mod util; -use crate::notifications::tauri; use crate::sql::sql::initialize_db; use crate::sql::sql::print_users; use crate::transport::omikron_connection; @@ -53,13 +51,6 @@ async fn main() { } }); - let tauri_handle = tokio::spawn(async move { - match tauri::start(9189).await { - Err(e) => log_err!(0, PrintType::General, "{:?}", e), - _ => {} - } - }); - log!("Started"); log!(" .env"); if let Err(e) = initialize_db().await { @@ -94,5 +85,4 @@ async fn main() { } omikron_handle.abort(); - tauri_handle.abort(); } diff --git a/src/notifications/mod.rs b/src/notifications/mod.rs deleted file mode 100644 index 62411f7..0000000 --- a/src/notifications/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod tauri; diff --git a/src/notifications/tauri.rs b/src/notifications/tauri.rs deleted file mode 100644 index 8195988..0000000 --- a/src/notifications/tauri.rs +++ /dev/null @@ -1,157 +0,0 @@ -use crate::log; -use crate::util::file_util::load_file_vec; -use crate::util::logger::PrintType; -use dashmap::DashMap; -use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; -use mtp::transport::{Host, Policy, Receiver, SendMode, Sender, host}; -use once_cell::sync::Lazy; -use std::net::{IpAddr, Ipv4Addr}; -use std::sync::Arc; -use std::time::Duration; - -pub struct TauriConnection { - pub user_id: i64, - pub sender: Arc, -} - -static TAURI_CONNECTIONS: Lazy>>> = Lazy::new(DashMap::new); - -pub async fn start(port: u16) -> Result<(), Box> { - let cert_pem = load_file_vec("certs", "transport_cert.pem").expect("Error loading Pemfile"); - let key_pem = load_file_vec("certs", "transport_key.pem").expect("Error loading Keyfile"); - - let mut host: Host = host( - IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), - port, - cert_pem, - key_pem, - 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, - }, - ) - .await?; - - log!("TauriServer listening on port {}", port); - - while let Some((sender, mut receiver)) = host.next().await { - tokio::spawn(async move { - handle_connection(sender, &mut receiver).await; - }); - } - - Ok(()) -} - -async fn handle_connection(sender: Sender, receiver: &mut Receiver) { - let mut current_user_id = 0; - let sender = Arc::new(sender); - - while let Ok(cv) = receiver.receive().await { - if cv.is_type(CommunicationType::TauriIdentification) { - let user_id = cv.get_data(DataType::UserId).as_number().unwrap_or(0) as i64; - if user_id != 0 { - current_user_id = user_id; - let conn = Arc::new(TauriConnection { - user_id, - sender: sender.clone(), - }); - - TAURI_CONNECTIONS - .entry(user_id) - .and_modify(|conns| { - conns.retain(|c| !Arc::ptr_eq(&c.sender.handle(), &sender.handle())); - conns.push(conn.clone()); - }) - .or_insert_with(|| vec![conn]); - - let response = - CommunicationValue::new(CommunicationType::Success).with_id(cv.get_id()); - if let Err(e) = sender.send(&response).await { - log!( - PrintType::General, - "Failed to send tauri success response: {}", - e - ); - } else { - log!(user_id, PrintType::Client, "Tauri device registered"); - } - } - } else if cv.is_type(CommunicationType::Ping) { - let response = - CommunicationValue::new(CommunicationType::Pong).with_id(cv.get_id()); - if let Err(_) = sender.send(&response).await { - break; - } - } - } - - // Cleanup - if current_user_id != 0 { - if let Some(mut conns) = TAURI_CONNECTIONS.get_mut(¤t_user_id) { - conns.retain(|c| !Arc::ptr_eq(c.sender.handle(), sender.handle())); - } - } -} - -pub async fn send_notification(user_id: i64, sender_id: i64) { - let conns_opt = { - let entry = TAURI_CONNECTIONS.get(&user_id); - entry.map(|e| e.clone()) - }; - - if let Some(conns) = conns_opt { - let cv = CommunicationValue::new(CommunicationType::PushNotification) - .add_typed_default(DataType::SenderId, DataValue::SignedNumber(sender_id as i128)); - - let mut remove_needed = false; - for conn in conns.iter() { - if conn.sender.send(&cv).await.is_err() { - remove_needed = true; - } - } - - if remove_needed { - if let Some(mut conns_mut) = TAURI_CONNECTIONS.get_mut(&user_id) { - conns_mut.retain(|conn| conn.sender.is_open()); - } - } - } -} - -pub async fn remove_notification(user_id: i64, sender_id: i64) { - let conns_opt = { - let entry = TAURI_CONNECTIONS.get(&user_id); - entry.map(|e| e.clone()) - }; - - if let Some(conns) = conns_opt { - let cv = CommunicationValue::new(CommunicationType::ReadNotification) - .add_typed_default(DataType::SenderId, DataValue::SignedNumber(sender_id as i128)); - - let mut remove_needed = false; - for conn in conns.iter() { - if conn.sender.send(&cv).await.is_err() { - remove_needed = true; - } - } - - if remove_needed { - if let Some(mut conns_mut) = TAURI_CONNECTIONS.get_mut(&user_id) { - conns_mut.retain(|conn| conn.sender.is_open()); - } - } - } -} diff --git a/src/server/api.rs b/src/server/api.rs index 219d5fc..93c594e 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -68,7 +68,7 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { let mut res = JsonValue::new_object(); res["status"] = "success".into(); res["id"] = id.into(); - res["public_key"] = public_key.into(); + res["public_key"] = public_key.to_base64().into(); res["ip_address"] = ip_address.into(); (StatusCode::OK, res.dump()) @@ -103,7 +103,7 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { let mut res = JsonValue::new_object(); res["status"] = "success".into(); res["id"] = id.into(); - res["public_key"] = public_key.into(); + res["public_key"] = public_key.to_base64().into(); res["ip_address"] = ip_address.into(); (StatusCode::OK, res.dump()) } else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) { @@ -111,7 +111,7 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { let mut res = JsonValue::new_object(); res["status"] = "success".into(); res["id"] = omikron_id.into(); - res["public_key"] = public_key.into(); + res["public_key"] = public_key.to_base64().into(); res["ip_address"] = ip_address.into(); (StatusCode::OK, res.dump()) } else { @@ -126,7 +126,7 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { let mut res = JsonValue::new_object(); res["status"] = "success".into(); res["id"] = omikron_id.into(); - res["public_key"] = public_key.into(); + res["public_key"] = public_key.to_base64().into(); res["ip_address"] = ip_address.into(); (StatusCode::OK, res.dump()) } else { @@ -172,7 +172,7 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { let mut res = JsonValue::new_object(); res["status"] = "success".into(); res["username"] = username.into(); - res["public_key"] = public_key.into(); + res["public_key"] = public_key.to_base64().into(); res["user_id"] = id.into(); res["iota_id"] = iota_id.into(); res["sub_level"] = sub_level.into(); @@ -227,7 +227,7 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { let mut res = JsonValue::new_object(); res["status"] = "success".into(); res["username"] = username.into(); - res["public_key"] = public_key.into(); + res["public_key"] = public_key.to_base64().into(); res["user_id"] = id.into(); res["iota_id"] = iota_id.into(); res["sub_level"] = sub_level.into(); diff --git a/src/sql/sql.rs b/src/sql/sql.rs index d9eb712..45b1899 100644 --- a/src/sql/sql.rs +++ b/src/sql/sql.rs @@ -1,4 +1,5 @@ use crate::log; +use mtp_crypto::PublicKeyBundle; use once_cell::sync::Lazy; use sqlx::{MySql, Pool, Row, mysql::MySqlPoolOptions}; use std::sync::atomic::{AtomicU64, Ordering}; @@ -61,7 +62,7 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> { avatar MEDIUMBLOB, sub_level INT(11) NOT NULL DEFAULT 0, sub_end BIGINT(20) NOT NULL DEFAULT 0, - public_key TEXT NOT NULL COLLATE utf8mb4_bin, + public_key BLOB NOT NULL, private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin DEFAULT '', iota_id BIGINT NOT NULL, token VARCHAR(256) NOT NULL UNIQUE COLLATE utf8mb4_bin @@ -73,7 +74,7 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> { "CREATE TABLE IF NOT EXISTS iotas ( id BIGINT NOT NULL PRIMARY KEY, - public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin + public_key BLOB NOT NULL )", ) .execute(&pool) @@ -82,7 +83,7 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> { "CREATE TABLE IF NOT EXISTS omikrons ( id BIGINT NOT NULL PRIMARY KEY, - public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin, + public_key BLOB NOT NULL, location VARCHAR(255) NOT NULL COLLATE utf8mb4_bin, ip_address VARCHAR(255) NOT NULL COLLATE utf8mb4_bin )", @@ -160,7 +161,7 @@ pub async fn get_by_username( Option>, i32, i64, - String, + PublicKeyBundle, String, String, ), @@ -192,7 +193,8 @@ pub async fn get_by_username( let avatar: Option> = row.get("avatar"); let sub_level: i32 = row.get("sub_level"); let sub_end: i64 = row.get("sub_end"); - let public_key: Vec = row.get("public_key"); + let public_key = PublicKeyBundle::from_bytes(&row.get::, _>("public_key")) + .map_err(|e| sqlx::Error::Decode(Box::new(e)))?; let private_key_hash: Vec = row.get("private_key_hash"); let token: Vec = row.get("token"); @@ -206,7 +208,7 @@ pub async fn get_by_username( avatar, sub_level, sub_end, - String::from_utf8_lossy(&public_key).to_string(), + public_key, String::from_utf8_lossy(&private_key_hash).to_string(), String::from_utf8_lossy(&token).to_string(), )) @@ -228,7 +230,7 @@ pub async fn get_by_user_id( Option>, i32, i64, - String, + PublicKeyBundle, String, String, ), @@ -260,7 +262,8 @@ pub async fn get_by_user_id( let avatar: Option> = row.get("avatar"); let sub_level: i32 = row.get("sub_level"); let sub_end: i64 = row.get("sub_end"); - let public_key: Vec = row.get("public_key"); + let public_key = PublicKeyBundle::from_bytes(&row.get::, _>("public_key")) + .map_err(|e| sqlx::Error::Decode(Box::new(e)))?; let private_key_hash: Vec = row.get("private_key_hash"); let token: Vec = row.get("token"); @@ -274,7 +277,7 @@ pub async fn get_by_user_id( avatar, sub_level, sub_end, - String::from_utf8_lossy(&public_key).to_string(), + public_key, String::from_utf8_lossy(&private_key_hash).to_string(), String::from_utf8_lossy(&token).to_string(), )) @@ -296,7 +299,7 @@ pub async fn get_users_by_iota_id( Option>, i32, i64, - String, + PublicKeyBundle, String, String, )>, @@ -328,7 +331,8 @@ pub async fn get_users_by_iota_id( let avatar: Option> = row.get("avatar"); let sub_level: i32 = row.get("sub_level"); let sub_end: i64 = row.get("sub_end"); - let public_key: Vec = row.get("public_key"); + let public_key = PublicKeyBundle::from_bytes(&row.get::, _>("public_key")) + .map_err(|e| sqlx::Error::Decode(Box::new(e)))?; let private_key_hash: Vec = row.get("private_key_hash"); let token: Vec = row.get("token"); @@ -342,7 +346,7 @@ pub async fn get_users_by_iota_id( avatar, sub_level, sub_end, - String::from_utf8_lossy(&public_key).to_string(), + public_key, String::from_utf8_lossy(&private_key_hash).to_string(), String::from_utf8_lossy(&token).to_string(), )); @@ -477,7 +481,7 @@ pub async fn change_iota_id(id: i64, new_iota_id: i64) -> Result<(), sqlx::Error pub async fn change_keys( id: i64, - new_public_key: String, + new_public_key: PublicKeyBundle, new_private_key_hash: String, ) -> Result<(), sqlx::Error> { let pool = { @@ -489,7 +493,7 @@ pub async fn change_keys( }; sqlx::query("UPDATE users SET public_key = ?, private_key_hash = ? WHERE id = ?") - .bind(new_public_key.as_bytes().to_vec()) + .bind(new_public_key.as_bytes()) .bind(new_private_key_hash.as_bytes().to_vec()) .bind(id) .execute(&pool) @@ -517,7 +521,7 @@ pub async fn change_token(id: i64, new_token: String) -> Result<(), sqlx::Error> pub async fn register_complete_user( id: i64, username: String, - public_key: String, + public_key: PublicKeyBundle, iota_id: i64, token: String, ) -> Result<(), sqlx::Error> { @@ -534,7 +538,7 @@ pub async fn register_complete_user( ) .bind(id) .bind(username.as_bytes().to_vec()) - .bind(public_key.as_bytes().to_vec()) + .bind(public_key.as_bytes()) .bind(iota_id) .bind(token.as_bytes().to_vec()) .execute(&pool) @@ -589,13 +593,13 @@ pub async fn print_users() -> Result<(), Box> { // ========================================================================================== // IOTA // ========================================================================================== -pub async fn create_new_iota(public_key: String) -> Result { +pub async fn create_new_iota(public_key: PublicKeyBundle) -> Result { let new_id = get_register_id().await as i64; register_complete_iota(new_id, public_key).await?; Ok(new_id) } -pub async fn register_complete_iota(id: i64, public_key: String) -> Result<(), sqlx::Error> { +pub async fn register_complete_iota(id: i64, public_key: PublicKeyBundle) -> Result<(), sqlx::Error> { let pool = { let db_lock = SQL_DB.read().await; db_lock @@ -606,14 +610,14 @@ pub async fn register_complete_iota(id: i64, public_key: String) -> Result<(), s sqlx::query("INSERT INTO iotas (id, public_key) VALUES (?, ?)") .bind(id) - .bind(public_key.as_bytes().to_vec()) + .bind(public_key.as_bytes()) .execute(&pool) .await?; Ok(()) } -pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> { +pub async fn get_iota_by_id(id: i64) -> Result<(i64, PublicKeyBundle), sqlx::Error> { let pool = { let db_lock = SQL_DB.read().await; db_lock @@ -631,7 +635,9 @@ pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> { match result { Ok(optional_row) => match optional_row { Some((id_i64, public_key)) => { - Ok((id_i64, String::from_utf8_lossy(&public_key).to_string())) + let bundle = PublicKeyBundle::from_bytes(&public_key) + .map_err(|e| sqlx::Error::Decode(Box::new(e)))?; + Ok((id_i64, bundle)) } _ => Err(sqlx::Error::RowNotFound), }, @@ -639,7 +645,7 @@ pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> { } } -pub async fn change_iota_key(id: i64, new_key: String) -> Result<(), sqlx::Error> { +pub async fn change_iota_key(id: i64, new_key: PublicKeyBundle) -> Result<(), sqlx::Error> { let pool = { let db_lock = SQL_DB.read().await; db_lock @@ -649,7 +655,7 @@ pub async fn change_iota_key(id: i64, new_key: String) -> Result<(), sqlx::Error }; sqlx::query("UPDATE iotas SET public_key = ? WHERE id = ?") - .bind(new_key.as_bytes().to_vec()) + .bind(new_key.as_bytes()) .bind(id) .execute(&pool) .await?; @@ -678,7 +684,7 @@ pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> { // OMIKRONS // ========================================================================================== -pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> { +pub async fn get_omikron_by_id(id: i64) -> Result<(PublicKeyBundle, String), sqlx::Error> { let pool = { let db_lock = SQL_DB.read().await; db_lock @@ -695,10 +701,11 @@ pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> .await?; match row { - Some((public_key, ip_address)) => Ok(( - String::from_utf8_lossy(&public_key).to_string(), - String::from_utf8_lossy(&ip_address).to_string(), - )), + Some((public_key, ip_address)) => { + let bundle = PublicKeyBundle::from_bytes(&public_key) + .map_err(|e| sqlx::Error::Decode(Box::new(e)))?; + Ok((bundle, String::from_utf8_lossy(&ip_address).to_string())) + } _ => Err(sqlx::Error::RowNotFound), } } diff --git a/src/transport/omikron_connection.rs b/src/transport/omikron_connection.rs index 982db94..cb11548 100644 --- a/src/transport/omikron_connection.rs +++ b/src/transport/omikron_connection.rs @@ -1,9 +1,9 @@ use crate::{ - get_keyring, get_public_key_bundle, log, log_cv_in, log_cv_out, log_err, log_in, + load_keyring, log, log_cv_in, log_cv_out, log_err, log_in, server::short_link::add_short_link, sql::{ connection_status::UserStatus, - sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get_omikron_by_id}, + sql::{self, get_by_user_id, get_by_username, get_iota_by_id}, user_online_tracker::{self}, }, transport::omikron_manager, @@ -11,9 +11,12 @@ use crate::{ }; use base64::{Engine as _, engine::general_purpose::STANDARD}; use dashmap::DashMap; -use mtp::{codec::{CommunicationType, CommunicationValue, DataType, DataValue}, host::{HostConfig, MTPHost}}; -use mtp::transport::{Host, Policy, Receiver, SendMode, Sender, host}; -use rand::{Rng, distributions::Alphanumeric}; +use mtp::host::{AuthenticationPolicy, Receiver, Sender}; +use mtp::{ + codec::{CommunicationType, CommunicationValue, DataType, DataValue}, + host::{Host, HostConfig, Policy, SendMode}, +}; +use mtp_crypto::PublicKeyBundle; use std::net::{IpAddr, Ipv4Addr}; use std::{ sync::Arc, @@ -23,8 +26,6 @@ use tokio::{ sync::{Mutex, RwLock}, time::interval, }; -use x448::PublicKey; - // ============================================================================ // Configuration // ============================================================================ @@ -63,35 +64,6 @@ pub struct WaitingTask { pub inserted_at: Instant, } -// ============================================================================ -// Connection State -// ============================================================================ - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum AuthState { - Unauthenticated, - Identified { omikron_id: i64 }, - Authenticated { omikron_id: i64 }, -} - -impl AuthState { - fn is_authenticated(&self) -> bool { - match self { - AuthState::Authenticated { omikron_id } => true, - _ => false, - } - } - - fn omikron_id(&self) -> Option { - match self { - AuthState::Identified { omikron_id } | AuthState::Authenticated { omikron_id } => { - Some(*omikron_id) - } - _ => None, - } - } -} - // ============================================================================ // Omikron Connection (mtp/QUIC-based) // ============================================================================ @@ -99,7 +71,6 @@ impl AuthState { pub struct OmikronConnection { id: u64, sender: Mutex>, - state: RwLock, challenge: RwLock, pub_key: RwLock>>, pub ping: RwLock, @@ -124,7 +95,6 @@ impl OmikronConnection { let conn = Arc::new(Self { id: rand::random(), sender: Mutex::new(Some(sender)), - state: RwLock::new(AuthState::Unauthenticated), challenge: RwLock::new(String::new()), pub_key: RwLock::new(None), ping: RwLock::new(-1), @@ -197,144 +167,9 @@ impl OmikronConnection { if cv.is_type(CommunicationType::Ping) { return self.handle_ping(cv).await; } - - let current_state = *self.state.read().await; - // Route based on authentication state - match current_state { - AuthState::Unauthenticated => self.clone().handle_unauthenticated(cv).await, - AuthState::Identified { .. } => self.clone().handle_identified(cv).await, - AuthState::Authenticated { omikron_id } => { - self.clone().handle_authenticated(cv, omikron_id).await - } - } + return Ok(()); } - // ------------------------------------------------------------------------- - // Authentication Handlers - // ------------------------------------------------------------------------- - - async fn handle_unauthenticated(self: Arc, cv: CommunicationValue) -> OmikronResult<()> { - if !cv.is_type(CommunicationType::Identification) { - let _ = self - .send_error_response(cv.get_id(), CommunicationType::ErrorNotAuthenticated) - .await; - return Err(OmikronError::NotAuthenticated); - } - - // Extract omikron ID - let omikron_id = cv - .get_data(DataType::OmikronId) - .as_number() - .ok_or(OmikronError::InvalidResponse)?; - log!("Omikron {:?} connected", omikron_id); - - // Lookup omikron in database - let (public_key, _) = get_omikron_by_id(omikron_id as i64) - .await - .map_err(|e| OmikronError::Sql(e.to_string()))?; - - log!("Got public Key"); - - let pub_key_bytes = STANDARD - .decode(&public_key) - .map_err(|_| OmikronError::AuthenticationFailed)?; - - let pub_key_bytes_clone = pub_key_bytes.clone(); - let omikron_pub_key = PublicKey::from_bytes(&pub_key_bytes_clone) - .ok_or(OmikronError::AuthenticationFailed)?; - - log!("Decoded public Key"); - - // Generate challenge - let challenge: String = rand::thread_rng() - .sample_iter(&Alphanumeric) - .take(32) - .map(char::from) - .collect(); - - log!("Generated Challenge"); - - *self.challenge.write().await = challenge.clone(); - - log!("Stored Challenge"); - - *self.pub_key.write().await = Some(pub_key_bytes); - - log!("Stored Pubkey"); - - *self.state.write().await = AuthState::Identified { - omikron_id: omikron_id as i64, - }; - - log!("Stored State"); - - let challenge_clone: String = challenge.clone(); - let private_key = get_keyring(); - let public_key_for_encrypt = omikron_pub_key; - - let encrypted = tokio::task::spawn_blocking(move || { - challenge_clone - encrypt(private_key, public_key_for_encrypt, &challenge_clone) - .map_err(|_| OmikronError::AuthenticationFailed) - }) - .await - .map_err(|_| OmikronError::AuthenticationFailed)??; - - log!("Encrypted Challenge"); - - // Send challenge response - let response = CommunicationValue::new(CommunicationType::Challenge) - .with_id(cv.get_id()) - .add_typed_default( - DataType::PublicKey, - DataValue::Str(STANDARD.encode(get_public_key_bundle().as_bytes())), - ) - .add_typed_default(DataType::Content, DataValue::Str(encrypted)); - - log!("Sending Challenge"); - self.send(&response).await - } - - async fn handle_identified(self: Arc, cv: CommunicationValue) -> OmikronResult<()> { - if !cv.is_type(CommunicationType::ChallengeResponse) { - let _ = self - .send_error_response(cv.get_id(), CommunicationType::ErrorNotAuthenticated) - .await; - return Err(OmikronError::NotAuthenticated); - } - - let client_response = cv - .get_data(DataType::Content) - .as_str() - .ok_or(OmikronError::InvalidResponse)?; - - let expected_challenge = self.challenge.read().await.clone(); - - if client_response == expected_challenge { - let omikron_id = self.state.read().await.omikron_id().unwrap_or(0); - *self.state.write().await = AuthState::Authenticated { omikron_id }; - - omikron_manager::add_omikron(self.clone()).await; - - let response = CommunicationValue::new(CommunicationType::IdentificationResponse) - .with_id(cv.get_id()) - .add_typed_default(DataType::Accepted, DataValue::Bool(true)); - - self.clone().send(&response).await?; - log_in!(omikron_id, PrintType::Omega, "Omikron authenticated"); - Ok(()) - } else { - let _ = self - .send_error_response(cv.get_id(), CommunicationType::ErrorInvalidChallenge) - .await; - Err(OmikronError::AuthenticationFailed) - } - } - - // ------------------------------------------------------------------------- - // Authenticated Message Handlers - // ------------------------------------------------------------------------- - async fn handle_authenticated( self: Arc, cv: CommunicationValue, @@ -554,7 +389,7 @@ impl OmikronConnection { Option>, i32, i64, - String, + PublicKeyBundle, String, String, ), @@ -577,7 +412,7 @@ impl OmikronConnection { let mut response = CommunicationValue::new(CommunicationType::GetUserData) .with_id(msg_id) .add_typed_default(DataType::Username, DataValue::Str(username.clone())) - .add_typed_default(DataType::PublicKey, DataValue::Str(public_key)) + .add_typed_default(DataType::PublicKey, DataValue::Str(public_key.to_base64())) .add_typed_default(DataType::UserId, DataValue::SignedNumber(id.into())) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())) .add_typed_default( @@ -704,13 +539,13 @@ impl OmikronConnection { self: Arc, msg_id: u32, iota_id: i64, - public_key: String, + public_key: PublicKeyBundle, user_id: Option, username: Option, ) -> CommunicationValue { let mut response = CommunicationValue::new(CommunicationType::GetIotaData) .with_id(msg_id) - .add_typed_default(DataType::PublicKey, DataValue::Str(public_key)) + .add_typed_default(DataType::PublicKey, DataValue::Str(public_key.to_base64())) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())); if let Some(uid) = user_id { @@ -752,10 +587,15 @@ impl OmikronConnection { ) -> OmikronResult<()> { let iota_id_opt = cv.get_data(DataType::IotaId).as_number().map(|n| n as i64); - if let Some(public_key) = cv.get_data(DataType::PublicKey).as_str() { + let public_key = cv + .get_data(DataType::PublicKey) + .as_str() + .and_then(|s| PublicKeyBundle::from_base64(s).ok()); + + if let Some(public_key) = public_key { if let Some(iota_id) = iota_id_opt { // Register existing IOTA - match sql::register_complete_iota(iota_id, public_key.to_string()).await { + match sql::register_complete_iota(iota_id, public_key).await { Ok(_) => { let response = CommunicationValue::new(CommunicationType::Success) .with_id(cv.get_id()); @@ -770,7 +610,7 @@ impl OmikronConnection { } } else { // Create new IOTA - match sql::create_new_iota(public_key.to_string()).await { + match sql::create_new_iota(public_key).await { Ok(new_iota_id) => { let response = CommunicationValue::new(CommunicationType::CompleteRegisterIota) @@ -807,7 +647,7 @@ impl OmikronConnection { let public_key = cv .get_data(DataType::PublicKey) .as_str() - .map(|s| s.to_string()); + .and_then(|s| PublicKeyBundle::from_base64(s).ok()); let iota_id = cv.get_sender(); let reset_token = cv .get_data(DataType::ResetToken) @@ -873,15 +713,13 @@ impl OmikronConnection { } } if let (Some(public_key), Some(private_key_hash)) = ( - cv.get_data(DataType::PublicKey).as_str(), + cv.get_data(DataType::PublicKey) + .as_str() + .and_then(|s| PublicKeyBundle::from_base64(s).ok()), cv.get_data(DataType::PrivateKeyHash).as_str(), ) { - if let Err(e) = sql::change_keys( - user_id, - public_key.to_string(), - private_key_hash.to_string(), - ) - .await + if let Err(e) = + sql::change_keys(user_id, public_key, private_key_hash.to_string()).await { success = false; error_message = e.to_string(); @@ -1051,9 +889,6 @@ impl OmikronConnection { .with_id(cv.get_id()); let _ = self.send(&response).await; - // Sync with Tauri - crate::notifications::tauri::remove_notification(receiver_id, other_id).await; - // Sync with other Omikron clients let sync_cv = CommunicationValue::new(CommunicationType::ReadNotification) .with_receiver(receiver_id as u64) @@ -1091,9 +926,6 @@ impl OmikronConnection { CommunicationValue::new(CommunicationType::PushNotification).with_id(cv.get_id()); let _ = self.send(&response).await; - // Sync with Tauri - crate::notifications::tauri::send_notification(receiver_id, sender_id).await; - // Sync with other Omikron clients let push_cv = CommunicationValue::new(CommunicationType::PushNotification) .with_receiver(receiver_id as u64) @@ -1189,12 +1021,10 @@ impl OmikronConnection { } async fn cleanup(self: Arc) { - if let Some(omikron_id) = self.state.read().await.omikron_id() { - if omikron_id != 0 { - log_in!(omikron_id, PrintType::Omega, "Omikron disconnected"); - omikron_manager::remove_omikron(omikron_id).await; - user_online_tracker::untrack_omikron(omikron_id).await; - } + if self.id != 0 { + log_in!(self.id as i64, PrintType::Omega, "Omikron disconnected"); + omikron_manager::remove_omikron(self.id as i64).await; + user_online_tracker::untrack_omikron(self.id as i64).await; } if let Some(handle) = self.cleanup_handle.lock().unwrap().take() { @@ -1202,13 +1032,8 @@ impl OmikronConnection { } } - // Public API for external use - pub async fn is_authenticated(self: Arc) -> bool { - self.state.read().await.is_authenticated() - } - pub async fn get_omikron_id(self: Arc) -> Option { - self.state.read().await.omikron_id() + Some(self.id as i64) } pub async fn send_message(self: Arc, cv: &CommunicationValue) -> OmikronResult<()> { @@ -1220,27 +1045,60 @@ impl OmikronConnection { // Server Startup // ============================================================================ +pub async fn get_by_omikron_id( + omikron_id: u64, + _description: Option, +) -> Option { + sql::get_omikron_by_id(omikron_id as i64) + .await + .ok() + .map(|(bundle, _ip_address)| bundle) +} +pub async fn complete_register(pub_key: PublicKeyBundle, description: Option) -> u64 { + 0 +} + pub async fn start(port: u16) -> Result<(), Box> { let cert_pem = load_file_vec("certs", "transport_cert.pem").expect("Error loading Pemfile"); let key_pem = load_file_vec("certs", "transport_key.pem").expect("Error loading Keyfile"); - let mut host: MTPHost = MTPHost::new( - HostConfig::new( - IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)), + let host_config = HostConfig::new( + IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)), port, cert_pem, key_pem, - ) ) - .await?; + .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_authentication( + load_keyring(), + Box::new(|user_id, description| Box::pin(get_by_omikron_id(user_id, description))), + Box::new(|pub_key, description| Box::pin(complete_register(pub_key, description))), + ) + .with_authentication_policy(AuthenticationPolicy::ForceAuthentication); + let mut host: Host = Host::new(host_config).await?; log!("OmikronServer listening on port {}", port); - while let Some((sender, mut receiver)) = host.next().await { + while let Ok(Some(mut connection)) = host.accept().await { tokio::spawn(async move { - let conn = OmikronConnection::new(sender); - conn.handle(&mut receiver).await; + let conn = OmikronConnection::new(connection.sender); + conn.handle(&mut connection.receiver).await; }); }