From e660dd3870247657d4b2445e7875df14973c6c1e Mon Sep 17 00:00:00 2001
From: Adam Shiervani <adam.shiervani@gmail.com>
Date: Mon, 24 Mar 2025 20:12:31 +0100
Subject: [PATCH] refactor(WebRTC): simplify WebRTCVideo component and enhance
 connection error handling

---
 ui/src/components/WebRTCVideo.tsx |  9 ++++---
 ui/src/routes/devices.$id.tsx     | 40 ++++++++++++++++++-------------
 2 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/ui/src/components/WebRTCVideo.tsx b/ui/src/components/WebRTCVideo.tsx
index 66d110e..505ce73 100644
--- a/ui/src/components/WebRTCVideo.tsx
+++ b/ui/src/components/WebRTCVideo.tsx
@@ -19,7 +19,7 @@ import { HDMIErrorOverlay } from "./VideoOverlay";
 import { ConnectionErrorOverlay } from "./VideoOverlay";
 import { LoadingOverlay } from "./VideoOverlay";
 
-export default function WebRTCVideo({ connectionFailed }: { connectionFailed: boolean }) {
+export default function WebRTCVideo() {
   // Video and stream related refs and states
   const videoElm = useRef<HTMLVideoElement>(null);
   const mediaStream = useRTCStore(state => state.mediaStream);
@@ -47,10 +47,9 @@ export default function WebRTCVideo({ connectionFailed }: { connectionFailed: bo
   const hdmiState = useVideoStore(state => state.hdmiState);
   const hdmiError = ["no_lock", "no_signal", "out_of_range"].includes(hdmiState);
   const isLoading = !hdmiError && !isPlaying;
-  const isConnectionError =
-    ["error", "failed", "disconnected"].includes(peerConnectionState || "") ||
-    // Connection failed is passed from the parent component and becomes true after multiple failed connection attempts, indicating we should stop trying
-    connectionFailed;
+  const isConnectionError = ["error", "failed", "disconnected", "closed"].includes(
+    peerConnectionState || "",
+  );
 
   // Keyboard related states
   const { setIsNumLockActive, setIsCapsLockActive, setIsScrollLockActive } =
diff --git a/ui/src/routes/devices.$id.tsx b/ui/src/routes/devices.$id.tsx
index 5139400..f8cbd7e 100644
--- a/ui/src/routes/devices.$id.tsx
+++ b/ui/src/routes/devices.$id.tsx
@@ -142,13 +142,27 @@ export default function KvmIdRoute() {
   const navigate = useNavigate();
   const { otaState, setOtaState, setModalView } = useUpdateStore();
 
+  const closePeerConnection = useCallback(
+    function closePeerConnection() {
+      peerConnection?.close();
+      // "closed" is a valid RTCPeerConnection state according to the WebRTC spec
+      // https://developer.mozilla.org/en-US/docs/Web/API/RTCPeerConnection/connectionState#closed
+      // However, the onconnectionstatechange event doesn't fire when close() is called manually
+      // So we need to explicitly update our state to maintain consistency
+      // I don't know why this is happening, but this is the best way I can think of to handle it
+      setPeerConnectionState("closed");
+    },
+    [peerConnection, setPeerConnectionState],
+  );
+
   useEffect(() => {
     const connectionAttemptsThreshold = 30;
     if (connectionAttempts > connectionAttemptsThreshold) {
       console.log(`Connection failed after ${connectionAttempts} attempts.`);
       setConnectionFailed(true);
+      closePeerConnection();
     }
-  }, [connectionAttempts]);
+  }, [connectionAttempts, closePeerConnection]);
 
   useEffect(() => {
     // Skip if already connected
@@ -168,21 +182,15 @@ export default function KvmIdRoute() {
       );
 
       // Fail connection if it's been over X seconds since we started connecting
-      if (elapsedTime > 30 * 1000) {
+      if (elapsedTime > 1 * 1000) {
         console.error(`Connection failed after ${elapsedTime} ms.`);
         setConnectionFailed(true);
+        closePeerConnection();
       }
     }, 1000);
 
     return () => clearInterval(interval);
-  }, [connectedAt, connectionFailed, startedConnectingAt]);
-
-  useEffect(() => {
-    if (connectionFailed) {
-      console.log("Closing peer connection");
-      peerConnection?.close();
-    }
-  }, [connectionFailed, peerConnection]);
+  }, [closePeerConnection, connectedAt, connectionFailed, startedConnectingAt]);
 
   const sdp = useCallback(
     async (event: RTCPeerConnectionIceEvent, pc: RTCPeerConnection) => {
@@ -219,7 +227,7 @@ export default function KvmIdRoute() {
           // - In device mode, the device api would timeout, the fetch would throw an error, therefore the catch block would be hit
           // Regardless, we should close the peer connection and let the useInterval handle reconnecting
           if (!res.ok) {
-            pc?.close();
+            closePeerConnection();
             console.error(`Error setting SDP - Status: ${res.status}}`, json);
             return;
           }
@@ -230,10 +238,10 @@ export default function KvmIdRoute() {
         ).catch(e => console.log(`Error setting remote description: ${e}`));
       } catch (error) {
         console.error(`Error setting SDP: ${error}`);
-        pc?.close();
+        closePeerConnection();
       }
     },
-    [navigate, params.id],
+    [closePeerConnection, navigate, params.id],
   );
 
   const connectWebRTC = useCallback(async () => {
@@ -310,9 +318,9 @@ export default function KvmIdRoute() {
       return;
     }
 
-    // This is final and we declare connection failure
+    // In certain cases, we want to never connect again. This happens when we've tried for a long time and failed
     if (connectionFailed) {
-      console.log("Connection failed. We're unable to establish a connection");
+      console.log("Connection failed. We won't attempt to connect again.");
       return;
     }
 
@@ -587,7 +595,7 @@ export default function KvmIdRoute() {
           />
 
           <div className="flex h-full overflow-hidden">
-            <WebRTCVideo connectionFailed={connectionFailed} />
+            <WebRTCVideo />
             <SidebarContainer sidebarView={sidebarView} />
           </div>
         </div>