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

Better Logging

Browse source
This commit is contained in:
Alex Emmet 2025-10-26 11:20:00 +01:00
commit ab707c1907
13 changed files with 176 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

10
Cargo.lock generated
View file

@ -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"

1
Cargo.toml
View file

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

2
src/auth/auth_connector.rs
View file

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

2
src/data/communication.rs
View file

@ -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,

34
src/main.rs
View file

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

16
src/omega/omega_connection.rs
View file

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

19
src/rho/client_connection.rs
View file

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

16
src/rho/iota_connection.rs
View file

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

26
src/rho/rho_connection.rs
View file

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

20
src/rho/rho_manager.rs
View file

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

11
src/util/config_util.rs
View file

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

1
src/util/mod.rs
View file

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

71
src/util/print.rs Normal file
View file

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