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

Creation

Browse source
This commit is contained in:
Alex Emmet 2026-01-03 21:40:21 +01:00
commit 8a93f91e69
19 changed files with 7090 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

10
.gitignore vendored
View file

@ -18,4 +18,12 @@ target
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
.idea/
.env
/logs
# Added by cargo
/target

4462
Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

42
Cargo.toml Normal file
View file

@ -0,0 +1,42 @@
[package]
name = "Omega"
version = "0.1.0"
edition = "2024"
[dependencies]
aes-gcm = "*"
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"
bytes = "1.11.0"
color-eyre = "0.6.5"
dashmap = "6.1.0"
dotenv = "0.15.0"
futures = "0.3.31"
futures-util = "0.3.31"
hkdf = "0.12.4"
http-body-util = "0.1.3"
hyper = { version = "1.8.1", features = ["full"] }
hyper-util = "0.1.19"
json = "0.12.4"
once_cell = "1.21.3"
pnet = "0.35.0"
rand = "0.8"
rand_core = { version = "0.6", features = ["getrandom", "std"] }
reqwest = "0.12.28"
rustls = "0.23.35"
rustls-pemfile = "2.2.0"
sha1 = "0.10.6"
sha2 = "0.10.9"
sqlx = { version = "0.8.6", features = ["mysql", "runtime-async-std"] }
sysinfo = "0.37.2"
tokio = { version = "*", features = ["full"] }
tokio-rustls = "0.26.4"
tokio-tungstenite = "0.28.0"
tokio-util = { version = "*", features = ["full"] }
tower = "0.5.2"
tungstenite = "0.28.0"
uuid = { version = "1.19.0", features = ["v4"] }
walkdir = "2.5.0"
x448 = "0.6.0"
zip = "6.0.0"

454
src/data/communication.rs Normal file
View file

