From 344a549f0b83b8fe289c5b6be8c67544bbf61ba3 Mon Sep 17 00:00:00 2001
From: Adam Shiervani <adam.shiervani@gmail.com>
Date: Mon, 24 Mar 2025 18:20:05 +0100
Subject: [PATCH] feat(WebRTC): enhance connection management with connection
 failures after X attempts or a certain time

---
 ui/src/components/WebRTCVideo.tsx |  9 +--
 ui/src/routes/devices.$id.tsx     | 98 +++++++++++++++++++++++++++----
 2 files changed, 92 insertions(+), 15 deletions(-)

diff --git a/ui/src/components/WebRTCVideo.tsx b/ui/src/components/WebRTCVideo.tsx
index de36e37..66d110e 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() {
+export default function WebRTCVideo({ connectionFailed }: { connectionFailed: boolean }) {
   // Video and stream related refs and states
   const videoElm = useRef<HTMLVideoElement>(null);
   const mediaStream = useRTCStore(state => state.mediaStream);
@@ -47,9 +47,10 @@ export default function WebRTCVideo() {
   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 || "",
-  );
+  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;
 
   // Keyboard related states
   const { setIsNumLockActive, setIsCapsLockActive, setIsScrollLockActive } =
diff --git a/ui/src/routes/devices.$id.tsx b/ui/src/routes/devices.$id.tsx
index d25b848..5139400 100644
--- a/ui/src/routes/devices.$id.tsx
+++ b/ui/src/routes/devices.$id.tsx
@@ -130,10 +130,60 @@ export default function KvmIdRoute() {
   const setDiskChannel = useRTCStore(state => state.setDiskChannel);
   const setRpcDataChannel = useRTCStore(state => state.setRpcDataChannel);
   const setTransceiver = useRTCStore(state => state.setTransceiver);
+  const location = useLocation();
+
+  const [connectionAttempts, setConnectionAttempts] = useState(0);
+
+  const [startedConnectingAt, setStartedConnectingAt] = useState<Date | null>(null);
+  const [connectedAt, setConnectedAt] = useState<Date | null>(null);
+
+  const [connectionFailed, setConnectionFailed] = useState(false);
 
   const navigate = useNavigate();
   const { otaState, setOtaState, setModalView } = useUpdateStore();
 
+  useEffect(() => {
+    const connectionAttemptsThreshold = 30;
+    if (connectionAttempts > connectionAttemptsThreshold) {
+      console.log(`Connection failed after ${connectionAttempts} attempts.`);
+      setConnectionFailed(true);
+    }
+  }, [connectionAttempts]);
+
+  useEffect(() => {
+    // Skip if already connected
+    if (connectedAt) return;
+
+    // Skip if connection is declared as failed
+    if (connectionFailed) return;
+
+    const interval = setInterval(() => {
+      console.log("Checking connection status");
+
+      // Skip if connection hasn't started
+      if (!startedConnectingAt) return;
+
+      const elapsedTime = Math.floor(
+        new Date().getTime() - startedConnectingAt.getTime(),
+      );
+
+      // Fail connection if it's been over X seconds since we started connecting
+      if (elapsedTime > 30 * 1000) {
+        console.error(`Connection failed after ${elapsedTime} ms.`);
+        setConnectionFailed(true);
+      }
+    }, 1000);
+
+    return () => clearInterval(interval);
+  }, [connectedAt, connectionFailed, startedConnectingAt]);
+
+  useEffect(() => {
+    if (connectionFailed) {
+      console.log("Closing peer connection");
+      peerConnection?.close();
+    }
+  }, [connectionFailed, peerConnection]);
+
   const sdp = useCallback(
     async (event: RTCPeerConnectionIceEvent, pc: RTCPeerConnection) => {
       if (!pc) return;
@@ -188,6 +238,12 @@ export default function KvmIdRoute() {
 
   const connectWebRTC = useCallback(async () => {
     console.log("Attempting to connect WebRTC");
+
+    // Track connection status to detect failures and show error overlay
+    setConnectionAttempts(x => x + 1);
+    setStartedConnectingAt(new Date());
+    setConnectedAt(null);
+
     const pc = new RTCPeerConnection({
       // We only use STUN or TURN servers if we're in the cloud
       ...(isInCloud && iceConfig?.iceServers
@@ -197,6 +253,11 @@ export default function KvmIdRoute() {
 
     // Set up event listeners and data channels
     pc.onconnectionstatechange = () => {
+      // If the connection state is connected, we reset the connection attempts.
+      if (pc.connectionState === "connected") {
+        setConnectionAttempts(0);
+        setConnectedAt(new Date());
+      }
       setPeerConnectionState(pc.connectionState);
     };
 
@@ -236,16 +297,35 @@ export default function KvmIdRoute() {
     setTransceiver,
   ]);
 
-  // WebRTC connection management
-  useInterval(() => {
+  useEffect(() => {
+    console.log("Attempting to connect WebRTC");
+
+    // If we're in an other session, we don't need to connect
+    if (location.pathname.includes("other-session")) return;
+
+    // If we're already connected or connecting, we don't need to connect
     if (
       ["connected", "connecting", "new"].includes(peerConnection?.connectionState ?? "")
     ) {
       return;
     }
-    if (location.pathname.includes("other-session")) return;
-    connectWebRTC();
-  }, 3000);
+
+    // This is final and we declare connection failure
+    if (connectionFailed) {
+      console.log("Connection failed. We're unable to establish a connection");
+      return;
+    }
+
+    const interval = setInterval(() => {
+      connectWebRTC();
+    }, 3000);
+    return () => clearInterval(interval);
+  }, [
+    connectWebRTC,
+    connectionFailed,
+    location.pathname,
+    peerConnection?.connectionState,
+  ]);
 
   // On boot, if the connection state is undefined, we connect to the WebRTC
   useEffect(() => {
@@ -431,7 +511,6 @@ export default function KvmIdRoute() {
   }, [kvmTerminal, peerConnection, serialConsole]);
 
   const outlet = useOutlet();
-  const location = useLocation();
   const onModalClose = useCallback(() => {
     if (location.pathname !== "/other-session") navigateTo("/");
   }, [navigateTo, location.pathname]);
@@ -508,7 +587,7 @@ export default function KvmIdRoute() {
           />
 
           <div className="flex h-full overflow-hidden">
-            <WebRTCVideo />
+            <WebRTCVideo connectionFailed={connectionFailed} />
             <SidebarContainer sidebarView={sidebarView} />
           </div>
         </div>
@@ -516,10 +595,6 @@ export default function KvmIdRoute() {
 
       <div
         className="isolate"
-        // onMouseMove={e => e.stopPropagation()}
-        // onMouseDown={e => e.stopPropagation()}
-        // onMouseUp={e => e.stopPropagation()}
-        // onPointerMove={e => e.stopPropagation()}
         onKeyUp={e => e.stopPropagation()}
         onKeyDown={e => {
           e.stopPropagation();
@@ -527,6 +602,7 @@ export default function KvmIdRoute() {
         }}
       >
         <Modal open={outlet !== null} onClose={onModalClose}>
+          {/* The 'used by other session' modal needs to have access to the connectWebRTC function */}
           <Outlet context={{ connectWebRTC }} />
         </Modal>
       </div>