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

Sync

Browse source
This commit is contained in:
Alois 2026-04-02 22:28:41 +02:00
commit b9b1f1fc1c
22 changed files with 8234 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

29
.gitignore vendored Normal file
View file

@ -0,0 +1,29 @@
# Generated by Cargo
# will have compiled files and executables
debug
target
# These are backup files generated by rustfmt
**/*.rs.bk
# MSVC Windows builds of rustc generate these, which store debugging information
*.pdb
# Generated by cargo mutants
# Contains mutation testing data
**/mutants.out*/
# RustRover
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# 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/
.env
/logs
# Added by cargo
/target

4547
Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

39
Cargo.toml Executable file
View file

@ -0,0 +1,39 @@
[package]
name = "Omega"
version = "0.1.0"
edition = "2024"
[dependencies]
epsilon-core = { git = "https://github.com/Tensamin/Epsilon.git", package = "epsilon-core" }
epsilon-native = { git = "https://github.com/Tensamin/Epsilon.git", package = "epsilon-native" }
actix-web = { version = "4.12.1", features = ["rustls-0_23"] }
aes-gcm = "*"
ansi_term = "0.12.1"
anyhow = "1.0.101"
base64 = "0.22.1"
dashmap = "6.1.0"
dotenv = "0.15.0"
hex = "0.4.3"
hkdf = "0.12.4"
json = "0.12.4"
once_cell = "1.21.3"
rand = "0.8"
rand_core = { version = "0.6", features = ["getrandom", "std"] }
reqwest = { version = "0.13.2" }
rustls = { version = "0.23.37", default-features = false, features = [
"std",
"tls12",
"aws-lc-rs",
"prefer-post-quantum",
] }
rustls-pemfile = "2.2.0"
sha2 = "0.10.9"
sqlx = { version = "0.8.6", features = ["mysql", "runtime-async-std"] }
strum = "0.27.2"
strum_macros = "0.27.2"
tokio = { version = "*", features = ["full"] }
uuid = { version = "1.19.0", features = ["v4"] }
x448 = "0.6.0"
zip = "6.0.0"
thiserror = "2.0.18"

16
LICENSE Normal file
View file

@ -0,0 +1,16 @@
Copyright (c) [2025] [Methanium]
All rights reserved.
This software is protected by copyright. Copying, editing,
distributing, publicly performing, or any other use of this software
or its components, in source or binary form, is strictly prohibited without the express
written permission of the copyright holder.
FUTURE LICENSE ACCEPTANCE:
It is the copyright holder's intention to release this software in the future
under a license yet to be defined, which will, among other things,
allow private, non-commercial use. This statement does not constitute
a current license grant and does not alter the above
prohibition on use, copying, or modification. Until the formal
publication of such a future license, all rights remain
reserved.

1
README.md Normal file
View file

@ -0,0 +1 @@
# ⚠️ Moved to [git.methanium.net](https://git.methanium.net/tensamin/omega) ⚠️

65
src/main.rs Normal file
View file

@ -0,0 +1,65 @@
mod server;
mod sql;
mod transport;
mod util;
use crate::sql::sql::initialize_db;
use crate::sql::sql::print_users;
use crate::transport::omikron_connection;
use crate::util::crypto_helper::load_public_key;
use crate::util::crypto_helper::load_secret_key;
use crate::util::logger::PrintType;
use crate::util::logger::startup;
use dotenv::dotenv;
use once_cell::sync::Lazy;
use rustls::crypto::aws_lc_rs::default_provider;
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() {
if let Err(_) = default_provider().install_default() {
println!("Error loading Provider");
return;
}
dotenv().ok();
startup();
log_in!("Incoming messages");
log_out!("Outgoing messages");
tokio::spawn(async move {
match omikron_connection::start(9187).await {
Err(e) => log_err!(0, PrintType::General, "{:?}", e),
_ => {}
}
});
log!("Started");
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;
} else {
log!(" DB");
}
if let Err(e) = print_users().await {
log!("[ERROR] Failed to print users: {}", e);
} else {
log!(" Users");
}
let _ = server::server::start(9188).await;
tokio::signal::ctrl_c().await.unwrap();
}

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

