From db2c795b5beec9dcb595307864bcc0b275797e1e Mon Sep 17 00:00:00 2001
From: Aurelien Franky <au.re@live.com>
Date: Sat, 30 Nov 2024 01:54:41 +0100
Subject: [PATCH] fix various issues

---
 docs/typings.d.ts | 11 ---------
 src/guest.ts      | 15 ++++++------
 src/helpers.ts    | 56 +++++++++++++++++++++++++++++++++++++++++++--
 src/host.ts       | 58 +++++++++++++++++++++++++++++++----------------
 4 files changed, 100 insertions(+), 40 deletions(-)
 delete mode 100644 docs/typings.d.ts

diff --git a/docs/typings.d.ts b/docs/typings.d.ts
deleted file mode 100644
index 4240ce8..0000000
--- a/docs/typings.d.ts
+++ /dev/null
@@ -1,11 +0,0 @@
-declare module "*.html?raw" {
-  const content: any;
-  export default content;
-}
-
-declare module "*?worker" {
-  const WorkerFactory: {
-    new (): Worker;
-  };
-  export default WorkerFactory;
-}
diff --git a/src/guest.ts b/src/guest.ts
index 32ca711..b8372cc 100644
--- a/src/guest.ts
+++ b/src/guest.ts
@@ -1,4 +1,4 @@
-import { extractMethods, isWorker } from "./helpers";
+import { extractMethods, getEventData, isWorker } from "./helpers";
 import { registerLocalMethods, registerRemoteMethods } from "./rpc";
 import { actions, EventHandlers, events, IConnection, ISchema } from "./types";
 
