From 504e5ff015db1c522d6c47c29f1006b93c420b3e Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Mon, 5 Jan 2026 02:17:12 +0100 Subject: [PATCH] Connection With Omega --- Cargo.lock | 74 +++++++ Cargo.toml | 4 +- src/auth/auth_connector.rs | 18 +- src/auth/crypto_helper.rs | 129 ++++++++++++ src/auth/mod.rs | 1 + src/calls/call_group.rs | 8 - src/calls/call_manager.rs | 54 ++--- src/calls/call_util.rs | 8 +- src/calls/caller.rs | 14 +- src/data/communication.rs | 18 +- src/main.rs | 44 ++-- src/omega/mod.rs | 1 + src/omega/omega_connection.rs | 373 ++++++++++++++++++++++++---------- src/omega/ping_pong_task.rs | 42 ++++ src/rho/client_connection.rs | 32 +-- src/rho/iota_connection.rs | 58 +++--- src/rho/rho_manager.rs | 20 +- src/util/logger.rs | 219 ++++++++++++++++++++ src/util/mod.rs | 2 +- src/util/print.rs | 85 -------- 20 files changed, 882 insertions(+), 322 deletions(-) create mode 100644 src/auth/crypto_helper.rs create mode 100644 src/omega/ping_pong_task.rs create mode 100644 src/util/logger.rs delete mode 100644 src/util/print.rs diff --git a/Cargo.lock b/Cargo.lock index 55de0ff..ad9b22e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6,6 +6,7 @@ version = 4 name = "Omikron" version = "0.1.0" dependencies = [ + "aes-gcm", "ansi_term", "async-tungstenite", "axum", @@ -38,6 +39,7 @@ dependencies = [ "sys-info", "sysinfo", "tokio", + "tokio-native-tls", "tokio-rustls", "tokio-stream", "tokio-util", @@ -59,6 +61,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common", + "generic-array", +] + [[package]] name = "aes" version = "0.8.4" @@ -70,6 +82,20 @@ dependencies = [ "cpufeatures", ] +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "ghash", + "subtle", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -759,9 +785,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "cxx" version = "1.0.189" @@ -1235,6 +1271,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", +] + [[package]] name = "gio" version = "0.21.4" @@ -2163,6 +2209,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "openssl" version = "0.10.75" @@ -2385,6 +2437,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -3741,6 +3805,16 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + [[package]] name = "untrusted" version = "0.9.0" diff --git a/Cargo.toml b/Cargo.toml index 38abf8f..e72c80e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ edition = "2024" [dependencies] ansi_term = "*" -async-tungstenite = { version = "0.32.0", features = ["futures-03-sink", "futures-util", "handshake", "__rustls-tls", "async-native-tls", "async-std", "async-std-runtime", "async-tls", "gio", "gio-runtime", "glib", "openssl", "real-async-native-tls", "real-async-tls", "real-native-tls", "real-tokio-native-tls", "real-tokio-openssl", "real-tokio-rustls", "rustls-native-certs", "rustls-pki-types", "tokio", "tokio-native-tls", "tokio-openssl", "tokio-runtime", "tokio-rustls-manual-roots", "tokio-rustls-native-certs", "tokio-rustls-webpki-roots", "url", "verbose-logging", "webpki-roots" ] } +async-tungstenite = { version = "0.32.0", features = ["futures-03-sink", "futures-util", "handshake", "__rustls-tls", "async-native-tls", "async-std", "async-std-runtime", "async-tls", "gio", "gio-runtime", "glib", "openssl", "real-async-native-tls", "real-async-tls", "real-native-tls", "real-tokio-native-tls", "real-tokio-openssl", "real-tokio-rustls", "rustls-native-certs", "rustls-pki-types", "tokio", "tokio-native-tls", "tokio-openssl", "tokio-runtime", "tokio-rustls-manual-roots", "tokio-rustls-native-certs", "tokio-rustls-webpki-roots", "url", "verbose-logging", "webpki-roots"] } axum = "*" base64 = "0.22.1" bytes = "*" @@ -49,3 +49,5 @@ log = "0.4" livekit = "0.7.25" livekit-api = "0.4.10" dotenv = "0.15.0" +aes-gcm = "0.10.3" +tokio-native-tls = "0.3.1" diff --git a/src/auth/auth_connector.rs b/src/auth/auth_connector.rs index c52a9fd..30da00e 100644 --- a/src/auth/auth_connector.rs +++ b/src/auth/auth_connector.rs @@ -1,6 +1,7 @@ use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; +use crate::log; use crate::util::config_util::CONFIG; -use crate::util::print::{PrintType, line}; +use crate::util::logger::PrintType; use json::number::Number; use reqwest::{Client, Response}; use std::time::Duration; @@ -26,6 +27,19 @@ fn client() -> Client { .build() .unwrap() } +pub async fn get_auth_public_key() -> Option { + let url = format!("https://auth.tensamin.net/api/get/public_key"); + let client = client(); + let res = client.get(&url).send().await.ok()?; + let json = res.text().await.ok()?; + + let cv = CommunicationValue::from_json(&json); + if cv.comm_type != CommunicationType::success { + return None; + } + + Some(cv.get_data(DataTypes::public_key).unwrap().to_string()) +} pub async fn get_user(user_id: Uuid) -> Option { let url = format!("https://auth.tensamin.net/api/get/{}", user_id); @@ -83,7 +97,7 @@ pub async fn get_iota_id(user_id: i64) -> Option { let cv = CommunicationValue::from_json(&json); if cv.comm_type != CommunicationType::success { - line(PrintType::IotaIn, &cv.to_json().to_string()); + log!(PrintType::Iota, "{}", &cv.to_json().to_string()); return None; } diff --git a/src/auth/crypto_helper.rs b/src/auth/crypto_helper.rs new file mode 100644 index 0000000..4f7d17a --- /dev/null +++ b/src/auth/crypto_helper.rs @@ -0,0 +1,129 @@ +use aes_gcm::{ + Aes256Gcm, Nonce, + aead::{Aead, KeyInit, OsRng}, +}; +use base64::{Engine as _, engine::general_purpose::STANDARD}; +use rand_core::RngCore; +use sha2::{Digest, Sha256}; +use x448::{PublicKey, Secret, SharedSecret}; + +/// Errors for crypto operations +#[derive(Debug)] +pub enum CryptoError { + Base64Decode(base64::DecodeError), + InvalidKey, + AgreementError, + EncryptionError(aes_gcm::Error), + DecryptionError(aes_gcm::Error), +} + +impl From for CryptoError { + fn from(err: base64::DecodeError) -> Self { + CryptoError::Base64Decode(err) + } +} + +pub struct KeyPair { + pub secret: Secret, + pub public: PublicKey, +} + +pub fn generate_keypair() -> KeyPair { + let mut buf = [0u8; 56]; + let mut rng = OsRng; + rng.fill_bytes(&mut buf); + let secret = Secret::from_bytes(&buf).unwrap(); + let public = PublicKey::from(&secret); + KeyPair { secret, public } +} + +pub fn public_key_to_base64(pubkey: &PublicKey) -> String { + STANDARD.encode(pubkey.as_bytes().as_ref()) +} + +pub fn secret_key_to_base64(secret: &Secret) -> String { + STANDARD.encode(secret.as_bytes().as_ref()) +} + +pub fn load_public_key(base64_pub: &str) -> Option { + let bytes = STANDARD.decode(base64_pub).unwrap(); + PublicKey::from_bytes(&bytes) +} + +pub fn load_secret_key(base64_secret: &str) -> Option { + let bytes = STANDARD.decode(base64_secret).unwrap(); + Secret::from_bytes(&bytes) +} + +fn derive_aes_key(shared: &SharedSecret) -> [u8; 32] { + let mut hasher = Sha256::new(); + hasher.update(shared.as_bytes()); + let result = hasher.finalize(); + let mut key = [0u8; 32]; + key.copy_from_slice(&result[..32]); + key +} + +pub fn encrypt( + base64_secret: &str, + base64_peer_pub: &str, + plaintext: &str, +) -> Result { + let secret = load_secret_key(base64_secret).unwrap(); + let peer_pub = load_public_key(base64_peer_pub).unwrap(); + let shared = secret + .to_diffie_hellman(&peer_pub) + .ok_or(CryptoError::AgreementError)?; + let key_bytes = derive_aes_key(&shared); + let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct"); + let mut nonce_bytes = [0u8; 12]; + OsRng.fill_bytes(&mut nonce_bytes); + let nonce = Nonce::from_slice(&nonce_bytes); + let ciphertext = cipher + .encrypt(nonce, plaintext.as_bytes()) + .map_err(CryptoError::EncryptionError)?; + // prefix nonce to ciphertext + let mut out = Vec::with_capacity(nonce_bytes.len() + ciphertext.len()); + out.extend_from_slice(&nonce_bytes); + out.extend_from_slice(&ciphertext); + Ok(STANDARD.encode(&out)) +} + +pub fn decrypt( + base64_secret: &str, + base64_peer_pub: &str, + encrypted_base64: &str, +) -> Result { + let secret = load_secret_key(base64_secret).unwrap(); + let peer_pub = load_public_key(base64_peer_pub).unwrap(); + let shared = secret + .to_diffie_hellman(&peer_pub) + .ok_or(CryptoError::AgreementError)?; + let key_bytes = derive_aes_key(&shared); + let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct"); + + let encrypted = STANDARD.decode(encrypted_base64)?; + if encrypted.len() < 12 { + return Err(CryptoError::DecryptionError(aes_gcm::Error)); + } + let nonce_bytes = &encrypted[..12]; + let ciphertext = &encrypted[12..]; + let nonce = Nonce::from_slice(nonce_bytes); + let plaintext_bytes = cipher + .decrypt(nonce, ciphertext) + .map_err(CryptoError::DecryptionError)?; + let plaintext = String::from_utf8(plaintext_bytes) + .map_err(|_| CryptoError::DecryptionError(aes_gcm::Error))?; + Ok(plaintext) +} + +pub fn hash_it(input: &str) -> Vec { + let mut hasher = Sha256::new(); + hasher.update(input.as_bytes()); + hasher.finalize().to_vec() +} + +pub fn hex_hash(input: &str) -> String { + let digest = hash_it(input); + digest.iter().map(|b| format!("{:02x}", b)).collect() +} diff --git a/src/auth/mod.rs b/src/auth/mod.rs index 7b384b2..8eef5c3 100644 --- a/src/auth/mod.rs +++ b/src/auth/mod.rs @@ -1 +1,2 @@ pub mod auth_connector; +pub mod crypto_helper; diff --git a/src/calls/call_group.rs b/src/calls/call_group.rs index 286e6f5..3880046 100644 --- a/src/calls/call_group.rs +++ b/src/calls/call_group.rs @@ -19,12 +19,4 @@ impl CallGroup { show: RwLock::new(true), } } - - pub async fn add_member(self: Arc, member: i64, inviter: i64) { - *self.show.write().await = true; - self.members - .write() - .await - .push(Arc::new(Caller::new(member, self.call_id, inviter))); - } } diff --git a/src/calls/call_manager.rs b/src/calls/call_manager.rs index ad6eb1e..622c108 100644 --- a/src/calls/call_manager.rs +++ b/src/calls/call_manager.rs @@ -6,24 +6,12 @@ use uuid::Uuid; use crate::{ calls::{call_group::CallGroup, caller::Caller}, - util::print::{PrintType, line}, + log, + util::logger::PrintType, }; static CALL_GROUPS: Lazy>>> = Lazy::new(|| RwLock::new(Vec::new())); -pub async fn get_call_invites(user_id: i64) -> Vec> { - let mut callers = Vec::new(); - for cg in CALL_GROUPS.read().await.iter() { - let members = cg.members.read().await; - for member in members.iter() { - if member.user_id == user_id { - callers.push(member.clone()); - } - } - } - callers -} - pub async fn get_call_groups(user_id: i64) -> Vec> { let mut call_groups = Vec::new(); for cg in CALL_GROUPS.read().await.iter() { @@ -45,38 +33,28 @@ pub async fn get_call_token(user_id: i64, call_id: Uuid) -> Option { call_groups.iter().find(|g| g.call_id == call_id).cloned() }; + // if the group exists if let Some(cg) = existing_group { - let mut members = cg.members.write().await; + let members = cg.members.write().await; + // if the user is already a member if let Some(member) = members.iter().find(|m| m.user_id == user_id) { return Some(member.create_token()); } - + return None; + /* let new_caller = Arc::new(Caller::new(user_id, call_id, user_id)); let token = new_caller.create_token(); members.push(new_caller); return Some(token); + */ } let mut call_groups = CALL_GROUPS.write().await; - if let Some(cg) = call_groups.iter().find(|g| g.call_id == call_id) { - let cg_clone = cg.clone(); - drop(call_groups); - - let mut members = cg_clone.members.write().await; - if let Some(member) = members.iter().find(|m| m.user_id == user_id) { - return Some(member.create_token()); - } - let new_caller = Arc::new(Caller::new(user_id, call_id, user_id)); - let token = new_caller.create_token(); - members.push(new_caller); - return Some(token); - } - - let caller = Arc::new(Caller::new(user_id, call_id, user_id)); + let caller = Arc::new(Caller::new(user_id, call_id, true)); let call_group = CallGroup::new(call_id, caller.clone()); call_groups.push(Arc::new(call_group)); @@ -99,7 +77,7 @@ pub async fn add_invite(call_id: Uuid, inviter_id: i64, invitee_id: i64) -> bool if is_inviter_member { if !members.iter().any(|m| m.user_id == invitee_id) { - members.push(Arc::new(Caller::new(invitee_id, call_id, inviter_id))); + members.push(Arc::new(Caller::new(invitee_id, call_id, false))); } return true; } @@ -138,13 +116,11 @@ pub async fn clean_calls() { let size_post = call_groups.len(); drop(call_groups); if size_pre - size_post != 0 { - line( - PrintType::CallIn, - &format!( - "Cleaned {} calls, {} remaining", - size_pre - size_post, - size_post - ), + log!( + PrintType::Call, + "Cleaned {} calls, {} remaining", + size_pre - size_post, + size_post ); } } diff --git a/src/calls/call_util.rs b/src/calls/call_util.rs index 7a07ccd..2d06718 100644 --- a/src/calls/call_util.rs +++ b/src/calls/call_util.rs @@ -2,7 +2,11 @@ use livekit_api::access_token; use std::env; use uuid::Uuid; -pub fn create_token(user_id: i64, call_id: Uuid) -> Result { +pub fn create_token( + user_id: i64, + call_id: Uuid, + has_admin: bool, +) -> Result { let api_key = env::var("LIVEKIT_API_KEY").expect("LIVEKIT_API_KEY is not set"); let api_secret = env::var("LIVEKIT_API_SECRET").expect("LIVEKIT_API_SECRET is not set"); @@ -11,9 +15,11 @@ pub fn create_token(user_id: i64, call_id: Uuid) -> Result, + pub has_admin: bool, } impl Caller { - pub fn new(user_id: i64, call_id: Uuid, inviter_id: i64) -> Self { + pub fn new(user_id: i64, call_id: Uuid, has_admin: bool) -> Self { Caller { user_id, call_id, - inviters: vec![inviter_id], + has_admin, } } + pub fn set_admin(&mut self, has_admin: bool) { + self.has_admin = has_admin; + } + pub fn has_admin(&self) -> bool { + self.has_admin + } pub fn create_token(&self) -> String { - if let Ok(token) = call_util::create_token(self.user_id, self.call_id) { + if let Ok(token) = call_util::create_token(self.user_id, self.call_id, self.has_admin()) { token } else { String::new() diff --git a/src/data/communication.rs b/src/data/communication.rs index 2f8c689..7af1340 100644 --- a/src/data/communication.rs +++ b/src/data/communication.rs @@ -36,6 +36,8 @@ pub enum DataTypes { shared_secret, call_id, call_token, + untill, + enable, start_date, end_date, receiver_id, @@ -110,6 +112,8 @@ impl DataTypes { "sharedsecret" => DataTypes::shared_secret, "callid" => DataTypes::call_id, "calltoken" => DataTypes::call_token, + "untill" => DataTypes::untill, + "enable" => DataTypes::enable, "startdate" => DataTypes::start_date, "enddate" => DataTypes::end_date, "receiverid" => DataTypes::receiver_id, @@ -159,11 +163,14 @@ impl DataTypes { pub enum CommunicationType { error, error_invalid_user_id, + error_invalid_omikron_id, error_not_found, + error_not_authenticated, error_no_iota, error_invalid_challenge, error_invalid_secret, error_invalid_private_key, + error_invalid_public_key, error_no_user_id, error_no_call_id, error_invalid_call_id, @@ -210,6 +217,9 @@ pub enum CommunicationType { watch_stream, call_token, call_invite, + call_disconnect_user, + call_timeout_user, + call_set_anonymous_joining, end_call, function, update, @@ -223,14 +233,20 @@ impl CommunicationType { "watchstream" => CommunicationType::watch_stream, "calltoken" => CommunicationType::call_token, "callinvite" => CommunicationType::call_invite, + "calldisconnectuser" => CommunicationType::call_disconnect_user, + "calltimeoutuser" => CommunicationType::call_timeout_user, + "callsetanonymousjoining" => CommunicationType::call_set_anonymous_joining, "endcall" => CommunicationType::end_call, "function" => CommunicationType::function, "update" => CommunicationType::update, "createuser" => CommunicationType::create_user, "errorinvaliduserid" => CommunicationType::error_invalid_user_id, + "errorinvalidomikronid" => CommunicationType::error_invalid_omikron_id, "errornotfound" => CommunicationType::error_not_found, + "errornotauthenticated" => CommunicationType::error_not_authenticated, "errornoiota" => CommunicationType::error_no_iota, "errorinvalidchallenge" => CommunicationType::error_invalid_challenge, + "errorinvalidpublickey" => CommunicationType::error_invalid_public_key, "errorinvalidsecret" => CommunicationType::error_invalid_secret, "errorinvalidprivatekey" => CommunicationType::error_invalid_private_key, "errornouserid" => CommunicationType::error_no_user_id, @@ -243,7 +259,7 @@ impl CommunicationType { "message" => CommunicationType::message, "messagesend" => CommunicationType::message_send, "messagelive" => CommunicationType::message_live, - "messageother_iota" => CommunicationType::message_other_iota, + "messageotheriota" => CommunicationType::message_other_iota, "messagechunk" => CommunicationType::message_chunk, "messagesget" => CommunicationType::messages_get, "changeconfirm" => CommunicationType::change_confirm, diff --git a/src/main.rs b/src/main.rs index 491cb3c..d641ef4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,33 +8,46 @@ mod util; use async_tungstenite::accept_hdr_async; use dotenv::dotenv; use futures::StreamExt; -use std::sync::Arc; +use once_cell::sync::Lazy; +use std::{env, sync::Arc}; use tokio::net::TcpListener; use tokio_util::compat::TokioAsyncReadCompatExt; use tungstenite::handshake::server::{Request, Response}; use crate::{ + auth::crypto_helper::{load_public_key, load_secret_key}, calls::call_manager::garbage_collect_calls, omega::omega_connection::OmegaConnection, rho::{client_connection::ClientConnection, iota_connection::IotaConnection}, util::{ config_util::CONFIG, - print::{PrintType, line, line_err}, + logger::{PrintType, startup}, }, }; +static PRIVATE_KEY: Lazy = Lazy::new(|| env::var("PRIVATE_KEY").unwrap()); +pub fn get_private_key() -> x448::Secret { + load_secret_key(&*PRIVATE_KEY).unwrap() +} +static PUBLIC_KEY: Lazy = Lazy::new(|| env::var("PUBLIC_KEY").unwrap()); +pub fn get_public_key() -> x448::PublicKey { + load_public_key(&*PUBLIC_KEY).unwrap() +} + #[tokio::main] async fn main() { dotenv().ok(); tokio::spawn(async move { - OmegaConnection::new().connect().await; + Arc::new(OmegaConnection::new()).connect(); }); + startup(); let address = format!("{}:{}", &CONFIG.read().await.ip, &CONFIG.read().await.port); let listener = TcpListener::bind(&address).await.unwrap(); - line( + log!( PrintType::General, - &format!("WebSocket server listening on {}", &address), + "WebSocket server listening on {}", + address, ); garbage_collect_calls(); @@ -50,16 +63,13 @@ async fn main() { let ws_stream = match accept_hdr_async(stream.compat(), callback).await { Ok(ws) => ws, Err(e) => { - line_err( - PrintType::General, - &format!("WebSocket upgrade failed: {}", e), - ); + log!(PrintType::General, "WebSocket upgrade failed: {}", e,); return; } }; let (sender, receiver) = ws_stream.split(); if path == "/ws/client/" { - line(PrintType::ClientIn, "New Client connection"); + log_in!(PrintType::Client, "New Client connection"); let client_conn: Arc = Arc::from(ClientConnection::new(sender, receiver)); loop { @@ -74,25 +84,25 @@ async fn main() { let text = msg.into_text().unwrap(); client_conn.clone().handle_message(text).await; } else if msg.is_close() { - line(PrintType::ClientIn, "Client disconnected"); + log_in!(PrintType::Client, "Client disconnected"); client_conn.handle_close().await; return; } } Some(Err(e)) => { - line_err(PrintType::ClientIn, &format!("WebSocket error: {}", e)); + log_err!(PrintType::Client, "WebSocket error: {}", e); client_conn.handle_close().await; return; } None => { - line(PrintType::ClientIn, "Client stream ended"); + log_in!(PrintType::Client, "Client stream ended"); client_conn.handle_close().await; return; } } } } else if path == "/ws/iota/" { - line(PrintType::IotaIn, "New Iota connection"); + log_in!(PrintType::Iota, "New Iota connection"); let iota_conn: Arc = Arc::from(IotaConnection::new(sender, receiver)); loop { @@ -107,19 +117,19 @@ async fn main() { let text = msg.into_text().unwrap(); iota_conn.clone().handle_message(text).await; } else if msg.is_close() { - line(PrintType::IotaIn, "Iota disconnected"); + log_in!(PrintType::Iota, "Iota disconnected"); iota_conn.handle_close().await; return; } } Some(Err(e)) => { - line_err(PrintType::IotaIn, &format!("WebSocket error: {}", e)); + log_err!(PrintType::Iota, "WebSocket error: {}", e); iota_conn.handle_close().await; return; } None => { // Stream ended - line(PrintType::IotaIn, "Iota stream ended"); + log_in!(PrintType::Iota, "Iota stream ended"); iota_conn.handle_close().await; return; } diff --git a/src/omega/mod.rs b/src/omega/mod.rs index cbe77d9..cb00438 100644 --- a/src/omega/mod.rs +++ b/src/omega/mod.rs @@ -1 +1,2 @@ pub mod omega_connection; +pub mod ping_pong_task; diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index 107f75f..0fdcaaf 100644 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -1,75 +1,251 @@ -use std::sync::Arc; -use std::time::Duration; - -use crate::data::communication::{CommunicationType, CommunicationValue}; -use crate::util::print::PrintType; -use crate::util::print::{line, line_err}; -use crate::{ - data::{ - communication::DataTypes, - user::{User, UserStatus}, - }, - rho::rho_manager, - util::config_util::CONFIG, +use async_tungstenite::{ + WebSocketReceiver, WebSocketSender, WebSocketStream, + stream::Stream, + tokio::{TokioAdapter, connect_async}, + tungstenite::protocol::Message, }; -use async_tungstenite::tungstenite::protocol::Message; use dashmap::DashMap; -use futures::StreamExt; -use json::JsonValue; +use futures::prelude::*; +use json::{JsonValue, number::Number}; use once_cell::sync::Lazy; -use tokio::sync::Mutex; -use tokio::time::sleep; -use tokio_util::compat::Compat; -use tungstenite::{Utf8Bytes, connect}; +use std::{collections::HashMap, env, sync::Arc, time::Duration}; +use tokio::{ + net::TcpStream, + sync::{Mutex, RwLock}, + time::{Instant, sleep}, +}; +use tokio_native_tls::TlsStream; use uuid::Uuid; -static WAITING_TASKS: Lazy bool + Send + Sync>>> = - Lazy::new(DashMap::new); +use crate::{ + auth::crypto_helper::decrypt, + data::{ + communication::{CommunicationType, CommunicationValue, DataTypes}, + user::UserStatus, + }, + get_private_key, log_in, log_out, + rho::rho_manager, + util::logger::PrintType, +}; +use crate::{auth::crypto_helper::secret_key_to_base64, log_err}; +static WAITING_TASKS: Lazy< + DashMap, CommunicationValue) -> bool + Send + Sync>>, +> = Lazy::new(DashMap::new); + +static GENERIC_TASK: Lazy< + Mutex, CommunicationValue) -> bool + Send + Sync>>>, +> = Lazy::new(|| Mutex::new(None)); + +static OMEGA_CONNECTION: Lazy> = Lazy::new(|| { + let conn = Arc::new(OmegaConnection::new()); + let conn_clone = conn.clone(); + tokio::spawn(async move { + conn_clone.connect_internal(0).await; + }); + conn +}); + +pub fn get_omega_connection() -> Arc { + OMEGA_CONNECTION.clone() +} + +#[derive(Clone)] pub struct OmegaConnection { - ws_stream: - Arc>>>>, + write: Arc< + RwLock< + Option< + WebSocketSender< + Stream, TokioAdapter>>, + >, + >, + >, + >, + read: Arc< + RwLock< + Option< + WebSocketReceiver< + Stream, TokioAdapter>>, + >, + >, + >, + >, + pingpong: Arc>>>, + pub last_ping: Arc>, + pub message_send_times: Arc>>, + pub is_connected: Arc>, } impl OmegaConnection { pub fn new() -> Self { OmegaConnection { - ws_stream: Arc::new(Mutex::new(None)), + read: Arc::new(RwLock::new(None)), + write: Arc::new(RwLock::new(None)), + pingpong: Arc::new(Mutex::new(None)), + last_ping: Arc::new(Mutex::new(-1)), + message_send_times: Arc::new(Mutex::new(HashMap::new())), + is_connected: Arc::new(RwLock::new(false)), } } - - pub async fn connect(&self) { - self.connect_internal(0).await; + pub fn connect(self: Arc) { + let cloned_self = self.clone(); + tokio::spawn(async move { + cloned_self.connect_internal(0).await; + }); } - async fn connect_internal(&self, mut retry: usize) { + async fn connect_internal(self: Arc, mut retry: usize) { loop { if retry > 5 { - line_err( - PrintType::OmegaIn, - &"Max retry attempts reached, giving up.", - ); + log_err!(PrintType::Omega, "Max retry attempts reached, giving up."); return; } - match connect("wss://tensamin.methanium.net/ws/omega") { - Ok((_, _)) => { + let url_str = + env::var("OMEGA_HOST").unwrap_or("wss://omega.tensamin.net/ws/omikron".to_string()); + match connect_async(&url_str).await { + Ok((ws_stream, _)) => { + *self.is_connected.write().await = true; + log_in!(PrintType::Omega, "WebSocket connected to {}", url_str); retry = 0; - let identify_msg = CommunicationValue::new(CommunicationType::identification) - .add_data( - DataTypes::uuid, - JsonValue::String(CONFIG.read().await.omikron_id.to_string()), - ); - self.send_message(&identify_msg).await; + let (write, read) = ws_stream.split(); + *self.read.write().await = Some(read); + *self.write.write().await = Some(write); - let ws_stream_clone = self.ws_stream.clone(); + let cloned_self = self.clone(); tokio::spawn(async move { - OmegaConnection::read_loop(ws_stream_clone).await; + cloned_self.clone().read_loop().await; }); + + let cloned_self = self.clone(); + tokio::spawn(async move { + let id = Uuid::new_v4(); + let identify_msg = + CommunicationValue::new(CommunicationType::identification) + .with_id(id) + .add_data( + DataTypes::omikron, + JsonValue::Number(Number::from( + env::var("ID") + .unwrap_or("0".to_string()) + .parse::() + .unwrap_or(0), + )), + ); + WAITING_TASKS.insert( + id, + Box::new(|selfc, cv| { + if cv.is_type(CommunicationType::error_not_found) { + log_err!( + PrintType::Omega, + "Identification failed: Omikron ID not found on Omega.", + ); + return false; + } + if !cv.is_type(CommunicationType::challenge) { + return false; + } + tokio::spawn(async move { + let task = async move { + let challenge = cv + .get_data(DataTypes::challenge) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + "Challenge not found or not a string".to_string() + })?; + + let server_pub_key = cv + .get_data(DataTypes::public_key) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + "Public key from server not found or not a string" + .to_string() + })?; + + let decrypted_challenge = decrypt( + &secret_key_to_base64(&get_private_key()), + server_pub_key, + challenge, + ) + .map_err(|e| { + format!("Failed to decrypt challenge: {:?}", e) + })?; + + let response_msg = CommunicationValue::new( + CommunicationType::challenge_response, + ) + .with_id(cv.get_id()) + .add_data( + DataTypes::challenge, + JsonValue::String(decrypted_challenge), + ); + + let response_id = response_msg.get_id(); + WAITING_TASKS.insert( + response_id, + Box::new(|_self, final_cv| { + if !final_cv + .is_type(CommunicationType::identification_response) + { + log_err!( + PrintType::Omega, + "Expected identification_response, got something else.", + ); + return false; + } + + log_err!( + PrintType::Omega, + "Successfully identified with Omega.", + ); + true + }), + ); + + selfc.send_message(&response_msg).await; + + Ok::<(), String>(()) + }; + + if let Err(e) = task.await { + log_err!(PrintType::Omega, "{}", &e); + } + }); + + true + }), + ); + cloned_self.send_message(&identify_msg).await + }); + + let cloned_self = self.clone(); + let handle = tokio::spawn(async move { + loop { + if *cloned_self.is_connected.read().await == false { + break; + } + cloned_self.send_ping().await; + sleep(Duration::from_secs(1)).await; + } + }); + + *self.is_connected.write().await = true; + *self.pingpong.lock().await = Some(handle); + + while *self.is_connected.read().await { + sleep(Duration::from_secs(2)).await; + } + *self.read.write().await = None; + *self.write.write().await = None; + log_err!(PrintType::Omega, "Connection lost. Retrying..."); + retry += 1; + sleep(Duration::from_secs(2)).await; } Err(e) => { - line_err( - PrintType::OmegaIn, - &format!("WebSocket connection failed (attempt {}): {}", retry, e), + log_err!( + PrintType::Omega, + "WebSocket connection failed (attempt {}): {}", + retry + 1, + e, ); retry += 1; sleep(Duration::from_secs(2)).await; @@ -79,72 +255,60 @@ impl OmegaConnection { } } - async fn read_loop( - ws_stream: Arc< - Mutex>>>, - >, - ) { + async fn read_loop(self: Arc) { loop { - let mut lock = ws_stream.lock().await; - let Some(ws) = lock.as_mut() else { - break; + let msg = { + let mut guard = self.read.write().await; + let ws = match guard.as_mut() { + Some(ws) => ws, + None => break, + }; + ws.next().await }; - match ws.next().await { + match msg { Some(Ok(Message::Text(msg))) => { let cv = CommunicationValue::from_json(&msg); + if cv.is_type(CommunicationType::pong) { + self.handle_pong(&cv, true).await; + continue; + } let msg_id = cv.get_id(); - + log_in!(PrintType::Omikron, "{}", &cv.to_json().to_string()); // Handle waiting tasks if let Some(task) = WAITING_TASKS.remove(&msg_id) { - if (task.1)(cv.clone()) { - continue; + if (task.1)(self.clone(), cv.clone()) { + // continue in the read_loop } - } - - // Handle CLIENT_CHANGED - if cv.is_type(CommunicationType::client_changed) { - let iota_id = cv - .get_data(DataTypes::iota_id) - .unwrap() - .as_i64() - .unwrap_or(0); - let user_id = cv - .get_data(DataTypes::user_id) - .unwrap() - .as_i64() - .unwrap_or(0); - let status_str = cv - .get_data(DataTypes::user_state) - .unwrap() - .as_str() - .unwrap(); - let status = UserStatus::from_string(&status_str) - .unwrap_or(UserStatus::iota_offline); - - let user = User::new(iota_id, user_id, status); - for rho_con in rho_manager::get_all_connections().await { - rho_con.are_they_interested(&user).await; + } else { + // Handle generic task + let generic_task_option = GENERIC_TASK.lock().await; + if let Some(generic_task) = generic_task_option.as_ref() { + if generic_task(self.clone(), cv.clone()) { + // continue in the read_loop + } } } } - Some(Ok(Message::Close(_))) | None => { - break; - } - Some(Err(_)) => { - break; - } + Some(Ok(Message::Close(_))) | None => break, + Some(Err(_)) => break, _ => {} } } + + *self.is_connected.write().await = false; + *self.read.write().await = None; + *self.write.write().await = None; } pub async fn send_message(&self, cv: &CommunicationValue) { - let mut guard = self.ws_stream.lock().await; + let mut guard = self.write.write().await; if let Some(ws) = guard.as_mut() { - line(PrintType::OmegaOut, &cv.to_json().to_string()); + if !cv.is_type(CommunicationType::ping) { + log_out!(PrintType::Omega, "{}", &cv.to_json().to_string()); + } let _ = ws - .send(Message::Text(Utf8Bytes::from(cv.to_json().to_string()))) + .send(Message::Text(cv.to_json().to_string().into())) .await; } } @@ -187,25 +351,26 @@ impl OmegaConnection { WAITING_TASKS.insert( msg_id, - Box::new(move |response: CommunicationValue| { - let _ = Box::pin(async move |_: CommunicationValue| { - let rho = rho_manager::get_rho_con_for_user(user_id).await; - if let Some(rho) = rho { - for client in rho.get_client_connections_for_user(user_id).await { - client.send_message(&response).await; + Box::new( + move |_: Arc, response: CommunicationValue| { + let _ = Box::pin(async move |_: CommunicationValue| { + let rho = rho_manager::get_rho_con_for_user(user_id).await; + if let Some(rho) = rho { + for client in rho.get_client_connections_for_user(user_id).await { + client.send_message(&response).await; + } } - } + true + }); true - }); - true - }), + }, + ), ); OmegaConnection::send_global(cv).await; } async fn send_global(cv: CommunicationValue) { - let conn = OmegaConnection::new(); - conn.send_message(&cv).await; + OMEGA_CONNECTION.send_message(&cv).await; } } diff --git a/src/omega/ping_pong_task.rs b/src/omega/ping_pong_task.rs new file mode 100644 index 0000000..95b63f0 --- /dev/null +++ b/src/omega/ping_pong_task.rs @@ -0,0 +1,42 @@ +use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; +use crate::omega::omega_connection::OmegaConnection; +use json::number::Number; +use tokio::time::Instant; +use uuid::Uuid; + +impl OmegaConnection { + pub async fn send_ping(&self) { + let uuid = Uuid::new_v4(); + let send_time = Instant::now(); + + self.message_send_times.lock().await.insert(uuid, send_time); + self.send_ping_message(uuid).await; + } + + pub async fn send_ping_message(&self, uuid: Uuid) { + let ping_message = CommunicationValue::new(CommunicationType::ping) + .with_id(uuid) + .add_data_num( + DataTypes::last_ping, + Number::from(*self.last_ping.lock().await), + ); + + self.send_message(&ping_message).await; + } + + /// Handles incoming pong and calculates latency + pub async fn handle_pong(&self, cv: &CommunicationValue, _log: bool) { + let id = cv.get_id(); + let send_time_opt = { + let queue = self.message_send_times.lock().await; + queue.get(&id).cloned() + }; + + if let Some(send_time) = send_time_opt { + let ping = Instant::now().duration_since(send_time).as_millis() as i64; + self.message_send_times.lock().await.remove(&id); + + *self.last_ping.lock().await = ping as i64; + } + } +} diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index 20f90a4..eb840cf 100644 --- a/src/rho/client_connection.rs +++ b/src/rho/client_connection.rs @@ -9,9 +9,7 @@ use uuid::Uuid; use super::{rho_connection::RhoConnection, rho_manager}; use crate::calls::call_manager; -use crate::util::print::PrintType; -use crate::util::print::line; -use crate::util::print::line_err; +use crate::util::logger::PrintType; use crate::{ auth::auth_connector, // calls::call_manager::CallManager, @@ -21,6 +19,7 @@ use crate::{ }, omega::omega_connection::OmegaConnection, }; +use crate::{log_in, log_out}; /// ClientConnection represents a WebSocket connection from a client device pub struct ClientConnection { @@ -94,17 +93,14 @@ impl ClientConnection { .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await { - line_err( - PrintType::ClientOut, - &format!("Failed to send message to client: {}", e), - ); + log_out!(PrintType::Client, "Failed to send message to client: {}", e,); } } /// Send a CommunicationValue to the client pub async fn send_message(&self, cv: &CommunicationValue) { if !cv.is_type(CommunicationType::pong) { - line(PrintType::ClientOut, &cv.to_json().to_string()); + log_out!(PrintType::Client, "{}", &cv.to_json().to_string()); } self.send_message_str(&cv.to_json().to_string()).await; } @@ -129,7 +125,7 @@ impl ClientConnection { self.handle_ping(cv).await; return; } - line(PrintType::ClientIn, &cv.to_json().to_string()); + log_in!(PrintType::Client, "{}", &cv.to_json().to_string()); // Handle client status changes if cv.is_type(CommunicationType::client_changed) { self.handle_client_changed(cv).await; @@ -181,7 +177,7 @@ impl ClientConnection { return; } } else { - line(PrintType::ClientIn, "Missing private key"); + log_in!(PrintType::Client, "Missing private key"); self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_private_key) .await; return; @@ -246,8 +242,7 @@ impl ClientConnection { async fn handle_client_changed(&self, cv: CommunicationValue) { let user_id = self.get_user_id().await; if let Some(_status_str) = cv.get_data(DataTypes::user_state) { - // Parse user status - this would need to be implemented properly - let user_status = UserStatus::online; // placeholder + let user_status = UserStatus::online; if let Some(rho_conn) = self.get_rho_connection().await { OmegaConnection::client_changed(rho_conn.get_iota_id().await, user_id, user_status) .await; @@ -354,6 +349,19 @@ impl ClientConnection { return; } } + async fn handle_call_timeout_user(&self, cv: CommunicationValue) { + let user_id = cv.get_data(DataTypes::call_id).unwrap(); + let call_id = cv.get_data(DataTypes::user_id).unwrap(); // JA man braucht CALL_ID + } + async fn handle_call_disconnect_user(&self, cv: CommunicationValue) { + let user_id = cv.get_data(DataTypes::call_id).unwrap(); + let call_id = cv.get_data(DataTypes::user_id).unwrap(); // JA man braucht CALL_ID + let untill = cv.get_data(DataTypes::untill).unwrap(); + } + async fn handle_call_set_anonymous_joining(&self, cv: CommunicationValue) { + let call_id = cv.get_data(DataTypes::user_id).unwrap(); + let enable = cv.get_data(DataTypes::enable).unwrap(); + } /// Forward message to Iota async fn forward_to_iota(&self, cv: CommunicationValue) { diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index eea650d..e34521d 100644 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -1,8 +1,9 @@ use crate::calls::call_group::CallGroup; use crate::calls::call_manager; -use crate::util::print::PrintType; -use crate::util::print::line; -use crate::util::print::line_err; +use crate::log_err; +use crate::log_in; +use crate::log_out; +use crate::util::logger::PrintType; use async_tungstenite::WebSocketReceiver; use async_tungstenite::WebSocketSender; use async_tungstenite::tungstenite::Message; @@ -94,17 +95,14 @@ impl IotaConnection { .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await { - line_err( - PrintType::IotaOut, - &format!("Failed to send WebSocket message: {:?}", e), - ); + log_err!(PrintType::Iota, "Failed to send WebSocket message: {:?}", e,); } } /// Send a CommunicationValue to the Iota pub async fn send_message(&self, cv: CommunicationValue) { if !cv.is_type(CommunicationType::pong) { - line(PrintType::IotaOut, &cv.to_json().to_string()); + log_out!(PrintType::Iota, "{}", cv.to_json().to_string()); } self.send_message_str(&cv.to_json().to_string()).await; } @@ -115,7 +113,7 @@ impl IotaConnection { // Handle identification if cv.is_type(CommunicationType::identification) && !self.is_identified().await { - line(PrintType::IotaIn, &cv.to_json().to_string()); + log_in!(PrintType::Iota, "{}", &cv.to_json().to_string()); self.handle_identification(cv).await; return; } @@ -129,7 +127,7 @@ impl IotaConnection { self.handle_ping(cv).await; return; } - line(PrintType::IotaIn, &cv.to_json().to_string()); + log_in!(PrintType::Iota, "{}", &cv.to_json().to_string()); // Handle forwarding to other Iotas or clients let receiver_id = cv.get_receiver(); if !self.get_user_ids().await.contains(&receiver_id) @@ -171,27 +169,26 @@ impl IotaConnection { match id_str.parse::() { Ok(user_id) => { if let Some(auth_iota_id) = auth_connector::get_iota_id(user_id).await { - line( - PrintType::IotaIn, - &format!( - "auth for {} should be {} is {}", - user_id, iota_id, auth_iota_id - ), + log_in!( + PrintType::Iota, + "auth for {} should be {} is {}", + user_id, + iota_id, + auth_iota_id ); if auth_iota_id == iota_id { validated_user_ids.push(user_id); } } else { - line( - PrintType::IotaIn, - &format!("User ID {} not parsed", id_str.trim()), - ); + log_in!(PrintType::Iota, "User ID {} not parsed", user_id); } } Err(e) => { - line( - PrintType::IotaIn, - &format!("Failed to parse '{}' as i64: {:?}", id_str, e), + log_in!( + PrintType::Iota, + "Failed to parse '{}' as i64: {:?}", + id_str, + e, ); } } @@ -273,16 +270,7 @@ impl IotaConnection { let receiver_id = cv.get_receiver(); let sender_id = cv.get_sender(); - if self.get_user_ids().await.contains(&receiver_id) { - if let Some(target_rho) = self.get_rho_connection().await { - target_rho.message_to_iota(cv).await; - } else { - let error = CommunicationValue::new(CommunicationType::error) - .with_id(cv.get_id()) - .with_sender(cv.get_sender()); - self.send_message(error).await; - } - } else if self.get_user_ids().await.contains(&sender_id) { + if self.get_user_ids().await.contains(&sender_id) { if let Some(target_rho) = rho_manager::get_rho_con_for_user(receiver_id).await { target_rho.message_to_iota(cv).await; } else { @@ -329,10 +317,10 @@ impl IotaConnection { // Process contacts and add call information let enriched_contacts = if *empty { if let Some(contacts_data) = cv.get_data(DataTypes::user_ids) { - line(PrintType::CallIn, "Call empty"); + log_in!(PrintType::Call, "Call empty"); contacts_data.clone() } else { - line(PrintType::CallIn, "Call empty No Data"); + log_in!(PrintType::Call, "Call empty No Data"); JsonValue::new_array() } } else { diff --git a/src/rho/rho_manager.rs b/src/rho/rho_manager.rs index 01afb2c..fb9d74a 100644 --- a/src/rho/rho_manager.rs +++ b/src/rho/rho_manager.rs @@ -1,6 +1,7 @@ use super::rho_connection::RhoConnection; -use crate::util::print::PrintType; -use crate::util::print::line; +use crate::log_in; +use crate::log_out; +use crate::util::logger::PrintType; use std::{ collections::HashMap, sync::{Arc, LazyLock}, @@ -12,17 +13,12 @@ pub static RHO_CONNECTIONS: LazyLock> pub async fn get_rho_con_for_user(user_id: i64) -> Option> { let connections = RHO_CONNECTIONS.read().await; - line( - PrintType::ClientIn, - &format!("Checking user ID: {:?}", user_id), - ); + log_in!(PrintType::Client, "Checking user ID: {:?}", user_id,); for rho_connection in connections.values() { - line( - PrintType::ClientIn, - &format!( - "Comparing user IDs: {:?}", - rho_connection.get_user_ids().to_vec() - ), + log_in!( + PrintType::Client, + "Comparing user IDs: {:?}", + rho_connection.get_user_ids().to_vec() ); if rho_connection.get_user_ids().contains(&user_id) { return Some(Arc::clone(rho_connection)); diff --git a/src/util/logger.rs b/src/util/logger.rs new file mode 100644 index 0000000..d537733 --- /dev/null +++ b/src/util/logger.rs @@ -0,0 +1,219 @@ +use std::{ + fs::{self, OpenOptions}, + io::Write, + path::Path, + sync::{OnceLock, mpsc}, + thread, + time::{SystemTime, UNIX_EPOCH}, +}; + +use ansi_term::Color; + +static LOGGER: OnceLock> = OnceLock::new(); + +#[derive(Clone, Copy)] +pub enum PrintType { + Call, + Client, + Iota, + Omikron, + Omega, + General, +} + +struct LogMessage { + timestamp_ms: u128, + sender: Option, + prefix: &'static str, + kind: PrintType, + is_error: bool, + message: String, +} + +/// Initialize the logging subsystem. +/// Must be called exactly once during startup. +pub fn startup() { + let (tx, rx) = mpsc::channel::(); + LOGGER.set(tx).expect("Logger already initialized"); + + thread::spawn(move || { + let log_dir = Path::new("logs"); + fs::create_dir_all(log_dir).expect("Failed to create log directory"); + + let start_ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + let path = log_dir.join(format!("log_{}.txt", start_ts)); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .expect("Failed to open log file"); + + for msg in rx { + let ts = fixed_box(&msg.timestamp_ms.to_string(), 13); + let sender = match msg.sender { + Some(id) => fixed_box(&id.to_string(), 19), + None => fixed_box("", 19), + }; + + let line = format!("{} {} {} {}", ts, sender, msg.prefix, msg.message); + + // Console (ANSI-colored) + println!("{}", colorize(msg.kind, msg.is_error).paint(&line)); + + // File (plain text) + let _ = writeln!(file, "{}", line); + } + }); +} + +fn colorize(kind: PrintType, is_error: bool) -> Color { + if is_error { + return Color::Red; + } + + match kind { + PrintType::Call => Color::Purple, + PrintType::Client => Color::Green, + PrintType::Iota => Color::Yellow, + PrintType::Omikron => Color::Blue, + PrintType::Omega => Color::Cyan, + PrintType::General => Color::White, + } +} + +fn fixed_box(content: &str, width: usize) -> String { + let s: String = content.chars().take(width).collect(); + let len = s.chars().count(); + if len < width { + format!("[{}{}]", " ".repeat(width - len), s) + } else { + s + } +} + +/** Internal async logging entry point. +* Not exposed publicly; all access goes through macros. +*/ +pub fn log_internal( + sender: Option, + kind: PrintType, + prefix: &'static str, + is_error: bool, + message: String, +) { + if let Some(tx) = LOGGER.get() { + let _ = tx.send(LogMessage { + timestamp_ms: SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis(), + sender, + prefix, + kind, + is_error, + message, + }); + } +} +/// Log a general informational message. +#[macro_export] +macro_rules! log { + + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, "", false, format!($($arg)*)) + }; + + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, "", false, format!($($arg)*)) + }; + + // plain + ($($arg:tt)*) => { + $crate::util::logger::log_internal( + None, + $crate::util::logger::PrintType::General, + "", + false, + format!($($arg)*) + ) + }; +} +/// Log an inbound message (`>`). +#[macro_export] +macro_rules! log_in { + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, ">", false, format!($($arg)*)) + }; + + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, ">", false, format!($($arg)*)) + }; + + // plain + ($($arg:tt)*) => { + $crate::util::logger::log_internal( + None, + $crate::util::logger::PrintType::General, + ">", + false, + format!($($arg)*) + ) + }; +} +/// Log an outbound message (`<`). +#[macro_export] +macro_rules! log_out { + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, "<", false, format!($($arg)*)) + }; + + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, "<", false, format!($($arg)*)) + }; + + // plain + ($($arg:tt)*) => { + $crate::util::logger::log_internal( + None, + $crate::util::logger::PrintType::General, + "<", + false, + format!($($arg)*) + ) + }; +} +/// Log an error message (`>>`). +#[macro_export] +macro_rules! log_err { + + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, ">>", true, format!($($arg)*)) + }; + + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, ">>", true, format!($($arg)*)) + }; + + // plain + ($($arg:tt)*) => { + $crate::util::logger::log_internal( + None, + $crate::util::logger::PrintType::General, + ">>", + true, + format!($($arg)*) + ) + }; +} diff --git a/src/util/mod.rs b/src/util/mod.rs index 611b7be..58403c4 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,3 +1,3 @@ pub mod config_util; pub mod file_util; -pub mod print; +pub mod logger; diff --git a/src/util/print.rs b/src/util/print.rs deleted file mode 100644 index 72e8630..0000000 --- a/src/util/print.rs +++ /dev/null @@ -1,85 +0,0 @@ -use ansi_term::Color; - -pub fn print_start_message() { - println!("{}", Color::Yellow.paint("> Iota inbound")); - println!("{}", Color::Purple.paint("< Iota outbound")); - println!("{}", Color::Green.paint("> Client inbound")); - println!("{}", Color::Blue.paint("< Client outbound")); - println!("{}", Color::Red.paint("> Call inbound")); - println!("{}", Color::Red.paint("< Call outbound")); - println!("{}", Color::Cyan.paint("> Omega inbound")); - println!("{}", Color::Cyan.paint("< Omega outbound")); - println!("{}", Color::White.paint("General info")); - println!("{}", Color::White.paint(">> Erros")); -} -pub enum PrintType { - IotaIn, - IotaOut, - OmegaIn, - OmegaOut, - ClientIn, - ClientOut, - CallIn, - CallOut, - General, -} -pub fn line(key: PrintType, message: &str) { - match key { - PrintType::IotaIn => println!( - "{}{}", - Color::Yellow.paint(">"), - Color::Yellow.paint(message) - ), - PrintType::IotaOut => println!( - "{}{}", - Color::Purple.paint("<"), - Color::Purple.paint(message) - ), - PrintType::OmegaIn => println!("{}{}", Color::Cyan.paint(">"), Color::Cyan.paint(message)), - PrintType::OmegaOut => println!("{}{}", Color::Cyan.paint("<"), Color::Cyan.paint(message)), - PrintType::ClientIn => { - println!("{}{}", Color::Green.paint(">"), Color::Green.paint(message)) - } - PrintType::ClientOut => { - println!("{}{}", Color::Blue.paint("<"), Color::Blue.paint(message)) - } - PrintType::CallIn => { - println!("{}{}", Color::Red.paint(">"), Color::Red.paint(message)) - } - PrintType::CallOut => { - println!("{}{}", Color::Red.paint("<"), Color::Red.paint(message)) - } - PrintType::General => println!("{}", Color::White.paint(message)), - } -} -pub fn line_err(key: PrintType, message: &str) { - match key { - PrintType::IotaIn => println!( - "{}{}", - Color::Yellow.paint(">>"), - Color::Yellow.paint(message) - ), - PrintType::IotaOut => println!( - "{}{}", - Color::Purple.paint("<<"), - Color::Purple.paint(message) - ), - PrintType::OmegaIn => println!("{}{}", Color::Cyan.paint(">>"), Color::Cyan.paint(message)), - PrintType::OmegaOut => { - println!("{}{}", Color::Cyan.paint("<<"), Color::Cyan.paint(message)) - } - PrintType::ClientIn => println!( - "{}{}", - Color::Green.paint(">>"), - Color::Green.paint(message) - ), - PrintType::ClientOut => { - println!("{}{}", Color::Blue.paint("<<"), Color::Blue.paint(message)) - } - PrintType::CallIn => println!("{}{}", Color::Red.paint(">>"), Color::Red.paint(message)), - PrintType::CallOut => { - println!("{}{}", Color::Red.paint("<<"), Color::Red.paint(message)) - } - PrintType::General => println!("{}", Color::Red.paint(message)), - } -}