@ -0,0 +1,454 @@
use json::number::Number;
use json::{Array, JsonValue, object, parse};
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;
#[derive(Eq, Hash, PartialEq, Clone, Debug)]
#[allow(non_camel_case_types, dead_code)]
pub enum DataTypes {
error_type,
accepted_ids,
uuid,
settings,
settings_name,
chat_partner_id,
iota_id,
user_id,
user_ids,
user_state,
user_states,
user_pings,
call_state,
screen_share,
private_key_hash,
accepted,
accepted_profiles,
denied_profiles,
content,
messages,
send_time,
get_time,
get_variant,
shared_secret_own,
shared_secret_other,
shared_secret_sign,
shared_secret,
call_id,
call_token,
start_date,
end_date,
receiver_id,
sender_id,
signature,
signed,
message,
last_ping,
ping_iota,
ping_clients,
matches,
omikron,
offset,
amount,
position,
name,
path,
codec,
function,
payload,
result,
interactables,
want_to_watch,
watcher,
created_at,
username,
display,
avatar,
about,
status,
public_key,
sub_level,
sub_end,
community_address,
challenge,
community_title,
communities,
user,
}
impl DataTypes {
pub fn parse(p0: String) -> DataTypes {
let normalized = p0.to_lowercase().replace('_', "");
match normalized.as_str() {
"errortype" => DataTypes::error_type,
"chatpartnerid" => DataTypes::chat_partner_id,
"uuid" => DataTypes::uuid,
"settings" => DataTypes::settings,
"settingsname" => DataTypes::settings_name,
"iotaid" => DataTypes::iota_id,
"userid" => DataTypes::user_id,
"userids" => DataTypes::user_ids,
"userstate" => DataTypes::user_state,
"userstates" => DataTypes::user_states,
"userpings" => DataTypes::user_pings,
"callstate" => DataTypes::call_state,
"screenshare" => DataTypes::screen_share,
"privatekeyhash" => DataTypes::private_key_hash,
"accepted" => DataTypes::accepted,
"acceptedprofiles" => DataTypes::accepted_profiles,
"deniedprofiles" => DataTypes::denied_profiles,
"content" => DataTypes::content,
"messages" => DataTypes::messages,
"sendtime" => DataTypes::send_time,
"gettime" => DataTypes::get_time,
"getvariant" => DataTypes::get_variant,
"sharedsecretown" => DataTypes::shared_secret_own,
"sharedsecretother" => DataTypes::shared_secret_other,
"sharedsecretsign" => DataTypes::shared_secret_sign,
"sharedsecret" => DataTypes::shared_secret,
"callid" => DataTypes::call_id,
"calltoken" => DataTypes::call_token,
"startdate" => DataTypes::start_date,
"enddate" => DataTypes::end_date,
"receiverid" => DataTypes::receiver_id,
"senderid" => DataTypes::sender_id,
"signature" => DataTypes::signature,
"signed" => DataTypes::signed,
"message" => DataTypes::message,
"lastping" => DataTypes::last_ping,
"pingiota" => DataTypes::ping_iota,
"pingclients" => DataTypes::ping_clients,
"matches" => DataTypes::matches,
"omikron" => DataTypes::omikron,
"offset" => DataTypes::offset,
"amount" => DataTypes::amount,
"position" => DataTypes::position,
"name" => DataTypes::name,
"path" => DataTypes::path,
"codec" => DataTypes::codec,
"function" => DataTypes::function,
"payload" => DataTypes::payload,
"result" => DataTypes::result,
"interactables" => DataTypes::interactables,
"wanttowatch" => DataTypes::want_to_watch,
"watcher" => DataTypes::watcher,
"createdat" => DataTypes::created_at,
"username" => DataTypes::username,
"display" => DataTypes::display,
"avatar" => DataTypes::avatar,
"about" => DataTypes::about,
"status" => DataTypes::status,
"publickey" => DataTypes::public_key,
"sublevel" => DataTypes::sub_level,
"subend" => DataTypes::sub_end,
"communityaddress" => DataTypes::community_address,
"challenge" => DataTypes::challenge,
"communitytitle" => DataTypes::community_title,
"communities" => DataTypes::communities,
"user" => DataTypes::user,
_ => DataTypes::error_type, // fallback if unknown
}
}
}
#[derive(PartialEq, Clone, Debug)]
#[allow(non_camel_case_types, dead_code)]
pub enum CommunicationType {
error,
error_invalid_user_id,
error_not_found,
error_no_iota,
error_invalid_challenge,
error_invalid_secret,
error_invalid_private_key,
error_no_user_id,
error_no_call_id,
error_invalid_call_id,
success,
settings_save,
settings_load,
settings_list,
message,
message_send,
message_live,
message_other_iota,
message_chunk,
messages_get,
change_confirm,
confirm_receive,
confirm_read,
get_chats,
get_states,
add_community,
remove_community,
get_communities,
challenge,
challenge_response,
register,
register_response,
identification,
identification_response,
ping,
pong,
add_chat,
send_chat,
iota_connected,
iota_closed,
client_changed,
client_connected,
client_disconnected,
client_closed,
public_key,
private_key,
webrtc_sdp,
webrtc_ice,
start_stream,
end_stream,
watch_stream,
call_token,
call_invite,
end_call,
function,
update,
create_user,
}
impl CommunicationType {
pub fn parse(p0: String) -> CommunicationType {
let normalized = p0.to_lowercase().replace('_', "");
match normalized.as_str() {
"watchstream" => CommunicationType::watch_stream,
"calltoken" => CommunicationType::call_token,
"callinvite" => CommunicationType::call_invite,
"endcall" => CommunicationType::end_call,
"function" => CommunicationType::function,
"update" => CommunicationType::update,
"createuser" => CommunicationType::create_user,
"errorinvaliduserid" => CommunicationType::error_invalid_user_id,
"errornotfound" => CommunicationType::error_not_found,
"errornoiota" => CommunicationType::error_no_iota,
"errorinvalidchallenge" => CommunicationType::error_invalid_challenge,
"errorinvalidsecret" => CommunicationType::error_invalid_secret,
"errorinvalidprivatekey" => CommunicationType::error_invalid_private_key,
"errornouserid" => CommunicationType::error_no_user_id,
"errornocallid" => CommunicationType::error_no_call_id,
"errorinvalidcallid" => CommunicationType::error_invalid_call_id,
"success" => CommunicationType::success,
"settingssave" => CommunicationType::settings_save,
"settingsload" => CommunicationType::settings_load,
"settingslist" => CommunicationType::settings_list,
"message" => CommunicationType::message,
"messagesend" => CommunicationType::message_send,
"messagelive" => CommunicationType::message_live,
"messageother_iota" => CommunicationType::message_other_iota,
"messagechunk" => CommunicationType::message_chunk,
"messagesget" => CommunicationType::messages_get,
"changeconfirm" => CommunicationType::change_confirm,
"confirmreceive" => CommunicationType::confirm_receive,
"confirmread" => CommunicationType::confirm_read,
"getchats" => CommunicationType::get_chats,
"getstates" => CommunicationType::get_states,
"addcommunity" => CommunicationType::add_community,
"removecommunity" => CommunicationType::remove_community,
"getcommunities" => CommunicationType::get_communities,
"challenge" => CommunicationType::challenge,
"challengeresponse" => CommunicationType::challenge_response,
"register" => CommunicationType::register,
"registerresponse" => CommunicationType::register_response,
"identification" => CommunicationType::identification,
"identificationresponse" => CommunicationType::identification_response,
"ping" => CommunicationType::ping,
"pong" => CommunicationType::pong,
"addchat" => CommunicationType::add_chat,
"sendchat" => CommunicationType::send_chat,
"iotaconnected" => CommunicationType::iota_connected,
"iotaclosed" => CommunicationType::iota_closed,
"clientchanged" => CommunicationType::client_changed,
"clientconnected" => CommunicationType::client_connected,
"clientdisconnected" => CommunicationType::client_disconnected,
"clientclosed" => CommunicationType::client_closed,
"publickey" => CommunicationType::public_key,
"privatekey" => CommunicationType::private_key,
"webrtcsdp" => CommunicationType::webrtc_sdp,
"webrtcice" => CommunicationType::webrtc_ice,
"startstream" => CommunicationType::start_stream,
"endstream" => CommunicationType::end_stream,
_ => CommunicationType::error,
}
}
}
#[derive(Debug, Clone)]
pub struct CommunicationValue {
pub id: Uuid,
pub comm_type: CommunicationType,
pub sender: i64,
pub receiver: i64,
pub data: HashMap<DataTypes, JsonValue>,
}
#[allow(dead_code)]
impl CommunicationValue {
pub fn new(comm_type: CommunicationType) -> Self {
Self {
id: Uuid::new_v4(),
comm_type,
sender: 0,
receiver: 0,
data: HashMap::new(),
}
}
pub fn with_id(mut self, p0: Uuid) -> Self {
self.id = p0;
self
}
pub fn get_id(&self) -> Uuid {
self.id.clone()
}
pub fn with_sender(mut self, sender: i64) -> Self {
self.sender = sender;
self
}
pub fn get_sender(&self) -> i64 {
self.sender.clone()
}
pub fn with_receiver(mut self, receiver: i64) -> Self {
self.receiver = receiver;
self
}
pub fn get_receiver(&self) -> i64 {
self.receiver.clone()
}
pub fn add_data_num(mut self, key: DataTypes, value: Number) -> Self {
self.data.insert(key, JsonValue::Number(value));
self
}
pub fn add_data_str(mut self, key: DataTypes, value: String) -> Self {
self.data.insert(key, JsonValue::String(value));
self
}
pub fn add_data(mut self, key: DataTypes, value: JsonValue) -> Self {
self.data.insert(key, value);
self
}
pub fn add_array(mut self, key: DataTypes, value: Array) -> Self {
self.data.insert(key, JsonValue::Array(value));
self
}
pub fn get_data(&self, key: DataTypes) -> Option<&JsonValue> {
self.data.get(&key)
}
pub(crate) fn is_type(&self, p0: CommunicationType) -> bool {
self.comm_type == p0
}
pub fn to_json(&self) -> JsonValue {
let mut jdata = object! {};
for (k, v) in &self.data {
jdata[&format!("{:?}", k)] = JsonValue::from(v.clone());
}
if self.sender > 0 && self.receiver > 0 {
object! {
id: self.id.to_string(),
type: format!("{:?}", self.comm_type),
sender: self.sender,
receiver: self.receiver,
data: jdata
}
} else if self.sender > 0 {
object! {
id: self.id.to_string(),
type: format!("{:?}", self.comm_type),
sender: self.sender,
data: jdata
}
} else if self.receiver > 0 {
object! {
id: self.id.to_string(),
type: format!("{:?}", self.comm_type),
receiver: self.receiver,
data: jdata
}
} else {
object! {
id: self.id.to_string(),
type: format!("{:?}", self.comm_type),
data: jdata
}
}
}
pub fn from_json(json_str: &str) -> Self {
if let Ok(parsed) = parse(json_str) {
let comm_type = CommunicationType::parse(parsed["type"].to_string());
let mut sender: i64 = 0;
if parsed.has_key("sender") {
sender = parsed["sender"].as_i64().unwrap_or(0);
}
let mut receiver: i64 = 0;
if parsed.has_key("receiver") {
receiver = parsed["receiver"].as_i64().unwrap_or(0);
}
let uuid =
Uuid::parse_str(parsed["id"].as_str().unwrap_or("")).unwrap_or(Uuid::new_v4());
let mut data = HashMap::new();
if parsed["data"].is_object() {
for (k, v) in parsed["data"].entries() {
data.insert(DataTypes::parse(k.to_string()), v.clone());
}
}
Self {
id: uuid,
comm_type,
sender,
receiver,
data,
}
} else {
Self {
id: Uuid::new_v4(),
comm_type: CommunicationType::error,
sender: 0,
receiver: 0,
data: HashMap::new(),
}
}
}
pub fn forward_to_other_iota(original: &mut CommunicationValue) -> CommunicationValue {
let receiver = original
.get_data(DataTypes::receiver_id)
.unwrap_or(&JsonValue::Number(Number::from(0)))
.as_i64()
.unwrap_or(0);
let now_ms = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as i64;
let sender = original.get_sender();
CommunicationValue::new(CommunicationType::message_other_iota)
.with_id(original.get_id())
.with_receiver(receiver)
.add_data(
DataTypes::receiver_id,
JsonValue::Number(Number::from(receiver)),
)
.with_sender(sender)
.add_data(DataTypes::send_time, JsonValue::String(now_ms.to_string()))
.add_data(
DataTypes::sender_id,
JsonValue::Number(Number::from(sender)),
)
.add_data(
DataTypes::content,
JsonValue::String(original.get_data(DataTypes::content).unwrap().to_string()),
)
}
}

1
src/data/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod communication;

44
src/main.rs Normal file
View file

@ -0,0 +1,44 @@
mod data;
mod server;
mod sql;
mod util;
use crate::sql::sql::initialize_db;
use crate::sql::sql::print_users;
use crate::util::crypto_helper::load_public_key;
use crate::util::crypto_helper::load_secret_key;
use crate::util::logger::startup;
use dotenv::dotenv;
use once_cell::sync::Lazy;
use std::env;
static PRIVATE_KEY: Lazy<String> = Lazy::new(|| env::var("PRIVATE_KEY").unwrap());
pub fn get_private_key() -> x448::Secret {
load_secret_key(&*PRIVATE_KEY).unwrap()
}
static PUBLIC_KEY: Lazy<String> = Lazy::new(|| env::var("PUBLIC_KEY").unwrap());
pub fn get_public_key() -> x448::PublicKey {
load_public_key(&*PUBLIC_KEY).unwrap()
}
#[tokio::main]
async fn main() {
dotenv().ok();
startup();
log!("Started");
log_in!("Incoming messages");
log_out!("Outgoing messages");
log!(" .env");
if let Err(e) = initialize_db().await {
log!("[FATAL] Database initialization failed: {}", e);
log!(
"[FATAL] Please ensure the database is running and the .env file is configured correctly."
);
return;
}
let _ = print_users().await;
log!(" DB");
server::server::start(9187).await;
log!(" Server");
loop {}
}