@ -0,0 +1,275 @@
use crate::get_public_key;
use crate::sql::sql;
use crate::sql::user_online_tracker::get_iota_primary_omikron_connection;
use crate::transport::omikron_manager::get_random_omikron;
use crate::util::file_util::get_directory;
use crate::{
sql::sql::{get_by_user_id, get_omikron_by_id},
util::crypto_helper::public_key_to_base64,
};
use actix_web::HttpResponse;
use actix_web::http::{StatusCode, header};
use base64::Engine as _;
use json::JsonValue;
pub async fn handle(path: &str, body_string: Option<String>) -> HttpResponse {
if path == "OPTIONS" {
return HttpResponse::Ok()
.insert_header(("Access-Control-Allow-Origin", "*"))
.insert_header(("Access-Control-Allow-Methods", "GET, POST, OPTIONS"))
.insert_header(("Access-Control-Allow-Headers", "*"))
.finish();
}
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
};
let (status, body_text) = match path_parts.as_slice() {
// ==================================================
// DOWNLOAD IOTA FRONTEND
// ==================================================
["api", "download", "iota_frontend"] => {
let file_path = format!("{}/downloads/iota_frontend.zip", get_directory());
match std::fs::read(file_path) {
Ok(file_bytes) => {
return HttpResponse::Ok()
.insert_header(("Access-Control-Allow-Origin", "*"))
.insert_header(("Content-Type", "application/zip"))
.insert_header((
"Content-Disposition",
"attachment; filename=\"iota_frontend.zip\"",
))
.body(file_bytes);
}
Err(_) => {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
return HttpResponse::NotFound()
.insert_header(("Access-Control-Allow-Origin", "*"))
.body(res.dump());
}
}
}
// ==================================================
// GET RANDOM OMIKRON
// ==================================================
["api", "get", "omikron"] => {
if let Ok(omikron_conn) = get_random_omikron().await {
if let Some(id) = omikron_conn.get_omikron_id().await {
if let Ok((public_key, ip_address)) = sql::get_omikron_by_id(id).await {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error".into();
(StatusCode::INTERNAL_SERVER_ERROR, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error".into();
(StatusCode::INTERNAL_SERVER_ERROR, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
}
// ==================================================
// GET OMIKRON BY ID
// ==================================================
["api", "get", "omikron", id] => {
let id = id.parse::<i64>().unwrap_or(0);
if id == 0 {
let mut res = JsonValue::new_object();
res["status"] = "error_bad_request".into();
(StatusCode::BAD_REQUEST, res.dump())
} else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) {
if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = omikron_id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = get_by_user_id(id).await
{
if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) {
if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["id"] = omikron_id.into();
res["public_key"] = public_key.into();
res["ip_address"] = ip_address.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::NOT_FOUND, res.dump())
}
}
// ==================================================
// GET ID BY USERNAME
// ==================================================
["api", "get", "id", username] => {
if username.is_empty() {
let mut res = JsonValue::new_object();
res["status"] = "error_bad_request".into();
(StatusCode::BAD_REQUEST, res.dump())
} else if let Ok((
id,
iota_id,
username,
_,
_,
_,
_,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_username(username).await
{
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["username"] = username.into();
res["public_key"] = public_key.into();
res["user_id"] = id.into();
res["iota_id"] = iota_id.into();
res["sub_level"] = sub_level.into();
res["sub_end"] = sub_end.into();
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::OK, res.dump())
}
}
// ==================================================
// GET SERVER PUBLIC KEY
// ==================================================
["api", "get", "public_key"] => {
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["public_key"] = public_key_to_base64(&get_public_key()).into();
(StatusCode::OK, res.dump())
}
// ==================================================
// GET USER BY ID
// ==================================================
["api", "get", "user", id] => {
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
let mut res = JsonValue::new_object();
res["status"] = "error_bad_request".into();
(StatusCode::BAD_REQUEST, res.dump())
} else if let Ok((
id,
iota_id,
username,
display,
status_msg,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_user_id(id).await
{
let mut res = JsonValue::new_object();
res["status"] = "success".into();
res["username"] = username.into();
res["public_key"] = public_key.into();
res["user_id"] = id.into();
res["iota_id"] = iota_id.into();
res["sub_level"] = sub_level.into();
res["sub_end"] = sub_end.into();
if let Some(display) = display {
res["display"] = display.into();
}
if let Some(status_msg) = status_msg {
res["status_message"] = status_msg.into();
}
if let Some(about) = about {
res["about"] = about.into();
}
if let Some(avatar) = avatar {
res["avatar"] = base64::engine::general_purpose::STANDARD
.encode(avatar)
.into();
}
(StatusCode::OK, res.dump())
} else {
let mut res = JsonValue::new_object();
res["status"] = "error_not_found".into();
(StatusCode::OK, res.dump())
}
}
// ==================================================
// DEFAULT
// ==================================================
_ => {
let mut res = JsonValue::new_object();
res["status"] = "error".into();
(StatusCode::INTERNAL_SERVER_ERROR, res.dump())
}
};
let body_bytes = body_text.into_bytes();
HttpResponse::build(status)
.insert_header((header::ACCESS_CONTROL_ALLOW_ORIGIN, "*"))
.insert_header((header::ACCESS_CONTROL_ALLOW_HEADERS, "*"))
.insert_header((header::ACCESS_CONTROL_ALLOW_METHODS, "GET, POST, OPTIONS"))
.body(body_bytes)
}

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

@ -0,0 +1,3 @@
pub mod api;
pub mod server;
pub mod short_link;

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

@ -0,0 +1,67 @@
use crate::{
log,
server::{api, short_link::get_short_link},
util::file_util::load_file_buf,
};
use actix_web::{App, HttpRequest, HttpResponse, HttpServer, Responder, http::header, web};
use rustls::ServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls_pemfile::{certs, pkcs8_private_keys};
pub async fn start(port: u16) -> anyhow::Result<()> {
let mut cert_reader = load_file_buf("certs", "server_cert.pem")?;
let mut key_reader = load_file_buf("certs", "server_key.pem")?;
let cert_chain: Vec<CertificateDer<'static>> =
certs(&mut cert_reader).collect::<Result<_, _>>()?;
let mut keys: Vec<PrivateKeyDer<'static>> = pkcs8_private_keys(&mut key_reader)
.map(|res| res.map(Into::into))
.collect::<Result<_, _>>()?;
let key = keys.remove(0);
let mut config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_chain, key)?;
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
let addr = format!("0.0.0.0:{port}");
log!(" Server on {}", addr);
HttpServer::new(move || {
App::new()
.route("/api/{path:.*}", web::to(api_handler))
.route("/direct/{path:.*}", web::to(direct_handler))
})
.bind_rustls_0_23(addr, config)?
.run()
.await?;
Ok(())
}
async fn direct_handler(req: HttpRequest) -> impl Responder {
let path = req.uri().path().to_string();
let short = path.replace("/direct/", "");
if let Ok(long) = get_short_link(&short).await {
HttpResponse::TemporaryRedirect()
.append_header((header::LOCATION, long))
.finish()
} else {
HttpResponse::TemporaryRedirect()
.append_header((header::LOCATION, "https://tensamin.net"))
.finish()
}
}
async fn api_handler(req: HttpRequest, body: web::Bytes) -> HttpResponse {
let path = req.uri().path().to_string();
let body_string = String::from_utf8_lossy(&body).to_string();
api::handle(&path, Some(body_string)).await
}

87
src/server/short_link.rs Normal file
View file

@ -0,0 +1,87 @@
use dashmap::DashMap;
use once_cell::sync::Lazy;
use rand::{Rng, thread_rng};
static LINKS: Lazy<DashMap<String, String>> = Lazy::new(DashMap::new);
const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHJKLMNPRSTUVWXYZ1234567890";
pub async fn add_short_link(long: &str) -> Result<String, ()> {
let raw = generate_unique_short_link().await;
LINKS.insert(raw.clone(), long.to_string());
Ok(format!(
"https://omega.tensamin.net/direct/{}",
format_with_dashes(&raw)
))
}
async fn generate_unique_short_link() -> String {
loop {
let short = generate_short_link().await;
if !LINKS.contains_key(&short) {
return short;
}
}
}
pub async fn generate_short_link() -> String {
let len = short_length();
let mut rng = thread_rng();
(0..len)
.map(|_| {
let idx = rng.gen_range(0..CHARSET.len());
CHARSET[idx] as char
})
.collect()
}
pub async fn get_short_link(short: &str) -> Result<String, ()> {
let key = if short.contains("/") {
short.split("/").nth(1).unwrap_or_default()
} else {
short
};
let frag = short.replace(key, "");
let normalized = normalize_short(&key);
if let Ok(t) = LINKS.get(&normalized).map(|v| v.value().clone()).ok_or(()) {
Ok(format!("{}{}", t, frag))
} else {
Err(())
}
}
/* ---------------- helpers ---------------- */
fn short_length() -> usize {
let count = LINKS.len();
match count {
0..=1_999 => 4,
2_000..=999_999 => 8,
_ => 12,
}
}
fn format_with_dashes(s: &str) -> String {
s.chars()
.collect::<Vec<_>>()
.chunks(4)
.map(|c| c.iter().collect::<String>())
.collect::<Vec<_>>()
.join("-")
}
fn normalize_short(input: &str) -> String {
input
.chars()
.filter(|c| *c != '-')
.map(|c| match c {
'Q' | 'O' => '0',
'I' => 'l',
_ => c,
})
.collect()
}

30
src/sql/connection_status.rs Normal file
View file

@ -0,0 +1,30 @@
use strum::IntoEnumIterator;
use strum_macros::EnumIter;
#[derive(Debug, Clone, PartialEq, EnumIter, Eq)]
#[allow(unused, non_camel_case_types)]
pub enum UserStatus {
user_offline,
user_online,
user_dnd,
user_idle,
user_wc,
user_borked,
iota_offline,
iota_online,
iota_borked,
}
#[allow(unused)]
impl UserStatus {
pub fn to_string(&self) -> String {
format!("{:?}", self)
}
pub fn from_str(s: &str) -> Option<UserStatus> {
for sel in UserStatus::iter() {
if &sel.to_string() == s {
return Some(sel);
}
}
None
}
}

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

@ -0,0 +1,3 @@
pub mod connection_status;
pub mod sql;
pub mod user_online_tracker;

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

@ -0,0 +1,765 @@
use crate::log;
use once_cell::sync::Lazy;
use sqlx::{MySql, Pool, Row, 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(200)
.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 DEFAULT 0,
public_key TEXT NOT NULL COLLATE utf8mb4_bin,
private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin DEFAULT '',
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 VARCHAR(255) NOT NULL COLLATE utf8mb4_bin
)",
)
.execute(&pool)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
omikrons (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY,
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)
.await;
let _ = sqlx::query(
"CREATE TABLE IF NOT EXISTS
notifications (
id BIGINT UNSIGNED NOT NULL PRIMARY KEY AUTO INCREMENT,
sender_id BIGINT UNSIGNED NOT NULL,
receiver_id BIGINT UNSIGNED NOT NULL,
amount BIGINT UNSIGNED NOT NULL DEFAULT 0
)",
)
.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,
Option<String>,
Option<String>,
Option<String>,
Option<Vec<u8>>,
i32,
i64,
String,
String,
String,
),
sqlx::Error,
> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
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?;
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: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let avatar: Option<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,
display.map(|d| String::from_utf8_lossy(&d).to_string()),
status.map(|s| String::from_utf8_lossy(&s).to_string()),
about.map(|a| String::from_utf8_lossy(&a).to_string()),
avatar,
sub_level,
sub_end,
public_key,
private_key_hash,
String::from_utf8_lossy(&token).to_string(),
))
}
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_by_user_id(
id: i64,
) -> Result<
(
i64,
i64,
String,
Option<String>,
Option<String>,
Option<String>,
Option<Vec<u8>>,
i32,
i64,
String,
String,
String,
),
sqlx::Error,
> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
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 = CAST(? AS UNSIGNED)",
)
.bind(id)
.fetch_optional(&pool)
.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: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let avatar: Option<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,
display.map(|d| String::from_utf8_lossy(&d).to_string()),
status.map(|s| String::from_utf8_lossy(&s).to_string()),
about.map(|a| String::from_utf8_lossy(&a).to_string()),
avatar,
sub_level,
sub_end,
public_key,
private_key_hash,
String::from_utf8_lossy(&token).to_string(),
))
}
_ => Err(sqlx::Error::RowNotFound),
}
}
pub async fn get_users_by_iota_id(
iota_id_param: i64,
) -> Result<
Vec<(
i64,
i64,
String,
Option<String>,
Option<String>,
Option<String>,
Option<Vec<u8>>,
i32,
i64,
String,
String,
String,
)>,
sqlx::Error,
> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let rows = sqlx::query(
"SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE iota_id = CAST(? AS UNSIGNED)",
)
.bind(iota_id_param)
.fetch_all(&pool)
.await?;
let mut users = Vec::new();
for row in rows {
let id: i64 = row.get("id");
let iota_id: i64 = row.get("iota_id");
let username: String = row.get("username");
let display: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let avatar: Option<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");
users.push((
id,
iota_id,
username,
display.map(|d| String::from_utf8_lossy(&d).to_string()),
status.map(|s| String::from_utf8_lossy(&s).to_string()),
about.map(|a| String::from_utf8_lossy(&a).to_string()),
avatar,
sub_level,
sub_end,
public_key,
private_key_hash,
String::from_utf8_lossy(&token).to_string(),
));
}
Ok(users)
}
pub async fn change_username(id: i64, new_username: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET username = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_username)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_display_name(id: i64, new_display: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET display = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_display)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_avatar(id: i64, new_avatar: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET avatar = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_avatar)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_about(id: i64, new_about: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET about = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_about)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_status(id: i64, new_status: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET status = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_status)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn delete_user(id: i64) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("DELETE FROM users WHERE id = CAST(? AS UNSIGNED)")
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_iota_id(id: i64, new_iota_id: i64) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET iota_id = CAST(? AS UNSIGNED) WHERE id = CAST(? AS UNSIGNED)")
.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 pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query(
"UPDATE users SET public_key = ?, private_key_hash = ? WHERE id = CAST(? AS UNSIGNED)",
)
.bind(new_public_key)
.bind(new_private_key_hash)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_token(id: i64, new_token: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE users SET token = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_token)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn register_complete_user(
id: i64,
username: String,
public_key: String,
iota_id: i64,
token: String,
) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query(
"INSERT INTO users (id, username, public_key, iota_id, token) VALUES (?, ?, ?, ?, ?)",
)
.bind(id)
.bind(username)
.bind(public_key)
.bind(iota_id)
.bind(token)
.execute(&pool)
.await?;
Ok(())
}
pub async fn print_users() -> Result<(), Box<dyn std::error::Error>> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
log!("Printing 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: Option<Vec<u8>> = row.get("display");
let status: Option<Vec<u8>> = row.get("status");
let about: Option<Vec<u8>> = row.get("about");
let sub_level: i32 = row.get("sub_level");
let sub_end: i64 = row.get("sub_end");
log!(
"User: {:?}",
(
id,
iota_id,
username,
display.map_or("".to_string(), |d| String::from_utf8_lossy(&d).to_string()),
status.map_or("".to_string(), |s| String::from_utf8_lossy(&s).to_string()),
about.map_or("".to_string(), |a| String::from_utf8_lossy(&a).to_string()),
sub_level,
sub_end
)
);
}
Ok(())
}
// ==========================================================================================
// IOTA
// ==========================================================================================
pub async fn create_new_iota(public_key: String) -> Result<i64, sqlx::Error> {
let new_id = get_register_id().await as i64;
register_complete_iota(new_id, public_key).await?;
Ok(new_id)
}
pub async fn register_complete_iota(id: i64, public_key: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("INSERT INTO iotas (id, public_key) VALUES (?, ?)")
.bind(id)
.bind(public_key)
.execute(&pool)
.await?;
Ok(())
}
pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let result = sqlx::query_as::<_, (u64, Vec<u8>)>(
"SELECT id, public_key FROM iotas WHERE id = CAST(? AS UNSIGNED)",
)
.bind(id)
.fetch_optional(&pool)
.await;
match result {
Ok(optional_row) => match optional_row {
Some((id_u64, public_key)) => Ok((
id_u64 as i64,
String::from_utf8_lossy(&public_key).to_string(),
)),
_ => Err(sqlx::Error::RowNotFound),
},
Err(e) => Err(e),
}
}
pub async fn change_iota_key(id: i64, new_key: String) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("UPDATE iotas SET public_key = ? WHERE id = CAST(? AS UNSIGNED)")
.bind(new_key)
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
sqlx::query("DELETE FROM iotas WHERE id = CAST(? AS UNSIGNED)")
.bind(id)
.execute(&pool)
.await?;
Ok(())
}
// ==========================================================================================
// OMIKRONS
// ==========================================================================================
pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> {
let pool = {
let db_lock = SQL_DB.read().await;
db_lock
.as_ref()
.cloned()
.expect("Database pool not initialized")
};
let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>)>(
"SELECT public_key, ip_address FROM omikrons WHERE id = CAST(? AS UNSIGNED)",
)
.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(),
)),
_ => Err(sqlx::Error::RowNotFound),
}
}
// ==========================================================================================
// PHI
// ==========================================================================================
pub async fn add_notification(sender_id: i64, receiver_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(
r#"
INSERT INTO notifications (sender_id, receiver_id, amount)
VALUES (?, ?, 1)
ON DUPLICATE KEY UPDATE amount = amount + 1
"#,
)
.bind(sender_id)
.bind(receiver_id)
.execute(pool)
.await?;
Ok(())
}
pub async fn read_notification(sender_id: i64, receiver_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(
r#"
DELETE FROM notifications WHERE sender_id = ? AND receiver_id = ?
"#,
)
.bind(sender_id)
.bind(receiver_id)
.execute(pool)
.await?;
Ok(())
}
pub async fn get_notifications(user_id: i64) -> Result<Vec<(i64, i64)>, 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)>(
r#"
SELECT sender_id, amount FROM notifications WHERE receiver_id = ?
"#,
)
.bind(user_id)
.fetch_all(pool)
.await
}

