From 8d343e5ccac18717b38cb07c3280ee5bd4904f01 Mon Sep 17 00:00:00 2001
From: Tsu-ba-me <ynho.li.aa.e@gmail.com>
Date: Fri, 16 Jun 2023 01:39:39 -0400
Subject: [PATCH] fix(striker-ui): attempt reconnect after unclean disconnect
 in VNC

---
 striker-ui/components/Display/FullSize.tsx   | 64 +++++++++++---------
 striker-ui/components/Display/VncDisplay.tsx |  6 ++
 2 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/striker-ui/components/Display/FullSize.tsx b/striker-ui/components/Display/FullSize.tsx
index a4ee7782..1b22b089 100644
--- a/striker-ui/components/Display/FullSize.tsx
+++ b/striker-ui/components/Display/FullSize.tsx
@@ -12,7 +12,7 @@ import {
   Typography,
 } from '@mui/material';
 import RFB from '@novnc/novnc/core/rfb';
-import { useState, useRef, useEffect, FC } from 'react';
+import { useState, useRef, useEffect, FC, useCallback } from 'react';
 import dynamic from 'next/dynamic';
 
 import API_BASE_URL from '../../lib/consts/API_BASE_URL';
@@ -131,42 +131,44 @@ const FullSize: FC<FullSizeProps> = ({
   const [vncConnecting, setVncConnecting] = useProtectedState<boolean>(false);
   const [isError, setIsError] = useProtectedState<boolean>(false);
 
-  useEffect(() => {
-    if (typeof window !== 'undefined') {
-      hostname.current = window.location.hostname;
-    }
+  const connectVnc = useCallback(async () => {
+    if (vncConnection || vncConnecting) return;
 
-    if (!vncConnection && !vncConnecting)
-      (async () => {
-        setVncConnecting(true);
+    setVncConnecting(true);
 
-        try {
-          const res = await putFetchWithTimeout(
-            CMD_VNC_PIPE_URL,
-            {
-              serverUuid: serverUUID,
-              open: true,
-            },
-            120000,
-          );
+    try {
+      const res = await putFetchWithTimeout(
+        CMD_VNC_PIPE_URL,
+        {
+          serverUuid: serverUUID,
+          open: true,
+        },
+        120000,
+      );
 
-          setVncConnection(await res.json());
-        } catch {
-          setIsError(true);
-        } finally {
-          setVncConnecting(false);
-        }
-      })();
+      setVncConnection(await res.json());
+    } catch {
+      setIsError(true);
+    } finally {
+      setVncConnecting(false);
+    }
   }, [
     serverUUID,
-    vncConnection,
-    isError,
-    setVncConnection,
     setIsError,
-    vncConnecting,
     setVncConnecting,
+    setVncConnection,
+    vncConnecting,
+    vncConnection,
   ]);
 
+  useEffect(() => {
+    if (typeof window !== 'undefined') {
+      hostname.current = window.location.hostname;
+    }
+
+    connectVnc();
+  }, [connectVnc]);
+
   const handleClick = (event: React.MouseEvent<HTMLButtonElement>): void => {
     setAnchorEl(event.currentTarget);
   };
@@ -214,6 +216,12 @@ const FullSize: FC<FullSizeProps> = ({
               background=""
               qualityLevel={6}
               compressionLevel={2}
+              onDisconnect={({ detail: { clean } }) => {
+                if (!clean) {
+                  setVncConnection(undefined);
+                  connectVnc();
+                }
+              }}
             />
             <Box>
               <Box className={classes.closeBox}>
diff --git a/striker-ui/components/Display/VncDisplay.tsx b/striker-ui/components/Display/VncDisplay.tsx
index 9afb222e..c9f743fc 100644
--- a/striker-ui/components/Display/VncDisplay.tsx
+++ b/striker-ui/components/Display/VncDisplay.tsx
@@ -14,6 +14,7 @@ type Props = {
   background: string;
   qualityLevel: number;
   compressionLevel: number;
+  onDisconnect: (event: { detail: { clean: boolean } }) => void;
 };
 
 const VncDisplay = (props: Props): JSX.Element => {
@@ -32,6 +33,7 @@ const VncDisplay = (props: Props): JSX.Element => {
     background,
     qualityLevel,
     compressionLevel,
+    onDisconnect,
   } = props;
 
   useEffect(() => {
@@ -59,6 +61,10 @@ const VncDisplay = (props: Props): JSX.Element => {
       rfb.current.background = background;
       rfb.current.qualityLevel = qualityLevel;
       rfb.current.compressionLevel = compressionLevel;
+
+      // RFB extends custom class EventTargetMixin;
+      // the usual .on or .once doesn't exist.
+      rfb.current.addEventListener('disconnect', onDisconnect);
     }
 
     /* eslint-disable consistent-return */