177
src/server/api.rs Normal file
View file

@ -0,0 +1,177 @@
use crate::{get_public_key, log};
use crate::{
sql::{
iota_omikron_tracker::{get_omikron_for_iota, track_iota_omikron, untrack_iota},
sql::{
change_about, change_avatar, change_display_name, change_iota_id, change_iota_key,
change_keys, change_status, change_username, get_by_id, get_by_username,
get_iota_by_id, get_omikron_by_id, get_random_omikron, get_register_id,
register_complete_iota, register_complete_user,
},
},
util::crypto_helper::public_key_to_base64,
};
use axum::http::HeaderValue;
use http_body_util::Full;
use hyper::body::Bytes;
use hyper::{HeaderMap, Response as HttpResponse, StatusCode};
use json::JsonValue;
pub async fn handle(
path: &str,
headers: HeaderMap<HeaderValue>,
body_string: Option<String>,
) -> HttpResponse<Full<Bytes>> {
let path_parts: Vec<&str> = path.split("/").filter(|s| !s.is_empty()).collect();
let body: Option<JsonValue> = if body_string.is_some() {
if let Ok(body_json) = json::parse(&body_string.unwrap()) {
Some(body_json)
} else {
None
}
} else {
None
};
// api/
// get/
// omikron/
// id/
// register/
// innit/
// complete
log!("{}", path);
log!("{:?} .len = {}", path_parts, path_parts.len());
let (status, content, body_text) = if path_parts.len() >= 2 {
match path_parts[1] {
"get" => match path_parts[2] {
// api/get/omikron -> any omikron
// api/get/omikron/<id> -> omikron for id (user / iota / omikron)
"omikron" => {
if path_parts.len() == 3 {
if let Ok((id, public_key, ip_address)) = get_random_omikron().await {
(
StatusCode::OK,
"application/json",
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
id, public_key, ip_address
),
)
} else {
not_found()
}
} else if path_parts.len() == 4 {
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
{
(
StatusCode::OK,
"application/json",
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_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)) =
get_omikron_by_id(omikron_id).await
{
(
StatusCode::OK,
"application/json",
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_id, public_key, ip_address
),
)
} else {
not_found()
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
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)) =
get_omikron_by_id(omikron_id).await
{
(
StatusCode::OK,
"application/json",
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_id, public_key, ip_address
),
)
} else {
not_found()
}
} else {
not_found()
}
} else {
not_found()
}
} else {
bad_request()
}
}
// get/id/<username>
"id" => {
let username = path_parts[3];
bad_request()
}
"public_key" => (
StatusCode::OK,
"application/json",
public_key_to_base64(&get_public_key()),
),
_ => {
let id = path_parts[2];
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
bad_request()
} else {
bad_request()
}
}
},
_ => not_found(),
}
} else {
not_found()
};
let body = Full::new(Bytes::from(body_text.to_string()));
HttpResponse::builder().status(status).body(body).unwrap()
}
pub fn bad_request() -> (StatusCode, &'static str, String) {
(
StatusCode::BAD_REQUEST,
"text/text",
"400 Bad Request".to_string(),
)
}
pub fn unauthorized() -> (StatusCode, &'static str, String) {
(
StatusCode::UNAUTHORIZED,
"text/text",
"401 Unauthorized".to_string(),
)
}
pub fn forbidden() -> (StatusCode, &'static str, String) {
(
StatusCode::FORBIDDEN,
"text/text",
"403 Forbidden".to_string(),
)
}
pub fn not_found() -> (StatusCode, &'static str, String) {
(
StatusCode::NOT_FOUND,
"text/text",
"404 Not Found".to_string(),
)
}

5
src/server/mod.rs Normal file
View file

@ -0,0 +1,5 @@
pub mod api;
pub mod omikron_connection;
pub mod omikron_manager;
pub mod server;
pub mod socket;

258
src/server/omikron_connection.rs Normal file
View file

@ -0,0 +1,258 @@
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 base64::{Engine as _, engine::general_purpose::STANDARD};
use dashmap::DashMap;
use futures::SinkExt;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use json::JsonValue;
use rand::Rng;
use rand::distributions::Alphanumeric;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio_tungstenite::WebSocketStream;
use tungstenite::Message;
use tungstenite::Utf8Bytes;
use uuid::Uuid;
use x448::PublicKey;
pub struct OmikronConnection {
pub sender: Arc<RwLock<SplitSink<WebSocketStream<TokioIo<Upgraded>>, Message>>>,
pub receiver: Arc<RwLock<SplitStream<WebSocketStream<TokioIo<Upgraded>>>>>,
pub omikron_id: Arc<RwLock<i64>>,
pub pub_key: Arc<RwLock<Option<Vec<u8>>>>,
identified: Arc<RwLock<bool>>,
challenged: Arc<RwLock<bool>>,
challenge: Arc<RwLock<String>>,
pub ping: Arc<RwLock<i64>>,
waiting_tasks: DashMap<
Uuid,
Box<dyn Fn(Arc<OmikronConnection>, CommunicationValue) -> bool + Send + Sync>,
>,
}
impl OmikronConnection {
pub fn new(
sender: SplitSink<WebSocketStream<TokioIo<Upgraded>>, Message>,
receiver: SplitStream<WebSocketStream<TokioIo<Upgraded>>>,
) -> Arc<Self> {
Arc::new(Self {
sender: Arc::new(RwLock::new(sender)),
receiver: Arc::new(RwLock::new(receiver)),
omikron_id: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)),
identified: Arc::new(RwLock::new(false)),
challenged: Arc::new(RwLock::new(false)),
challenge: Arc::new(RwLock::new(String::new())),
ping: Arc::new(RwLock::new(-1)),
waiting_tasks: DashMap::new(),
})
}
pub async fn send_message(&self, message: &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);
sender.send(message_text).await.unwrap();
}
pub async fn get_user_id(&self) -> i64 {
*self.omikron_id.read().await
}
pub async fn is_identified(&self) -> bool {
*self.identified.read().await && *self.challenged.read().await
}
pub async fn get_public_key(&self) -> PublicKey {
PublicKey::from_bytes(self.pub_key.read().await.as_ref().unwrap()).unwrap()
}
pub async fn handle_message(self: Arc<Self>, message: String) {
let cv = CommunicationValue::from_json(&message);
if cv.is_type(CommunicationType::ping) {
self.handle_ping(cv).await;
return;
}
log_in_from!(*self.omikron_id.read().await, "{}", message);
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;
let challenge_str: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
*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,
)
.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_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()),
)
.add_data_str(DataTypes::challenge, encrypted_challenge)
.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)
.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)
.await;
return;
}
}
async fn send_error_response(&self, message_id: &Uuid, error_type: CommunicationType) {
let error = CommunicationValue::new(error_type).with_id(*message_id);
self.send_message(&error).await;
}
pub async fn close(&self) {
let mut sender = self.sender.write().await;
let _ = sender.close().await;
}
pub async fn handle_close(self: Arc<Self>) {
if self.is_identified().await {
if self.get_user_id().await != 0 {}
}
}
async fn handle_ping(&self, cv: CommunicationValue) {
if let Some(last_ping) = cv.get_data(DataTypes::last_ping) {
if let Ok(ping_val) = last_ping.to_string().parse::<i64>() {
let mut ping_guard = self.ping.write().await;
*ping_guard = ping_val;
}
}
let response = CommunicationValue::new(CommunicationType::pong).with_id(cv.get_id());
self.send_message(&response).await;
}
}

