Better Logging

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"

Cargo.toml

@@ -42,3 +42,4 @@ warp = { version = "*" }
 chrono = "*"
 tokio-stream = "*"
 loom = "0.7.2"
+ansi_term = "*"

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;
     }

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,

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<ClientConnection> =
                     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<IotaConnection> = 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;
                         }

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<DashMap<Uuid, Box<dyn Fn(CommunicationValue) -> 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;

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
+        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;

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");
                     }
                 }
             }

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<IotaConnection>,
@@ -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() {

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<Arc<RwLock<HashMap<Uuid, Arc<RhoConnection>>>>> =
     LazyLock::new(|| Arc::new(RwLock::new(HashMap::new())));
 
 pub async fn get_rho_con_for_user(user_id: Uuid) -> Option<Arc<RhoConnection>> {
     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!(
+        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));

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);
+        line(PrintType::ClientIn, &format!("{:?}", json));
-        println!(
+        line(PrintType::ClientIn, &format!("{:?}", json["omikron_id"]));
-            "{:?}",
-            Uuid::parse_str(json["omikron_id"].as_str().unwrap_or_default()).unwrap_or_default()
-        );
         Self {
             omega_server: json["omega_server"].as_str().unwrap_or_default().into(),
             auth_server: json["auth_server"].as_str().unwrap_or_default().into(),

src/util/mod.rs

@@ -1,2 +1,3 @@
 pub mod config_util;
 pub mod file_util;
+pub mod print;

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)),
+    }
+}