From 8cb4c66fec4981347ae366b882a7e6d514c069ed Mon Sep 17 00:00:00 2001
From: "irvingouj @ Devolutions" <jou@devolutions.net>
Date: Thu, 9 Jan 2025 11:47:57 -0500
Subject: [PATCH] fix(dgw): fallback to recording playback when shadowing is
 not possible (#1181)

---
 webapp/player-project/src/main.ts            | 18 +++++++-
 webapp/player-project/src/streamers/index.ts | 11 ++++-
 webapp/player-project/src/terminal.ts        |  7 +++
 webapp/player-project/src/ws-proxy.ts        | 45 ++++++++++++++++++++
 4 files changed, 79 insertions(+), 2 deletions(-)
 create mode 100644 webapp/player-project/src/ws-proxy.ts

diff --git a/webapp/player-project/src/main.ts b/webapp/player-project/src/main.ts
index 01d5d958b..3c26ff3f2 100644
--- a/webapp/player-project/src/main.ts
+++ b/webapp/player-project/src/main.ts
@@ -1,6 +1,8 @@
 import { GatewayAccessApi } from './gateway';
 import { getPlayer } from './players/index.js';
-import { getShadowPlayer } from './streamers/index.js';
+import { cleanUpStreamers, getShadowPlayer } from './streamers/index.js';
+import './ws-proxy.ts';
+import { OnBeforeClose as BeforeWebsocketClose } from './ws-proxy.ts';
 
 async function main() {
   const { sessionId, token, gatewayAccessUrl, isActive } = getSessionDetails();
@@ -23,6 +25,20 @@ async function playSessionShadowing(gatewayAccessApi) {
   try {
     const recordingInfo = await gatewayAccessApi.fetchRecordingInfo();
     const fileType = getFileType(recordingInfo);
+    BeforeWebsocketClose((closeEvent) => {
+      if (closeEvent.code !== 1000) {
+        // The session playback failed; attempt to play the recording as usual as a fallback.
+        cleanUpStreamers();
+        playStaticRecording(gatewayAccessApi);
+        return {
+          ...closeEvent,
+          // This prevents extra handling by other listeners, particularly for asciinema-player in this scenario.
+          // For more details, see the asciinema-player WebSocket driver’s socket close handler.
+          // https://github.com/asciinema/asciinema-player/blob/c09e1d2625450a32e9e76063cdc315fd54ecdd9d/src/driver/websocket.js#L219
+          code: 1000,
+        };
+      }
+    });
 
     getShadowPlayer(fileType).play(gatewayAccessApi);
   } catch (error) {
diff --git a/webapp/player-project/src/streamers/index.ts b/webapp/player-project/src/streamers/index.ts
index 9732c6e6b..320c222eb 100644
--- a/webapp/player-project/src/streamers/index.ts
+++ b/webapp/player-project/src/streamers/index.ts
@@ -1,7 +1,7 @@
 import { GatewayAccessApi } from '../gateway';
+import { removeTerminal } from '../terminal';
 import { handleCast } from './cast';
 import { handleWebm } from './webm';
-
 export const getShadowPlayer = (fileType) => {
   const player = {
     play: (_: GatewayAccessApi) => {},
@@ -17,3 +17,12 @@ export const getShadowPlayer = (fileType) => {
 
   return player;
 };
+
+export const cleanUpStreamers = () => {
+  // Remove all shadow-player elements.
+  const shadowPlayers = document.querySelectorAll('shadow-player');
+  for (const shadowPlayer of shadowPlayers) {
+    shadowPlayer.remove();
+  }
+  removeTerminal();
+};
diff --git a/webapp/player-project/src/terminal.ts b/webapp/player-project/src/terminal.ts
index dd375734c..7be2501d9 100644
--- a/webapp/player-project/src/terminal.ts
+++ b/webapp/player-project/src/terminal.ts
@@ -18,3 +18,10 @@ export function createTerminalDiv() {
   document.body.appendChild(terminalDiv);
   return terminalDiv;
 }
+
+export function removeTerminal() {
+  const terminalDiv = document.getElementById('terminal');
+  if (terminalDiv) {
+    terminalDiv.remove();
+  }
+}
diff --git a/webapp/player-project/src/ws-proxy.ts b/webapp/player-project/src/ws-proxy.ts
new file mode 100644
index 000000000..ee6add51d
--- /dev/null
+++ b/webapp/player-project/src/ws-proxy.ts
@@ -0,0 +1,45 @@
+let beforeClose = (args: CloseEvent): CloseEvent => {
+  return args;
+};
+
+export const OnBeforeClose = (callback: (args: CloseEvent) => CloseEvent) => {
+  beforeClose = callback;
+};
+
+const WebSocketProxy = new Proxy(window.WebSocket, {
+  construct(target, args: [url: string | URL, protocols?: string | string[]]) {
+    console.log('Proxying WebSocket connection', ...args);
+    const ws = new target(...args); // Create the actual WebSocket instance
+
+    // Proxy for intercepting `addEventListener`
+    ws.addEventListener = new Proxy(ws.addEventListener, {
+      apply(target, thisArg, args) {
+        if (args[0] === 'close') {
+          console.log('Intercepted addEventListener for close event');
+          const transformedArgs = beforeClose(args as unknown as CloseEvent);
+          return target.apply(thisArg, transformedArgs);
+        }
+        return target.apply(thisArg, args);
+      },
+    });
+
+    // Proxy for intercepting `onclose`
+    return new Proxy(ws, {
+      set(target, prop, value) {
+        if (prop === 'onclose') {
+          console.log('Intercepted setting of onclose');
+          const transformedValue = (...args) => {
+            const transformedArgs = beforeClose(args as unknown as CloseEvent);
+            if (typeof value === 'function') {
+              value(transformedArgs); // Call the original handler
+            }
+          };
+          return Reflect.set(target, prop, transformedValue);
+        }
+        return Reflect.set(target, prop, value);
+      },
+    });
+  },
+});
+
+window.WebSocket = WebSocketProxy;