From cb443135d9149af85c22fbf3177207a38648c338 Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Tue, 9 Jul 2024 10:28:36 -0700
Subject: [PATCH] remove broadcastchannel for workerPath, let the client pass
 in a messageport, use snapshot for more stuff

---
 rollup.config.js  |  2 +-
 src/client.ts     |  8 +++++--
 src/connection.ts | 53 ++++++++++++++++++++---------------------------
 src/snapshot.ts   |  2 ++
 4 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/rollup.config.js b/rollup.config.js
index 6babc32..fb53aa2 100644
--- a/rollup.config.js
+++ b/rollup.config.js
@@ -6,7 +6,7 @@ const commonPlugins = () => [
 	typescript(),
 	inject(
 		Object.fromEntries(
-			['fetch', 'Request', 'Response', 'WebSocket', 'XMLHttpRequest', 'SharedWorker'].map(
+			['fetch', 'Request', 'Response', 'WebSocket', 'XMLHttpRequest', 'SharedWorker', 'localStorage', 'serviceWorker'].map(
 				(name) => [
 					name,
 					[fileURLToPath(new URL('./src/snapshot.ts', import.meta.url)), name],
diff --git a/src/client.ts b/src/client.ts
index 89a6ebc..eb221f2 100644
--- a/src/client.ts
+++ b/src/client.ts
@@ -124,6 +124,10 @@ export class BareMuxConnection {
 			client: functionBody,
 		});
 	}
+
+	async getInnerPort(): Promise<MessagePort> {
+		return await this.worker.port;
+	}
 }
 
 export class BareClient {
@@ -132,8 +136,8 @@ export class BareClient {
 	/**
 	 * Create a BareClient. Calls to fetch and connect will wait for an implementation to be ready.
 	 */
-	constructor(workerPath?: string) {
-		this.worker = new WorkerConnection(workerPath);
+	constructor(worker?: string | MessagePort) {
+		this.worker = new WorkerConnection(worker);
 	}
 
 	createWebSocket(
diff --git a/src/connection.ts b/src/connection.ts
index fdff6de..8dc6e4b 100644
--- a/src/connection.ts
+++ b/src/connection.ts
@@ -33,8 +33,7 @@ export type WorkerResponse = {
 }
 
 export type BroadcastMessage = {
-	type: "getPath" | "path" | "refreshPort",
-	path?: string,
+	type: "refreshPort",
 }
 
 async function searchForPort(): Promise<MessagePort> {
@@ -59,25 +58,17 @@ function tryGetPort(client: SWClient): Promise<MessagePort> {
 	});
 }
 
-function createPort(path: string, channel: BroadcastChannel, registerHandlers: boolean): MessagePort {
+function createPort(path: string, registerHandlers: boolean): MessagePort {
 	const worker = new SharedWorker(path, "bare-mux-worker");
 	if (registerHandlers) {
-		// uv removes navigator.serviceWorker so this errors
-		if (navigator.serviceWorker) {
-			navigator.serviceWorker.addEventListener("message", event => {
-				if (event.data.type === "getPort" && event.data.port) {
-					console.debug("bare-mux: recieved request for port from sw");
-					const worker = new SharedWorker(path, "bare-mux-worker");
-					event.data.port.postMessage(worker.port, [worker.port]);
-				}
-			});
-		}
-		channel.onmessage = (event: MessageEvent) => {
-			if (event.data.type === "getPath") {
-				console.debug("bare-mux: recieved request for worker path from broadcast channel");
-				channel.postMessage(<BroadcastMessage>{ type: "path", path: path });
+		// @ts-expect-error we are using snapshot.ts
+		serviceWorker.addEventListener("message", (event: MessageEvent) => {
+			if (event.data.type === "getPort" && event.data.port) {
+				console.debug("bare-mux: recieved request for port from sw");
+				const worker = new SharedWorker(path, "bare-mux-worker");
+				event.data.port.postMessage(worker.port, [worker.port]);
 			}
-		};
+		});
 	}
 	return worker.port;
 }
@@ -87,9 +78,13 @@ export class WorkerConnection {
 	port: MessagePort | Promise<MessagePort>;
 	workerPath: string;
 
-	constructor(workerPath?: string) {
+	constructor(worker?: string | MessagePort) {
 		this.channel = new BroadcastChannel("bare-mux");
-		this.createChannel(workerPath, true);
+		if (worker instanceof MessagePort) {
+			this.port = worker;
+		} else {
+			this.createChannel(worker, true);
+		}
 	}
 
 	createChannel(workerPath?: string, inInit?: boolean) {
@@ -108,18 +103,16 @@ export class WorkerConnection {
 			// create the SharedWorker and help other bare-mux clients get the workerPath
 
 			if (!workerPath.startsWith("/") && !workerPath.includes("://")) throw new Error("Invalid URL. Must be absolute or start at the root.");
-			this.port = createPort(workerPath, this.channel, inInit);
+			this.port = createPort(workerPath, inInit);
+			console.debug("bare-mux: setting localStorage bare-mux-path to", workerPath);
+			localStorage["bare-mux-path"] = workerPath;
 		} else if (SharedWorker) {
 			// running in a window, was not passed a workerPath
-			// ask other bare-mux clients for the workerPath
-			this.port = new Promise(resolve => {
-				this.channel.onmessage = (event: MessageEvent) => {
-					if (event.data.type === "path") {
-						resolve(createPort(event.data.path, this.channel, inInit));
-					}
-				}
-				this.channel.postMessage(<BroadcastMessage>{ type: "getPath" });
-			});
+			// use sessionStorage for the workerPath
+			const path = localStorage["bare-mux-path"];
+			console.debug("bare-mux: got localStorage bare-mux-path:", path);
+			if (!path) throw new Error("Unable to get bare-mux workerPath from localStorage.");
+			this.port = createPort(path, inInit);
 		} else {
 			// SharedWorker does not exist
 			throw new Error("Unable to get a channel to the SharedWorker.");
diff --git a/src/snapshot.ts b/src/snapshot.ts
index 0cee3cd..2d5d582 100644
--- a/src/snapshot.ts
+++ b/src/snapshot.ts
@@ -7,6 +7,8 @@ export const Request = globalThis.Request;
 export const Response = globalThis.Response;
 export const XMLHttpRequest = globalThis.XMLHttpRequest;
 export const SharedWorker = globalThis.SharedWorker;
+export const localStorage = globalThis.localStorage;
+export const serviceWorker = globalThis.navigator.serviceWorker;
 
 export const WebSocketFields = {
   prototype: {