fix: hidRPC handshake packet should be only sent once (#969)

Author: ym

ui/package-lock.json

@@ -56,6 +56,7 @@
     "react-xtermjs": "^1.0.10",
     "recharts": "^3.3.0",
     "tailwind-merge": "^3.3.1",
+    "tslog": "^4.10.2",
     "usehooks-ts": "^3.1.1",
     "validator": "^13.15.20",
     "zustand": "^4.5.2"

ui/package.json

@@ -1,4 +1,5 @@
 import { useCallback, useEffect, useMemo } from "react";
+import { Logger } from "tslog";
 
 import { useRTCStore } from "@hooks/stores";
 
@@ -25,6 +26,128 @@ interface sendMessageParams {
   requireOrdered?: boolean;
 }
 
+const HANDSHAKE_TIMEOUT = 30 * 1000; // 30 seconds
+const HANDSHAKE_MAX_ATTEMPTS = 10;
+const logger = new Logger({ name: "hidrpc" });
+
+export function doRpcHidHandshake(rpcHidChannel: RTCDataChannel, setRpcHidProtocolVersion: (version: number | null) => void) {
+  let attempts = 0;
+  let lastConnectedTime: Date | undefined;
+  let lastSendTime: Date | undefined;
+  let handshakeCompleted = false;
+  let handshakeInterval: ReturnType<typeof setInterval> | null = null;
+
+  const shouldGiveUp = () => {
+    if (attempts > HANDSHAKE_MAX_ATTEMPTS) {
+      logger.error(`Failed to send handshake message after ${HANDSHAKE_MAX_ATTEMPTS} attempts`);
+      return true;
+    }
+
+    const timeSinceConnected = lastConnectedTime ? Date.now() - lastConnectedTime.getTime() : 0;
+    if (timeSinceConnected > HANDSHAKE_TIMEOUT) {
+      logger.error(`Handshake timed out after ${timeSinceConnected}ms`);
+      return true;
+    }
+
+    return false;
+  }
+
+  const sendHandshake = (initial: boolean) => {
+    if (handshakeCompleted) return;
+
+    attempts++;
+    lastSendTime = new Date();
+
+    if (!initial && shouldGiveUp()) {
+      if (handshakeInterval) {
+        clearInterval(handshakeInterval);
+        handshakeInterval = null;
+      }
+      return;
+    }
+
+    let data: Uint8Array | undefined;
+    try {
+      const message = new HandshakeMessage(HID_RPC_VERSION);
+      data = message.marshal();
+    } catch (e) {
+      logger.error("Failed to marshal message", e);
+      return;
+    }
+    if (!data) return;
+    rpcHidChannel.send(data as unknown as ArrayBuffer);
+
+    if (initial) {
+      handshakeInterval = setInterval(() => {
+        sendHandshake(false);
+      }, 1000);
+    }
+  };
+
+  const onMessage = (ev: MessageEvent) => {
+    const message = unmarshalHidRpcMessage(new Uint8Array(ev.data));
+    if (!message || !(message instanceof HandshakeMessage)) return;
+
+    if (!message.version) {
+      logger.error("Received handshake message without version", message);
+      return;
+    }
+
+    if (message.version > HID_RPC_VERSION) {
+      // we assume that the UI is always using the latest version of the HID RPC protocol
+      // so we can't support this
+      // TODO: use capabilities to determine rather than version number
+      logger.error("Server is using a newer version than the client", message);
+      return;
+    }
+
+    setRpcHidProtocolVersion(message.version);
+
+    const timeUsed = lastSendTime ? Date.now() - lastSendTime.getTime() : 0;
+    logger.info(`Handshake completed in ${timeUsed}ms after ${attempts} attempts (Version: ${message.version} / ${HID_RPC_VERSION})`);
+
+    // clean up 
+    rpcHidChannel.removeEventListener("message", onMessage);
+    resetHandshake({ completed: true });
+  };
+
+  const resetHandshake = ({ lastConnectedTime: newLastConnectedTime, completed }: { lastConnectedTime?: Date | undefined, completed?: boolean }) => {
+    if (newLastConnectedTime) lastConnectedTime = newLastConnectedTime;
+    lastSendTime = undefined;
+    attempts = 0;
+    if (completed !== undefined) handshakeCompleted = completed;
+    if (handshakeInterval) {
+      clearInterval(handshakeInterval);
+      handshakeInterval = null;
+    }
+  };
+
+  const onConnected = () => {
+    resetHandshake({ lastConnectedTime: new Date() });
+    logger.info("Channel connected");
+
+    sendHandshake(true);
+    rpcHidChannel.addEventListener("message", onMessage);
+  };
+
+  const onClose = () => {
+    resetHandshake({ lastConnectedTime: undefined, completed: false });
+
+    logger.info("Channel closed");
+    setRpcHidProtocolVersion(null);
+
+    rpcHidChannel.removeEventListener("message", onMessage);
+  };
+
+  rpcHidChannel.addEventListener("open", onConnected);
+  rpcHidChannel.addEventListener("close", onClose);
+
+  // handle case where channel is already open when the hook is mounted
+  if (rpcHidChannel.readyState === "open") {
+    onConnected();
+  }
+}
+
 export function useHidRpc(onHidRpcMessage?: (payload: RpcMessage) => void) {
   const {
     rpcHidChannel,
@@ -78,7 +201,7 @@ export function useHidRpc(onHidRpcMessage?: (payload: RpcMessage) => void) {
       try {
         data = message.marshal();
       } catch (e) {
-        console.error("Failed to marshal HID RPC message", e);
+        logger.error("Failed to marshal message", e);
       }
       if (!data) return;
 
@@ -151,99 +274,46 @@ export function useHidRpc(onHidRpcMessage?: (payload: RpcMessage) => void) {
     sendMessage(KEEPALIVE_MESSAGE);
   }, [sendMessage]);
 
-  const sendHandshake = useCallback(() => {
-    if (hidRpcDisabled) return;
-    if (rpcHidProtocolVersion) return;
-    if (!rpcHidChannel) return;
-
-    sendMessage(new HandshakeMessage(HID_RPC_VERSION), { ignoreHandshakeState: true });
-  }, [rpcHidChannel, rpcHidProtocolVersion, sendMessage, hidRpcDisabled]);
-
-  const handleHandshake = useCallback(
-    (message: HandshakeMessage) => {
-      if (hidRpcDisabled) return;
-
-      if (!message.version) {
-        console.error("Received handshake message without version", message);
-        return;
-      }
-
-      if (message.version > HID_RPC_VERSION) {
-        // we assume that the UI is always using the latest version of the HID RPC protocol
-        // so we can't support this
-        // TODO: use capabilities to determine rather than version number
-        console.error("Server is using a newer HID RPC version than the client", message);
-        return;
-      }
-
-      setRpcHidProtocolVersion(message.version);
-    },
-    [setRpcHidProtocolVersion, hidRpcDisabled],
-  );
-
   useEffect(() => {
     if (!rpcHidChannel) return;
     if (hidRpcDisabled) return;
 
-    // send handshake message
-    sendHandshake();
-
     const messageHandler = (e: MessageEvent) => {
       if (typeof e.data === "string") {
-        console.warn("Received string data in HID RPC message handler", e.data);
+        logger.warn("Received string data in message handler", e.data);
         return;
       }
 
       const message = unmarshalHidRpcMessage(new Uint8Array(e.data));
       if (!message) {
-        console.warn("Received invalid HID RPC message", e.data);
+        logger.warn("Received invalid message", e.data);
         return;
       }
 
-      console.debug("Received HID RPC message", message);
-      switch (message.constructor) {
-        case HandshakeMessage:
-          handleHandshake(message as HandshakeMessage);
-          break;
-        default:
-          // not all events are handled here, the rest are handled by the onHidRpcMessage callback
-          break;
-      }
+      if (message instanceof HandshakeMessage) return; // handshake message is handled by the doRpcHidHandshake function
 
-      onHidRpcMessage?.(message);
-    };
+      // to remove it from the production build, we need to use the /* @__PURE__ */ comment here
+      // setting `esbuild.pure` doesn't work
+      /* @__PURE__ */ logger.debug("Received message", message);
 
-    const openHandler = () => {
-      console.info("HID RPC channel opened");
-      sendHandshake();
-    };
-
-    const closeHandler = () => {
-      console.info("HID RPC channel closed");
-      setRpcHidProtocolVersion(null);
+      onHidRpcMessage?.(message);
     };
 
     const errorHandler = (e: Event) => {
-      console.error(`Error on rpcHidChannel '${rpcHidChannel.label}': ${e}`)
+      logger.error(`Error on channel '${rpcHidChannel.label}'`, e);
     };
 
     rpcHidChannel.addEventListener("message", messageHandler);
-    rpcHidChannel.addEventListener("close", closeHandler);
     rpcHidChannel.addEventListener("error", errorHandler);
-    rpcHidChannel.addEventListener("open", openHandler);
 
     return () => {
       rpcHidChannel.removeEventListener("message", messageHandler);
-      rpcHidChannel.removeEventListener("close", closeHandler);
       rpcHidChannel.removeEventListener("error", errorHandler);
-      rpcHidChannel.removeEventListener("open", openHandler);
     };
   }, [
     rpcHidChannel,
     onHidRpcMessage,
     setRpcHidProtocolVersion,
-    sendHandshake,
-    handleHandshake,
     hidRpcDisabled,
   ]);
 

ui/src/hooks/useHidRpc.ts

@@ -53,6 +53,7 @@ import {
 } from "@components/VideoOverlay";
 import { FeatureFlagProvider } from "@providers/FeatureFlagProvider";
 import { m } from "@localizations/messages.js";
+import { doRpcHidHandshake } from "@hooks/useHidRpc";
 
 export type AuthMode = "password" | "noPassword" | null;
 
@@ -127,6 +128,7 @@ export default function KvmIdRoute() {
     setRpcHidChannel,
     setRpcHidUnreliableNonOrderedChannel,
     setRpcHidUnreliableChannel,
+    setRpcHidProtocolVersion,
   } = useRTCStore();
 
   const location = useLocation();
@@ -498,6 +500,7 @@ export default function KvmIdRoute() {
     rpcHidChannel.onopen = () => {
       setRpcHidChannel(rpcHidChannel);
     };
+    doRpcHidHandshake(rpcHidChannel, setRpcHidProtocolVersion);
 
     const rpcHidUnreliableChannel = pc.createDataChannel("hidrpc-unreliable-ordered", {
       ordered: true,
@@ -534,6 +537,7 @@ export default function KvmIdRoute() {
     setRpcHidChannel,
     setRpcHidUnreliableNonOrderedChannel,
     setRpcHidUnreliableChannel,
+    setRpcHidProtocolVersion,
     setTransceiver,
   ]);
 

ui/src/routes/devices.$id.tsx

@@ -39,7 +39,7 @@ export default defineConfig(({ mode, command }) => {
   return {
     plugins,
     esbuild: {
-      pure: ["console.debug"],
+      pure: command === "build" ? ["console.debug"]: [],
     },
     assetsInclude: ["**/*.woff2"],
     build: {