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

Merge branch 'main' of github.com:Tensamin/Omega

Browse source
This commit is contained in:
Alex Emmet 2026-02-11 22:45:16 +01:00
commit d509cfb09c
8 changed files with 676 additions and 185 deletions
Download patch file Download diff file Expand all files Collapse all files

13
Cargo.lock generated
View file

@ -519,9 +519,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
[[package]]
name = "bytes"
version = "1.11.0"
version = "1.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3"
checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33"
[[package]]
name = "bzip2"
@ -1487,14 +1487,13 @@ dependencies = [
[[package]]
name = "hyper-util"
version = "0.1.19"
version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "727805d60e7938b76b826a6ef209eb70eaa1812794f9424d4a4e2d740662df5f"
checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"http",
"http-body",
@ -3213,9 +3212,9 @@ dependencies = [
[[package]]
name = "system-configuration"
version = "0.6.1"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3c879d448e9d986b661742763247d3693ed13609438cf3d006f51f5368a5ba6b"
checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b"
dependencies = [
"bitflags 2.10.0",
"core-foundation 0.9.4",

39
Cargo.toml Normal file → Executable file
View file

@ -7,10 +7,41 @@ edition = "2024"
aes-gcm = "*"
ansi_term = "0.12.1"
async-trait = "0.1.89"
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"
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.8"
base64 = "0.22.1"
bytes = "1.11.0"
bytes = "1.11.1"
color-eyre = "0.6.5"
dashmap = "6.1.0"
dotenv = "0.15.0"
@ -20,7 +51,7 @@ hex = "0.4.3"
hkdf = "0.12.4"
http-body-util = "0.1.3"
hyper = { version = "1.8.1", features = ["full"] }
hyper-util = "0.1.19"
hyper-util = "0.1.20"
json = "0.12.4"
once_cell = "1.21.3"
pnet = "0.35.0"

35
src/server/api.rs
View file

@ -3,17 +3,21 @@ use crate::get_public_key;
use crate::server::omikron_manager::get_random_omikron;
use crate::sql::sql;
use crate::sql::user_online_tracker::get_iota_primary_omikron_connection;
use crate::util::file_util::load_file_vec;
use crate::{
sql::sql::{get_by_user_id, get_omikron_by_id},
util::crypto_helper::public_key_to_base64,
};
use axum::http::HeaderValue;
use base64::Engine as _;
use http_body_util::Full;
use hyper::body::Bytes;
use http_body_util::{Full, StreamBody};
use hyper::body::{Body, Bytes, Frame};
use hyper::header::{CONTENT_DISPOSITION, CONTENT_LENGTH, CONTENT_TYPE};
use hyper::{HeaderMap, Response as HttpResponse, StatusCode};
use json::JsonValue;
use json::number::Number;
use tokio::fs::File;
use tokio_util::io::ReaderStream;
pub async fn handle(
path: &str,
@ -31,12 +35,31 @@ pub async fn handle(
} else {
None
};
// api/
// get/
// omikron/
// id/
let (status, body_text) = if path_parts.len() >= 2 {
match path_parts[1] {
"download" => {
if path_parts.len() == 3 && path_parts[2] == "iota_frontend" {
let file: Bytes = load_file_vec("downloads", "iota_frontend.zip")
.unwrap()
.into();
let len = file.len();
let body = Full::new(file);
let response = HttpResponse::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "application/zip")
.header(
CONTENT_DISPOSITION,
"attachment; filename=\"iota_frontend.zip\"",
)
.header(CONTENT_LENGTH, len)
.body(body)
.unwrap();
return response;
} else {
bad_request()
}
}
"get" => match path_parts[2] {
// api/get/omikron -> any omikron
// api/get/omikron/<id> -> omikron for id (user / iota / omikron)

221
src/server/server.rs Normal file → Executable file
View file

@ -1,18 +1,16 @@
use crate::log;
use crate::server::api;
use crate::server::short_link::get_short_link;
use crate::server::socket;
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::server::conn::http2;
use hyper::{Method, StatusCode};
use hyper::{
Request as HttpRequest, Response as HttpResponse, body::Incoming, server::conn::http1, upgrade,
};
use hyper_util::rt::TokioExecutor;
use hyper_util::rt::tokio::TokioIo;
use hyper_util::service::TowerToHyperService;
use pnet::datalink::NetworkInterface;
@ -20,19 +18,22 @@ use rustls::ServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use sha1::{Digest, Sha1};
use std::error::Error;
use std::fs::{self, File};
use std::io::{self, BufReader, ErrorKind};
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::io::ErrorKind;
use std::io::{self, BufReader};
use std::net::{IpAddr, 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 tower::Service;
use crate::log;
use crate::server::api;
use crate::server::short_link::get_short_link;
use crate::server::socket;
use crate::util::file_util::load_file_buf;
#[derive(Clone)]
struct HttpService {
peer_addr: SocketAddr,
@ -51,7 +52,7 @@ impl Service<HttpRequest<Incoming>> for HttpService {
}
fn call(&mut self, req: HttpRequest<Incoming>) -> Self::Future {
let _peer_ip = self.peer_addr.ip();
let peer_ip = self.peer_addr.ip();
let (parts, body) = req.into_parts();
@ -64,20 +65,12 @@ impl Service<HttpRequest<Incoming>> for HttpService {
&& method == Method::GET
&& headers
.get("connection")
.map(|v| {
v.to_str()
.unwrap_or("")
.split(',')
.any(|s| s.trim().eq_ignore_ascii_case("upgrade"))
})
.map(|v| v.to_str().unwrap_or("").contains("Upgrade"))
.unwrap_or(false)
&& headers
.get("upgrade")
.map(|v| v.to_str().unwrap_or("").eq_ignore_ascii_case("websocket"))
.unwrap_or(false);
&& headers.get("upgrade").map(|v| v.to_str().unwrap_or("")) == Some("websocket");
if is_websocket_upgrade {
log!("Attempting WebSocket upgrade on {}", path);
log!("Attempting WebSocket upgrade on /ws");
if let Some(sec_websocket_key) = headers.get("sec-websocket-key") {
let sec_websocket_key = sec_websocket_key.to_str().unwrap_or("").to_string();
@ -92,11 +85,8 @@ impl Service<HttpRequest<Incoming>> for HttpService {
.unwrap();
let req_for_upgrade = HttpRequest::from_parts(parts, body);
let upgrades = upgrade::on(req_for_upgrade);
log!("Handling WebSocket upgrade");
socket::handle(path, upgrades);
log!("Handled WebSocket connection initiation");
Ok(response)
} else {
log!("No Sec-WebSocket-Key found in request headers");
@ -107,19 +97,8 @@ impl Service<HttpRequest<Incoming>> for HttpService {
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 whole_body =
tokio::time::timeout(Duration::from_secs(10), body.collect()).await??;
let bytes = whole_body.to_bytes();
let body_string: Option<String> = match String::from_utf8(bytes.to_vec()) {
@ -145,11 +124,27 @@ impl Service<HttpRequest<Incoming>> for HttpService {
Ok(response)
}
} else {
let response = HttpResponse::builder()
.status(StatusCode::BAD_REQUEST)
.body(Full::new(Bytes::from("No path provided")))
.unwrap();
Ok(response)
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, body_string).await)
}
};
@ -162,8 +157,58 @@ impl Service<HttpRequest<Incoming>> for HttpService {
}
}
pub fn is_local_network(addr: IpAddr) -> bool {
match addr {
IpAddr::V4(v4) => {
let octets = v4.octets();
if octets[0] == 10 {
return true;
}
if octets[0] == 172 && (16..=31).contains(&octets[1]) {
return true;
}
if octets[0] == 192 && octets[1] == 168 {
return true;
}
if octets[0] == 127 {
return true;
}
if octets[0] == 169 && octets[1] == 254 {
return true;
}
false
}
IpAddr::V6(v6) => {
let segments = v6.segments();
if (segments[0] & 0xfe00) == 0xfc00 {
return true;
}
if (segments[0] & 0xffc0) == 0xfe80 {
return true;
}
if v6.is_loopback() {
return true;
}
false
}
}
}
async fn run_http_server(port: u16) -> bool {
let ip = "0.0.0.0".to_string();
let mut ip = "0.0.0.0".to_string();
for iface in pnet::datalink::interfaces() {
let iface: NetworkInterface = iface;
if iface.ips.len() > 0 {
let ipsv = format!("{}", iface.ips[0]);
let ips: &str = ipsv.split('/').next().unwrap();
if format!("{}", ips).starts_with("10.") || format!("{}", ips).starts_with("192.") {
ip = ips.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);
@ -176,28 +221,42 @@ async fn run_http_server(port: u16) -> bool {
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!("Standard 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 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();
use hyper::server::conn::http2;
let conn = http2::Builder::new(TokioExecutor::new())
.serve_connection(io, TowerToHyperService::new(service));
// Wait for either the connection to finish naturally OR the shutdown signal
tokio::select! {
res = conn => {
if let Err(err) = res {
@ -225,6 +284,8 @@ async fn run_http_server(port: u16) -> bool {
}
}
}
log!("Standard Server shutdown complete.");
});
true
@ -237,7 +298,7 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
let iface: NetworkInterface = iface;
let ipsv = format!("{}", iface.ips[0]);
let ips: &str = ipsv.split('/').next().unwrap();
log!("{}", ips.to_string());
log!("{}", ips);
if format!("{}", ips).starts_with("10.") {
ip = ips.to_string();
}
@ -257,11 +318,13 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
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;
@ -273,6 +336,7 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
break;
}
// Accept new connections
accepted = listener.accept() => {
match accepted {
std::result::Result::Ok((stream, addr)) => {
@ -295,12 +359,14 @@ async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
};
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 {
@ -352,37 +418,7 @@ fn calculate_accept_key(key: &str) -> String {
sha1.update(key.as_bytes());
sha1.update(websocket_guid.as_bytes());
let result = sha1.finalize();
STANDARD.encode(result)
}
pub fn load_file_buf(path: &str, name: &str) -> io::Result<BufReader<File>> {
let exe = std::env::current_exe().unwrap_or_else(|_| PathBuf::from("."));
let dir = exe
.parent()
.unwrap_or(Path::new("."))
.to_string_lossy()
.to_string();
let dir = Path::new(&dir).join(path);
let file_path = dir.join(name);
if !dir.exists() {
if let Err(_) = fs::create_dir_all(&dir) {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"Directory creation failed",
));
}
}
if !file_path.exists() {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"File creation failed",
));
}
let file = File::open(&file_path)?;
Ok(BufReader::new(file))
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.
@ -397,7 +433,7 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
log!("TLS certificate 'certs/cert.pem' not found.");
return Ok(None);
}
Err(e) => return Err(e.into()),
Err(e) => return Err(e.into()), // Other IO error
};
let key_file_buf = match key_file_res {
@ -406,7 +442,7 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
log!("TLS key 'certs/cert.key' not found.");
return Ok(None);
}
Err(e) => return Err(e.into()),
Err(e) => return Err(e.into()), // Other IO error
};
// Continue with configuration if both files were found
@ -417,12 +453,12 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn 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))
.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")?);
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>>()?;
@ -430,20 +466,21 @@ fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
if key_ders.is_empty() {
// EC
key_reader = BufReader::new(load_file_buf("certs", "cert.key")?);
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 valid private keys found in key file (Tried PKCS8, RSA and EC).".into());
return Err("No private keys found in key file. (Tried PKCS8, RSA, and EC)".into());
}
let config = rustls::ServerConfig::builder()
let mut 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()))?;
.with_single_cert(cert_ders, key_ders.remove(0))?;
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
Ok(Some(Arc::new(config)))
}

52
src/server/socket.rs
View file

@ -12,9 +12,13 @@ use crate::server::omikron_connection::OmikronConnection;
pub fn handle(path: String, upgrades: OnUpgrade) {
tokio::spawn(async move {
log!(
"[ws] Spawning new task to handle WebSocket upgrade for path: {}",
path
);
match upgrades.await {
Ok(upgraded_stream) => {
log!("Valid WebSocket upgrade");
log!("[ws] WebSocket upgrade successful for path: {}", path);
let raw_stream = TokioIo::new(upgraded_stream);
let ws_stream = WebSocketStream::from_raw_socket(
@ -24,20 +28,28 @@ pub fn handle(path: String, upgrades: OnUpgrade) {
)
.await;
log!(
"WebSocket handshake successful, handling connection for {}",
"[ws] WebSocket handshake successful, handling connection for {}",
path
);
let (writer, reader) = ws_stream.split();
if path == "/ws/omikron" {
let connection = OmikronConnection::new(writer, reader);
start_connecteable_handler(connection).await;
tokio::spawn(start_connecteable_handler(connection));
}
}
Err(e) => {
log!("WebSocket upgrade failed after response: {:?}", e);
log!(
"[ERROR] WebSocket upgrade failed for path {}: {:?}",
path,
e
);
}
}
log!(
"[ws] WebSocket handling task for path: {} is finished.",
path
);
});
}
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
@ -46,14 +58,13 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
log!("[ws_handler] Starting connection handler loop.");
loop {
let mut receiver = connection.receiver.write().await;
let mut receiver_guard = connection.receiver.write().await;
match tokio::time::timeout(IDLE_TIMEOUT, receiver.next()).await {
match tokio::time::timeout(IDLE_TIMEOUT, receiver_guard.next()).await {
Ok(Some(Ok(msg))) => {
// Drop the lock so other tasks can use the receiver if needed,
// and so we can handle the message without holding the lock.
drop(receiver);
drop(receiver_guard);
match msg {
Message::Text(text) => {
@ -63,28 +74,30 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
});
}
Message::Close(_) => {
log!("[ws_handler] Received 'Close' message. Breaking loop.");
break;
}
Message::Pong(_) => {
// Received a pong, connection is alive.
log!("[ws_handler] Received 'Pong'. Connection is alive.");
}
_ => {
log!("[ws_handler] Received unhandled message type.");
}
_ => {}
}
}
Ok(Some(Err(e))) => {
log!("WS Error: {}", e);
log!("[ERROR] WS Error: {}. Breaking loop.", e);
break;
}
Ok(None) => {
// Stream is closed
log!("[ws_handler] WebSocket stream closed by peer. Breaking loop.");
break;
}
Err(_) => {
// Timeout, we need to send a ping.
// Drop receiver lock before acquiring sender lock to avoid deadlock.
drop(receiver);
log!("WebSocket connection is idle. Sending a ping.");
drop(receiver_guard);
log!("[ws_handler] Timeout: Dropped receiver lock. Sending a ping.");
let mut sender = connection.sender.write().await;
log!("[ws_handler] Acquired sender lock for ping.");
if let Err(e) = sender
.send(Message::Text(Utf8Bytes::from(
CommunicationValue::new(CommunicationType::ping)
@ -93,11 +106,14 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
)))
.await
{
log!("Failed to send ping: {}. Closing connection.", e);
log!("[ERROR] Failed to send ping: {}. Closing connection.", e);
break;
}
log!("[ws_handler] Ping sent successfully.");
}
}
}
log!("[ws_handler] Connection handler loop finished.");
connection.handle_close().await;
log!("[ws_handler] Connection closed.");
}

211
src/sql/sql.rs
View file

@ -28,7 +28,7 @@ pub async fn connect() -> Result<Pool<MySql>, sqlx::Error> {
let table = env::var("DB_TABLE").expect("DB_TABLE is not set");
MySqlPoolOptions::new()
.max_connections(5)
.max_connections(200)
.connect(&format!(
"mysql://{}:{}@127.0.0.1:3306/{}",
user, passwd, table
@ -171,14 +171,19 @@ pub async fn get_by_username(
),
sqlx::Error,
> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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)
.fetch_optional(&pool)
.await?;
match row {
@ -234,14 +239,19 @@ pub async fn get_by_user_id(
),
sqlx::Error,
> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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)
.fetch_optional(&pool)
.await?;
match row {
@ -297,14 +307,19 @@ pub async fn get_users_by_iota_id(
)>,
sqlx::Error,
> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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)
.fetch_all(&pool)
.await?;
let mut users = Vec::new();
@ -342,89 +357,124 @@ pub async fn get_users_by_iota_id(
}
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");
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)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_display_name(id: i64, new_display: String) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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)
.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");
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)
.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");
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)
.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");
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)
.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");
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)
.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");
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)
.execute(&pool)
.await?;
Ok(())
@ -435,8 +485,13 @@ pub async fn change_keys(
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");
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)",
@ -444,19 +499,24 @@ pub async fn change_keys(
.bind(new_public_key)
.bind(new_private_key_hash)
.bind(id)
.execute(pool)
.execute(&pool)
.await?;
Ok(())
}
pub async fn change_token(id: i64, new_token: String) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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)
.execute(&pool)
.await?;
Ok(())
@ -468,8 +528,13 @@ pub async fn register_complete_user(
iota_id: i64,
token: String,
) -> Result<(), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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 (?, ?, ?, ?, ?)",
@ -479,20 +544,25 @@ pub async fn register_complete_user(
.bind(public_key)
.bind(iota_id)
.bind(token)
.execute(pool)
.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");
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)
.fetch_all(&pool)
.await?
.iter()
{
@ -533,27 +603,37 @@ pub async fn create_new_iota(public_key: String) -> Result<i64, sqlx::Error> {
}
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");
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)
.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");
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)
.fetch_optional(&pool)
.await;
match result {
@ -569,25 +649,35 @@ pub async fn get_iota_by_id(id: i64) -> Result<(i64, String), sqlx::Error> {
}
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");
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)
.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");
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)
.execute(&pool)
.await?;
Ok(())
@ -598,14 +688,19 @@ pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> {
// ==========================================================================================
pub async fn get_omikron_by_id(id: i64) -> Result<(String, String), sqlx::Error> {
let db_lock = SQL_DB.read().await;
let pool = db_lock.as_ref().expect("Database pool is not initialized");
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)
.fetch_optional(&pool)
.await?;
match row {

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

1
src/util/mod.rs
View file

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