0
src/server/omikron_manager.rs Normal file
View file

475
src/server/server.rs Normal file
View file

@ -0,0 +1,475 @@
use crate::log;
use crate::server::api;
use crate::server::omikron_connection::OmikronConnection;
use crate::util::file_util::load_file_buf;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use bytes::Bytes;
use futures::StreamExt;
use futures_util::TryFutureExt;
use http_body_util::BodyExt;
use http_body_util::Full;
use hyper::{Method, StatusCode};
use hyper::{
Request as HttpRequest, Response as HttpResponse, body::Incoming, server::conn::http1, upgrade,
};
use hyper_util::rt::tokio::TokioIo;
use hyper_util::service::TowerToHyperService;
use pnet::datalink::NetworkInterface;
use rustls::ServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use sha1::{Digest, Sha1};
use std::error::Error;
use std::io::ErrorKind;
use std::io::{self, BufReader};
use std::net::SocketAddr;
use std::result::Result::Ok;
use std::sync::Arc;
use std::{future::Future, pin::Pin, time::Duration};
use tokio::net::TcpListener;
use tokio::sync::broadcast;
use tokio_rustls::TlsAcceptor;
use tokio_tungstenite::{WebSocketStream, accept_async};
use tower::Service;
use tungstenite::Message;
#[derive(Clone)]
struct HttpService {
peer_addr: SocketAddr,
}
impl Service<HttpRequest<Incoming>> for HttpService {
type Response = HttpResponse<Full<Bytes>>;
type Error = io::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(
&mut self,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(std::io::Result::Ok(()))
}
fn call(&mut self, req: HttpRequest<Incoming>) -> Self::Future {
let _peer_ip = self.peer_addr.ip();
let (parts, body) = req.into_parts();
let method = parts.method.clone();
let path = parts.uri.path().to_string();
let headers = parts.headers.clone();
let fut = async move {
let is_websocket_upgrade = path == "/ws/omikron"
&& method == Method::GET
&& headers
.get("connection")
.map(|v| {
v.to_str()
.unwrap_or("")
.split(',')
.any(|s| s.trim().eq_ignore_ascii_case("upgrade"))
})
.unwrap_or(false)
&& headers
.get("upgrade")
.map(|v| v.to_str().unwrap_or("").eq_ignore_ascii_case("websocket"))
.unwrap_or(false);
if is_websocket_upgrade {
log!("Attempting WebSocket upgrade on {}", path);
if let Some(sec_websocket_key) = headers.get("sec-websocket-key") {
let sec_websocket_key = sec_websocket_key.to_str().unwrap_or("").to_string();
let sec_websocket_accept = calculate_accept_key(&sec_websocket_key);
let response = HttpResponse::builder()
.status(StatusCode::SWITCHING_PROTOCOLS)
.header("Upgrade", "websocket")
.header("Connection", "Upgrade")
.header("Sec-WebSocket-Accept", sec_websocket_accept)
.body(Full::new(Bytes::from("")))
.unwrap();
let req_for_upgrade = HttpRequest::from_parts(parts, body);
let upgrades = upgrade::on(req_for_upgrade);
log!("Handling WebSocket upgrade");
// Spawn upgrade handling to avoid blocking the service call
tokio::spawn(async move {
match upgrades.await {
Ok(upgraded_stream) => {
log!("Valid WebSocket upgrade");
let raw_stream = TokioIo::new(upgraded_stream);
let ws_stream = WebSocketStream::from_raw_socket(
raw_stream,
tungstenite::protocol::Role::Server,
None,
)
.await;
log!(
"WebSocket handshake successful, handling connection for Omikron"
);
// Split stream for OmikronConnection
let (writer, reader) = ws_stream.split();
// INTEGRATION START
// Erstelle die Connection und starte den Handler
let connection = OmikronConnection::new(writer, reader);
// Handler in separatem Task starten
start_omikron_handler(connection).await;
// INTEGRATION END
}
Err(e) => {
log!("WebSocket upgrade failed after response: {:?}", e);
}
}
});
log!("Handled WebSocket connection initiation");
Ok(response)
} else {
log!("No Sec-WebSocket-Key found in request headers");
let response = HttpResponse::builder()
.status(StatusCode::BAD_REQUEST)
.body(Full::new(Bytes::from("Missing Sec-WebSocket-Key")))
.unwrap();
Ok(response)
}
} else if path.starts_with("/api") {
let whole_body = match body.collect().await {
Ok(collected) => collected,
Err(e) => {
log!("Error collecting body: {}", e);
return Ok(HttpResponse::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Full::new(Bytes::from(format!(
"Failed to read body: {}",
e
))))
.unwrap());
}
};
let bytes = whole_body.to_bytes();
let body_string: Option<String> = match String::from_utf8(bytes.to_vec()) {
Ok(s) => Some(s),
Err(_) => None,
};
Ok(api::handle(&path, headers.clone(), body_string).await)
} else {
let response = HttpResponse::builder()
.status(StatusCode::BAD_REQUEST)
.body(Full::new(Bytes::from("No path provided")))
.unwrap();
Ok(response)
}
};
Box::pin(fut.map_err(|err: color_eyre::eyre::ErrReport| {
io::Error::new(
io::ErrorKind::Other,
format!("Error in request handling: {}", err),
)
}))
}
}
async fn run_http_server(port: u16) -> bool {
let ip = "0.0.0.0".to_string();
let listener = TcpListener::bind(format!("0.0.0.0:{}", port)).await;
if let Err(e) = listener {
log!("Failed to bind to port {}: {:?}", port, e);
return false;
}
let listener = listener.unwrap();
log!(
"Standard Server listening for HTTP and WS on {}:{}",
ip,
port
);
// Create a broadcast channel for graceful shutdown signal
let (shutdown_tx, _) = broadcast::channel::<()>(1);
tokio::spawn(async move {
loop {
tokio::select! {
accepted = listener.accept() => {
match accepted {
std::result::Result::Ok((stream, addr)) => {
let service = HttpService { peer_addr: addr };
let io = TokioIo::new(stream);
// Subscribe to the shutdown signal for this specific connection
let mut rx = shutdown_tx.subscribe();
tokio::spawn(async move {
// Prepare the connection future
let conn = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.serve_connection(io, TowerToHyperService::new(service))
.with_upgrades();
// Wait for either the connection to finish naturally OR the shutdown signal
tokio::select! {
res = conn => {
if let Err(err) = res {
if let Some(io_err) = err.source().and_then(|e| e.downcast_ref::<io::Error>()) {
if io_err.kind() != io::ErrorKind::ConnectionReset
&& io_err.kind() != io::ErrorKind::BrokenPipe
{
log!("Error serving connection: {:?}", err);
}
}
}
}
_ = rx.recv() => {
// Shutdown signal received.
// Dropping the 'conn' future here closes the socket immediately.
}
}
});
}
Err(e) => {
log!("Error accepting connection: {:?}", e);
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
}
}
}
});
true
}
/// Runs the encrypted HTTPS/WSS server loop using the provided TLS config.
async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
let mut ip = "0.0.0.0".to_string();
for iface in pnet::datalink::interfaces() {
let iface: NetworkInterface = iface;
let ipsv = format!("{}", iface.ips[0]);
let ips: &str = ipsv.split('/').next().unwrap();
log!("{}", ips.to_string());
if format!("{}", ips).starts_with("10.") {
ip = ips.to_string();
}
}
let acceptor = TlsAcceptor::from(tls_config);
let listener = TcpListener::bind(format!("0.0.0.0:{}", port)).await;
if let Err(e) = listener {
log!("Failed to bind to port {}: {:?}", port, e);
return false;
}
let listener = listener.unwrap();
log!(
"Encrypted Server listening for HTTPS and WSS on {}:{}",
ip,
port
);
// Create a broadcast channel for graceful shutdown signal
let (shutdown_tx, _) = broadcast::channel::<()>(1);
tokio::spawn(async move {
loop {
tokio::select! {
// Monitor for shutdown signal
_ = async {
loop {
tokio::time::sleep(Duration::from_millis(100)).await;
}
} => {
log!("Encrypted Server received shutdown signal.");
// Send kill signal to all active connection tasks
let _ = shutdown_tx.send(());
break;
}
// Accept new connections
accepted = listener.accept() => {
match accepted {
std::result::Result::Ok((stream, addr)) => {
let service = HttpService { peer_addr: addr };
let acceptor = acceptor.clone();
// Subscribe to the shutdown signal for this specific connection
let mut rx = shutdown_tx.subscribe();
tokio::spawn(async move {
// Perform TLS handshake
let tls_stream = match acceptor.accept(stream).await {
Ok(s) => s,
Err(e) => {
if e.kind() != io::ErrorKind::Interrupted {
log!("TLS Handshake error: {:?}", e);
}
return;
}
};
let io = TokioIo::new(tls_stream);
// Prepare connection future
let conn = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.serve_connection(io, TowerToHyperService::new(service))
.with_upgrades();
// Wait for either the connection to finish naturally OR the shutdown signal
tokio::select! {
res = conn => {
if let Err(err) = res {
if let Some(io_err) = err.source().and_then(|e| e.downcast_ref::<io::Error>()) {
if io_err.kind() != io::ErrorKind::ConnectionReset
&& io_err.kind() != io::ErrorKind::BrokenPipe
{
log!("Error serving connection: {:?}", err);
}
}
}
}
_ = rx.recv() => {
// Shutdown signal received.
// Dropping the 'conn' future here closes the socket immediately.
}
}
});
}
Err(e) => {
log!("Error accepting connection: {:?}", e);
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
}
}
}
log!("Encrypted Server shutdown complete.");
});
true
}
pub async fn start(port: u16) -> bool {
let tls_result = load_tls_config();
match tls_result {
Ok(Some(tls_config)) => run_tls_server(port, tls_config).await,
Ok(None) => run_http_server(port).await,
Err(e) => {
log!("Fatal error during TLS config load: {}", e);
false
}
}
}
fn calculate_accept_key(key: &str) -> String {
let websocket_guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
let mut sha1 = Sha1::new();
sha1.update(key.as_bytes());
sha1.update(websocket_guid.as_bytes());
let result = sha1.finalize();
STANDARD.encode(result) // Base64 encode the result
}
/// Loads TLS config. Returns Ok(None) if cert files are not found, and an error if parsing fails.
fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
let cert_file_res = load_file_buf("certs", "cert.pem");
let key_file_res = load_file_buf("certs", "cert.key");
// Check if certificate files are present. If not, return None.
let cert_file_buf = match cert_file_res {
Ok(b) => b,
Err(e) if e.kind() == ErrorKind::NotFound => {
log!("TLS certificate 'certs/cert.pem' not found.");
return Ok(None);
}
Err(e) => return Err(e.into()), // Other IO error
};
let key_file_buf = match key_file_res {
Ok(b) => b,
Err(e) if e.kind() == ErrorKind::NotFound => {
log!("TLS key 'certs/cert.key' not found.");
return Ok(None);
}
Err(e) => return Err(e.into()), // Other IO error
};
// Continue with configuration if both files were found
let mut cert_reader = BufReader::new(cert_file_buf);
let cert_ders = rustls_pemfile::certs(&mut cert_reader)
.collect::<Result<Vec<CertificateDer>, io::Error>>()?;
// PKCS8
let mut key_reader = BufReader::new(key_file_buf);
let mut key_ders = rustls_pemfile::pkcs8_private_keys(&mut key_reader)
.map(|r| r.map(Into::into)) // Explicit conversion
.collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
if key_ders.is_empty() {
// RSA
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?); // Re-read key file
key_ders = rustls_pemfile::rsa_private_keys(&mut key_reader)
.map(|r| r.map(Into::into))
.collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
}
if key_ders.is_empty() {
// EC
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?); // Re-read key file
key_ders = rustls_pemfile::ec_private_keys(&mut key_reader)
.map(|r| r.map(Into::into))
.collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
}
if key_ders.is_empty() {
return Err("No private keys found in key file. (Tried PKCS8, RSA, and EC)".into());
}
let config = rustls::ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_ders, key_ders.remove(0))
.map_err(|e| io::Error::new(ErrorKind::Other, e.to_string()))?;
Ok(Some(Arc::new(config)))
}
// In deiner Server-Logik, wo OmikronConnection initialisiert wird:
pub async fn start_omikron_handler(connection: Arc<OmikronConnection>) {
loop {
let msg = match {
let mut receiver = connection.receiver.write().await;
receiver.next().await
} {
Some(Ok(msg)) => msg,
Some(Err(e)) => {
log!("WS Error: {}", e);
break;
}
None => break, // Stream ended
};
match msg {
Message::Text(text) => {
let conn_clone = connection.clone();
tokio::spawn(async move {
conn_clone.handle_message(text.to_string()).await;
});
}
Message::Close(_) => {
break;
}
// Other message types like Binary, Ping, Pong are ignored.
_ => {}
}
}
connection.handle_close().await;
}

