From 45a9e3734280e19d300a2a149d10ac56c394cc8a Mon Sep 17 00:00:00 2001
From: Dave Richer <dave@imexsystems.ca>
Date: Wed, 25 Sep 2024 14:12:31 -0400
Subject: [PATCH] IO-2935-Add-Enhanced-Websocket-Provider - Add Firebase token
 refresh

Signed-off-by: Dave Richer <dave@imexsystems.ca>
---
 client/src/contexts/SocketIO/useSocket.js | 36 +++++++++++++++--------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/client/src/contexts/SocketIO/useSocket.js b/client/src/contexts/SocketIO/useSocket.js
index f75fd3fea..c5688cf3c 100644
--- a/client/src/contexts/SocketIO/useSocket.js
+++ b/client/src/contexts/SocketIO/useSocket.js
@@ -4,22 +4,35 @@ import { auth } from "../../firebase/firebase.utils";
 
 const useSocket = (bodyshop) => {
   const [socket, setSocket] = useState(null);
-  const [clientId, setClientId] = useState(null); // State to store unique identifier
+  const [clientId, setClientId] = useState(null);
+  const [token, setToken] = useState(null);
 
   useEffect(() => {
-    if (bodyshop && bodyshop.id) {
+    // Listener for token changes
+    const unsubscribe = auth.onIdTokenChanged(async (user) => {
+      if (user) {
+        const newToken = await user.getIdToken();
+        setToken(newToken);
+      } else {
+        setToken(null);
+      }
+    });
+
+    // Clean up the listener on unmount
+    return () => unsubscribe();
+  }, []);
+
+  useEffect(() => {
+    if (bodyshop && bodyshop.id && token) {
       const endpoint = import.meta.env.PROD ? import.meta.env.VITE_APP_AXIOS_BASE_API_URL : "";
 
       const socketInstance = SocketIO(endpoint, {
-        path: "/wss", // Ensure this matches the Vite proxy and backend path
+        path: "/wss",
         withCredentials: true,
-        auth: async (callback) => {
-          const token = auth.currentUser && (await auth.currentUser.getIdToken());
-          callback({ token });
-        },
-        reconnectionAttempts: Infinity, // Try reconnecting forever
-        reconnectionDelay: 2000, // How long to wait between reconnection attempts
-        reconnectionDelayMax: 10000 // Maximum delay between attempts
+        auth: { token }, // Use the current token
+        reconnectionAttempts: Infinity,
+        reconnectionDelay: 2000,
+        reconnectionDelayMax: 10000
       });
 
       setSocket(socketInstance);
@@ -62,9 +75,8 @@ const useSocket = (bodyshop) => {
         socketInstance.disconnect();
       };
     }
-  }, [bodyshop]);
+  }, [bodyshop, token]); // Include 'token' in dependencies to re-run effect when it changes
 
-  // Return both socket and clientId
   return { socket, clientId };
 };