From b00782671098e6b4425beb001636c5966e0b042a Mon Sep 17 00:00:00 2001
From: Pierre Hubert <pierre.git@communiquons.org>
Date: Wed, 18 Oct 2023 18:36:24 +0200
Subject: [PATCH] Fix VNC connection

---
 virtweb_backend/src/actors/vnc_actor.rs  | 24 ++++++--
 virtweb_frontend/src/routes/VNCRoute.tsx | 70 +++++++++++-------------
 2 files changed, 51 insertions(+), 43 deletions(-)

diff --git a/virtweb_backend/src/actors/vnc_actor.rs b/virtweb_backend/src/actors/vnc_actor.rs
index 9c859ff..fd4814c 100644
--- a/virtweb_backend/src/actors/vnc_actor.rs
+++ b/virtweb_backend/src/actors/vnc_actor.rs
@@ -84,11 +84,27 @@ impl VNCActor {
             let mut buff: [u8; 5000] = [0; 5000];
             loop {
                 match read_half.read(&mut buff).await {
-                    Ok(l) => {
+                    Ok(mut l) => {
                         if l == 0 {
-                            log::error!("Got empty read. Closing read end...");
-                            addr.do_send(CloseWebSocketReq);
-                            return;
+                            log::warn!("Got empty read!");
+
+                            // Ugly hack made to wait for next byte
+                            let mut one_byte_buff: [u8; 1] = [0; 1];
+                            match read_half.read_exact(&mut one_byte_buff).await {
+                                Ok(b) => {
+                                    if b == 0 {
+                                        log::error!("Did not get a byte !");
+                                        let _ = addr.send(CloseWebSocketReq).await;
+                                        break;
+                                    }
+                                    buff[0] = one_byte_buff[0];
+                                    l = 1;
+                                }
+                                Err(e) => {
+                                    log::error!("Failed to read 1 BYTE from remote socket. Stopping now... {:?}", e);
+                                    break;
+                                }
+                            }
                         }
 
                         let to_send = SendBytesReq(Vec::from(&buff[0..l]));
diff --git a/virtweb_frontend/src/routes/VNCRoute.tsx b/virtweb_frontend/src/routes/VNCRoute.tsx
index 20fe5bf..ef76976 100644
--- a/virtweb_frontend/src/routes/VNCRoute.tsx
+++ b/virtweb_frontend/src/routes/VNCRoute.tsx
@@ -1,4 +1,3 @@
-import { CircularProgress } from "@mui/material";
 import React from "react";
 import { useParams } from "react-router-dom";
 import { VncScreen } from "react-vnc";
@@ -28,15 +27,15 @@ export function VNCRoute(): React.ReactElement {
 function VNCInner(p: { vm: VMInfo }): React.ReactElement {
   const snackbar = useSnackbar();
 
+  const counter = React.useRef(false);
   const [url, setURL] = React.useState<string | undefined>();
-  const urlRef = React.useRef<string | undefined>();
 
   const load = async () => {
     try {
       const u = await VMApi.OneShotVNCURL(p.vm);
       console.info(u);
-      if (urlRef.current === undefined) {
-        urlRef.current = u;
+      if (counter.current === false) {
+        counter.current = true;
         setURL(u);
       }
     } catch (e) {
@@ -45,42 +44,35 @@ function VNCInner(p: { vm: VMInfo }): React.ReactElement {
     }
   };
 
-  React.useEffect(() => {
-    load();
-  });
-
-  if (urlRef.current === undefined)
-    return (
-      <div
-        style={{
-          display: "flex",
-          alignItems: "center",
-          justifyContent: "center",
-          height: "100%",
-        }}
-      >
-        <CircularProgress />
-      </div>
-    );
+  const reconnect = () => {
+    counter.current = false;
+    setURL(undefined);
+  };
 
   return (
-    <div
-      style={{
-        display: "flex",
-        alignItems: "center",
-        justifyContent: "center",
-        height: "100%",
-      }}
-    >
-      <VncScreen
-        url={urlRef.current}
-        onDisconnect={() => {
-          console.info("VNC disconnected " + urlRef.current);
-          urlRef.current = undefined;
-          load();
-          setURL(undefined);
-        }}
-      />
-    </div>
+    <AsyncWidget
+      loadKey={counter.current}
+      load={load}
+      ready={url !== undefined && counter.current}
+      errMsg="Failed to get a token to initialize VNC connection!"
+      build={() => (
+        <div
+          style={{
+            display: "flex",
+            alignItems: "center",
+            justifyContent: "center",
+            height: "100%",
+          }}
+        >
+          <VncScreen
+            url={url!}
+            onDisconnect={() => {
+              console.info("VNC disconnected " + url);
+              reconnect();
+            }}
+          />
+        </div>
+      )}
+    />
   );
 }