From bc39909f3c4206ac92f08d067b970da745c4b73c Mon Sep 17 00:00:00 2001
From: Xevion <xevion@xevion.dev>
Date: Mon, 23 Dec 2024 19:53:26 -0600
Subject: [PATCH] Apply Session middleware to all but /notify with separate
 router, improve /notify key error handling

---
 src/main.rs | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index aaf5fc3..58aecc3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -203,7 +203,7 @@ pub async fn download(req: &mut Request, res: &mut Response, depot: &mut Depot)
 }
 
 #[handler]
-pub async fn notify(req: &mut Request, res: &mut Response, depot: &mut Depot) {
+pub async fn notify(req: &mut Request, res: &mut Response) {
     let key = req.query::<String>("key");
 
     match key {
@@ -219,14 +219,23 @@ pub async fn notify(req: &mut Request, res: &mut Response, depot: &mut Depot) {
             };
 
             let store = &mut *STORE.lock().await;
-            let session = store.sessions.get_mut(&key).expect("Session not found");
+            let session = store.sessions.get_mut(&key);
 
-            let message = OutgoingMessage::TokenAlert { token: key };
-            session
-                .send_message(message)
-                .expect("Failed to buffer token alert message");
+            match session {
+                Some(session) => {
+                    let message = OutgoingMessage::TokenAlert { token: key };
+                    session
+                        .send_message(message)
+                        .expect("Failed to buffer token alert message");
 
-            res.render("Notification sent");
+                    res.render("Notification sent");
+                }
+                None => {
+                    tracing::warn!("Session not found for key while attempting notify: {}", key);
+                    res.status_code(StatusCode::UNAUTHORIZED);
+                    return;
+                }
+            }
         }
         None => {
             res.status_code(StatusCode::BAD_REQUEST);
@@ -337,12 +346,18 @@ async fn main() {
     let router = Router::new()
         .hoop(CatchPanic::new())
         .hoop(cors)
-        .hoop(session_middleware)
-        .push(Router::with_path("download/<id>").get(download))
-        .push(Router::with_path("session").get(get_session))
-        .push(Router::with_path("ws").goal(connect))
+        // /notify does not need a session, nor should it have one
         .push(Router::with_path("notify").post(notify))
-        .push(Router::with_path("<**path>").get(static_dir));
+        .push(
+            Router::new()
+                .hoop(session_middleware)
+                .push(Router::with_path("download/<id>").get(download))
+                .push(Router::with_path("session").get(get_session))
+                // websocket /ws
+                .push(Router::with_path("ws").goal(connect))
+                // static files
+                .push(Router::with_path("<**path>").get(static_dir)),
+        );
 
     let service = Service::new(router).hoop(Logger::new());