144
src/sql/user_online_tracker.rs Normal file
View file

@ -0,0 +1,144 @@
use crate::sql;
use crate::sql::connection_status::UserStatus;
use dashmap::DashMap;
use once_cell::sync::Lazy;
#[derive(Debug, Clone)]
pub struct UserConnection {
pub connection_type: UserStatus,
pub omikron_id: i64,
}
// IotaID -> Primary OmikronID
static IOTA_PRIMARY_OMIKRON_CONNECTION: Lazy<DashMap<i64, i64>> = Lazy::new(DashMap::new);
// IotaID -> Vec<OmikronID>
static IOTA_OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Vec<i64>>> = Lazy::new(DashMap::new);
// UserID -> UserStatus
static USER_STATUS_MAP: Lazy<DashMap<i64, UserConnection>> = Lazy::new(DashMap::new);
pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) {
let mut entry = IOTA_OMIKRON_CONNECTIONS
.entry(iota_id)
.or_insert_with(Vec::new);
if !entry.contains(&omikron_id) {
entry.push(omikron_id);
}
if primary {
IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, omikron_id);
}
}
pub fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
let connections_empty = if let Some(r) = IOTA_OMIKRON_CONNECTIONS.get(&iota_id) {
let mut vec = r.value().clone();
vec.retain(|&id| id != omikron_id);
let empty = vec.is_empty();
drop(r);
IOTA_OMIKRON_CONNECTIONS.insert(iota_id, vec);
empty
} else {
false
};
if let Some(primary_ref) = IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id) {
let primary_id = *primary_ref.value();
drop(primary_ref);
if primary_id == omikron_id {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
}
}
if connections_empty {
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
}
connections_empty
}
pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id).map(|v| *v)
}
pub fn get_iota_omikron_connections(iota_id: i64) -> Option<Vec<i64>> {
IOTA_OMIKRON_CONNECTIONS.get(&iota_id).map(|v| v.clone())
}
pub fn track_user_status(user_id: i64, status: UserStatus, omikron_id: i64) {
USER_STATUS_MAP.insert(
user_id,
UserConnection {
connection_type: status,
omikron_id,
},
);
}
pub fn get_user_status(user_id: i64) -> Option<UserConnection> {
USER_STATUS_MAP.get(&user_id).map(|v| v.clone())
}
pub fn untrack_many_users(user_ids: &[i64]) {
for user_id in user_ids {
USER_STATUS_MAP.remove(user_id);
}
}
pub async fn untrack_omikron(omikron_id: i64) {
let primary_keys_to_remove: Vec<i64> = IOTA_PRIMARY_OMIKRON_CONNECTION
.iter()
.filter(|entry| *entry.value() == omikron_id)
.map(|entry| *entry.key())
.collect();
for key in primary_keys_to_remove {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&key);
}
let mut offline_iotas = Vec::new();
let mut primary_to_remove = Vec::new();
// Collect iotas and primary info first
for r in IOTA_OMIKRON_CONNECTIONS.iter() {
let iota_id = *r.key();
let mut connections = r.value().clone();
connections.retain(|&id| id != omikron_id);
if connections.is_empty() {
offline_iotas.push(iota_id);
}
if IOTA_PRIMARY_OMIKRON_CONNECTION
.get(&iota_id)
.map(|p| *p == omikron_id)
.unwrap_or(false)
{
primary_to_remove.push(iota_id);
}
// Update the connections vector after filtering
IOTA_OMIKRON_CONNECTIONS.insert(iota_id, connections);
}
// Step 2: Remove primary connections safely
for iota_id in primary_to_remove {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
}
// Step 3: Remove users that were on this omikron
USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
// Step 4: For offline iotas, remove associated users from USER_STATUS_MAP
for iota_id in offline_iotas {
if let Ok(users) = sql::sql::get_users_by_iota_id(iota_id).await {
for user in users {
USER_STATUS_MAP.remove(&user.0);
}
}
// Finally remove the empty connections vector
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
}
}

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

