diff --git a/Cargo.lock b/Cargo.lock index 6bf40cc..45d3a1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7,6 +7,7 @@ name = "Omega" version = "0.1.0" dependencies = [ "aes-gcm", + "ansi_term", "async-tungstenite", "axum", "base64 0.22.1", @@ -109,6 +110,15 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "ansi_term" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2" +dependencies = [ + "winapi", +] + [[package]] name = "arbitrary" version = "1.4.2" diff --git a/Cargo.toml b/Cargo.toml index a1cb5fa..1331a56 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] aes-gcm = "*" +ansi_term = "0.12.1" 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 = "0.8.7" base64 = "0.22.1" 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 01b51b8..bf945e4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -25,9 +25,9 @@ pub fn get_public_key() -> x448::PublicKey { async fn main() { dotenv().ok(); startup(); - log!("Started"); log_in!("Incoming messages"); log_out!("Outgoing messages"); + log!("Started"); log!(" .env"); if let Err(e) = initialize_db().await { log!("[FATAL] Database initialization failed: {}", e); @@ -35,9 +35,14 @@ async fn main() { "[FATAL] Please ensure the database is running and the .env file is configured correctly." ); return; + } else { + log!(" DB"); + } + if let Err(e) = print_users().await { + log!("[ERROR] Failed to print users: {}", e); + } else { + log!(" Users"); } - let _ = print_users().await; - log!(" DB"); server::server::start(9187).await; log!(" Server"); loop {} diff --git a/src/server/api.rs b/src/server/api.rs index db464a2..3b5fbcc 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -65,19 +65,17 @@ pub async fn handle( let id = path_parts[3].parse::().unwrap_or(0); if id == 0 { not_found() - } else if let Ok((omikron_id, public_key, ip_address)) = - get_omikron_by_id(id).await - { + } else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await { ( StatusCode::OK, "application/json", format!( "{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}", - omikron_id, public_key, ip_address + id, public_key, ip_address ), ) } else if let Some(omikron_id) = get_omikron_for_iota(id).await { - if let Ok((omikron_id, public_key, ip_address)) = + if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await { ( @@ -95,7 +93,7 @@ pub async fn handle( get_by_id(id).await { if let Some(omikron_id) = get_omikron_for_iota(iota_id).await { - if let Ok((omikron_id, public_key, ip_address)) = + if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await { ( diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs index 5aafd16..ea2fa85 100644 --- a/src/server/omikron_connection.rs +++ b/src/server/omikron_connection.rs @@ -1,8 +1,9 @@ use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; -use crate::get_public_key; use crate::sql::sql::get_omikron_by_id; use crate::util::crypto_helper::encrypt; -use crate::{get_private_key, log_in_from, log_out_from}; +use crate::util::logger::PrintType; +use crate::{get_private_key, log_out}; +use crate::{get_public_key, log_in}; use base64::{Engine as _, engine::general_purpose::STANDARD}; use dashmap::DashMap; use futures::SinkExt; @@ -53,10 +54,17 @@ impl OmikronConnection { waiting_tasks: DashMap::new(), }) } - pub async fn send_message(&self, message: &CommunicationValue) { + pub async fn send_message(&self, cv: &CommunicationValue) { let mut sender = self.sender.write().await; - let message_text = Message::Text(Utf8Bytes::from(message.to_json().to_string())); - log_out_from!(*self.omikron_id.read().await, "{}", message_text); + let message_text = Message::Text(Utf8Bytes::from(cv.to_json().to_string())); + if !cv.is_type(CommunicationType::pong) { + log_out!( + *self.omikron_id.read().await, + PrintType::Omikron, + "{}", + message_text + ); + } sender.send(message_text).await.unwrap(); } pub async fn get_user_id(&self) -> i64 { @@ -76,155 +84,124 @@ impl OmikronConnection { return; } - log_in_from!(*self.omikron_id.read().await, "{}", message); + log_in!( + *self.omikron_id.read().await, + PrintType::Omikron, + "{}", + message + ); + if let Some((_, task)) = self.waiting_tasks.remove(&cv.get_id()) { + let _ = task(self.clone(), cv.clone()); + return; + } + + // Handle identification if !*self.identified.read().await && cv.is_type(CommunicationType::identification) { let omikron_id = cv .get_data(DataTypes::omikron) .unwrap_or(&JsonValue::Null) .as_i64() .unwrap_or(0); - if let Ok((_, public_key, _)) = get_omikron_by_id(omikron_id).await { - // Generate Challenge, encrypt it and send it to the omikron - *self.omikron_id.write().await = omikron_id; - *self.identified.write().await = true; + match get_omikron_by_id(omikron_id).await { + Ok((public_key, _)) => { + // Generate Challenge, encrypt it and send it to the omikron + *self.omikron_id.write().await = omikron_id; - let challenge_str: String = rand::thread_rng() - .sample_iter(&Alphanumeric) - .take(32) - .map(char::from) - .collect(); + let challenge_str: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(32) + .map(char::from) + .collect(); - *self.challenge.write().await = challenge_str.clone(); + *self.challenge.write().await = challenge_str.clone(); - let user_public_key_bytes = match STANDARD.decode(&public_key) { - Ok(bytes) => bytes, - Err(_) => { - self.send_error_response( - &cv.get_id(), - CommunicationType::error_invalid_user_id, + let user_public_key_bytes = match STANDARD.decode(&public_key) { + Ok(bytes) => bytes, + Err(_) => { + self.send_error_response( + &cv.get_id(), + CommunicationType::error_invalid_omikron_id, + ) + .await; + return; + } + }; + *self.pub_key.write().await = Some(user_public_key_bytes.clone()); + + let omikron_pub_key: PublicKey = + match PublicKey::from_bytes(&user_public_key_bytes) { + Some(key) => key, + None => { + self.send_error_response( + &cv.get_id(), + CommunicationType::error_invalid_public_key, + ) + .await; + return; + } + }; + + let encrypted_challenge = + encrypt(get_private_key(), omikron_pub_key, &challenge_str) + .unwrap_or("".to_string()); + + let response = CommunicationValue::new(CommunicationType::challenge) + .add_data_str( + DataTypes::public_key, + STANDARD.encode(get_public_key().as_bytes()), ) - .await; - return; - } - }; - *self.pub_key.write().await = Some(user_public_key_bytes.clone()); + .add_data_str(DataTypes::challenge, encrypted_challenge) + .with_id(cv.get_id()); - let omikron_pub_key: PublicKey = match PublicKey::from_bytes(&user_public_key_bytes) - { - Some(key) => key, - None => { - self.send_error_response( - &cv.get_id(), - CommunicationType::error_invalid_user_id, - ) - .await; - return; - } - }; - - let encrypted_challenge = - encrypt(get_private_key(), omikron_pub_key, &challenge_str) - .unwrap_or("".to_string()); - - let response = CommunicationValue::new(CommunicationType::challenge) - .add_data_str( - DataTypes::public_key, - STANDARD.encode(get_public_key().as_bytes()), + self.send_message(&response).await; + *self.identified.write().await = true; + return; + } + Err(e) => { + self.send_message( + &CommunicationValue::new(CommunicationType::error_not_authenticated) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, e.to_string()), ) - .add_data_str(DataTypes::challenge, encrypted_challenge) + .await; + + return; + } + } + } + + // Handle challenge response + if *self.identified.read().await + && !*self.challenged.read().await + && cv.is_type(CommunicationType::challenge_response) + { + let client_response = cv + .get_data(DataTypes::challenge) + .unwrap_or(&JsonValue::Null) + .as_str() + .unwrap_or(""); + let expected_challenge = self.challenge.read().await.clone(); + + if client_response == expected_challenge { + *self.challenged.write().await = true; + + let response = CommunicationValue::new(CommunicationType::identification_response) .with_id(cv.get_id()); self.send_message(&response).await; - // prepare Challenge Response handling - self.waiting_tasks.insert( - cv.get_id(), - Box::new( - |selfc: Arc, cv: CommunicationValue| -> bool { - tokio::spawn(async move { - let client_challenge_response_b64 = - match cv.get_data(DataTypes::challenge) { - Some(data) => data.to_string(), - None => { - selfc - .send_error_response( - &cv.get_id(), - CommunicationType::error, - ) - .await; - return; - } - }; - - let challenge_response_bytes = - match STANDARD.decode(&client_challenge_response_b64) { - Ok(bytes) => bytes, - Err(_) => { - selfc - .send_error_response( - &cv.get_id(), - CommunicationType::error, - ) - .await; - return; - } - }; - - if challenge_response_bytes.len() < 12 { - selfc - .send_error_response(&cv.get_id(), CommunicationType::error) - .await; - return; - } - - let client_response = cv - .get_data(DataTypes::challenge) - .unwrap_or(&JsonValue::Null) - .as_str() - .unwrap_or(""); - - let expected_challenge = selfc.challenge.read().await.clone(); - - if client_response != expected_challenge { - selfc - .send_error_response( - &cv.get_id(), - CommunicationType::error_invalid_challenge, - ) - .await; - selfc.close().await; - return; - } - - *selfc.challenged.write().await = true; - - let response = CommunicationValue::new( - CommunicationType::identification_response, - ) - .with_id(cv.get_id()); - - selfc.send_message(&response).await; - return; - }); - return true; - }, - ), - ); } else { - self.send_error_response(&cv.get_id(), CommunicationType::error_not_found) + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_challenge) .await; + self.close().await; } return; } - if self.waiting_tasks.contains_key(&cv.get_id()) { - let (_, task) = self.waiting_tasks.remove(&cv.get_id()).unwrap(); - let _ = task(self.clone(), cv.clone()); - } - if !self.is_identified().await { - self.send_error_response(&cv.get_id(), CommunicationType::error_not_found) + self.send_error_response(&cv.get_id(), CommunicationType::error_not_authenticated) .await; + self.close().await; return; } } diff --git a/src/server/server.rs b/src/server/server.rs index 7dc19ab..263ff88 100644 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -464,10 +464,19 @@ pub async fn start_omikron_handler(connection: Arc) { conn_clone.handle_message(text.to_string()).await; }); } + Message::Ping(_) => { + let pong_response = crate::data::communication::CommunicationValue::new( + crate::data::communication::CommunicationType::pong, + ); + let conn_clone = connection.clone(); + tokio::spawn(async move { + conn_clone.send_message(&pong_response).await; + }); + } Message::Close(_) => { break; } - // Other message types like Binary, Ping, Pong are ignored. + // Other message types like Binary, Pong are ignored. _ => {} } } diff --git a/src/server/socket.rs b/src/server/socket.rs index 2ff5966..6cc825d 100644 --- a/src/server/socket.rs +++ b/src/server/socket.rs @@ -35,6 +35,11 @@ pub fn handle( .clone() .handle_message(text.to_string()) .await; + } else if msg.is_ping() { + let pong_response = crate::data::communication::CommunicationValue::new( + crate::data::communication::CommunicationType::pong, + ); + community_conn.send_message(&pong_response).await; } else if msg.is_close() { log!("Closing: {}", msg); community_conn.handle_close().await; diff --git a/src/sql/sql.rs b/src/sql/sql.rs index de749fa..a5463d7 100644 --- a/src/sql/sql.rs +++ b/src/sql/sql.rs @@ -1,6 +1,6 @@ use crate::log; use once_cell::sync::Lazy; -use sqlx::{MySql, Pool, mysql::MySqlPoolOptions}; +use sqlx::{MySql, Pool, Row, mysql::MySqlPoolOptions}; use std::sync::atomic::{AtomicU64, Ordering}; use std::{ env, @@ -78,7 +78,7 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> { "CREATE TABLE IF NOT EXISTS iotas ( id BIGINT UNSIGNED NOT NULL PRIMARY KEY, - public_key TEXT NOT NULL COLLATE utf8mb4_bin + public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin )", ) .execute(&pool) @@ -88,9 +88,9 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> { omikrons ( id BIGINT UNSIGNED NOT NULL PRIMARY KEY, is_active INT(1) NOT NULL DEFAULT 0, - public_key TEXT NOT NULL COLLATE utf8mb4_bin, - location TEXT NOT NULL COLLATE utf8mb4_bin, - ip_address TEXT NOT NULL COLLATE utf8mb4_bin + public_key VARCHAR(255) NOT NULL COLLATE utf8mb4_bin, + location VARCHAR(255) NOT NULL COLLATE utf8mb4_bin, + ip_address VARCHAR(255) NOT NULL COLLATE utf8mb4_bin )", ) .execute(&pool) @@ -165,12 +165,45 @@ pub async fn get_by_username( let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - sqlx::query_as::<_, (i64, i64, String, String, String, String, String, i32, i64, String, String, String)>( - "SELECT id, iota_id, username, display_name, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE username = ?", + let row = sqlx::query( + "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE username = ?", ) .bind(username) .fetch_optional(pool) - .await?.ok_or_else(|| sqlx::Error::RowNotFound) + .await?; + + match row { + Some(row) => { + let id: i64 = row.get("id"); + let iota_id: i64 = row.get("iota_id"); + let username: String = row.get("username"); + let display: Vec = row.get("display"); + let status: Vec = row.get("status"); + let about: Vec = row.get("about"); + let avatar: Vec = row.get("avatar"); + let sub_level: i32 = row.get("sub_level"); + let sub_end: i64 = row.get("sub_end"); + let public_key: String = row.get("public_key"); + let private_key_hash: String = row.get("private_key_hash"); + let token: Vec = row.get("token"); + + Ok(( + id, + iota_id, + username, + String::from_utf8_lossy(&display).to_string(), + String::from_utf8_lossy(&status).to_string(), + String::from_utf8_lossy(&about).to_string(), + String::from_utf8_lossy(&avatar).to_string(), + sub_level, + sub_end, + public_key, + private_key_hash, + String::from_utf8_lossy(&token).to_string(), + )) + } + None => Err(sqlx::Error::RowNotFound), + } } pub async fn get_by_id( @@ -195,12 +228,45 @@ pub async fn get_by_id( let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - sqlx::query_as::<_, (i64, i64, String, String, String, String, String, i32, i64, String, String, String)>( - "SELECT id, iota_id, username, display_name, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE id = ?", + let row = sqlx::query( + "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE id = ?", ) .bind(id) .fetch_optional(pool) - .await?.ok_or_else(|| sqlx::Error::RowNotFound) + .await?; + + match row { + Some(row) => { + let id: i64 = row.get("id"); + let iota_id: i64 = row.get("iota_id"); + let username: String = row.get("username"); + let display: Vec = row.get("display"); + let status: Vec = row.get("status"); + let about: Vec = row.get("about"); + let avatar: Vec = row.get("avatar"); + let sub_level: i32 = row.get("sub_level"); + let sub_end: i64 = row.get("sub_end"); + let public_key: String = row.get("public_key"); + let private_key_hash: String = row.get("private_key_hash"); + let token: Vec = row.get("token"); + + Ok(( + id, + iota_id, + username, + String::from_utf8_lossy(&display).to_string(), + String::from_utf8_lossy(&status).to_string(), + String::from_utf8_lossy(&about).to_string(), + String::from_utf8_lossy(&avatar).to_string(), + sub_level, + sub_end, + public_key, + private_key_hash, + String::from_utf8_lossy(&token).to_string(), + )) + } + None => Err(sqlx::Error::RowNotFound), + } } pub async fn change_username(id: i64, new_username: String) -> Result<(), sqlx::Error> { @@ -216,12 +282,12 @@ pub async fn change_username(id: i64, new_username: String) -> Result<(), sqlx:: Ok(()) } -pub async fn change_display_name(id: i64, new_display_name: String) -> Result<(), sqlx::Error> { +pub async fn change_display_name(id: i64, new_display: String) -> Result<(), sqlx::Error> { let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - sqlx::query("UPDATE users SET display_name = ? WHERE id = ?") - .bind(new_display_name) + sqlx::query("UPDATE users SET display = ? WHERE id = ?") + .bind(new_display) .bind(id) .execute(pool) .await?; @@ -304,18 +370,18 @@ pub async fn register_complete_user( public_key: String, private_key_hash: String, iota_id: i64, - reset_token: String, + token: String, ) -> Result<(), sqlx::Error> { let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - sqlx::query("INSERT INTO users (id, username, public_key, private_key_hash, iota_id, reset_token) VALUES (?, ?, ?, ?, ?, ?)") + sqlx::query("INSERT INTO users (id, username, public_key, private_key_hash, iota_id, token) VALUES (?, ?, ?, ?, ?, ?)") .bind(id) .bind(username) .bind(public_key) .bind(private_key_hash) .bind(iota_id) - .bind(reset_token) + .bind(token) .execute(pool) .await?; @@ -326,17 +392,40 @@ pub async fn print_users() -> Result<(), Box> { let pool = db_lock.as_ref().expect("Database pool is not initialized"); log!("Printing users..."); - for (id, iota_id, username, display_name, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token) - in sqlx::query_as::<_, (i64, i64, String, String, String, String, String, i32, i64, String, String, String)>( - "SELECT id, iota_id, username, display_name, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users", + for row in sqlx::query( + "SELECT id, iota_id, username, display, status, about, sub_level, sub_end, public_key, private_key_hash, token FROM users", ) .fetch_all(pool) .await? .iter() { + let id: i64 = row.get("id"); + let iota_id: i64 = row.get("iota_id"); + let username: String = row.get("username"); + let display: Vec = row.get("display"); + let status: Vec = row.get("status"); + let about: Vec = row.get("about"); + let sub_level: i32 = row.get("sub_level"); + let sub_end: i64 = row.get("sub_end"); + let public_key: String = row.get("public_key"); + let private_key_hash: String = row.get("private_key_hash"); + let token: Vec = row.get("token"); + log!( "User: {:?}", - (id, iota_id, username, display_name, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token) + ( + id, + iota_id, + username, + String::from_utf8_lossy(&display), + String::from_utf8_lossy(&status), + String::from_utf8_lossy(&about), + sub_level, + sub_end, + public_key, + private_key_hash, + String::from_utf8_lossy(&token) + ) ); } @@ -375,10 +464,11 @@ pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> { let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - sqlx::query_as::<_, (i64, String)>("SELECT id, public_key FROM iotas WHERE id = ?") + sqlx::query_as::<_, (i64, Vec)>("SELECT id, public_key FROM iotas WHERE id = ?") .bind(id) .fetch_optional(pool) .await? + .map(|(id, public_key)| (id, String::from_utf8_lossy(&public_key).to_string())) .ok_or_else(|| sqlx::Error::RowNotFound) } @@ -414,24 +504,38 @@ pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> { pub async fn get_random_omikron() -> Result<(i64, String, String), sqlx::Error> { let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - let a = sqlx::query_as("SELECT id, public_key, ip_address FROM omikrons WHERE is_active = 1 ORDER BY RAND() LIMIT 1") + let row = sqlx::query_as::<_, (i64, Vec, Vec)>("SELECT id, public_key, ip_address FROM omikrons WHERE is_active = 1 ORDER BY RAND() LIMIT 1") .fetch_optional(pool) - .await? - .ok_or_else(|| sqlx::Error::RowNotFound); - // this log isn't printing - log!("Random Omikron: {:?}", a); - a + .await?; + + match row { + Some((id, public_key, ip_address)) => Ok(( + id, + String::from_utf8_lossy(&public_key).to_string(), + String::from_utf8_lossy(&ip_address).to_string(), + )), + None => Err(sqlx::Error::RowNotFound), + } } -pub async fn get_omikron_by_id(id: i64) -> Result<(i64, String, String), sqlx::Error> { +pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> { let db_lock = SQL_DB.read().await; let pool = db_lock.as_ref().expect("Database pool is not initialized"); - sqlx::query_as("SELECT id, public_key, ip_address FROM omikrons WHERE id = ?") - .bind(id) - .fetch_optional(pool) - .await? - .ok_or_else(|| sqlx::Error::RowNotFound) + let row = sqlx::query_as::<_, (Vec, Vec)>( + "SELECT public_key, ip_address FROM omikrons WHERE id = ?", + ) + .bind(id) + .fetch_optional(pool) + .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(), + )), + None => Err(sqlx::Error::RowNotFound), + } } pub async fn set_omikron_active(id: i64, active: bool) -> Result<(), sqlx::Error> { diff --git a/src/util/logger.rs b/src/util/logger.rs index 47173e9..dede826 100644 --- a/src/util/logger.rs +++ b/src/util/logger.rs @@ -6,22 +6,37 @@ use std::{ 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 logger (call once) +/// 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 || { - // Prepare log directory let log_dir = Path::new("logs"); fs::create_dir_all(log_dir).expect("Failed to create log directory"); @@ -31,42 +46,65 @@ pub fn startup() { .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"); - // Dedicated logging loop for msg in rx { - let timestamp_box = fixed_box(&msg.timestamp_ms.to_string(), 13); - - let sender_box = match msg.sender { - Some(id) => fixed_box(&format!("{}", id), 19), + 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!("{} {} {}", timestamp_box, sender_box, msg.message); + let line = format!("{} {} {} {}", ts, sender, msg.prefix, msg.message); - println!("{}", line); + // 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 = content.chars().take(width).collect::(); + let s: String = content.chars().take(width).collect(); let len = s.chars().count(); if len < width { - let mut a = " ".repeat(width - len); - a.push_str(&s); - format!("[{}]", a) + format!("[{}{}]", " ".repeat(width - len), s) } else { s } } -/// Internal function (sync + async safe) -pub fn log_internal(sender: Option, message: String) { + +/** 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() @@ -74,62 +112,108 @@ pub fn log_internal(sender: Option, message: String) { .unwrap() .as_millis(), sender, + prefix, + kind, + is_error, message, }); - } else { - println!("{}", message); } } - +/// Log a general informational message. #[macro_export] macro_rules! log { + // plain ($($arg:tt)*) => { - $crate::util::logger::log_internal(None, format!($($arg)*)) + $crate::util::logger::log_internal( + None, + $crate::util::logger::PrintType::General, + "", + false, + format!($($arg)*) + ) }; -} -#[macro_export] -macro_rules! log_from { - ($sender:expr, $($arg:tt)*) => { - $crate::util::logger::log_internal(Some($sender), format!($($arg)*)) + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, "", false, format!($($arg)*)) + }; + + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, "", false, format!($($arg)*)) }; } +/// Log an inbound message (`>`). #[macro_export] macro_rules! log_in { + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, ">", false, format!($($arg)*)) + }; + + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, ">", false, format!($($arg)*)) + }; + + // plain ($($arg:tt)*) => { $crate::util::logger::log_internal( None, - format!("> {}", format!($($arg)*)) + $crate::util::logger::PrintType::General, + ">", + false, + format!($($arg)*) ) }; } - +/// Log an outbound message (`<`). #[macro_export] macro_rules! log_out { + + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, "<", false, format!($($arg)*)) + }; + + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, "<", false, format!($($arg)*)) + }; + + // plain ($($arg:tt)*) => { $crate::util::logger::log_internal( None, - format!("< {}", format!($($arg)*)) + $crate::util::logger::PrintType::General, + "<", + false, + format!($($arg)*) ) }; } - +/// Log an error message (`>>`). #[macro_export] -macro_rules! log_in_from { - ($sender:expr, $($arg:tt)*) => { - $crate::util::logger::log_internal( - Some($sender), - format!("> {}", format!($($arg)*)) - ) - }; -} +macro_rules! log_err { -#[macro_export] -macro_rules! log_out_from { - ($sender:expr, $($arg:tt)*) => { + // sender + actor + ($sender:expr, $kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(Some($sender), $kind, ">>", true, format!($($arg)*)) + }; + + // actor only + ($kind:expr, $($arg:tt)*) => { + $crate::util::logger::log_internal(None, $kind, ">>", true, format!($($arg)*)) + }; + + // plain + ($($arg:tt)*) => { $crate::util::logger::log_internal( - Some($sender), - format!("< {}", format!($($arg)*)) + None, + $crate::util::logger::PrintType::General, + ">>", + true, + format!($($arg)*) ) }; }