fix(ws): enable ws proxy via MessageChannel

Author: jourdain

vue3-app/src/main.js

@@ -4,6 +4,7 @@ import wslink from "./core/wslink";
 import { handlePageResources, loadScript } from "./core/trame/setup";
 import { createTrameInstance } from "./core/trame";
 import { registerUserScripts } from "./user_script_handler";
+import { createMessageChannelWSFactory } from "./messageChannel";
 import TrameUse from "./use";
 
 const { createApp } = window.Vue;
@@ -12,7 +13,7 @@ async function start() {
   // Check if we need to override websocket
   try {
     // Can throw exception if parent is cross-origin
-    if (window?.parent?.trameJupyter?.init) {
+    if (!window.WSLINK && window?.parent?.trameJupyter?.init) {
       window.WSLINK = window.parent.trameJupyter.init(window);
     }
   } catch (e) {
@@ -123,11 +124,31 @@ async function start() {
   registerUserScripts();
 }
 
+const urlParams = vtkURLExtract.extractURLParameters();
 // Initialize service worker to override headers for SharedArrayBuffer
 // > Cross-Origin-Opener-Policy: same-origin
 // > Cross-Origin-Embedder-Policy: require-corp
-if (vtkURLExtract.extractURLParameters().enableSharedArrayBufferServiceWorker) {
+if (urlParams.enableSharedArrayBufferServiceWorker) {
   loadScript("coi-serviceworker.min.js");
 }
 
-start();
+if (urlParams.wsChannel) {
+  window.addEventListener("message", (event) => {
+    // Waiting for message initializing channel communication
+    // for websocket proxying
+    if (event.data !== "trame-ws-channel-init") {
+      // not for us! skip
+      return;
+    }
+
+    // Grab communication port and register fake WebSocket
+    window.WSLINK = {
+      createWebSocket: createMessageChannelWSFactory(event.ports[0]),
+    };
+
+    // Start trame client
+    start();
+  });
+} else {
+  start();
+}

vue3-app/src/messageChannel.js

@@ -0,0 +1,73 @@
+class MessageChannelWebSocket {
+  constructor(port) {
+    console.log("create MessageChannelWebSocket");
+    this._channel_port = port;
+    this.onmessage = null;
+    this.onclose = null;
+    this.onerror = null;
+
+    // Listen to channel messages
+    this._channel_port.onmessage = (e) => {
+      if (!this.onmessage) return;
+      const { a, b, t } = e.data;
+
+      if (a === "s") {
+        this.onmessage({ data: b || t });
+        return;
+      }
+      if (a === "c" && this.onclose) {
+        this.onclose(e.data.c);
+        return;
+      }
+      if (a === "e" && this.onerror) {
+        this.onerror(e.data.e);
+        return;
+      }
+    };
+  }
+
+  set onopen(callback) {
+    // Auto open but allow current listeners registration to complete before triggering
+    setTimeout(callback, 0);
+  }
+
+  send(data) {
+    console.log("send", data);
+    if (data.buffer) {
+      // may need a copy
+      if (data.buffer.byteLength !== data.length) {
+        const tmp = new Uint8Array(data.length);
+        tmp.set(data);
+        self._channel_port.postMessage(
+          {
+            a: "s",
+            b: tmp.buffer,
+          },
+          [tmp.buffer]
+        );
+        return;
+      }
+      self._channel_port.postMessage(
+        {
+          a: "s",
+          b: data.buffer,
+        },
+        [data.buffer]
+      );
+      return;
+    }
+    self._channel_port.postMessage({
+      a: "s",
+      t: data,
+    });
+  }
+
+  close() {
+    self._channel_port.postMessage({ a: "c", c: "Closing from client" });
+    self._channel_port.close();
+  }
+}
+
+export function createMessageChannelWSFactory(port) {
+  return () => new MessageChannelWebSocket(port);
+}