@ -0,0 +1,2 @@
pub mod omikron_connection;
pub mod omikron_manager;

1139
src/transport/omikron_connection.rs Normal file
View file

File diff suppressed because it is too large Load diff

39
src/transport/omikron_manager.rs Normal file
View file

@ -0,0 +1,39 @@
use crate::transport::omikron_connection::OmikronConnection;
use dashmap::DashMap;
use once_cell::sync::Lazy;
use rand::prelude::IteratorRandom;
use std::sync::Arc;
pub static OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Arc<OmikronConnection>>> =
Lazy::new(|| DashMap::new());
pub async fn add_omikron(conn: Arc<OmikronConnection>) {
let id = match conn.clone().get_omikron_id().await {
Some(id) => id,
_ => {
conn.close().await;
return;
}
};
if let Some(old) = OMIKRON_CONNECTIONS.insert(id, conn.clone()) {
old.close().await;
}
}
pub async fn remove_omikron(omikron_id: i64) {
OMIKRON_CONNECTIONS.remove(&omikron_id);
}
pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
let mut rng = rand::thread_rng();
let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
if let Some(key) = keys.into_iter().choose(&mut rng) {
if let Some(entry) = OMIKRON_CONNECTIONS.get(&key) {
return Ok(entry.clone());
}
}
Err(())
}

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

