From a69755936d28da7c8784fb78333c72b6682f6ed9 Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Fri, 21 Feb 2025 17:24:09 -0800
Subject: [PATCH] use rewriter wasm directly

---
 rewriter/wasm/build.sh     | 12 ------------
 src/shared/rewriters/js.ts | 19 ++++++++++++++-----
 src/types.d.ts             |  1 +
 src/worker/fetch.ts        | 22 ++++++++++++++++++++++
 src/worker/index.ts        |  7 ++++++-
 static/sw.js               |  6 +-----
 static/ui.js               |  2 +-
 7 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/rewriter/wasm/build.sh b/rewriter/wasm/build.sh
index dd340af..5eef277 100644
--- a/rewriter/wasm/build.sh
+++ b/rewriter/wasm/build.sh
@@ -104,16 +104,4 @@ fi
 mkdir -p dist/
 
 cp rewriter/wasm/out/optimized.wasm dist/scramjet.wasm.wasm
-{
-
-cat <<EOF
-if ("document" in self && document?.currentScript) {
-	document.currentScript.remove();
-}
-EOF
-echo -n "self.WASM = '"
-base64 -w0 < "rewriter/wasm/out/optimized.wasm"
-echo -n "';"
-
-} > dist/scramjet.wasm.js
 echo "Rewriter Build Complete!"
diff --git a/src/shared/rewriters/js.ts b/src/shared/rewriters/js.ts
index 291eac3..2ce19ed 100644
--- a/src/shared/rewriters/js.ts
+++ b/src/shared/rewriters/js.ts
@@ -9,11 +9,16 @@ import {
 } from "../../../rewriter/wasm/out/wasm.js";
 import { $scramjet, flagEnabled } from "../../scramjet";
 
-initSync({
-	module: new WebAssembly.Module(
-		Uint8Array.from(atob(self.WASM), (c) => c.charCodeAt(0))
-	),
-});
+if (self.WASM)
+	self.REAL_WASM = Uint8Array.from(atob(self.WASM), (c) => c.charCodeAt(0));
+
+// only use in sw
+export async function asyncInitRewriter() {
+	const buf = await fetch($scramjet.config.files.wasm).then((r) =>
+		r.arrayBuffer()
+	);
+	self.REAL_WASM = new Uint8Array(buf);
+}
 
 Error.stackTraceLimit = 50;
 
@@ -25,6 +30,10 @@ function rewriteJsWrapper(
 	meta: URLMeta,
 	module: boolean
 ): string | ArrayBuffer {
+	initSync({
+		module: new WebAssembly.Module(self.REAL_WASM),
+	});
+
 	let out: RewriterOutput;
 	const before = performance.now();
 	try {
diff --git a/src/types.d.ts b/src/types.d.ts
index f89bfc1..9f20355 100644
--- a/src/types.d.ts
+++ b/src/types.d.ts
@@ -103,6 +103,7 @@ declare global {
 		};
 		COOKIE: string;
 		WASM: string;
+		REAL_WASM: Uint8Array;
 		ScramjetController: typeof ScramjetController;
 
 		// the scramjet client belonging to a window
diff --git a/src/worker/fetch.ts b/src/worker/fetch.ts
index fde3301..56a464c 100644
--- a/src/worker/fetch.ts
+++ b/src/worker/fetch.ts
@@ -49,6 +49,28 @@ export async function handleFetch(
 			requesturl.searchParams.delete("dest");
 		}
 
+		if (requesturl.pathname == this.config.files.wasm) {
+			return fetch(this.config.files.wasm).then(async (x) => {
+				const buf = await x.arrayBuffer();
+				const b64 = btoa(
+					new Uint8Array(buf)
+						.reduce(
+							(data, byte) => (data.push(String.fromCharCode(byte)), data),
+							[]
+						)
+						.join("")
+				);
+
+				let payload = "";
+				payload += `if ("document" in self && document.currentScript) { document.currentScript.remove(); }\n`;
+				payload += `self.WASM = '${b64}';`;
+
+				return new Response(payload, {
+					headers: { "content-type": "text/javascript" },
+				});
+			});
+		}
+
 		if (
 			requesturl.pathname.startsWith(this.config.prefix + "blob:") ||
 			requesturl.pathname.startsWith(this.config.prefix + "data:")
diff --git a/src/worker/index.ts b/src/worker/index.ts
index 3a451f0..5d1f37b 100644
--- a/src/worker/index.ts
+++ b/src/worker/index.ts
@@ -3,6 +3,7 @@ import { handleFetch } from "./fetch";
 import type BareClient from "@mercuryworkshop/bare-mux";
 import { ScramjetConfig } from "../types";
 import { $scramjet, loadCodecs } from "../scramjet";
+import { asyncInitRewriter } from "../shared/rewriters/js";
 
 export class ScramjetServiceWorker extends EventTarget {
 	client: BareClient;
@@ -69,12 +70,14 @@ export class ScramjetServiceWorker extends EventTarget {
 				const store = tx.objectStore("config");
 				const config = store.get("config");
 
-				config.onsuccess = () => {
+				config.onsuccess = async () => {
 					this.config = config.result;
 					$scramjet.config = config.result;
 
 					loadCodecs();
 
+					await asyncInitRewriter();
+
 					resolve();
 				};
 				config.onerror = () => reject(config.error);
@@ -87,6 +90,8 @@ export class ScramjetServiceWorker extends EventTarget {
 	route({ request }: FetchEvent) {
 		if (request.url.startsWith(location.origin + this.config.prefix))
 			return true;
+		else if (request.url.startsWith(location.origin + this.config.files.wasm))
+			return true;
 		else return false;
 	}
 
diff --git a/static/sw.js b/static/sw.js
index 8daf1f4..83ea084 100644
--- a/static/sw.js
+++ b/static/sw.js
@@ -6,11 +6,7 @@ if (navigator.userAgent.includes("Firefox")) {
 	});
 }
 
-importScripts(
-	"/scram/scramjet.wasm.js",
-	"/scram/scramjet.shared.js",
-	"/scram/scramjet.worker.js"
-);
+importScripts("/scram/scramjet.shared.js", "/scram/scramjet.worker.js");
 
 const scramjet = new ScramjetServiceWorker();
 
diff --git a/static/ui.js b/static/ui.js
index a3e7e1c..d023bf7 100644
--- a/static/ui.js
+++ b/static/ui.js
@@ -1,6 +1,6 @@
 const scramjet = new ScramjetController({
 	files: {
-		wasm: "/scram/scramjet.wasm.js",
+		wasm: "/scram/scramjet.wasm.wasm",
 		worker: "/scram/scramjet.worker.js",
 		client: "/scram/scramjet.client.js",
 		shared: "/scram/scramjet.shared.js",