From b9e44e1fa3703967aa8ce9784a43590afd091fd0 Mon Sep 17 00:00:00 2001
From: Riz Ashraf <reazul.ashraf@inseinc.com>
Date: Sun, 20 Sep 2026 23:08:40 +0100
Subject: [PATCH] fix(server): Fix memory leak in handle_socket on abrupt
 disconnects

- Used a SessionCleanup Drop guard (RAII pattern) to guarantee that the websocket session is evicted from the global clients map and background read/write tasks are forcibly aborted even when Axum terminates the handler mid-execution due to a TCP connection loss.
---
 server/src/main.rs | 51 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/server/src/main.rs b/server/src/main.rs
index 9d50188..1b93a5f 100644
--- a/server/src/main.rs
+++ b/server/src/main.rs
@@ -557,22 +557,45 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
         );
     });
 
-    tokio::select! {
-        _ = (&mut send_task) => {
-            tracing::info!("Websocket send task finished for session {}", session_id);
-            recv_task.abort();
-        },
-        _ = (&mut recv_task) => {
-            tracing::info!("Websocket recv task finished for session {}", session_id);
-            send_task.abort();
-        },
+    struct SessionCleanup {
+        session_id: String,
+        state: Arc<AppState>,
+        send_task: Option<tokio::task::JoinHandle<()>>,
+        recv_task: Option<tokio::task::JoinHandle<()>>,
+    }
+
+    impl Drop for SessionCleanup {
+        fn drop(&mut self) {
+            self.state.clients.write().unwrap().remove(&self.session_id);
+            if let Some(task) = self.send_task.take() {
+                task.abort();
+            }
+            if let Some(task) = self.recv_task.take() {
+                task.abort();
+            }
+            tracing::info!(
+                "Websocket session {} closed and cleaned up",
+                self.session_id
+            );
+        }
+    }
+
+    let mut cleanup = SessionCleanup {
+        session_id: session_id.clone(),
+        state: Arc::clone(&state),
+        send_task: Some(send_task),
+        recv_task: Some(recv_task),
     };
 
-    state.clients.write().unwrap().remove(&session_id);
-    tracing::info!(
-        "Websocket session {} closed and removed from state",
-        session_id
-    );
+    tokio::select! {
+        _ = cleanup.send_task.as_mut().unwrap() => {
+            tracing::info!("Websocket send task finished for session {}", session_id);
+        },
+        _ = cleanup.recv_task.as_mut().unwrap() => {
+            tracing::info!("Websocket recv task finished for session {}", session_id);
+        },
+    };
+    // Drop guard automatically handles removal and aborts the other task.
 }
 
 #[derive(serde::Deserialize, serde::Serialize, Debug)]