58
src/server/socket.rs Normal file
View file

@ -0,0 +1,58 @@
use std::sync::Arc;
use futures::StreamExt;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use tungstenite::Message;
use crate::log;
use crate::server::omikron_connection::OmikronConnection;
pub fn handle(
path: String,
writer: SplitSink<tokio_tungstenite::WebSocketStream<TokioIo<Upgraded>>, Message>,
reader: SplitStream<tokio_tungstenite::WebSocketStream<TokioIo<Upgraded>>>,
) {
log!("handling");
tokio::spawn(async move {
if path.starts_with("/ws/phi/") {
} else if path.starts_with("/ws/omikron/") {
let community_conn: Arc<OmikronConnection> =
Arc::from(OmikronConnection::new(writer, reader));
loop {
let msg_result: Option<Result<_, _>> = {
let mut session_lock = community_conn.receiver.write().await;
session_lock.next().await
};
match msg_result {
Some(Ok(msg)) => {
if msg.is_text() {
let text = msg.into_text().unwrap();
community_conn
.clone()
.handle_message(text.to_string())
.await;
} else if msg.is_close() {
log!("Closing: {}", msg);
community_conn.handle_close().await;
return;
}
}
Some(Err(e)) => {
log!("Closing ERR: {}", e);
community_conn.handle_close().await;
return;
}
None => {
log!("Closed Session me!");
community_conn.handle_close().await;
return;
}
}
}
}
});
}

22
src/sql/iota_omikron_tracker.rs Normal file
View file

@ -0,0 +1,22 @@
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
static IOTA_OMIKRON_MAP: Lazy<Arc<RwLock<HashMap<i64, i64>>>> =
Lazy::new(|| Arc::new(RwLock::new(HashMap::new())));
pub async fn track_iota_omikron(iota: i64, omikron: i64) {
let mut c = IOTA_OMIKRON_MAP.write().await;
c.insert(iota, omikron);
}
pub async fn get_omikron_for_iota(iota: i64) -> Option<i64> {
let c = IOTA_OMIKRON_MAP.read().await;
c.get(&iota).cloned()
}
pub async fn untrack_iota(iota: i64) {
let mut c = IOTA_OMIKRON_MAP.write().await;
c.remove(&iota);
}