@ -0,0 +1,137 @@
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,
AgreementError,
EncryptionError(aes_gcm::Error),
DecryptionError(aes_gcm::Error),
}
impl From<base64::DecodeError> for CryptoError {
fn from(_: base64::DecodeError) -> Self {
CryptoError::Base64Decode
}
}
pub fn generate_keypair() -> (Secret, PublicKey) {
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);
(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()
}

201
src/util/crypto_util.rs Normal file
View file

@ -0,0 +1,201 @@
use aes_gcm::{
Aes256Gcm, Nonce,
aead::{Aead, KeyInit, Payload},
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STD};
use hkdf::Hkdf;
use sha2::{Digest, Sha256};
use x448::{PublicKey, Secret};
// --- Custom Errors ---
#[derive(Debug)]
pub enum SecurePayloadError {
InvalidBase64,
InvalidHex,
EncryptionError,
DecryptionError,
InvalidKeyLength,
}
// --- Data Format Enum ---
#[derive(Clone, Copy, Debug)]
pub enum DataFormat {
Raw,
Base64,
Hex,
}
// --- Main Class Structure ---
pub struct SecurePayload {
/// The internal canonical representation is always raw bytes.
inner_data: Vec<u8>,
/// The private key of the user associated with this payload instance.
private_key: Secret,
}
impl Clone for SecurePayload {
fn clone(&self) -> Self {
Self {
inner_data: self.inner_data.clone(),
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
}
}
}
impl SecurePayload {
/// Clear Constructor: Takes data in any format and the user's private key.
pub fn new<S, T: AsRef<[u8]>>(
data: T,
format: DataFormat,
private_key: S,
) -> Result<Self, SecurePayloadError>
where
S: Into<Secret>,
{
let raw_data = match format {
DataFormat::Raw => data.as_ref().to_vec(),
DataFormat::Base64 => BASE64_STD
.decode(data.as_ref())
.map_err(|_| SecurePayloadError::InvalidBase64)?,
DataFormat::Hex => {
hex::decode(data.as_ref()).map_err(|_| SecurePayloadError::InvalidHex)?
}
};
Ok(Self {
inner_data: raw_data,
private_key: private_key.into(),
})
}
/// Helper to get the public key associated with this instance's private key.
pub fn get_public_key(&self) -> [u8; 56] {
*PublicKey::from(&self.private_key).as_bytes()
}
/// Exports the internal data to the requested format
pub fn export(&self, format: DataFormat) -> String {
match format.into() {
DataFormat::Raw => String::from_utf8_lossy(&self.inner_data).to_string(),
DataFormat::Base64 => BASE64_STD.encode(&self.inner_data),
DataFormat::Hex => hex::encode(&self.inner_data),
}
}
/// Access raw bytes directly
pub fn get_bytes(&self) -> &[u8] {
&self.inner_data
}
/// Returns the SHA-256 Hash of the data in the requested format
pub fn get_hash(&self, format: DataFormat) -> String {
let mut hasher = Sha256::new();
hasher.update(&self.inner_data);
let result = hasher.finalize();
match format {
DataFormat::Raw => String::from_utf8_lossy(&result).to_string(),
DataFormat::Base64 => BASE64_STD.encode(result),
DataFormat::Hex => hex::encode(result),
}
}
/// Encrypts the held data for a specific recipient using AES-256-GCM.
/// The message will contain ONLY the ciphertext.
pub fn encrypt_x448<S>(&self, public_key: S) -> Result<SecurePayload, SecurePayloadError>
where
S: Into<PublicKey>,
{
let peer_pub = public_key.into();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
println!(
"Encryption Shared Secret (Hex): {}",
hex::encode(shared_secret.as_bytes())
);
// 3. Key & Nonce Derivation (HKDF)
// We derive 32 bytes for the key and 12 bytes for a deterministic nonce.
let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44]; // 32 (Key) + 12 (Nonce)
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::EncryptionError)?;
let key = &okm[..32];
let nonce_bytes = &okm[32..];
// 4. Encrypt with AES-256-GCM
let cipher = Aes256Gcm::new(key.into());
let nonce = Nonce::from_slice(nonce_bytes);
let ciphertext = cipher
.encrypt(
nonce,
Payload {
msg: &self.inner_data,
aad: &[],
},
)
.map_err(|_| SecurePayloadError::EncryptionError)?;
// 5. Result is ONLY the ciphertext. No key or nonce is packed.
Ok(SecurePayload {
inner_data: ciphertext,
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
})
}
/// Decrypts the held data providing the sender's public key manually.
pub fn decrypt_to_format(
&self,
peer_public_key_bytes: &[u8; 56],
output_format: DataFormat,
) -> Result<String, SecurePayloadError> {
let decrypted_instance = self.decrypt_x448(peer_public_key_bytes)?;
Ok(decrypted_instance.export(output_format))
}
/// Decrypts the held data using the internal Private Key and the provided Peer Public Key.
pub fn decrypt_x448(
&self,
peer_public_key_bytes: &[u8; 56],
) -> Result<SecurePayload, SecurePayloadError> {
// 1. Perform Exchange
let peer_pub = PublicKey::from_bytes(peer_public_key_bytes).unwrap();
let shared_secret = self.private_key.as_diffie_hellman(&peer_pub).unwrap();
// LOGGING: Shared Secret
println!(
"Decryption Shared Secret (Hex): {}",
hex::encode(shared_secret.as_bytes())
);
// 2. Key & Nonce Derivation (Must match encryption exactly)
let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
let mut okm = [0u8; 44];
hkdf.expand(b"x448-aes-gcm-no-overhead", &mut okm)
.map_err(|_| SecurePayloadError::DecryptionError)?;
let key = &okm[..32];
let nonce_bytes = &okm[32..];
// 3. Decrypt with AES-256-GCM
let cipher = Aes256Gcm::new(key.into());
let nonce = Nonce::from_slice(nonce_bytes);
let plaintext = cipher
.decrypt(
nonce,
Payload {
msg: &self.inner_data,
aad: &[],
},
)
.map_err(|_| SecurePayloadError::DecryptionError)?;
Ok(SecurePayload {
inner_data: plaintext,
private_key: Secret::from_bytes(self.private_key.as_bytes()).unwrap(),
})
}
}

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

