From eebefbc0707d0b9759a97c66adcfc2ada60d706e Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Fri, 26 Jul 2024 22:24:25 -0700
Subject: [PATCH] use js function to encode

---
 rewriter/src/lib.rs     | 38 ++++++++++++++++++++++++++------------
 rewriter/src/rewrite.rs |  2 +-
 server.js               |  5 +++++
 src/scramjet.config.ts  |  2 +-
 src/thread/thread.ts    |  3 ++-
 src/worker/fetch.ts     |  6 +++++-
 static/ui.js            |  2 +-
 7 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/rewriter/src/lib.rs b/rewriter/src/lib.rs
index 12919a9..ef66100 100644
--- a/rewriter/src/lib.rs
+++ b/rewriter/src/lib.rs
@@ -2,8 +2,8 @@ pub mod rewrite;
 
 use std::{panic, str::FromStr};
 
-use js_sys::encode_uri_component;
-use rewrite::rewrite;
+use js_sys::Function;
+use rewrite::{rewrite, EncodeFn};
 use url::Url;
 use wasm_bindgen::prelude::*;
 
@@ -13,26 +13,40 @@ extern "C" {
     fn log(s: &str);
 }
 
-// import the SCRAM!!! jet encoder here later
-fn encode(s: String) -> String {
-    encode_uri_component(&s).into()
-}
-
 #[wasm_bindgen]
 pub fn init() {
     panic::set_hook(Box::new(console_error_panic_hook::hook));
 }
 
-#[wasm_bindgen]
-pub fn rewrite_js(js: &str, url: &str) -> Vec<u8> {
-    rewrite(js, Url::from_str(url).unwrap(), Box::new(encode))
+fn create_encode_function(encode: Function) -> EncodeFn {
+    Box::new(move |str| {
+        encode
+            .call1(&JsValue::NULL, &str.into())
+            .unwrap()
+            .as_string()
+            .unwrap()
+            .to_string()
+    })
 }
 
 #[wasm_bindgen]
-pub fn rewrite_js_from_arraybuffer(js: &[u8], url: &str) -> Vec<u8> {
+pub fn rewrite_js(js: &str, url: &str, encode: Function) -> Vec<u8> {
+    rewrite(
+        js,
+        Url::from_str(url).unwrap(),
+        create_encode_function(encode),
+    )
+}
+
+#[wasm_bindgen]
+pub fn rewrite_js_from_arraybuffer(js: &[u8], url: &str, encode: Function) -> Vec<u8> {
     // technically slower than the c++ string conversion but it will create *less copies*
 
     let js = unsafe { std::str::from_utf8_unchecked(js) };
 
-    rewrite(js, Url::from_str(url).unwrap(), Box::new(encode))
+    rewrite(
+        js,
+        Url::from_str(url).unwrap(),
+        create_encode_function(encode),
+    )
 }
diff --git a/rewriter/src/rewrite.rs b/rewriter/src/rewrite.rs
index e086b70..27c72ab 100644
--- a/rewriter/src/rewrite.rs
+++ b/rewriter/src/rewrite.rs
@@ -27,7 +27,7 @@ enum JsChange {
     },
 }
 
-type EncodeFn = Box<dyn Fn(String) -> String>;
+pub type EncodeFn = Box<dyn Fn(String) -> String>;
 struct Rewriter {
     jschanges: Vec<JsChange>,
     base: Url,
diff --git a/server.js b/server.js
index 35745f5..164ec03 100644
--- a/server.js
+++ b/server.js
@@ -49,6 +49,11 @@ fastify.register(fastifyStatic, {
 	prefix: "/scram/",
 	decorateReply: false,
 });
+fastify.register(fastifyStatic, {
+	root: join(fileURLToPath(new URL(".", import.meta.url)), "./assets"),
+	prefix: "/assets/",
+	decorateReply: false,
+});
 fastify.register(fastifyStatic, {
 	root: baremuxPath,
 	prefix: "/baremux/",
diff --git a/src/scramjet.config.ts b/src/scramjet.config.ts
index 736622b..468fb34 100644
--- a/src/scramjet.config.ts
+++ b/src/scramjet.config.ts
@@ -4,7 +4,7 @@ if (!self.$scramjet) {
 }
 self.$scramjet.config = {
 	prefix: "/scramjet/",
-	codec: self.$scramjet.codecs.plain,
+	codec: self.$scramjet.codecs.base64,
 	config: "/scram/scramjet.config.js",
 	shared: "/scram/scramjet.shared.js",
 	worker: "/scram/scramjet.worker.js",
diff --git a/src/thread/thread.ts b/src/thread/thread.ts
index 5cb44f8..5b9345d 100644
--- a/src/thread/thread.ts
+++ b/src/thread/thread.ts
@@ -36,5 +36,6 @@ const tasks = {
 };
 
 function taskRewriteJs(js: ArrayBuffer, origin: string): string {
-	return rewriteJs(js, new URL(origin));
+	// idk how to get the codec from here
+	return rewriteJs(js, new URL(origin), () => {});
 }
diff --git a/src/worker/fetch.ts b/src/worker/fetch.ts
index 94cef3a..90a7f6e 100644
--- a/src/worker/fetch.ts
+++ b/src/worker/fetch.ts
@@ -80,7 +80,11 @@ async function handleResponse(
 				}
 				break;
 			case "script":
-				responseBody = rewriteJs(await response.arrayBuffer(), url);
+				responseBody = rewriteJs(
+					await response.arrayBuffer(),
+					url,
+					self.$scramjet.config.codec.encode
+				);
 				// Disable threading for now, it's causing issues.
 				// responseBody = await this.threadpool.rewriteJs(await responseBody.arrayBuffer(), url.toString());
 				break;
diff --git a/static/ui.js b/static/ui.js
index 2346d2d..c328832 100644
--- a/static/ui.js
+++ b/static/ui.js
@@ -157,7 +157,7 @@ window.addEventListener("load", async () => {
 
 		return btoa(binary);
 	}
-	const arraybuffer = await (await fetch("/scramjet.png")).arrayBuffer();
+	const arraybuffer = await (await fetch("/assets/scramjet.png")).arrayBuffer();
 	console.log(
 		"%cb",
 		`