2
src/sql/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod iota_omikron_tracker;
pub mod sql;

450
src/sql/sql.rs Normal file
View file

@ -0,0 +1,450 @@
use crate::log;
use once_cell::sync::Lazy;
use sqlx::{MySql, Pool, mysql::MySqlPoolOptions};
use std::sync::atomic::{AtomicU64, Ordering};
use std::{
env,
sync::Arc,
time::{SystemTime, UNIX_EPOCH},
};
use tokio::sync::RwLock;
/*
use crate::sql::{
iota_omikron_tracker::{get_omikron_for_iota, track_iota_omikron, untrack_iota},
sql::{
change_about, change_avatar, change_display_name, change_iota_id, change_iota_key,
change_keys, change_status, change_username, get_by_id, get_by_username, get_iota_by_id,
get_register_id, register_complete_iota, register_complete_user,
},
};
*/
static SQL_DB: Lazy<Arc<RwLock<Option<Pool<MySql>>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
pub async fn connect() -> Result<Pool<MySql>, sqlx::Error> {
let user = env::var("DB_USERNAME").expect("DB_USERNAME is not set");
let passwd = env::var("DB_PASSWD").expect("DB_PASSWD is not set");
let table = env::var("DB_TABLE").expect("DB_TABLE is not set");
MySqlPoolOptions::new()
.max_connections(5)
.connect(&format!(
"mysql://{}:{}@127.0.0.1:3306/{}",
user, passwd, table
))
.await
}
// Omega
// - Omikron
// - Iota
// - User
// - User
// - Iota
// - User
// - User
// - Omikron
// - Iota
// - User
// - User
// - Iota
// - User
// - User
pub async fn initialize_db() -> Result<(), sqlx::Error> {
let pool = connect().await?;
let mut db_lock = SQL_DB.write().await;
// create tables
// with indexes
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
users (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
username VARCHAR(15) NOT NULL UNIQUE COLLATE utf8mb4_bin,
display VARCHAR(15) COLLATE utf8mb4_bin,
status VARCHAR(15) COLLATE utf8mb4_bin,
about VARCHAR(200) COLLATE utf8mb4_bin,
avatar MEDIUMBLOB,
sub_level INT(11) NOT NULL DEFAULT 0,
sub_end BIGINT(20) NOT NULL,
public_key TEXT NOT NULL COLLATE utf8mb4_bin,
private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin,
iota_id BIGINT UNSIGNED NOT NULL,
token VARCHAR(255) NOT NULL UNIQUE COLLATE utf8mb4_bin
)",
)
.execute(&pool)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
iotas (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
public_key TEXT NOT NULL COLLATE utf8mb4_bin
)",
)
.execute(&pool)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
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
)",
)
.execute(&pool)
.await;
*db_lock = Some(pool);
Ok(())
}
// ==========================================================================================
// REGISTER
// ==========================================================================================
pub static CURRENT_MILLI_USED: Lazy<Arc<AtomicU64>> = Lazy::new(|| Arc::new(AtomicU64::new(0)));
pub static CURRENT_REGISTER_PROCESS: Lazy<Arc<RwLock<Vec<u64>>>> =
Lazy::new(|| Arc::new(RwLock::new(Vec::new())));
pub async fn get_register_id() -> u64 {
let mut current_time = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
loop {
let current_locked = CURRENT_MILLI_USED.load(Ordering::SeqCst);
if current_locked < current_time {
let result = CURRENT_MILLI_USED.compare_exchange(
current_locked, // expected value
current_time, // new value
Ordering::SeqCst, // acquire/release ordering
Ordering::SeqCst, // failure ordering
);
match result {
Ok(_) => {
CURRENT_REGISTER_PROCESS.write().await.push(current_time);
return current_time;
}
Err(_) => {
continue;
}
}
} else {
current_time = current_locked + 1;
}
}
}
// ==========================================================================================
// USERS
// ==========================================================================================
pub async fn get_by_username(
username: &str,
) -> Result<
(
i64,
i64,
String,
String,
String,
String,
String,
i32,
i64,
String,
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::<_, (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 = ?",
)
.bind(username)
.fetch_optional(pool)
.await?.ok_or_else(|| sqlx::Error::RowNotFound)
}
pub async fn get_by_id(
id: i64,
) -> Result<
(
i64,
i64,
String,
String,
String,
String,
String,
i32,
i64,
String,
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::<_, (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 = ?",
)
.bind(id)
.fetch_optional(pool)
.await?.ok_or_else(|| sqlx::Error::RowNotFound)
}
pub async fn change_username(id: i64, new_username: 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 username = ? WHERE id = ?")
.bind(new_username)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn change_display_name(id: i64, new_display_name: 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)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn change_avatar(id: i64, new_avatar: 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 avatar = ? WHERE id = ?")
.bind(new_avatar)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn change_about(id: i64, new_about: 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 about = ? WHERE id = ?")
.bind(new_about)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn change_status(id: i64, new_status: 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 status = ? WHERE id = ?")
.bind(new_status)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn change_iota_id(id: i64, new_iota_id: i64) -> 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 iota_id = ? WHERE id = ?")
.bind(new_iota_id)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn change_keys(
id: i64,
new_public_key: String,
new_private_key_hash: 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 public_key = ?, private_key_hash = ? WHERE id = ?")
.bind(new_public_key)
.bind(new_private_key_hash)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn register_complete_user(
id: i64,
username: String,
public_key: String,
private_key_hash: String,
iota_id: i64,
reset_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 (?, ?, ?, ?, ?, ?)")
.bind(id)
.bind(username)
.bind(public_key)
.bind(private_key_hash)
.bind(iota_id)
.bind(reset_token)
.execute(pool)
.await?;
Ok(())
}
pub async fn print_users() -> Result<(), Box<dyn std::error::Error>> {
let db_lock = SQL_DB.read().await;
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",
)
.fetch_all(pool)
.await?
.iter()
{
log!(
"User: {:?}",
(id, iota_id, username, display_name, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token)
);
}
Ok(())
}
// ==========================================================================================
// IOTA
// ==========================================================================================
pub async fn register_complete_iota(id: i64, public_key: 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 iotas (id, public_key) VALUES (?, ?)")
.bind(id)
.bind(public_key)
.execute(pool)
.await?;
Ok(())
}
pub async fn delete_user(id: i64) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
sqlx::query("DELETE FROM users WHERE id = ?")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
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 = ?")
.bind(id)
.fetch_optional(pool)
.await?
.ok_or_else(|| sqlx::Error::RowNotFound)
}
pub async fn change_iota_key(id: i64, new_key: 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 iotas SET public_key = ? WHERE id = ?")
.bind(new_key)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
sqlx::query("DELETE FROM iotas WHERE id = ?")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
// ==========================================================================================
// OMIKRONS
// ==========================================================================================
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")
.fetch_optional(pool)
.await?
.ok_or_else(|| sqlx::Error::RowNotFound);
// this log isn't printing
log!("Random Omikron: {:?}", a);
a
}
pub async fn get_omikron_by_id(id: i64) -> 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");
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)
}
pub async fn set_omikron_active(id: i64, active: bool) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
let active = if active { 1 } else { 0 };
sqlx::query("UPDATE omikrons SET active = ? WHERE id = ?")
.bind(active)
.bind(id)
.execute(pool)
.await?;
Ok(())
}

143
src/util/crypto_helper.rs Normal file
View file

