From ab707c19076c6cef54328f312feca94efdc52d38 Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Sun, 26 Oct 2025 11:20:00 +0100 Subject: [PATCH] Better Logging --- Cargo.lock | 10 +++++ Cargo.toml | 1 + src/auth/auth_connector.rs | 2 - src/data/communication.rs | 2 +- src/main.rs | 34 ++++++++++++----- src/omega/omega_connection.rs | 16 ++++++-- src/rho/client_connection.rs | 19 ++++++++-- src/rho/iota_connection.rs | 16 ++++++-- src/rho/rho_connection.rs | 26 ++++--------- src/rho/rho_manager.rs | 20 +++++++--- src/util/config_util.rs | 11 +++--- src/util/mod.rs | 1 + src/util/print.rs | 71 +++++++++++++++++++++++++++++++++++ 13 files changed, 176 insertions(+), 53 deletions(-) create mode 100644 src/util/print.rs diff --git a/Cargo.lock b/Cargo.lock index cc4af11..3d6c4f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6,6 +6,7 @@ version = 4 name = "Omikron" version = "0.1.0" dependencies = [ + "ansi_term", "axum", "base64", "bytes", @@ -79,6 +80,15 @@ dependencies = [ "libc", ] +[[package]] +name = "ansi_term" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2" +dependencies = [ + "winapi", +] + [[package]] name = "atomic-waker" version = "1.1.2" diff --git a/Cargo.toml b/Cargo.toml index 434fe00..7c97bdb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,3 +42,4 @@ warp = { version = "*" } chrono = "*" tokio-stream = "*" loom = "0.7.2" +ansi_term = "*" diff --git a/src/auth/auth_connector.rs b/src/auth/auth_connector.rs index 6855a7a..4c8f14f 100644 --- a/src/auth/auth_connector.rs +++ b/src/auth/auth_connector.rs @@ -97,7 +97,6 @@ pub async fn is_private_key_valid(user_id: Uuid, pk_hash: &str) -> bool { ); let client = client(); - println!("Auth: {}", CONFIG.lock().await.omikron_id); let res = client .get(&url) .header("Authorization", CONFIG.lock().await.omikron_id.to_string()) @@ -115,7 +114,6 @@ pub async fn is_private_key_valid(user_id: Uuid, pk_hash: &str) -> bool { }; let cv = CommunicationValue::from_json(&body); - println!("Auth: {}", &body); if cv.comm_type != CommunicationType::success { return false; } diff --git a/src/data/communication.rs b/src/data/communication.rs index 2a771da..58727a0 100644 --- a/src/data/communication.rs +++ b/src/data/communication.rs @@ -214,7 +214,7 @@ impl CommunicationType { "messageotheriota" => CommunicationType::message_other_iota, "messagechunk" => CommunicationType::message_chunk, "messagesget" => CommunicationType::messages_get, - "message_send" => CommunicationType::message_send, + "messagesend" => CommunicationType::message_send, "changeconfirm" => CommunicationType::change_confirm, "confirmreceive" => CommunicationType::confirm_receive, "confirmread" => CommunicationType::confirm_read, diff --git a/src/main.rs b/src/main.rs index 64288cd..0e03ce6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,8 @@ mod omega; mod rho; mod util; +use ansi_term::Color; +use crossterm::style::PrintStyledContent; use futures::StreamExt; use std::sync::Arc; use tokio::net::TcpListener; @@ -14,13 +16,23 @@ use tungstenite::handshake::server::{Request, Response}; use crate::{ omega::omega_connection::OmegaConnection, rho::{client_connection::ClientConnection, iota_connection::IotaConnection}, + util::print::PrintType, + util::print::line, + util::print::line_err, + util::print::print_start_message, }; #[tokio::main] + async fn main() { + print_start_message(); + OmegaConnection::new().connect().await; let listener = TcpListener::bind("0.0.0.0:959").await.unwrap(); - println!("WebSocket server listening on 0.0.0.0:959"); + line( + PrintType::OmegaIn, + "WebSocket server listening on 0.0.0.0:959", + ); while let Ok((stream, _)) = listener.accept().await { tokio::spawn(async move { @@ -32,12 +44,15 @@ async fn main() { let ws_stream = match accept_hdr_async(stream, callback).await { Ok(ws) => ws, Err(e) => { - eprintln!("WebSocket upgrade failed: {}", e); + line_err( + PrintType::General, + &format!("WebSocket upgrade failed: {}", e), + ); return; } }; - println!("New {} connection", path); if path == "/ws/client/" { + line(PrintType::ClientIn, "New Client connection"); let client_conn: Arc = Arc::from(ClientConnection::new(ws_stream)); loop { @@ -52,24 +67,25 @@ async fn main() { let text = msg.into_text().unwrap(); client_conn.clone().handle_message(text).await; } else if msg.is_close() { - println!("Client disconnected"); + line(PrintType::ClientIn, "Client disconnected"); client_conn.handle_close().await; return; } } Some(Err(e)) => { - eprintln!("WebSocket error: {}", e); + line_err(PrintType::ClientIn, &format!("WebSocket error: {}", e)); client_conn.handle_close().await; return; } None => { - println!("Client stream ended"); + line(PrintType::ClientIn, "Client stream ended"); client_conn.handle_close().await; return; } } } } else if path == "/ws/iota/" { + line(PrintType::IotaIn, "New Iota connection"); let iota_conn: Arc = Arc::from(IotaConnection::new(ws_stream)); loop { let msg_result = { @@ -83,19 +99,19 @@ async fn main() { let text = msg.into_text().unwrap(); iota_conn.clone().handle_message(text).await; } else if msg.is_close() { - println!("Iota disconnected"); + line(PrintType::IotaIn, "Iota disconnected"); iota_conn.handle_close().await; return; } } Some(Err(e)) => { - eprintln!("WebSocket error: {}", e); + line_err(PrintType::IotaIn, &format!("WebSocket error: {}", e)); iota_conn.handle_close().await; return; } None => { // Stream ended - println!("Iota stream ended"); + line(PrintType::IotaIn, "Iota stream ended"); iota_conn.handle_close().await; return; } diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs index 8a5b18a..923b865 100644 --- a/src/omega/omega_connection.rs +++ b/src/omega/omega_connection.rs @@ -1,6 +1,10 @@ use std::sync::Arc; use std::time::Duration; +use crate::data::communication::{CommunicationType, CommunicationValue}; +use crate::util::print::PrintType; +use crate::util::print::line; +use crate::util::print::line_err; use crate::{ data::{ communication::DataTypes, @@ -19,8 +23,6 @@ use tokio_tungstenite::{connect_async, tungstenite::protocol::Message}; use tungstenite::Utf8Bytes; use uuid::Uuid; -use crate::data::communication::{CommunicationType, CommunicationValue}; - static WAITING_TASKS: Lazy bool + Send + Sync>>> = Lazy::new(DashMap::new); @@ -41,7 +43,10 @@ impl OmegaConnection { async fn connect_internal(&self, mut retry: usize) { loop { if retry > 5 { - eprintln!("Max retry attempts reached, giving up."); + line_err( + PrintType::OmegaIn, + &"Max retry attempts reached, giving up.", + ); return; } @@ -61,7 +66,10 @@ impl OmegaConnection { }); } Err(e) => { - eprintln!("WebSocket connection failed (attempt {}): {}", retry, e); + line_err( + PrintType::OmegaIn, + &format!("WebSocket connection failed (attempt {}): {}", retry, e), + ); retry += 1; sleep(Duration::from_secs(2)).await; continue; diff --git a/src/rho/client_connection.rs b/src/rho/client_connection.rs index be5d4f6..8a54d3f 100644 --- a/src/rho/client_connection.rs +++ b/src/rho/client_connection.rs @@ -7,6 +7,9 @@ use tungstenite::Utf8Bytes; use uuid::Uuid; use super::{rho_connection::RhoConnection, rho_manager}; +use crate::util::print::PrintType; +use crate::util::print::line; +use crate::util::print::line_err; use crate::{ auth::auth_connector, // calls::call_manager::CallManager, @@ -84,12 +87,18 @@ impl ClientConnection { .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await { - eprintln!("Failed to send message to client: {}", e); + line_err( + PrintType::ClientOut, + &format!("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()); + } self.send_message_str(&cv.to_json().to_string()).await; } @@ -112,7 +121,7 @@ impl ClientConnection { self.handle_ping(cv).await; return; } - + line(PrintType::ClientIn, &cv.to_json().to_string()); // Handle client status changes if cv.is_type(CommunicationType::client_changed) { self.handle_client_changed(cv).await; @@ -132,7 +141,9 @@ impl ClientConnection { } // Forward other messages to Iota - self.forward_to_iota(cv).await; + tokio::spawn(async move { + self.forward_to_iota(cv).await; + }); } /// Handle identification message @@ -173,7 +184,7 @@ impl ClientConnection { return; } } else { - println!("Missing private key"); + line(PrintType::ClientIn, "Missing private key"); self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_private_key) .await; return; diff --git a/src/rho/iota_connection.rs b/src/rho/iota_connection.rs index ce6ccda..d1319fa 100644 --- a/src/rho/iota_connection.rs +++ b/src/rho/iota_connection.rs @@ -1,3 +1,8 @@ +use crate::util::print::PrintType; +use crate::util::print::line; +use crate::util::print::line_err; +use ansi_term::Color; +use futures::FutureExt; use futures::SinkExt; use json::JsonValue; use std::{ @@ -83,14 +88,17 @@ impl IotaConnection { .send(Message::Text(Utf8Bytes::from(message.to_string()))) .await { - eprintln!("Failed to send WebSocket message: {:?}", e); + line_err( + PrintType::IotaOut, + &format!("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) { - println!("{}", cv.to_json().to_string()); + line(PrintType::IotaOut, &cv.to_json().to_string()); } self.send_message_str(&cv.to_json().to_string()).await; } @@ -114,7 +122,7 @@ impl IotaConnection { self.handle_ping(cv).await; return; } - + line(PrintType::IotaIn, &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) @@ -164,6 +172,8 @@ impl IotaConnection { if auth_iota_id == iota_id { validated_user_ids.push(user_id); } + } else { + line(PrintType::IotaIn, "User ID not found"); } } } diff --git a/src/rho/rho_connection.rs b/src/rho/rho_connection.rs index 35d7448..53ec5b4 100644 --- a/src/rho/rho_connection.rs +++ b/src/rho/rho_connection.rs @@ -1,14 +1,16 @@ -use std::collections::HashMap; -use std::sync::Arc; -use tokio::sync::RwLock; -use uuid::Uuid; - use super::{client_connection::ClientConnection, iota_connection::IotaConnection, rho_manager}; use crate::data::{ communication::{CommunicationType, CommunicationValue, DataTypes}, user::UserStatus, }; use crate::omega::omega_connection::OmegaConnection; +use crate::util::print::PrintType; +use crate::util::print::line; +use crate::util::print::line_err; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::RwLock; +use uuid::Uuid; pub struct RhoConnection { iota_connection: Arc, @@ -130,20 +132,8 @@ impl RhoConnection { OmegaConnection::close_iota(self.get_iota_id().await).await; } - /// Send message from Iota to specific client by user ID - pub async fn message_iota_to_client_by_user(&self, user_id: Uuid, message: &str) { - let connections = self.client_connections.read().await; - for connection in connections.iter() { - if let Some(conn_user_id) = connection.get_user_id().await { - if conn_user_id == user_id { - connection.send_message_str(message).await; - } - } - } - } - /// Send message from Iota to specific client - pub async fn message_iota_to_client(&self, cv: CommunicationValue) { + pub async fn message_to_client(&self, cv: CommunicationValue) { if let Some(receiver_id) = Some(cv.get_receiver()) { let connections = self.client_connections.read().await; for connection in connections.iter() { diff --git a/src/rho/rho_manager.rs b/src/rho/rho_manager.rs index c832273..7a26603 100644 --- a/src/rho/rho_manager.rs +++ b/src/rho/rho_manager.rs @@ -1,3 +1,7 @@ +use super::rho_connection::RhoConnection; +use crate::util::print::PrintType; +use crate::util::print::line; +use crate::util::print::line_err; use std::{ collections::HashMap, sync::{Arc, LazyLock}, @@ -5,18 +9,22 @@ use std::{ use tokio::sync::RwLock; use uuid::Uuid; -use super::rho_connection::RhoConnection; - pub static RHO_CONNECTIONS: LazyLock>>>> = LazyLock::new(|| Arc::new(RwLock::new(HashMap::new()))); pub async fn get_rho_con_for_user(user_id: Uuid) -> Option> { let connections = RHO_CONNECTIONS.read().await; - println!("Checking user ID: {:?}", user_id); + line( + PrintType::ClientIn, + &format!("Checking user ID: {:?}", user_id), + ); for rho_connection in connections.values() { - println!( - "Comparing user IDs: {:?}", - rho_connection.get_user_ids().to_vec() + line( + PrintType::ClientIn, + &format!( + "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/config_util.rs b/src/util/config_util.rs index a13bdd6..c792257 100644 --- a/src/util/config_util.rs +++ b/src/util/config_util.rs @@ -1,9 +1,11 @@ use crate::util::file_util::load_file; +use crate::util::print::PrintType; +use crate::util::print::line; +use crate::util::print::line_err; use futures::lock::Mutex; use once_cell::sync::Lazy; use serde::{Deserialize, Serialize}; use uuid::Uuid; - #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Config { pub omega_server: String, @@ -39,11 +41,8 @@ impl Config { } let json = json::parse(&content).unwrap(); - println!("{:?}", json); - println!( - "{:?}", - Uuid::parse_str(json["omikron_id"].as_str().unwrap_or_default()).unwrap_or_default() - ); + line(PrintType::ClientIn, &format!("{:?}", json)); + line(PrintType::ClientIn, &format!("{:?}", json["omikron_id"])); Self { omega_server: json["omega_server"].as_str().unwrap_or_default().into(), auth_server: json["auth_server"].as_str().unwrap_or_default().into(), diff --git a/src/util/mod.rs b/src/util/mod.rs index 5e461b2..611b7be 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,2 +1,3 @@ pub mod config_util; pub mod file_util; +pub mod print; diff --git a/src/util/print.rs b/src/util/print.rs new file mode 100644 index 0000000..4ca77c1 --- /dev/null +++ b/src/util/print.rs @@ -0,0 +1,71 @@ +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::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, + 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::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::General => println!("{}", Color::Red.paint(message)), + } +}