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

Connection With Omikron

Browse source
This commit is contained in:
Alex Emmet 2026-01-05 02:15:55 +01:00
commit 0122d5b342
10 changed files with 426 additions and 217 deletions
Download patch file Download diff file Expand all files Collapse all files

10
Cargo.lock generated
View file

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

1
Cargo.toml
View file

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

18
src/data/communication.rs
View file

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

11
src/main.rs
View file

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

10
src/server/api.rs
View file

@ -65,19 +65,17 @@ pub async fn handle(
let id = path_parts[3].parse::<i64>().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
{
(

237
src/server/omikron_connection.rs
View file

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

11
src/server/server.rs
View file

@ -464,10 +464,19 @@ pub async fn start_omikron_handler(connection: Arc<OmikronConnection>) {
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.
_ => {}
}
}

5
src/server/socket.rs
View file

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

172
src/sql/sql.rs
View file

@ -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<u8> = row.get("display");
let status: Vec<u8> = row.get("status");
let about: Vec<u8> = row.get("about");
let avatar: Vec<u8> = 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<u8> = 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<u8> = row.get("display");
let status: Vec<u8> = row.get("status");
let about: Vec<u8> = row.get("about");
let avatar: Vec<u8> = 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<u8> = 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<dyn std::error::Error>> {
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<u8> = row.get("display");
let status: Vec<u8> = row.get("status");
let about: Vec<u8> = 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<u8> = 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<u8>)>("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<u8>, Vec<u8>)>("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<u8>, Vec<u8>)>(
"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> {

170
src/util/logger.rs
View file

@ -6,22 +6,37 @@ use std::{
thread,
time::{SystemTime, UNIX_EPOCH},
};
use ansi_term::Color;
static LOGGER: OnceLock<mpsc::Sender<LogMessage>> = OnceLock::new();
#[derive(Clone, Copy)]
pub enum PrintType {
Call,
Client,
Iota,
Omikron,
Omega,
General,
}
struct LogMessage {
timestamp_ms: u128,
sender: Option<i64>,
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::<LogMessage>();
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::<String>();
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<i64>, message: String) {
/** Internal async logging entry point.
* Not exposed publicly; all access goes through macros.
*/
pub fn log_internal(
sender: Option<i64>,
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<i64>, 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)*)
)
};
}