From 8fc98e66c85a5c61be26d1b278b1cd8486013a00 Mon Sep 17 00:00:00 2001
From: velzie <velzie@velzie.rip>
Date: Sun, 13 Oct 2024 15:05:55 -0400
Subject: [PATCH] fix siteflags

---
 rewriter/src/lib.rs        | 33 +++++++++++++++++++++++++--------
 src/client/client.ts       |  2 +-
 src/scramjet.ts            |  1 +
 src/shared/rewriters/js.ts |  5 +++--
 static/ui.js               |  7 ++++++-
 5 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/rewriter/src/lib.rs b/rewriter/src/lib.rs
index 42049c6..73d2ab8 100644
--- a/rewriter/src/lib.rs
+++ b/rewriter/src/lib.rs
@@ -2,7 +2,10 @@ pub mod rewrite;
 
 use std::{panic, str::FromStr};
 
-use js_sys::{Function, Object, Reflect};
+use js_sys::{
+	global, Array, Function, Object,
+	Reflect::{self, construct},
+};
 use rewrite::{rewrite, Config, EncodeFn};
 use url::Url;
 use wasm_bindgen::{prelude::*, throw_str};
@@ -45,10 +48,24 @@ fn get_str(obj: &JsValue, k: &str) -> String {
 	Reflect::get(obj, &k.into()).unwrap().as_string().unwrap()
 }
 
-fn get_config(scramjet: &Object) -> Config {
+fn get_flag(scramjet: &Object, url: &str, flag: &str) -> bool {
+	let urlconstructor = get_obj(&global(), "URL");
+	let args = Array::new();
+	args.push(&JsValue::from_str(url));
+	let url = construct(&urlconstructor.dyn_into::<Function>().unwrap(), &args).unwrap();
+	let fenabled = get_obj(scramjet, "flagEnabled")
+		.dyn_into::<Function>()
+		.unwrap();
+	return fenabled
+		.call2(&JsValue::NULL, &JsValue::from_str(flag), &url)
+		.unwrap()
+		.as_bool()
+		.unwrap();
+}
+
+fn get_config(scramjet: &Object, url: &str) -> Config {
 	let codec = &get_obj(scramjet, "codec");
 	let config = &get_obj(scramjet, "config");
-	let flags = &get_obj(config, "flags");
 	let globals = &get_obj(config, "globals");
 
 	Config {
@@ -62,9 +79,9 @@ fn get_config(scramjet: &Object) -> Config {
 		setrealmfn: get_str(globals, "setrealmfn"),
 		pushsourcemapfn: get_str(globals, "pushsourcemapfn"),
 
-		do_sourcemaps: get_bool(flags, "sourcemaps"),
-		capture_errors: get_bool(flags, "captureErrors"),
-		scramitize: get_bool(flags, "scramitize"),
+		do_sourcemaps: get_flag(scramjet, url, "sourcemaps"),
+		capture_errors: get_flag(scramjet, url, "captureErrors"),
+		scramitize: get_flag(scramjet, url, "scramitize"),
 	}
 }
 
@@ -85,7 +102,7 @@ pub fn rewrite_js(js: &str, url: &str, scramjet: &Object) -> Vec<u8> {
 		return Vec::new();
 	}
 
-	rewrite(js, Url::from_str(url).unwrap(), get_config(scramjet))
+	rewrite(js, Url::from_str(url).unwrap(), get_config(scramjet, url))
 }
 
 #[wasm_bindgen]
@@ -98,5 +115,5 @@ pub fn rewrite_js_from_arraybuffer(js: &[u8], url: &str, scramjet: &Object) -> V
 	// we know that this is a valid utf-8 string
 	let js = unsafe { std::str::from_utf8_unchecked(js) };
 
-	rewrite(js, Url::from_str(url).unwrap(), get_config(scramjet))
+	rewrite(js, Url::from_str(url).unwrap(), get_config(scramjet, url))
 }
diff --git a/src/client/client.ts b/src/client/client.ts
index cef1ed4..332b757 100644
--- a/src/client/client.ts
+++ b/src/client/client.ts
@@ -192,7 +192,7 @@ export class ScramjetClient {
 		});
 
 		for (const module of modules) {
-			if (!module.enabled || module.enabled())
+			if (!module.enabled || module.enabled(this))
 				module.default(this, this.global);
 			else if (module.disabled) module.disabled(this, this.global);
 		}
diff --git a/src/scramjet.ts b/src/scramjet.ts
index e168734..f49e89c 100644
--- a/src/scramjet.ts
+++ b/src/scramjet.ts
@@ -8,6 +8,7 @@ if (!("$scramjet" in self)) {
 			version: VERSION,
 		},
 		codec: {},
+		flagEnabled,
 	};
 }
 
diff --git a/src/shared/rewriters/js.ts b/src/shared/rewriters/js.ts
index cc48d99..2591be8 100644
--- a/src/shared/rewriters/js.ts
+++ b/src/shared/rewriters/js.ts
@@ -7,7 +7,7 @@ import {
 	rewrite_js,
 	rewrite_js_from_arraybuffer,
 } from "../../../rewriter/out/rewriter.js";
-import { $scramjet } from "../../scramjet";
+import { $scramjet, flagEnabled } from "../../scramjet";
 
 initSync({
 	module: new WebAssembly.Module(
@@ -20,9 +20,10 @@ init();
 Error.stackTraceLimit = 50;
 
 export function rewriteJs(js: string | ArrayBuffer, meta: URLMeta) {
-	if ($scramjet.config.flags.naiiveRewriter) {
+	if (flagEnabled("naiiveRewriter", meta.origin)) {
 		const text = typeof js === "string" ? js : new TextDecoder().decode(js);
 
+		console.log("naiive");
 		return rewriteJsNaiive(text);
 	}
 
diff --git a/static/ui.js b/static/ui.js
index dd784f5..5f8915b 100644
--- a/static/ui.js
+++ b/static/ui.js
@@ -6,6 +6,11 @@ const scramjet = new ScramjetController({
 		shared: "/scram/scramjet.shared.js",
 		sync: "/scram/scramjet.sync.js",
 	},
+	siteFlags: {
+		"https://discord.com/.*": {
+			naiiveRewriter: true,
+		},
+	},
 });
 
 scramjet.init("./sw.js");
@@ -183,7 +188,7 @@ function App() {
 		if (!url) return;
 		if (url === "about:blank") return;
 
-		this.url = $scramjet.codecs.plain.decode(
+		this.url = $scramjet.codec.decode(
 			url.substring((location.href + "/scramjet").length)
 		);
 	});