@@ -8,16 +8,17 @@ function connect(schema: ISchema = {}, eventHandlers?: EventHandlers): Promise<I
 
     // on handshake response
     async function handleHandshakeResponse(event: any) {
-      if (event.data.action !== actions.HANDSHAKE_REPLY) return;
+      const eventData = getEventData(event);
+      if (eventData?.action !== actions.HANDSHAKE_REPLY) return;
 
       // register local methods
-      const unregisterLocal = registerLocalMethods(schema, localMethods, event.data.connectionID);
+      const unregisterLocal = registerLocalMethods(schema, localMethods, eventData.connectionID);
 
       // register remote methods
       const { remote, unregisterRemote } = registerRemoteMethods(
-        event.data.schema,
-        event.data.methods,
-        event.data.connectionID,
+        eventData.schema,
+        eventData.methods,
+        eventData.connectionID,
         event
       );
 
@@ -26,7 +27,7 @@ function connect(schema: ISchema = {}, eventHandlers?: EventHandlers): Promise<I
       // send a HANDSHAKE REPLY to the host
       const payload = {
         action: actions.HANDSHAKE_REPLY,
-        connectionID: event.data.connectionID,
+        connectionID: eventData.connectionID,
       };
 
       if (isWorker()) self.postMessage(payload);
diff --git a/src/helpers.ts b/src/helpers.ts
index d4cd628..bc1a3ad 100644
--- a/src/helpers.ts
+++ b/src/helpers.ts
@@ -1,5 +1,3 @@
-export const CONNECTION_TIMEOUT = 1000;
-
 /**
  * check if run in a webworker
  *
@@ -109,3 +107,57 @@ export function generateId(length: number = 10): string {
   }
   return result;
 }
+
+export interface NodeWorker {
+  on(event: string, handler: any): void;
+  off(event: string, handler: any): void;
+  postMessage(message: any): void;
+  terminate(): void;
+}
+
+// Type that captures common properties between Web Workers and Node Workers
+export type WorkerLike = Worker | NodeWorker;
+
+let NodeWorkerClass: any = null;
+
+if (isNodeEnv()) {
+  try {
+    const workerThreads = require('worker_threads');
+    NodeWorkerClass = workerThreads.Worker;
+  } catch {
+  }
+}
+
+export function isNodeWorker(target: any): target is NodeWorker {
+  return NodeWorkerClass !== null && target instanceof NodeWorkerClass;
+}
+
+export function isWorkerLike(target: any): target is WorkerLike {
+  return isNodeWorker(target) || target instanceof Worker;
+}
+
+export function addEventListener(target: Window | WorkerLike | HTMLIFrameElement, event: string, handler: any) {
+  if (isNodeWorker(target)) {
+    target.on(event, handler);
+  } else if ('addEventListener' in target) {
+    target.addEventListener(event, handler);
+  }
+}
+
+export function removeEventListener(target: Window | WorkerLike | HTMLIFrameElement, event: string, handler: any) {
+  if (isNodeWorker(target)) {
+    target.off(event, handler);
+  } else if ('removeEventListener' in target) {
+    target.removeEventListener(event, handler);
+  }
+}
+
+/**
+ * Normalize message event data across Web and Node.js environments
+ * In web, data is in event.data
+ * In Node.js, the event itself contains the data
+ */
+export function getEventData(event: any): any {
+  return event.data || event;
+}
+
diff --git a/src/host.ts b/src/host.ts
index e23a143..3fa8fcd 100644
--- a/src/host.ts
+++ b/src/host.ts
@@ -1,16 +1,28 @@
-import { extractMethods, generateId, getOriginFromURL, isNodeEnv } from "./helpers";
+import { extractMethods, generateId, getOriginFromURL, isNodeEnv, addEventListener, removeEventListener, isNodeWorker, NodeWorker, getEventData } from "./helpers";
 import { registerLocalMethods, registerRemoteMethods } from "./rpc";
 import { actions, events, IConnection, IConnections, ISchema } from "./types";
 
 const connections: IConnections = {};
 
-function isValidTarget(iframe: HTMLIFrameElement, event: any) {
-  const childURL = iframe.getAttribute("src");
-  const childOrigin = getOriginFromURL(childURL);
-  const hasProperOrigin = event.origin === childOrigin;
-  const hasProperSource = event.source === iframe.contentWindow;
-
-  return hasProperOrigin && hasProperSource;
+function isValidTarget(guest: HTMLIFrameElement | Worker | NodeWorker, event: any) {
+  // If it's a worker, we don't need to validate origin
+  if (isNodeWorker(guest) || (typeof Worker !== 'undefined' && guest instanceof Worker)) {
+    return true;
+  }
+
+  // For iframes, check origin and source
+  const iframe = guest as HTMLIFrameElement;
+  try {
+    const childURL = iframe.src;
+    const childOrigin = getOriginFromURL(childURL);
+    const hasProperOrigin = event.origin === childOrigin;
+    const hasProperSource = event.source === iframe.contentWindow;
+
+    return (hasProperOrigin && hasProperSource) || !childURL;
+  } catch (e) {
+    console.warn('Error checking iframe target:', e);
+    return false;
+  }
 }
 
 /**
@@ -21,10 +33,10 @@ function isValidTarget(iframe: HTMLIFrameElement, event: any) {
  * @param schema
  * @returns Promise
  */
-function connect(guest: HTMLIFrameElement | Worker, schema: ISchema = {}): Promise<IConnection> {
+function connect(guest: HTMLIFrameElement | Worker | NodeWorker, schema: ISchema = {}): Promise<IConnection> {
   if (!guest) throw new Error("a target is required");
 
-  const guestIsWorker = (guest as Worker).onerror !== undefined && (guest as Worker).onmessage !== undefined;
+  const guestIsWorker = isNodeWorker(guest) || ((guest as Worker).onerror !== undefined && (guest as Worker).onmessage !== undefined);
   const listeners = guestIsWorker || isNodeEnv() ? guest : window;
 
   return new Promise((resolve) => {
@@ -32,8 +44,11 @@ function connect(guest: HTMLIFrameElement | Worker, schema: ISchema = {}): Promi
 
     // on handshake request
     function handleHandshake(event: any) {
-      if (!guestIsWorker && !isNodeEnv() && !isValidTarget(guest as HTMLIFrameElement, event)) return;
-      if (event.data.action !== actions.HANDSHAKE_REQUEST) return;
+
+      if (!guestIsWorker && !isNodeEnv() && !isValidTarget(guest, event)) return;
+      
+      const eventData = getEventData(event);
+      if (eventData?.action !== actions.HANDSHAKE_REQUEST) return;
 
       // register local methods
       const localMethods = extractMethods(schema);
@@ -46,8 +61,8 @@ function connect(guest: HTMLIFrameElement | Worker, schema: ISchema = {}): Promi
 
       // register remote methods
       const { remote, unregisterRemote } = registerRemoteMethods(
-        event.data.schema,
-        event.data.methods,
+        eventData.schema,
+        eventData.methods,
         connectionID,
         event,
         guestIsWorker || isNodeEnv() ? (guest as Worker) : undefined
@@ -66,10 +81,12 @@ function connect(guest: HTMLIFrameElement | Worker, schema: ISchema = {}): Promi
 
       // close the connection and all listeners when called
       const close = () => {
-        listeners.removeEventListener(events.MESSAGE, handleHandshake);
+        removeEventListener(listeners, events.MESSAGE, handleHandshake);
         unregisterRemote();
         unregisterLocal();
-        if (guestIsWorker) (guest as Worker).terminate();
+        if (guestIsWorker) {
+          (guest as Worker).terminate();
+        }
       };
 
       const connection: IConnection = { remote, close };
@@ -77,15 +94,16 @@ function connect(guest: HTMLIFrameElement | Worker, schema: ISchema = {}): Promi
     }
 
     // subscribe to HANDSHAKE MESSAGES
-    listeners.addEventListener(events.MESSAGE, handleHandshake);
+    addEventListener(listeners, events.MESSAGE, handleHandshake);
 
     // on handshake reply
     function handleHandshakeReply(event: any) {
-      if (event.data.action !== actions.HANDSHAKE_REPLY) return;
-      return resolve(connections[event.data.connectionID]);
+      const eventData = getEventData(event);
+      if (eventData?.action !== actions.HANDSHAKE_REPLY) return;
+      return resolve(connections[eventData.connectionID]);
     }
 
-    listeners.addEventListener(events.MESSAGE, handleHandshakeReply);
+    addEventListener(listeners, events.MESSAGE, handleHandshakeReply);
   });
 }