From 53f0ecd9e5fb26db7ab60204ae1f0b57d5aeed87 Mon Sep 17 00:00:00 2001
From: Percs <83934299+Percslol@users.noreply.github.com>
Date: Mon, 3 Feb 2025 13:42:26 -0600
Subject: [PATCH] add module rewriting more deeply and clean things up

---
 src/client/client.ts           | 10 +++++-----
 src/client/dom/attr.ts         |  8 ++++----
 src/client/dom/css.ts          |  4 ++--
 src/client/location.ts         |  8 ++++----
 src/client/shared/err.ts       |  6 +++---
 src/client/shared/event.ts     | 10 +++++-----
 src/client/shared/worker.ts    |  5 +++--
 src/shared/rewriters/html.ts   |  7 +++++++
 src/shared/rewriters/worker.ts |  8 +++-----
 src/worker/fetch.ts            |  7 ++++++-
 10 files changed, 42 insertions(+), 31 deletions(-)

diff --git a/src/client/client.ts b/src/client/client.ts
index 300824e..b84cfaa 100644
--- a/src/client/client.ts
+++ b/src/client/client.ts
@@ -345,7 +345,7 @@ export class ScramjetClient {
 		if (handler.construct) {
 			h.construct = function (
 				constructor: any,
-				argArray: any[],
+				args: any[],
 				newTarget: AnyFunction
 			) {
 				let returnValue: any = undefined;
@@ -354,7 +354,7 @@ export class ScramjetClient {
 				const ctx: ProxyCtx = {
 					fn: constructor,
 					this: null,
-					args: argArray,
+					args,
 					newTarget: newTarget,
 					return: (r: any) => {
 						earlyreturn = true;
@@ -379,14 +379,14 @@ export class ScramjetClient {
 		}
 
 		if (handler.apply) {
-			h.apply = function (fn: any, thisArg: any, argArray: any[]) {
+			h.apply = function (fn: any, that: any, args: any[]) {
 				let returnValue: any = undefined;
 				let earlyreturn = false;
 
 				const ctx: ProxyCtx = {
 					fn,
-					this: thisArg,
-					args: argArray,
+					this: that,
+					args,
 					newTarget: null,
 					return: (r: any) => {
 						earlyreturn = true;
diff --git a/src/client/dom/attr.ts b/src/client/dom/attr.ts
index 6505e08..dd6f5fe 100644
--- a/src/client/dom/attr.ts
+++ b/src/client/dom/attr.ts
@@ -22,12 +22,12 @@ export default function (client: ScramjetClient, _self: typeof window) {
 
 					if (prop in NamedNodeMap.prototype && typeof value === "function") {
 						return new Proxy(value, {
-							apply(target, thisArg, argArray) {
-								if (thisArg === proxy) {
-									return Reflect.apply(target, map, argArray);
+							apply(target, that, args) {
+								if (that === proxy) {
+									return Reflect.apply(target, map, args);
 								}
 
-								return Reflect.apply(target, thisArg, argArray);
+								return Reflect.apply(target, that, args);
 							},
 						});
 					}
diff --git a/src/client/dom/css.ts b/src/client/dom/css.ts
index 9150a9d..55c4fff 100644
--- a/src/client/dom/css.ts
+++ b/src/client/dom/css.ts
@@ -39,8 +39,8 @@ export default function (client: ScramjetClient) {
 
 					if (typeof v === "function") {
 						return new Proxy(v, {
-							apply(target, thisArg, argArray) {
-								return Reflect.apply(target, style, argArray);
+							apply(target, that, args) {
+								return Reflect.apply(target, style, args);
 							},
 						});
 					}
diff --git a/src/client/location.ts b/src/client/location.ts
index 633bc1e..50143c8 100644
--- a/src/client/location.ts
+++ b/src/client/location.ts
@@ -45,7 +45,7 @@ export function createLocationProxy(
 		}
 		if (native.set) {
 			desc.set = new Proxy(native.set, {
-				apply(target, thisArg, args) {
+				apply(target, that, args) {
 					if (prop === "href") {
 						// special case
 						client.url = args[0];
@@ -83,20 +83,20 @@ export function createLocationProxy(
 		});
 	if (self.location.assign)
 		fakeLocation.assign = new Proxy(self.location.assign, {
-			apply(target, thisArg, args) {
+			apply(target, that, args) {
 				args[0] = rewriteUrl(args[0], client.meta);
 				Reflect.apply(target, self.location, args);
 			},
 		});
 	if (self.location.reload)
 		fakeLocation.reload = new Proxy(self.location.reload, {
-			apply(target, thisArg, args) {
+			apply(target, that, args) {
 				Reflect.apply(target, self.location, args);
 			},
 		});
 	if (self.location.replace)
 		fakeLocation.replace = new Proxy(self.location.replace, {
-			apply(target, thisArg, args) {
+			apply(target, that, args) {
 				args[0] = rewriteUrl(args[0], client.meta);
 				Reflect.apply(target, self.location, args);
 			},
diff --git a/src/client/shared/err.ts b/src/client/shared/err.ts
index 7d5ceef..03b4a78 100644
--- a/src/client/shared/err.ts
+++ b/src/client/shared/err.ts
@@ -43,9 +43,9 @@ export default function (client: ScramjetClient, self: typeof globalThis) {
 		apply(ctx) {
 			if (ctx.args[0])
 				ctx.args[0] = new Proxy(ctx.args[0], {
-					apply(target, thisArg, argArray) {
-						// console.warn("CAUGHT PROMISE REJECTION", argArray);
-						Reflect.apply(target, thisArg, argArray);
+					apply(target, that, args) {
+						// console.warn("CAUGHT PROMISE REJECTION", args);
+						Reflect.apply(target, that, args);
 					},
 				});
 		},
diff --git a/src/client/shared/event.ts b/src/client/shared/event.ts
index ba8951c..1190a3e 100644
--- a/src/client/shared/event.ts
+++ b/src/client/shared/event.ts
@@ -48,8 +48,8 @@ export default function (client: ScramjetClient, self: Self) {
 
 	function wraplistener(listener: (...args: any) => any) {
 		return new Proxy(listener, {
-			apply(target, thisArg, argArray) {
-				const realEvent: Event = argArray[0];
+			apply(target, that, args) {
+				const realEvent: Event = args[0];
 
 				// we only need to handle events dispatched from the browser
 				if (realEvent.isTrusted) {
@@ -62,7 +62,7 @@ export default function (client: ScramjetClient, self: Self) {
 							if (handler._init.call(realEvent) === false) return;
 						}
 
-						argArray[0] = new Proxy(realEvent, {
+						args[0] = new Proxy(realEvent, {
 							get(_target, prop, reciever) {
 								if (prop in handler) {
 									return handler[prop].call(_target);
@@ -78,13 +78,13 @@ export default function (client: ScramjetClient, self: Self) {
 				if (!self.event) {
 					Object.defineProperty(self, "event", {
 						get() {
-							return argArray[0];
+							return args[0];
 						},
 						configurable: true,
 					});
 				}
 
-				const rv = Reflect.apply(target, thisArg, argArray);
+				const rv = Reflect.apply(target, that, args);
 
 				return rv;
 			},
diff --git a/src/client/shared/worker.ts b/src/client/shared/worker.ts
index ffa6a88..fe07e27 100644
--- a/src/client/shared/worker.ts
+++ b/src/client/shared/worker.ts
@@ -32,7 +32,7 @@ export default function (client: ScramjetClient, _self: typeof globalThis) {
 	// sharedworkers can only be constructed from window
 	client.Proxy("SharedWorker", {
 		construct({ args, call }) {
-			args[0] = rewriteUrl(args[0], client.meta) + "?dest=worker";
+			args[0] = rewriteUrl(args[0], client.meta) + "?dest=sharedworker";
 
 			if (args[1] && typeof args[1] === "string")
 				args[1] = `${client.url.origin}@${args[1]}`;
@@ -67,7 +67,8 @@ export default function (client: ScramjetClient, _self: typeof globalThis) {
 
 	client.Proxy("Worklet.prototype.addModule", {
 		apply(ctx) {
-			if (ctx.args[0]) ctx.args[0] = rewriteUrl(ctx.args[0], client.meta);
+			if (ctx.args[0])
+				ctx.args[0] = rewriteUrl(ctx.args[0], client.meta) + "?dest=worklet";
 		},
 	});
 }
diff --git a/src/shared/rewriters/html.ts b/src/shared/rewriters/html.ts
index 80b9aa1..6c13ead 100644
--- a/src/shared/rewriters/html.ts
+++ b/src/shared/rewriters/html.ts
@@ -240,6 +240,13 @@ function traverseParsedHtml(
 	if (node.name === "style" && node.children[0] !== undefined)
 		node.children[0].data = rewriteCss(node.children[0].data, meta);
 
+	if (
+		node.name === "script" &&
+		node.attribs.type === "module" &&
+		node.attribs.src
+	)
+		node.attribs.src = node.attribs.src + "?type=module";
+
 	if (
 		node.name === "script" &&
 		/(application|text)\/javascript|module|importmap|undefined/.test(
diff --git a/src/shared/rewriters/worker.ts b/src/shared/rewriters/worker.ts
index d4b486b..912fe1d 100644
--- a/src/shared/rewriters/worker.ts
+++ b/src/shared/rewriters/worker.ts
@@ -9,9 +9,9 @@ export function rewriteWorkers(
 	meta: URLMeta
 ) {
 	let str = "";
-
+	const module = type === "module";
 	const script = (script) => {
-		if (type === "module") {
+		if (module) {
 			str += `import "${$scramjet.config.files[script]}"\n`;
 		} else {
 			str += `importScripts("${$scramjet.config.files[script]}");\n`;
@@ -23,14 +23,12 @@ export function rewriteWorkers(
 	str += `self.$scramjet.config = ${JSON.stringify($scramjet.config)};`;
 	script("client");
 
-	let rewritten = rewriteJs(js, url, meta);
+	let rewritten = rewriteJs(js, url, meta, module);
 	if (rewritten instanceof Uint8Array) {
 		rewritten = new TextDecoder().decode(rewritten);
 	}
 
 	str += rewritten;
 
-	// dbg.log("Rewrite", type, str);
-
 	return str;
 }
diff --git a/src/worker/fetch.ts b/src/worker/fetch.ts
index ab8d1f3..fde3301 100644
--- a/src/worker/fetch.ts
+++ b/src/worker/fetch.ts
@@ -341,7 +341,12 @@ async function rewriteBody(
 				return response.body;
 			}
 		case "script":
-			return rewriteJs(await response.arrayBuffer(), response.finalURL, meta);
+			return rewriteJs(
+				await response.arrayBuffer(),
+				response.finalURL,
+				meta,
+				workertype === "module"
+			);
 		case "style":
 			return rewriteCss(await response.text(), meta);
 		case "sharedworker":