From e8fe3e8ec97fbae6ff5310b7a5b6d58359003ad4 Mon Sep 17 00:00:00 2001
From: Xevion <xevion@xevion.dev>
Date: Wed, 1 Jan 2025 11:21:32 -0600
Subject: [PATCH] Setup websocket reconnection logic

---
 frontend/src/components/useSocket.ts | 69 +++++++++++++++++-----------
 1 file changed, 43 insertions(+), 26 deletions(-)

diff --git a/frontend/src/components/useSocket.ts b/frontend/src/components/useSocket.ts
index 7815f90..4e7e1a5 100644
--- a/frontend/src/components/useSocket.ts
+++ b/frontend/src/components/useSocket.ts
@@ -25,6 +25,7 @@ function useSocket(): UseSocketResult {
   const [downloads, setDownloads] = useState<Download[] | null>(null);
   const [executables, setExecutables] = useState<Executable[] | null>(null);
   const socketRef = useRef<WebSocket | null>(null);
+  const allowReconnectRef = useRef<boolean>(true);
 
   function deleteDownload(download_token: number) {
     if (socketRef.current == null) {
@@ -43,40 +44,56 @@ function useSocket(): UseSocketResult {
   }
 
   useEffect(() => {
-    const socket = new WebSocket(
-      (window.location.protocol === "https:" ? "wss://" : "ws://") +
-        (import.meta.env.DEV ? "localhost:5800" : window.location.host) +
-        "/ws"
-    );
-    socketRef.current = socket;
+    function connect() {
+      const socket = new WebSocket(
+        (window.location.protocol === "https:" ? "wss://" : "ws://") +
+          (import.meta.env.DEV ? "localhost:5800" : window.location.host) +
+          "/ws"
+      );
+      socketRef.current = socket;
 
-    socket.onmessage = (event) => {
-      const data = JSON.parse(event.data);
+      socket.onmessage = (event) => {
+        const data = JSON.parse(event.data);
 
-      if (data.type == undefined)
-        throw new Error("Received message without type");
+        if (data.type == undefined)
+          throw new Error("Received message without type");
 
-      switch (data.type) {
-        case "state":
-          setId(data.session.id as number);
-          setDownloads(data.session.downloads as Download[]);
-          break;
-        case "executables":
-          setExecutables(data.executables as Executable[]);
-          break;
-        default:
-          console.warn("Received unknown message type", data.type);
-      }
-    };
+        switch (data.type) {
+          case "state":
+            setId(data.session.id as number);
+            setDownloads(data.session.downloads as Download[]);
+            break;
+          case "executables":
+            setExecutables(data.executables as Executable[]);
+            break;
+          default:
+            console.warn("Received unknown message type", data.type);
+        }
+      };
 
-    socket.onclose = (event) => {
-      console.log("WebSocket connection closed", event);
-    };
+      socket.onclose = (event) => {
+        console.log("WebSocket connection closed", event);
+
+        socketRef.current = null;
+        if (allowReconnectRef.current) {
+          setId(null);
+          setDownloads(null);
+          setExecutables(null);
+
+          setTimeout(() => {
+            connect();
+          }, 3000);
+        }
+      };
+    }
+
+    connect();
 
     return () => {
       // Close the socket when the component is unmounted
       console.log("Unmounting, closing WebSocket connection");
-      socket.close();
+      socketRef.current?.close();
+      allowReconnectRef.current = false;
     };
   }, []);