From de3dacda771a02d290c54bea8370020b4b049398 Mon Sep 17 00:00:00 2001
From: Xevion <xevion@xevion.dev>
Date: Mon, 23 Dec 2024 19:36:51 -0600
Subject: [PATCH] Fix notify sending TokenAlert to incorrect session id, use
 'notify' id for TokenAlert name

---
 src/main.rs   | 31 ++++++++++++++++---------------
 src/models.rs | 18 +++++++++++++-----
 2 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index c273900..aaf5fc3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -106,7 +106,9 @@ async fn handle_socket(session_id: u32, websocket: WebSocket) {
     session.tx = Some(tx_channel);
 
     session.send_state();
-    session.send_message(executable_message);
+    session
+        .send_message(executable_message)
+        .expect("Failed to buffer executables message");
 
     // Handle incoming messages
     let fut = async move {
@@ -207,23 +209,22 @@ pub async fn notify(req: &mut Request, res: &mut Response, depot: &mut Depot) {
     match key {
         Some(key) => {
             // Parse key into u32
-            let key = key.parse::<u32>();
-            if key.is_err() {
-                res.status_code(StatusCode::BAD_REQUEST);
-                return;
-            }
+            let key = match key.parse::<u32>() {
+                Ok(k) => k,
+                Err(e) => {
+                    tracing::error!("Error parsing key: {}", e);
+                    res.status_code(StatusCode::BAD_REQUEST);
+                    return;
+                }
+            };
 
             let store = &mut *STORE.lock().await;
-            let session_id = get_session_id(req, depot).unwrap();
-            let session = store
-                .sessions
-                .get_mut(&session_id)
-                .expect("Session not found");
+            let session = store.sessions.get_mut(&key).expect("Session not found");
 
-            let message = OutgoingMessage::TokenAlert {
-                token: key.unwrap(),
-            };
-            session.send_message(message);
+            let message = OutgoingMessage::TokenAlert { token: key };
+            session
+                .send_message(message)
+                .expect("Failed to buffer token alert message");
 
             res.render("Notification sent");
         }
diff --git a/src/models.rs b/src/models.rs
index 1f89c8c..e527884 100644
--- a/src/models.rs
+++ b/src/models.rs
@@ -51,9 +51,10 @@ impl Session {
         return self.downloads.last().unwrap();
     }
 
+    // This function's failure is not a failure to transmit the message, but a failure to buffer it into the channel (or any preceding steps).
     pub fn send_message(&mut self, message: OutgoingMessage) -> Result<(), anyhow::Error> {
         if self.tx.is_none() {
-            return Err(anyhow::anyhow!("Session has no sender"));
+            return Err(anyhow::anyhow!("Session {} has no sender", self.id));
         }
 
         // TODO: Error handling
@@ -210,13 +211,20 @@ pub enum IncomingMessage {
 }
 
 #[derive(Debug, Serialize)]
-#[serde(tag = "type", rename_all = "lowercase")]
+#[serde(tag = "type", rename_all = "kebab-case")]
 pub enum OutgoingMessage {
     // An alert to the client that a session download has been used.
-    TokenAlert { token: u32 },
+    #[serde(rename = "notify")]
+    TokenAlert {
+        token: u32,
+    },
     // A message describing the current session state
-    State { session: Session },
-    Executables { executables: Vec<ExecutableJson> },
+    State {
+        session: Session,
+    },
+    Executables {
+        executables: Vec<ExecutableJson>,
+    },
 }
 
 #[derive(Debug, Serialize)]