diff --git a/Cargo.lock b/Cargo.lock index 22b7dc0..fec9aac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -61,6 +61,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "android-tzdata" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "async-trait" version = "0.1.83" @@ -165,6 +180,21 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chrono" +version = "0.4.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e36cc9d416881d2e24f9a963be5fb1cd90966419ac844274161d10488b3e825" +dependencies = [ + "android-tzdata", + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-targets", +] + [[package]] name = "cipher" version = "0.4.4" @@ -283,6 +313,7 @@ dependencies = [ name = "dynamic-preauth" version = "0.1.0" dependencies = [ + "chrono", "rand", "salvo", "serde", @@ -667,6 +698,29 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.61" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "1.5.0" @@ -2473,6 +2527,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-registry" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index ec21290..12b6321 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition = "2021" [dependencies] +chrono = { version = "0.4.39", features = ["serde"] } rand = "0.8.5" salvo = { version = "0.74.3", features = ["affix-state", "logging", "serve-static"] } serde = { version = "1.0.216", features = ["derive"] } diff --git a/src/main.rs b/src/main.rs index 59a92f4..26858f5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,12 @@ use std::sync::LazyLock; -use salvo::http::HeaderValue; +use salvo::http::{HeaderValue, StatusCode}; use salvo::logging::Logger; use salvo::prelude::{ handler, Listener, Request, Response, Router, Server, Service, StaticDir, TcpListener, }; use salvo::writing::Json; +use salvo::Depot; use tokio::sync::Mutex; use crate::models::State; @@ -16,7 +17,7 @@ mod models; mod utility; #[handler] -async fn session_middleware(req: &mut Request, res: &mut Response) { +async fn session_middleware(req: &mut Request, res: &mut Response, depot: &mut Depot) { match req.cookie("Session") { Some(cookie) => { // Check if the session exists @@ -24,18 +25,23 @@ async fn session_middleware(req: &mut Request, res: &mut Response) { Ok(session_id) => { let mut store = STORE.lock().await; if !store.sessions.contains_key(&session_id) { - store.new_session(res).await; + let id = store.new_session(res).await; + depot.insert("session_id", id); } } Err(_) => { let mut store = STORE.lock().await; - store.new_session(res).await; + let id = store.new_session(res).await; + + depot.insert("session_id", id); } } } None => { let mut store = STORE.lock().await; - store.new_session(res).await; + let id = store.new_session(res).await; + + depot.insert("session_id", id); } } } @@ -64,18 +70,34 @@ pub async fn download(req: &mut Request, res: &mut Response) { } #[handler] -pub async fn get_session(req: &mut Request, res: &mut Response) { +pub async fn get_session(req: &mut Request, res: &mut Response, depot: &mut Depot) { let store = STORE.lock().await; - let session_id = req - .cookie("Session") - .unwrap() - .value() - .parse::() - .unwrap(); - let session = store.sessions.get(&session_id).unwrap(); + let session_id = match req.cookie("Session") { + Some(cookie) => match cookie.value().parse::() { + Ok(id) => id, + _ => { + res.status_code(StatusCode::BAD_REQUEST); + return; + } + }, + None => match depot.get::("session_id") { + Ok(id) => *id, + _ => { + res.status_code(StatusCode::BAD_REQUEST); + return; + } + }, + }; - res.render(Json(&session)); + match store.sessions.get(&session_id) { + Some(session) => { + res.render(Json(&session)); + } + None => { + res.status_code(StatusCode::BAD_REQUEST); + } + } } #[tokio::main] diff --git a/src/models.rs b/src/models.rs index 9832879..84006f5 100644 --- a/src/models.rs +++ b/src/models.rs @@ -9,10 +9,8 @@ use crate::utility::search; #[derive(Clone, Debug, Serialize)] pub struct Session { pub tokens: Vec, - #[serde(skip_serializing)] - pub last_seen: std::time::Instant, - #[serde(skip_serializing)] - pub first_seen: std::time::Instant, + pub last_seen: chrono::DateTime, + pub first_seen: chrono::DateTime, } #[derive(Default, Clone, Debug)] @@ -56,12 +54,13 @@ impl<'a> State<'a> { let mut rng = rand::thread_rng(); let id: usize = rng.gen(); + let now = chrono::Utc::now(); self.sessions.insert( id, Session { tokens: vec![], - last_seen: std::time::Instant::now(), - first_seen: std::time::Instant::now(), + last_seen: now, + first_seen: now, }, );