@ -0,0 +1,289 @@
use std::fs::{self, File};
use std::io::{self, BufReader, Read};
use std::path::{Path, PathBuf};
use tokio::io::AsyncWriteExt;
use uuid::Uuid;
use zip::ZipArchive;
use crate::log;
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()
}
#[allow(dead_code)]
pub fn delete_directory(path: &str) -> bool {
let dir = Path::new(&get_directory()).join(path);
delete_dir_recursive(&dir)
}
#[allow(dead_code)]
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
}
#[allow(dead_code)]
pub fn delete_user_directory(user_id: Uuid) {
let user_dir = Path::new(&get_directory())
.join("users")
.join(user_id.to_string());
let _ = delete_dir_recursive(&user_dir);
}
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) -> Result<Vec<u8>, std::io::Error> {
let dir = Path::new(&get_directory()).join(path);
let file_path = dir.join(name);
std::fs::read(file_path)
}
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()
}
// Helper to download the zip file content to a file on disk
#[allow(dead_code)]
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 = tokio::fs::File::create(as_name).await?;
let body = response.bytes().await?;
zip_file.write_all(&body).await?;
Ok(())
}
#[allow(dead_code, deprecated)]
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(())
}
#[allow(dead_code)]
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());
let zip_path = base_dir.join(&zip_filename);
let target_dir = base_dir.join(as_name);
if let Err(e) = download_zip(url, &zip_path).await {
log!("Error downloading file: {}", e);
return;
}
let zip_path_clone = zip_path.clone();
let target_dir_clone = target_dir.clone();
let extract_result = extract_zip_contents_to_folder(&zip_path_clone, &target_dir_clone);
if let Err(e) = extract_result {
log!("Panic during ZIP extraction: {}", e);
}
if let Err(e) = tokio::fs::remove_file(&zip_path).await {
log!("Error cleaning up ZIP file {}: {}", zip_path.display(), e);
}
}

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