@ -0,0 +1,143 @@
use aes_gcm::{
Aes256Gcm, Nonce,
aead::{Aead, KeyInit, OsRng},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use rand_core::RngCore;
use sha2::{Digest, Sha256};
use x448::{PublicKey, Secret, SharedSecret};
/// Errors for crypto operations
#[derive(Debug)]
pub enum CryptoError {
Base64Decode(base64::DecodeError),
InvalidKey,
AgreementError,
EncryptionError(aes_gcm::Error),
DecryptionError(aes_gcm::Error),
}
impl From<base64::DecodeError> for CryptoError {
fn from(err: base64::DecodeError) -> Self {
CryptoError::Base64Decode(err)
}
}
pub struct KeyPair {
pub secret: Secret,
pub public: PublicKey,
}
pub fn generate_keypair() -> KeyPair {
let mut buf = [0u8; 56];
let mut rng = OsRng;
rng.fill_bytes(&mut buf);
let secret = Secret::from_bytes(&buf).unwrap();
let public = PublicKey::from(&secret);
KeyPair { secret, public }
}
pub fn public_key_to_base64(pubkey: &PublicKey) -> String {
STANDARD.encode(pubkey.as_bytes().as_ref())
}
pub fn secret_key_to_base64(secret: &Secret) -> String {
STANDARD.encode(secret.as_bytes().as_ref())
}
pub fn load_public_key(base64_pub: &str) -> Option<PublicKey> {
let bytes = STANDARD.decode(base64_pub).unwrap();
PublicKey::from_bytes(&bytes)
}
pub fn load_secret_key(base64_secret: &str) -> Option<Secret> {
let bytes = STANDARD.decode(base64_secret).unwrap();
Secret::from_bytes(&bytes)
}
fn derive_aes_key(shared: &SharedSecret) -> [u8; 32] {
let mut hasher = Sha256::new();
hasher.update(shared.as_bytes());
let result = hasher.finalize();
let mut key = [0u8; 32];
key.copy_from_slice(&result[..32]);
key
}
pub fn encrypt_b64(
base64_secret: &str,
base64_peer_pub: &str,
plaintext: &str,
) -> Result<String, CryptoError> {
let secret = load_secret_key(base64_secret).unwrap();
let peer_pub = load_public_key(base64_peer_pub).unwrap();
encrypt(secret, peer_pub, plaintext)
}
pub fn encrypt(
secret: Secret,
peer_pub: PublicKey,
plaintext: &str,
) -> Result<String, CryptoError> {
let shared = secret
.to_diffie_hellman(&peer_pub)
.ok_or(CryptoError::AgreementError)?;
let key_bytes = derive_aes_key(&shared);
let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct");
let mut nonce_bytes = [0u8; 12];
OsRng.fill_bytes(&mut nonce_bytes);
let nonce = Nonce::from_slice(&nonce_bytes);
let ciphertext = cipher
.encrypt(nonce, plaintext.as_bytes())
.map_err(CryptoError::EncryptionError)?;
// prefix nonce to ciphertext
let mut out = Vec::with_capacity(nonce_bytes.len() + ciphertext.len());
out.extend_from_slice(&nonce_bytes);
out.extend_from_slice(&ciphertext);
Ok(STANDARD.encode(&out))
}
pub fn decrypt_b64(
base64_secret: &str,
base64_peer_pub: &str,
encrypted_base64: &str,
) -> Result<String, CryptoError> {
let secret = load_secret_key(base64_secret).unwrap();
let peer_pub = load_public_key(base64_peer_pub).unwrap();
decrypt(secret, peer_pub, encrypted_base64)
}
pub fn decrypt(
secret: Secret,
peer_pub: PublicKey,
encrypted_base64: &str,
) -> Result<String, CryptoError> {
let shared = secret
.to_diffie_hellman(&peer_pub)
.ok_or(CryptoError::AgreementError)?;
let key_bytes = derive_aes_key(&shared);
let cipher = Aes256Gcm::new_from_slice(&key_bytes).expect("Key length should be correct");
let encrypted = STANDARD.decode(encrypted_base64)?;
if encrypted.len() < 12 {
return Err(CryptoError::DecryptionError(aes_gcm::Error));
}
let nonce_bytes = &encrypted[..12];
let ciphertext = &encrypted[12..];
let nonce = Nonce::from_slice(nonce_bytes);
let plaintext_bytes = cipher
.decrypt(nonce, ciphertext)
.map_err(CryptoError::DecryptionError)?;
let plaintext = String::from_utf8(plaintext_bytes)
.map_err(|_| CryptoError::DecryptionError(aes_gcm::Error))?;
Ok(plaintext)
}
pub fn hash_it(input: &str) -> Vec<u8> {
let mut hasher = Sha256::new();
hasher.update(input.as_bytes());
hasher.finalize().to_vec()
}
pub fn hex_hash(input: &str) -> String {
let digest = hash_it(input);
digest.iter().map(|b| format!("{:02x}", b)).collect()
}

350
src/util/file_util.rs Normal file
View file

@ -0,0 +1,350 @@
use crate::log;
use std::ffi::OsStr;
use std::fs::{self, File};
use std::io::{self, BufReader, Read};
use std::path::{Path, PathBuf};
use sysinfo::System;
use uuid::Uuid;
use walkdir::WalkDir;
use zip::ZipArchive;
pub fn delete_file(path: &str, name: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
let file = dir.join(name);
if !file.exists() {
return false;
}
fs::remove_file(file).is_ok()
}
pub fn delete_directory(path: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
delete_dir_recursive(&dir)
}
fn delete_dir_recursive(directory: &Path) -> bool {
if !directory.exists() {
return false;
}
if let Err(e) = fs::remove_dir_all(directory) {
log!(
"[IMPORTANT] Couldn't delete directory {}: {}",
directory.display(),
e
);
return false;
}
true
}
pub fn load_file_buf(path: &str, name: &str) -> io::Result<BufReader<File>> {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
// Ensure the directory exists, create if necessary
if !dir.exists() {
if let Err(_) = fs::create_dir_all(&dir) {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"Directory creation failed",
));
}
}
// Create the file if it doesn't exist
if !file_path.exists() {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"File creation failed",
));
}
// Open the file and return a BufReader for efficient reading
let file = File::open(&file_path)?;
Ok(BufReader::new(file))
}
pub fn has_file(path: &str, name: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
if !dir.exists() {
return false;
}
if !file_path.exists() {
return false;
}
true
}
pub fn has_dir(path: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
if !dir.exists() {
return false;
}
true
}
pub fn load_file(path: &str, name: &str) -> String {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
if !dir.exists() {
if let Err(e) = fs::create_dir_all(&dir) {
log!("[IMPORTANT] Couldn't create directories: {}", e);
return String::new();
}
return String::new();
}
if !file_path.exists() {
if let Err(e) = File::create(&file_path) {
log!("[IMPORTANT] Couldn't create file: {}", e);
}
return String::new();
}
let mut content = String::new();
if let Ok(mut f) = File::open(&file_path) {
let _ = f.read_to_string(&mut content);
}
content
}
pub fn load_file_vec(path: &str, name: &str) -> Vec<u8> {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
if !dir.exists() {
if let Err(e) = fs::create_dir_all(&dir) {
log!("[IMPORTANT] Couldn't create directories: {}", e);
return Vec::new();
}
return Vec::new();
}
if !file_path.exists() {
if let Err(e) = File::create(&file_path) {
log!("[IMPORTANT] Couldn't create file: {}", e);
}
return Vec::new();
}
let mut content = Vec::new();
if let Ok(mut f) = File::open(&file_path) {
let _ = f.read_to_end(&mut content);
}
content
}
pub fn save_file(path: &str, name: &str, value: &str) {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
if !dir.exists() {
if let Err(e) = fs::create_dir_all(&dir) {
log!("[IMPORTANT] Couldn't create directories: {}", e);
return;
}
}
if let Err(e) = fs::write(&file_path, value) {
log!(
"[IMPORTANT] Couldn't write file {}: {}",
file_path.display(),
e
);
}
}
pub fn get_children(path: &str) -> Vec<String> {
let dir = Path::new(&get_directory()).join(path);
let mut children = Vec::new();
if let Ok(entries) = fs::read_dir(&dir) {
for entry in entries {
if let Ok(entry) = entry {
children.push(entry.file_name().to_string_lossy().to_string());
}
}
}
children
}
pub fn get_directory() -> String {
let exe = std::env::current_exe().unwrap_or_else(|_| PathBuf::from("."));
exe.parent()
.unwrap_or(Path::new("."))
.to_string_lossy()
.to_string()
}
pub fn used_space() -> u64 {
get_directory_size(&PathBuf::from(get_directory()))
}
pub fn used_dir_space(path: &str) -> u64 {
get_directory_size(&PathBuf::from(format!("{}/{}", get_directory(), path)))
}
pub fn get_directory_size(directory: &Path) -> u64 {
let mut size = 0;
for entry in WalkDir::new(directory).into_iter().filter_map(|e| e.ok()) {
let path = entry.path();
if path.is_file() {
if let Ok(metadata) = path.metadata() {
size += path.file_name().unwrap_or(OsStr::new("")).len() as u64;
size += metadata.len();
}
}
}
size
}
pub fn get_designed_storage(user_id: Uuid) -> String {
let user_dir = Path::new(&get_directory())
.join("users")
.join(user_id.to_string());
design_byte(get_directory_size(&user_dir))
}
pub fn design_byte(bytes: u64) -> String {
let mut hr_size = format!("{:.2}B", bytes as f64);
let k = bytes as f64 / 1024.0;
let m = k / 1024.0;
let g = m / 1024.0;
let t = g / 1024.0;
if t >= 1.0 {
hr_size = format!("{:.2}TB", t);
} else if g >= 1.0 {
hr_size = format!("{:.2}GB", g);
} else if m >= 1.0 {
hr_size = format!("{:.2}MB", m);
} else if k >= 1.0 {
hr_size = format!("{:.2}KB", k);
}
hr_size
}
pub fn get_used_ram() -> String {
let mut sys = System::new_all();
sys.refresh_all();
let used = sys.used_memory() * 1024; // kB to bytes
let total = sys.total_memory() * 1024;
format!("{}/{}", design_byte(used), design_byte(total))
}
// Helper to download the zip file content to a file on disk
async fn download_zip(url: &str, as_name: &Path) -> Result<(), Box<dyn std::error::Error>> {
let response = reqwest::get(url).await?;
// Check for successful response status
if !response.status().is_success() {
return Err(format!("Failed to download file: Status {}", response.status()).into());
}
let mut zip_file = File::create(as_name)?;
let body = response.bytes().await?;
io::copy(&mut &*body, &mut zip_file)?;
Ok(())
}
fn extract_zip_contents_to_folder(
zip_path: &Path,
target_dir: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
let file = File::open(zip_path)?;
let mut archive = ZipArchive::new(file)?;
let staging_dir = target_dir.with_extension("staging");
let _ = fs::remove_dir_all(&staging_dir);
fs::create_dir_all(&staging_dir)?;
let mut first_item_name: Option<PathBuf> = None;
for i in 0..archive.len() {
let mut file = archive.by_index(i)?;
let entry_path = staging_dir.join(file.sanitized_name());
if i == 0 {
if file.name().ends_with('/') || file.sanitized_name().components().count() == 1 {
first_item_name = Some(file.sanitized_name());
}
}
if file.name().ends_with('/') {
fs::create_dir_all(&entry_path)?;
} else {
if let Some(parent) = entry_path.parent() {
fs::create_dir_all(parent)?;
}
let mut out_file = File::create(entry_path)?;
io::copy(&mut file, &mut out_file)?;
}
}
if let Some(root_path) = first_item_name {
let root_dir = staging_dir.join(&root_path);
if root_dir.is_dir() {
let root_contents_count = fs::read_dir(&staging_dir)?.count();
if root_contents_count == 1
|| (root_contents_count > 1 && fs::metadata(&root_dir).is_ok())
{
let _ = fs::remove_dir_all(target_dir);
fs::create_dir_all(target_dir)?;
for entry in fs::read_dir(root_dir)? {
let entry = entry?;
let src = entry.path();
let dest = target_dir.join(entry.file_name());
if let Err(_) = fs::rename(&src, &dest) {
if src.is_file() {
fs::copy(&src, &dest)?;
} else {
if entry.path().is_dir() {
fs::rename(&src, &dest)?;
}
}
}
}
let _ = fs::remove_dir_all(&staging_dir);
return Ok(());
}
}
}
log!("Extracting directly (no single root folder detected).");
let _ = fs::remove_dir_all(target_dir);
fs::rename(&staging_dir, target_dir)?;
Ok(())
}
pub async fn download_and_extract_zip(url: &str, as_name: &str) {
let base_dir = PathBuf::from(get_directory());
let zip_filename = format!("{}.zip", Uuid::new_v4()); // Use a unique name for the downloaded ZIP file
let zip_path = base_dir.join(&zip_filename);
let target_dir = base_dir.join(as_name);
// Step 1: Download the ZIP file
if let Err(e) = download_zip(url, &zip_path).await {
log!("Error downloading file: {}", e);
return;
}
// Step 2: Extract and flatten the ZIP file contents into the target directory
if let Err(e) = extract_zip_contents_to_folder(&zip_path, &target_dir) {
log!("Error extracting ZIP file contents: {}", e);
}
// Step 3: Clean up the downloaded ZIP file
if let Err(e) = fs::remove_file(&zip_path) {
log!("Error cleaning up ZIP file {}: {}", zip_path.display(), e);
}
}

