diff --git a/Cargo.lock b/Cargo.lock index 9830798..22b7dc0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -283,7 +283,9 @@ dependencies = [ name = "dynamic-preauth" version = "0.1.0" dependencies = [ + "rand", "salvo", + "serde", "tokio", "tracing", "tracing-subscriber", diff --git a/Cargo.toml b/Cargo.toml index 817449e..ec21290 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,9 @@ version = "0.1.0" edition = "2021" [dependencies] +rand = "0.8.5" salvo = { version = "0.74.3", features = ["affix-state", "logging", "serve-static"] } +serde = { version = "1.0.216", features = ["derive"] } tokio = { version = "1", features = ["macros"] } tracing = "0.1" tracing-subscriber = "0.3" diff --git a/src/main.rs b/src/main.rs index 33210a8..59a92f4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ use salvo::logging::Logger; use salvo::prelude::{ handler, Listener, Request, Response, Router, Server, Service, StaticDir, TcpListener, }; +use salvo::writing::Json; use tokio::sync::Mutex; use crate::models::State; @@ -14,6 +15,31 @@ static STORE: LazyLock> = LazyLock::new(State::new); mod models; mod utility; +#[handler] +async fn session_middleware(req: &mut Request, res: &mut Response) { + match req.cookie("Session") { + Some(cookie) => { + // Check if the session exists + match cookie.value().parse::() { + Ok(session_id) => { + let mut store = STORE.lock().await; + if !store.sessions.contains_key(&session_id) { + store.new_session(res).await; + } + } + Err(_) => { + let mut store = STORE.lock().await; + store.new_session(res).await; + } + } + } + None => { + let mut store = STORE.lock().await; + store.new_session(res).await; + } + } +} + #[handler] pub async fn download(req: &mut Request, res: &mut Response) { let article_id = req.param::("id").unwrap(); @@ -37,6 +63,21 @@ pub async fn download(req: &mut Request, res: &mut Response) { ); } +#[handler] +pub async fn get_session(req: &mut Request, res: &mut Response) { + let store = STORE.lock().await; + + let session_id = req + .cookie("Session") + .unwrap() + .value() + .parse::() + .unwrap(); + let session = store.sessions.get(&session_id).unwrap(); + + res.render(Json(&session)); +} + #[tokio::main] async fn main() { let port = std::env::var("PORT").unwrap_or_else(|_| "5800".to_string()); @@ -51,7 +92,9 @@ async fn main() { let static_dir = StaticDir::new(["./public"]).defaults("index.html"); let router = Router::new() + .hoop(session_middleware) .push(Router::with_path("download/").get(download)) + .push(Router::with_path("session").get(get_session)) .push(Router::with_path("<**path>").get(static_dir)); let service = Service::new(router).hoop(Logger::new()); diff --git a/src/models.rs b/src/models.rs index f816751..9832879 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1,21 +1,35 @@ +use rand::{distributions::Alphanumeric, Rng}; +use salvo::{http::cookie::Cookie, Response}; +use serde::Serialize; use std::{collections::HashMap, path}; use tokio::sync::Mutex; 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, +} + #[derive(Default, Clone, Debug)] -pub(crate) struct State<'a> { +pub struct State<'a> { pub executables: HashMap<&'a str, Executable>, + pub sessions: HashMap, } impl<'a> State<'a> { - pub(crate) fn new() -> Mutex { + pub fn new() -> Mutex { Mutex::new(Self { executables: HashMap::new(), + sessions: HashMap::new(), }) } - pub(crate) fn add_executable(&mut self, exe_type: &'a str, exe_path: &str) { + pub fn add_executable(&mut self, exe_type: &'a str, exe_path: &str) { let data = std::fs::read(&exe_path).expect("Unable to read file"); let pattern = "a".repeat(1024); @@ -37,10 +51,32 @@ impl<'a> State<'a> { self.executables.insert(exe_type, exe); } + + pub async fn new_session(&mut self, res: &mut Response) -> usize { + let mut rng = rand::thread_rng(); + let id: usize = rng.gen(); + + self.sessions.insert( + id, + Session { + tokens: vec![], + last_seen: std::time::Instant::now(), + first_seen: std::time::Instant::now(), + }, + ); + + res.add_cookie( + Cookie::build(("Session", id.to_string())) + .permanent() + .build(), + ); + + return id; + } } #[derive(Default, Clone, Debug)] -pub(crate) struct Executable { +pub struct Executable { pub data: Vec, pub filename: String, pub key_start: usize, @@ -48,7 +84,7 @@ pub(crate) struct Executable { } impl Executable { - pub(crate) fn with_key(&self, new_key: &[u8]) -> Vec { + pub fn with_key(&self, new_key: &[u8]) -> Vec { let mut data = self.data.clone(); // Copy the key into the data