finish porting session middleware handler, Session model & creator

Cargo.lock

@@ -283,7 +283,9 @@ dependencies = [
 name = "dynamic-preauth"
 version = "0.1.0"
 dependencies = [
+ "rand",
  "salvo",
+ "serde",
  "tokio",
  "tracing",
  "tracing-subscriber",

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"

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<Mutex<State>> = 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::<usize>() {
+                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::<String>("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::<usize>()
+        .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/<id>").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());

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<String>,
+    #[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<usize, Session>,
 }
 
 impl<'a> State<'a> {
-    pub(crate) fn new() -> Mutex<Self> {
+    pub fn new() -> Mutex<Self> {
         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<u8>,
     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<u8> {
+    pub fn with_key(&self, new_key: &[u8]) -> Vec<u8> {
         let mut data = self.data.clone();
 
         // Copy the key into the data