diff --git a/Cargo.lock b/Cargo.lock index 972464a..4b58524 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", diff --git a/Cargo.toml b/Cargo.toml old mode 100644 new mode 100755 index 344dbc0..22121da --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/src/server/api.rs b/src/server/api.rs index 10be9fd..35ee9ec 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -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/ -> omikron for id (user / iota / omikron) diff --git a/src/server/server.rs b/src/server/server.rs old mode 100644 new mode 100755 index fe88e3c..3926f08 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -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> for HttpService { } fn call(&mut self, req: HttpRequest) -> 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> 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> 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> 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 = match String::from_utf8(bytes.to_vec()) { @@ -145,11 +124,27 @@ impl Service> 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 = 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> 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) -> 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) -> 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) -> 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) -> 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> { - 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>, Box> { 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>, Box> { 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>, Box> { // 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::, 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::, io::Error>>()?; @@ -430,20 +466,21 @@ fn load_tls_config() -> Result>, Box> { 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::, 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))) } diff --git a/src/server/socket.rs b/src/server/socket.rs index f18f525..5a8b815 100644 --- a/src/server/socket.rs +++ b/src/server/socket.rs @@ -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) { @@ -46,14 +58,13 @@ pub async fn start_connecteable_handler(connection: Arc) { 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) { }); } 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) { ))) .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."); } diff --git a/src/sql/sql.rs b/src/sql/sql.rs index bb79893..1994551 100644 --- a/src/sql/sql.rs +++ b/src/sql/sql.rs @@ -28,7 +28,7 @@ pub async fn connect() -> Result, 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> { - 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 { } 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)>( "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, Vec)>( "SELECT public_key, ip_address FROM omikrons WHERE id = CAST(? AS UNSIGNED)", ) .bind(id) - .fetch_optional(pool) + .fetch_optional(&pool) .await?; match row { diff --git a/src/util/file_util.rs b/src/util/file_util.rs new file mode 100644 index 0000000..b108810 --- /dev/null +++ b/src/util/file_util.rs @@ -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> { + 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, 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 { + 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> { + 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> { + 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 = 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); + } +} diff --git a/src/util/mod.rs b/src/util/mod.rs index 3472249..1897223 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,3 +1,4 @@ pub mod crypto_helper; pub mod crypto_util; +pub mod file_util; pub mod logger;