@ -0,0 +1,352 @@
use std::{
collections::{BTreeMap, HashMap},
fs::{self, OpenOptions},
io::Write,
path::Path,
sync::{OnceLock, mpsc},
thread,
time::{SystemTime, UNIX_EPOCH},
};
use ansi_term::Color;
use epsilon_core::{CommunicationValue, DataTypes, DataValue};
use json::JsonValue;
static LOGGER: OnceLock<mpsc::Sender<LogMessage>> = OnceLock::new();
#[derive(Clone, Copy)]
#[allow(unused)]
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,
}
pub fn startup() {
let (tx, rx) = mpsc::channel::<LogMessage>();
LOGGER.set(tx).expect("Logger already initialized");
thread::spawn(move || {
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");
for msg in rx {
let ts = fixed_box(&msg.timestamp_ms.to_string(), 13);
let sender = match msg.sender {
Some(id) => fixed_box(&id.to_string(), 19),
_ => fixed_box("", 19),
};
let line = format!("{} {} {} {}", ts, sender, msg.prefix, msg.message);
println!("{}", colorize(msg.kind, msg.is_error).paint(&line));
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: String = content.chars().take(width).collect();
let len = s.chars().count();
if len < width {
format!("[{}{}]", " ".repeat(width - len), s)
} else {
s
}
}
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()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis(),
sender,
prefix,
kind,
is_error,
message,
});
}
}
#[macro_export]
macro_rules! log {
// plain
($($arg:tt)*) => {
$crate::util::logger::log_internal(
None,
$crate::util::logger::PrintType::General,
"",
false,
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)*))
};
}
#[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,
$crate::util::logger::PrintType::General,
">",
false,
format!($($arg)*)
)
};
}
#[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,
$crate::util::logger::PrintType::General,
"<",
false,
format!($($arg)*)
)
};
}
#[macro_export]
macro_rules! log_err {
// 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(
None,
$crate::util::logger::PrintType::General,
">>",
true,
format!($($arg)*)
)
};
}
// ******** COMMUNICATION VALUES ********
pub fn log_cv_internal(
prefix: &'static str,
cv: &CommunicationValue,
print_type: Option<PrintType>,
) {
let formatted = format_cv(cv);
log_internal(
Some(cv.get_sender() as i64),
print_type.unwrap_or(PrintType::General),
prefix,
false,
formatted,
);
}
pub fn format_cv(cv: &CommunicationValue) -> String {
let mut parts = Vec::new();
let sender = cv.get_sender();
let receiver = cv.get_receiver();
if sender > 0 && receiver > 0 {
parts.push(format!("{} > {}", sender, receiver));
} else if sender > 0 {
parts.push(format!("{}", sender));
} else if receiver > 0 {
parts.push(format!("> {}", receiver));
}
let comm_type = cv.get_type().to_string();
parts.push(format!("{}", comm_type));
let data: &BTreeMap<DataTypes, DataValue> = cv.get_data_container();
let formated_data =
format_data_container(data.iter().map(|(k, v)| (k.clone(), v.clone())).collect());
parts.push(format!("{}", formated_data));
parts.join(": ")
}
fn format_data_container(data: Vec<(DataTypes, DataValue)>) -> String {
let parts: Vec<String> = data
.into_iter()
.map(|(key, value)| {
let key_str = key.to_string();
match value {
DataValue::Str(s) => format!("{}=\"{}\"", key_str, s),
DataValue::Container(inner) => {
let inner_formatted = format_data_container(inner);
format!("{}={{ {} }}", key_str, inner_formatted)
}
DataValue::Array(arr) => {
let arr_formatted = format_array(arr);
format!("{}=[{}]", key_str, arr_formatted)
}
DataValue::Bool(b) => format!("{}={}", key_str, b),
DataValue::BoolTrue => format!("{}=true", key_str),
DataValue::BoolFalse => format!("{}=false", key_str),
DataValue::Number(num) => format!("{}={}", key_str, num),
_ => "".to_string(),
}
})
.collect();
parts.join(", ")
}
fn format_array(arr: Vec<DataValue>) -> String {
let parts: Vec<String> = arr
.into_iter()
.map(|value| match value {
DataValue::Str(s) => format!("\"{}\"", s),
DataValue::Container(inner) => {
let inner_formatted = format_data_container(inner);
format!("{{ {} }}", inner_formatted)
}
DataValue::Array(inner_arr) => {
let formatted = format_array(inner_arr);
format!("[{}]", formatted)
}
DataValue::Bool(b) => b.to_string(),
DataValue::BoolTrue => "true".to_string(),
DataValue::BoolFalse => "false".to_string(),
DataValue::Number(num) => num.to_string(),
_ => String::new(),
})
.collect();
parts.join(", ")
}
#[macro_export]
macro_rules! log_cv {
($kind:expr, $cv:expr) => {
$crate::util::logger::log_cv_internal("", &$cv, Some($kind))
};
($cv:expr) => {
$crate::util::logger::log_cv_internal("", &$cv, None)
};
}
#[macro_export]
macro_rules! log_cv_in {
($kind:expr, $cv:expr) => {
$crate::util::logger::log_cv_internal("> ", &$cv, Some($kind))
};
($cv:expr) => {
$crate::util::logger::log_cv_internal("> ", &$cv, None)
};
}
#[macro_export]
macro_rules! log_cv_out {
($kind:expr, $cv:expr) => {
$crate::util::logger::log_cv_internal("< ", &$cv, Some($kind))
};
($cv:expr) => {
$crate::util::logger::log_cv_internal("< ", &$cv, None)
};
}

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

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