135
src/util/logger.rs Normal file
View file

@ -0,0 +1,135 @@
use std::{
fs::{self, OpenOptions},
io::Write,
path::Path,
sync::{OnceLock, mpsc},
thread,
time::{SystemTime, UNIX_EPOCH},
};
static LOGGER: OnceLock<mpsc::Sender<LogMessage>> = OnceLock::new();
struct LogMessage {
timestamp_ms: u128,
sender: Option<i64>,
message: String,
}
/// Initialize logger (call once)
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");
let start_ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.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),
None => fixed_box("", 19),
};
let line = format!("{} {} {}", timestamp_box, sender_box, msg.message);
println!("{}", line);
let _ = writeln!(file, "{}", line);
}
});
}
fn fixed_box(content: &str, width: usize) -> String {
let s = content.chars().take(width).collect::<String>();
let len = s.chars().count();
if len < width {
let mut a = " ".repeat(width - len);
a.push_str(&s);
format!("[{}]", a)
} else {
s
}
}
/// Internal function (sync + async safe)
pub fn log_internal(sender: Option<i64>, message: String) {
if let Some(tx) = LOGGER.get() {
let _ = tx.send(LogMessage {
timestamp_ms: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis(),
sender,
message,
});
} else {
println!("{}", message);
}
}
#[macro_export]
macro_rules! log {
($($arg:tt)*) => {
$crate::util::logger::log_internal(None, format!($($arg)*))
};
}
#[macro_export]
macro_rules! log_from {
($sender:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(Some($sender), format!($($arg)*))
};
}
#[macro_export]
macro_rules! log_in {
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
format!("> {}", format!($($arg)*))
)
};
}
#[macro_export]
macro_rules! log_out {
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
format!("< {}", format!($($arg)*))
)
};
}
#[macro_export]
macro_rules! log_in_from {
($sender:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(
Some($sender),
format!("> {}", format!($($arg)*))
)
};
}
#[macro_export]
macro_rules! log_out_from {
($sender:expr, $($arg:tt)*) => {
$crate::util::logger::log_internal(
Some($sender),
format!("< {}", format!($($arg)*))
)
};
}

3
src/util/mod.rs Normal file
View file

@ -0,0 +1,3 @@
pub mod crypto_helper;
pub mod file_util;
pub mod logger;