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

MTP migration

Browse source
This commit is contained in:
Alex Emmet 2026-07-03 17:57:31 +02:00
commit 3e7b5121d4
7 changed files with 123 additions and 425 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/main.rs
View file

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

1
src/notifications/mod.rs
View file

@ -1 +0,0 @@
pub mod tauri;

157
src/notifications/tauri.rs
View file

@ -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<Sender>,
}
static TAURI_CONNECTIONS: Lazy<DashMap<i64, Vec<Arc<TauriConnection>>>> = Lazy::new(DashMap::new);
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
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(&current_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());
}
}
}
}

12
src/server/api.rs
View file

@ -68,7 +68,7 @@ pub async fn handle(path: &str, body_string: Option<String>) -> 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<String>) -> 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<String>) -> 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<String>) -> 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<String>) -> 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<String>) -> 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();

63
src/sql/sql.rs
View file

@ -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<Vec<u8>>,
i32,
i64,
String,
PublicKeyBundle,
String,
String,
),
@ -192,7 +193,8 @@ pub async fn get_by_username(
let avatar: Option<Vec<u8>> = row.get("avatar");
let sub_level: i32 = row.get("sub_level");
let sub_end: i64 = row.get("sub_end");
let public_key: Vec<u8> = row.get("public_key");
let public_key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
.map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
let private_key_hash: Vec<u8> = row.get("private_key_hash");
let token: Vec<u8> = 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<Vec<u8>>,
i32,
i64,
String,
PublicKeyBundle,
String,
String,
),
@ -260,7 +262,8 @@ pub async fn get_by_user_id(
let avatar: Option<Vec<u8>> = row.get("avatar");
let sub_level: i32 = row.get("sub_level");
let sub_end: i64 = row.get("sub_end");
let public_key: Vec<u8> = row.get("public_key");
let public_key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
.map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
let private_key_hash: Vec<u8> = row.get("private_key_hash");
let token: Vec<u8> = 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<Vec<u8>>,
i32,
i64,
String,
PublicKeyBundle,
String,
String,
)>,
@ -328,7 +331,8 @@ pub async fn get_users_by_iota_id(
let avatar: Option<Vec<u8>> = row.get("avatar");
let sub_level: i32 = row.get("sub_level");
let sub_end: i64 = row.get("sub_end");
let public_key: Vec<u8> = row.get("public_key");
let public_key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
.map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
let private_key_hash: Vec<u8> = row.get("private_key_hash");
let token: Vec<u8> = 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<dyn std::error::Error>> {
// ==========================================================================================
// IOTA
// ==========================================================================================
pub async fn create_new_iota(public_key: String) -> Result<i64, sqlx::Error> {
pub async fn create_new_iota(public_key: PublicKeyBundle) -> Result<i64, sqlx::Error> {
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),
}
}

288
src/transport/omikron_connection.rs
View file

@ -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<i64> {
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<Option<Sender>>,
state: RwLock<AuthState>,
challenge: RwLock<String>,
pub_key: RwLock<Option<Vec<u8>>>,
pub ping: RwLock<i64>,
@ -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<Self>, 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<Self>, 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<Self>,
cv: CommunicationValue,
@ -554,7 +389,7 @@ impl OmikronConnection {
Option<Vec<u8>>,
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<Self>,
msg_id: u32,
iota_id: i64,
public_key: String,
public_key: PublicKeyBundle,
user_id: Option<i64>,
username: Option<String>,
) -> 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<Self>) {
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<Self>) -> bool {
self.state.read().await.is_authenticated()
}
pub async fn get_omikron_id(self: Arc<Self>) -> Option<i64> {
self.state.read().await.omikron_id()
Some(self.id as i64)
}
pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
@ -1220,27 +1045,60 @@ impl OmikronConnection {
// Server Startup
// ============================================================================
pub async fn get_by_omikron_id(
omikron_id: u64,
_description: Option<String>,
) -> Option<PublicKeyBundle> {
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<String>) -> u64 {
0
}
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
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;
});
}