From e0a16f1c607f0867856488c6cb30a59e5366eefc Mon Sep 17 00:00:00 2001
From: velzie <velzie@velzie.rip>
Date: Mon, 28 Oct 2024 19:42:34 -0400
Subject: [PATCH] skull emoji

---
 rewriter/src/lib.rs            |   6 +-
 rewriter/src/main.rs           |   1 +
 rewriter/src/rewrite.rs        | 122 +++++++++++++++++++--------------
 src/client/dom/element.ts      |   2 +-
 src/client/shared/eval.ts      |   4 +-
 src/client/shared/function.ts  |   6 +-
 src/client/shared/import.ts    |   2 +-
 src/client/shared/wrap.ts      |  15 ++++
 src/shared/rewriters/html.ts   |  42 ++++++------
 src/shared/rewriters/js.ts     |  33 ++++++---
 src/shared/rewriters/url.ts    |   4 +-
 src/shared/rewriters/worker.ts |   2 +-
 src/worker/fetch.ts            |   6 +-
 13 files changed, 153 insertions(+), 92 deletions(-)

diff --git a/rewriter/src/lib.rs b/rewriter/src/lib.rs
index da833ec..ce60bfe 100644
--- a/rewriter/src/lib.rs
+++ b/rewriter/src/lib.rs
@@ -139,6 +139,7 @@ fn create_rewriter_output(
 #[wasm_bindgen]
 pub fn rewrite_js(
 	js: String,
+	module: bool,
 	url: &str,
 	script_url: String,
 	scramjet: &Object,
@@ -149,7 +150,7 @@ pub fn rewrite_js(
 	}
 
 	let before = Instant::now();
-	let out = rewrite(&js, Url::from_str(url)?, get_config(scramjet, url)?)?;
+	let out = rewrite(&js, module, Url::from_str(url)?, get_config(scramjet, url)?)?;
 	let after = Instant::now();
 
 	create_rewriter_output(out, script_url, js, after - before)
@@ -158,6 +159,7 @@ pub fn rewrite_js(
 #[wasm_bindgen]
 pub fn rewrite_js_from_arraybuffer(
 	js: Vec<u8>,
+	module: bool,
 	url: &str,
 	script_url: String,
 	scramjet: &Object,
@@ -171,7 +173,7 @@ pub fn rewrite_js_from_arraybuffer(
 	let js = unsafe { String::from_utf8_unchecked(js) };
 
 	let before = Instant::now();
-	let out = rewrite(&js, Url::from_str(url)?, get_config(scramjet, url)?)?;
+	let out = rewrite(&js, module, Url::from_str(url)?, get_config(scramjet, url)?)?;
 	let after = Instant::now();
 
 	create_rewriter_output(out, script_url, js, after - before)
diff --git a/rewriter/src/main.rs b/rewriter/src/main.rs
index 5948d34..a9ce0e0 100644
--- a/rewriter/src/main.rs
+++ b/rewriter/src/main.rs
@@ -113,6 +113,7 @@ fn dorewrite(source_text: &str) -> Result<String> {
 	Ok(from_utf8(
 		rewrite(
 			source_text,
+			false,
 			Url::from_str("https://google.com/glorngle/si.js").unwrap(),
 			Config {
 				prefix: "/scrammedjet/".to_string(),
diff --git a/rewriter/src/rewrite.rs b/rewriter/src/rewrite.rs
index d32288d..2ac0d5b 100644
--- a/rewriter/src/rewrite.rs
+++ b/rewriter/src/rewrite.rs
@@ -5,6 +5,7 @@ use oxc_allocator::Allocator;
 use oxc_ast::{
 	ast::{
 		AssignmentTarget, Expression, IdentifierReference, MemberExpression, ObjectPropertyKind,
+		Statement, VariableDeclaration,
 	},
 	visit::walk,
 	Visit,
@@ -65,7 +66,7 @@ impl Rewriter {
 
 		let urlencoded = (self.config.encode)(url.to_string());
 
-		format!("\"{}{}\"", self.config.prefix, urlencoded)
+		format!("\"{}{}?type=module\"", self.config.prefix, urlencoded)
 	}
 
 	fn rewrite_ident(&mut self, name: &Atom, span: Span) {
@@ -106,20 +107,23 @@ impl<'a> Visit<'a> for Rewriter {
 		// 		),
 		// 	});
 		// } else {
-		if UNSAFE_GLOBALS.contains(&it.name.to_string().as_str()) {
-			self.jschanges.push(JsChange::GenericChange {
-				span: it.span,
-				text: format!("{}({})", self.config.wrapfn, it.name),
-			});
-		}
+		// if UNSAFE_GLOBALS.contains(&it.name.to_string().as_str()) {
+		// 	self.jschanges.push(JsChange::GenericChange {
+		// 		span: it.span,
+		// 		text: format!("{}({})", self.config.wrapfn, it.name),
+		// 	});
+		// }
 		// }
 	}
 
+	// fn visit_program(&mut self, it: &oxc_ast::ast::Program<'a>) {
+	//     it.directives[0].
+	// }
 	// we need to rewrite `new Something` to `new (wrapfn(Something))` instead of `new wrapfn(Something)`, that's why there's weird extra code here
-	fn visit_new_expression(&mut self, it: &oxc_ast::ast::NewExpression<'a>) {
-		self.walk_member_expression(&it.callee);
-		walk::walk_arguments(self, &it.arguments);
-	}
+	// fn visit_new_expression(&mut self, it: &oxc_ast::ast::NewExpression<'a>) {
+	// 	// self.walk_member_expression(&it.callee);
+	// 	walk::walk_arguments(self, &it.arguments);
+	// }
 	fn visit_member_expression(&mut self, it: &oxc_ast::ast::MemberExpression<'a>) {
 		match it {
 			MemberExpression::StaticMemberExpression(s) => {
@@ -133,18 +137,6 @@ impl<'a> Visit<'a> for Rewriter {
 					return; // unwise to walk the rest of the tree
 				}
 
-				if !self.config.strict_rewrites
-					&& !UNSAFE_GLOBALS.contains(&s.property.name.as_str())
-				{
-					if let Expression::Identifier(_) = &s.object {
-						// cull tree - this should be safe
-						return;
-					}
-					if let Expression::ThisExpression(_) = &s.object {
-						return;
-					}
-				}
-
 				if self.config.scramitize
 					&& !matches!(s.object, Expression::MetaProperty(_))
 					&& !matches!(s.object, Expression::Super(_))
@@ -170,12 +162,12 @@ impl<'a> Visit<'a> for Rewriter {
 
 		walk::walk_member_expression(self, it);
 	}
-	fn visit_this_expression(&mut self, it: &oxc_ast::ast::ThisExpression) {
-		self.jschanges.push(JsChange::GenericChange {
-			span: it.span,
-			text: format!("{}(this)", self.config.wrapthisfn),
-		});
-	}
+	// fn visit_this_expression(&mut self, it: &oxc_ast::ast::ThisExpression) {
+	// 	self.jschanges.push(JsChange::GenericChange {
+	// 		span: it.span,
+	// 		text: format!("{}(this)", self.config.wrapthisfn),
+	// 	});
+	// }
 
 	fn visit_debugger_statement(&mut self, it: &oxc_ast::ast::DebuggerStatement) {
 		// delete debugger statements entirely. some sites will spam debugger as an anti-debugging measure, and we don't want that!
@@ -276,28 +268,28 @@ impl<'a> Visit<'a> for Rewriter {
 		walk::walk_try_statement(self, it);
 	}
 
-	fn visit_object_expression(&mut self, it: &oxc_ast::ast::ObjectExpression<'a>) {
-		for prop in &it.properties {
-			#[allow(clippy::single_match)]
-			match prop {
-				ObjectPropertyKind::ObjectProperty(p) => match &p.value {
-					Expression::Identifier(s) => {
-						if UNSAFE_GLOBALS.contains(&s.name.to_string().as_str()) && p.shorthand {
-							self.jschanges.push(JsChange::GenericChange {
-								span: s.span,
-								text: format!("{}: ({}({}))", s.name, self.config.wrapfn, s.name),
-							});
-							return;
-						}
-					}
-					_ => {}
-				},
-				_ => {}
-			}
-		}
-
-		walk::walk_object_expression(self, it);
-	}
+	// fn visit_object_expression(&mut self, it: &oxc_ast::ast::ObjectExpression<'a>) {
+	// 	for prop in &it.properties {
+	// 		#[allow(clippy::single_match)]
+	// 		match prop {
+	// 			ObjectPropertyKind::ObjectProperty(p) => match &p.value {
+	// 				Expression::Identifier(s) => {
+	// 					if UNSAFE_GLOBALS.contains(&s.name.to_string().as_str()) && p.shorthand {
+	// 						self.jschanges.push(JsChange::GenericChange {
+	// 							span: s.span,
+	// 							text: format!("{}: ({}({}))", s.name, self.config.wrapfn, s.name),
+	// 						});
+	// 						return;
+	// 					}
+	// 				}
+	// 				_ => {}
+	// 			},
+	// 			_ => {}
+	// 		}
+	// 	}
+	//
+	// 	walk::walk_object_expression(self, it);
+	// }
 
 	fn visit_function_body(&mut self, it: &oxc_ast::ast::FunctionBody<'a>) {
 		// tag function for use in sourcemaps
@@ -411,7 +403,12 @@ fn random_string() -> String {
 	.to_string()
 }
 
-pub fn rewrite(js: &str, url: Url, config: Config) -> Result<(Vec<u8>, Vec<OxcDiagnostic>)> {
+pub fn rewrite(
+	js: &str,
+	module: bool,
+	url: Url,
+	config: Config,
+) -> Result<(Vec<u8>, Vec<OxcDiagnostic>)> {
 	let allocator = Allocator::default();
 	let source_type = SourceType::default();
 	let ret = Parser::new(&allocator, js, source_type).parse();
@@ -426,8 +423,29 @@ pub fn rewrite(js: &str, url: Url, config: Config) -> Result<(Vec<u8>, Vec<OxcDi
 		config,
 	};
 
+	// dbg!(&program);
 	ast_pass.visit_program(&program);
 
+	if !module {
+		for s in &program.body {
+			if let Statement::VariableDeclaration(it) = s {
+				if !it.kind.is_var() {
+					continue;
+				}
+				ast_pass.jschanges.push(JsChange::GenericChange {
+					span: Span::new(it.span.start, it.declarations[0].span.start),
+					text: String::new(),
+				});
+				for d in &it.declarations {
+					ast_pass.jschanges.push(JsChange::GenericChange {
+						span: Span::new(d.id.span().start, d.id.span().start),
+						text: "fakeVars.".to_string(),
+					});
+				}
+			}
+		}
+	}
+
 	// sorrt changse
 	ast_pass.jschanges.sort_by(|a, b| {
 		let a = match a {
diff --git a/src/client/dom/element.ts b/src/client/dom/element.ts
index 1b51655..84b73ba 100644
--- a/src/client/dom/element.ts
+++ b/src/client/dom/element.ts
@@ -202,7 +202,7 @@ export default function (client: ScramjetClient, self: typeof window) {
 		set(ctx, value: string) {
 			let newval;
 			if (ctx.this instanceof self.HTMLScriptElement) {
-				newval = rewriteJs(value, client.meta);
+				newval = rewriteJs(value, false, client.meta);
 			} else if (ctx.this instanceof self.HTMLStyleElement) {
 				newval = rewriteCss(value, client.meta);
 			} else {
diff --git a/src/client/shared/eval.ts b/src/client/shared/eval.ts
index 011ab21..19808a8 100644
--- a/src/client/shared/eval.ts
+++ b/src/client/shared/eval.ts
@@ -8,7 +8,7 @@ export default function (client: ScramjetClient, self: Self) {
 		value: function (js: any) {
 			if (typeof js !== "string") return js;
 
-			const rewritten = rewriteJs(js, client.meta);
+			const rewritten = rewriteJs(js, false, client.meta);
 
 			return rewritten;
 		},
@@ -23,5 +23,5 @@ export function indirectEval(this: ScramjetClient, js: any) {
 
 	const indirection = this.global.eval;
 
-	return indirection(rewriteJs(js, this.meta) as string);
+	return indirection(rewriteJs(js, false, this.meta) as string);
 }
diff --git a/src/client/shared/function.ts b/src/client/shared/function.ts
index a9a1b45..40bc2a1 100644
--- a/src/client/shared/function.ts
+++ b/src/client/shared/function.ts
@@ -4,7 +4,11 @@ import { rewriteJs } from "../../shared";
 function rewriteFunction(ctx: ProxyCtx, client: ScramjetClient) {
 	const stringifiedFunction = ctx.call().toString();
 
-	const content = rewriteJs(`return ${stringifiedFunction}`, client.meta);
+	const content = rewriteJs(
+		`return ${stringifiedFunction}`,
+		false,
+		client.meta
+	);
 	ctx.return(ctx.fn(content)());
 }
 
diff --git a/src/client/shared/import.ts b/src/client/shared/import.ts
index 6f1811e..f4b6b5e 100644
--- a/src/client/shared/import.ts
+++ b/src/client/shared/import.ts
@@ -10,7 +10,7 @@ export default function (client: ScramjetClient, self: Self) {
 			const resolved = new URL(url, base).href;
 
 			return Function(
-				`return import("${rewriteUrl(resolved, client.meta)}")`
+				`return import("${rewriteUrl(resolved, client.meta)}?type=module")`
 			)();
 		};
 	};
diff --git a/src/client/shared/wrap.ts b/src/client/shared/wrap.ts
index b4dcab5..106e406 100644
--- a/src/client/shared/wrap.ts
+++ b/src/client/shared/wrap.ts
@@ -55,6 +55,21 @@ export default function (client: ScramjetClient, self: typeof globalThis) {
 		writable: false,
 		configurable: false,
 	});
+
+	Object.defineProperty(self, "fakeVars", {
+		writable: false,
+		configurable: false,
+		value: new Proxy(
+			{},
+			{
+				set(t, p, v) {
+					console.log(p, v);
+					self[p] = v;
+					return true;
+				},
+			}
+		),
+	});
 	Object.defineProperty(self, config.globals.wrapthisfn, {
 		value: function (i) {
 			if (i === self) return client.globalProxy;
diff --git a/src/shared/rewriters/html.ts b/src/shared/rewriters/html.ts
index 3cf5929..d14b5f3 100644
--- a/src/shared/rewriters/html.ts
+++ b/src/shared/rewriters/html.ts
@@ -113,16 +113,7 @@ export const htmlRules: {
 		},
 
 		// url rewrites
-		src: [
-			"embed",
-			"script",
-			"img",
-			"image",
-			"iframe",
-			"source",
-			"input",
-			"track",
-		],
+		src: ["embed", "img", "image", "iframe", "source", "input", "track"],
 		href: ["a", "link", "area"],
 		data: ["object"],
 		action: ["form"],
@@ -227,17 +218,28 @@ function traverseParsedHtml(
 		node.name === "script" &&
 		/(application|text)\/javascript|module|importmap|undefined/.test(
 			node.attribs.type
-		) &&
-		node.children[0] !== undefined
+		)
 	) {
-		let js = node.children[0].data;
-		// node.attribs[`data-scramjet-script-source-src`] = btoa(js);
-		node.attribs["data-scramjet-script-source-src"] = bytesToBase64(
-			new TextEncoder().encode(js)
-		);
-		const htmlcomment = /<!--[\s\S]*?-->/g;
-		js = js.replace(htmlcomment, "");
-		node.children[0].data = rewriteJs(js, meta);
+		if (node.children[0] !== undefined) {
+			let js = node.children[0].data;
+			// node.attribs[`data-scramjet-script-source-src`] = btoa(js);
+			node.attribs["data-scramjet-script-source-src"] = bytesToBase64(
+				new TextEncoder().encode(js)
+			);
+			const htmlcomment = /<!--[\s\S]*?-->/g;
+			js = js.replace(htmlcomment, "");
+			node.children[0].data = rewriteJs(
+				js,
+				node.attribs["type"] === "module",
+				meta
+			);
+		} else if (node.attribs["src"]) {
+			let url = rewriteUrl(node.attribs["src"], meta);
+			if (node.attribs["type"] === "module") url += "?type=module";
+
+			node.attribs["data-scramjet-src"] = node.attribs["src"];
+			node.attribs["src"] = url;
+		}
 	}
 
 	if (node.name === "meta" && node.attribs["http-equiv"] != undefined) {
diff --git a/src/shared/rewriters/js.ts b/src/shared/rewriters/js.ts
index 498c554..790edc8 100644
--- a/src/shared/rewriters/js.ts
+++ b/src/shared/rewriters/js.ts
@@ -24,12 +24,14 @@ const decoder = new TextDecoder();
 
 function rewriteJsWrapper(
 	input: string | ArrayBuffer,
+	module: boolean,
 	meta: URLMeta
 ): string | ArrayBuffer {
 	let out: RewriterOutput;
 	if (typeof input === "string") {
 		out = rewrite_js(
 			input,
+			module,
 			meta.base.href,
 			"PERCS_PLEASE_FILL_THIS_IN.js",
 			$scramjet
@@ -37,6 +39,7 @@ function rewriteJsWrapper(
 	} else {
 		out = rewrite_js_from_arraybuffer(
 			new Uint8Array(input),
+			module,
 			meta.base.href,
 			"PERCS_PLEASE_FILL_THIS_IN.js",
 			$scramjet
@@ -67,17 +70,27 @@ function rewriteJsWrapper(
 	return typeof input === "string" ? decoder.decode(js) : js;
 }
 
-export function rewriteJs(js: string | ArrayBuffer, meta: URLMeta) {
-	if (flagEnabled("naiiveRewriter", meta.origin)) {
-		const text = typeof js === "string" ? js : new TextDecoder().decode(js);
+export function rewriteJs(
+	js: string | ArrayBuffer,
+	module: boolean,
+	meta: URLMeta
+) {
+	const j = rewriteJsWrapper(js, module, meta);
 
-		console.log("naiive");
+	const text = typeof j === "string" ? j : new TextDecoder().decode(j);
+	if (module) {
+		js = `
+			let location = $scramjet$wrap(self.location);
+			let window = $scramjet$wrap(self.window);
+			let globalThis = $scramjet$wrap(self.globalThis);
 
-		return rewriteJsNaiive(text);
+			${text}
+		`;
+	} else {
+		js = `(function(window, location, globalThis) {
+			${text}
+		}).call($scramjet$wrap(this), $scramjet$wrap(window), $scramjet$wrap(location), $scramjet$wrap(globalThis))`;
 	}
-
-	js = rewriteJsWrapper(js, meta);
-
 	return js;
 }
 
@@ -93,10 +106,10 @@ export function rewriteJsNaiive(js: string | ArrayBuffer) {
 	}
 
 	return `
-		with (${$scramjet.config.globals.wrapfn}(globalThis)) {
+	with (${$scramjet.config.globals.wrapfn} (globalThis)) {
 
 			${js}
 
-		}
+	}
 	`;
 }
diff --git a/src/shared/rewriters/url.ts b/src/shared/rewriters/url.ts
index bb575a7..4760614 100644
--- a/src/shared/rewriters/url.ts
+++ b/src/shared/rewriters/url.ts
@@ -32,7 +32,9 @@ export function rewriteUrl(url: string | URL, meta: URLMeta) {
 	}
 
 	if (url.startsWith("javascript:")) {
-		return "javascript:" + rewriteJs(url.slice("javascript:".length), meta);
+		return (
+			"javascript:" + rewriteJs(url.slice("javascript:".length), false, meta)
+		);
 	} else if (url.startsWith("blob:")) {
 		return location.origin + $scramjet.config.prefix + url;
 	} else if (url.startsWith("data:")) {
diff --git a/src/shared/rewriters/worker.ts b/src/shared/rewriters/worker.ts
index 543d369..261b3ae 100644
--- a/src/shared/rewriters/worker.ts
+++ b/src/shared/rewriters/worker.ts
@@ -24,7 +24,7 @@ export function rewriteWorkers(
 	script("client");
 	console.log(str);
 
-	let rewritten = rewriteJs(js, meta);
+	let rewritten = rewriteJs(js, type === "module", meta);
 	if (rewritten instanceof Uint8Array) {
 		rewritten = new TextDecoder().decode(rewritten);
 	}
diff --git a/src/worker/fetch.ts b/src/worker/fetch.ts
index c1075a1..fc2d3e3 100644
--- a/src/worker/fetch.ts
+++ b/src/worker/fetch.ts
@@ -300,7 +300,11 @@ async function rewriteBody(
 				return response.body;
 			}
 		case "script":
-			return rewriteJs(await response.arrayBuffer(), meta);
+			return rewriteJs(
+				await response.arrayBuffer(),
+				workertype === "module",
+				meta
+			);
 		case "style":
 			return rewriteCss(await response.text(), meta);
 		case "sharedworker":