diff --git a/Cargo.toml b/Cargo.toml index 4d9de10..a5e3049 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ path = "src/main.rs" rustdoc-args = ["--document-private-items"] [dependencies] +actix = "0.13.5" actix-rt = "2.10.0" actix-web = { version = "4.9.0", features = ["openssl"] } actix-ws = "0.3.0" @@ -46,7 +47,6 @@ reqwest = { version = "0.11", features = ["blocking", "json"] } minijinja = { version = "2.3.1", features = ["loader"] } url = "2.5.2" regex = "1.10.6" -Inflector = "0.11" openssl = "0.10.64" dotenv = "0.15.0" futures-util = "0.3.30" diff --git a/src/resources/disks.rs b/src/resources/disks.rs index 9643039..6bfe057 100644 --- a/src/resources/disks.rs +++ b/src/resources/disks.rs @@ -4,8 +4,6 @@ use regex::Regex; use serde_json::Value; use crate::squire; -// todo: tested only on macOS - /// Function to convert byte size to human-readable format. /// /// # Arguments diff --git a/src/resources/info.rs b/src/resources/info.rs index d966f94..eb45869 100644 --- a/src/resources/info.rs +++ b/src/resources/info.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; use chrono::Utc; -use inflector::cases::titlecase::to_title_case; use sysinfo::{DiskExt, System, SystemExt}; use crate::{squire, resources}; @@ -44,7 +43,7 @@ pub fn get_sys_info() -> HashMap<&'static str, HashMap<&'static str, String>> { let os_arch = resources::system::os_arch(); let basic = HashMap::from_iter(vec![ ("node", sys.host_name().unwrap_or("Unknown".to_string())), - ("system", to_title_case(&os_arch.name)), + ("system", squire::util::capwords(&os_arch.name, None)), ("architecture", os_arch.architecture), ("uptime", uptime), ("CPU_cores_raw", sys.cpus().len().to_string() diff --git a/src/routes/websocket.rs b/src/routes/websocket.rs index df6bde1..8baa216 100644 --- a/src/routes/websocket.rs +++ b/src/routes/websocket.rs @@ -1,5 +1,5 @@ use crate::{constant, resources, routes, squire}; -use actix_web::rt::time::sleep; +use actix; use actix_web::{rt, web, Error, HttpRequest, HttpResponse}; use actix_ws::AggregatedMessage; use fernet::Fernet; @@ -13,19 +13,20 @@ use std::time::Duration; /// # Arguments /// /// * `request` - A reference to the Actix web `HttpRequest` object. -async fn send_system_resources(mut session: actix_ws::Session) { +async fn send_system_resources(request: HttpRequest, mut session: actix_ws::Session) { + let host = request.connection_info().host().to_string(); loop { let system_resources = resources::stream::system_resources(); let serialized = serde_json::to_string(&system_resources).unwrap(); match session.text(serialized).await { Ok(_) => (), Err(err) => { - log::warn!("Connection {} by the client!", err); + log::info!("Connection from '{}' has been {}", host, err.to_string().to_lowercase()); break; } } // 500ms / 0.5s delay - sleep(Duration::from_secs(1)).await; + rt::time::sleep(Duration::from_secs(1)).await; } } @@ -49,7 +50,7 @@ async fn send_system_resources(mut session: actix_ws::Session) { /// * `stream` - A stream of `AggregatedMessage` objects. async fn receive_messages( mut session: actix_ws::Session, - mut stream: impl futures::Stream> + Unpin + mut stream: impl futures::Stream> + Unpin, ) { while let Some(msg) = stream.next().await { match msg { @@ -70,6 +71,20 @@ async fn receive_messages( } } +/// Handles the session by closing it after a certain duration. +/// +/// # Arguments +/// +/// * `session` - A reference to the Actix web `Session` object. +/// * `duration` - Duration in seconds to keep the session alive. +async fn session_handler(session: actix_ws::Session, duration: i64) { + let session = session.clone(); + actix::spawn(async move { + rt::time::sleep(Duration::from_secs(duration as u64)).await; + let _ = session.close(None).await; + }); +} + /// Handles the WebSocket endpoint for system resources. /// /// # Arguments @@ -100,17 +115,16 @@ async fn echo( Ok(result) => result, Err(_) => { return Ok(HttpResponse::ServiceUnavailable().finish()); - }, + } }; - // todo: implement a session timeout here let stream = stream .aggregate_continuations(); rt::spawn(async move { log::warn!("Connection established"); - let send_task = send_system_resources(session.clone()); - let receive_task = receive_messages(session, stream); - future::join(send_task, receive_task).await; + let send_task = send_system_resources(request.clone(), session.clone()); + let receive_task = receive_messages(session.clone(), stream); + let session_task = session_handler(session.clone(), config.session_duration); + future::join3(send_task, receive_task, session_task).await; }); - // respond immediately with response connected to WS session Ok(response) } diff --git a/src/squire/util.rs b/src/squire/util.rs index 459949b..17e45b9 100644 --- a/src/squire/util.rs +++ b/src/squire/util.rs @@ -120,3 +120,30 @@ pub fn run_command(command: &str, args: &[&str]) -> Result { } } } + +/// Function to capitalize the first letter of each word in a string. +/// +/// # Arguments +/// +/// * `s` - The string to capitalize +/// * `sep` - The separator to split the string +/// +/// # Returns +/// +/// A `String` containing the capitalized string +pub fn capwords(string: &str, sep: Option<&str>) -> String { + let separator = sep.unwrap_or(" "); + let mut result = Vec::new(); + + for word in string.split(separator) { + let mut chars = word.chars(); + if let Some(first) = chars.next() { + let capitalized_word = first.to_uppercase().collect::() + chars.as_str(); + result.push(capitalized_word); + } else { + result.push(String::new()); + } + } + + result.join(separator) +} diff --git a/src/templates/mod.rs b/src/templates/mod.rs index 7f3b747..180259e 100644 --- a/src/templates/mod.rs +++ b/src/templates/mod.rs @@ -1,6 +1,6 @@ use std::sync::Arc; use minijinja::{value::Value}; -use inflector::cases::titlecase::to_title_case; +use crate::squire; /// Index page template that is served as HTML response for the root endpoint. mod index; @@ -44,7 +44,7 @@ fn capwords_filter(value: Value) -> Result { Ok(Value::from(result)) } else { let result = val.replace("_", " "); - Ok(Value::from(to_title_case(&result))) + Ok(Value::from(squire::util::capwords(&result, None))) } } else { panic!("capwords filter only works with strings");