test cases and drm

rewriter/Cargo.toml

@@ -9,6 +9,7 @@ crate-type = ["cdylib"]
 [features]
 default = ["debug"]
 debug = []
+drm = []
 
 [profile.speed]
 inherits = "release"
@@ -20,6 +21,7 @@ opt-level = "z"
 [dependencies]
 console_error_panic_hook = "0.1.7"
 js-sys = "0.3.69"
+obfstr = "0.4.3"
 oxc_allocator = "0.20.0"
 oxc_ast = { version = "0.20.0", features = ["serialize"]}
 oxc_parser = { version = "0.20.0" }
@@ -29,3 +31,6 @@ serde = "1.0.204"
 serde-wasm-bindgen = "0.6.5"
 url = "2.5.2"
 wasm-bindgen =  "0.2.92"
+
+[dev-dependencies]
+boa_engine = "0.19.0"

rewriter/src/lib.rs

@@ -38,10 +38,7 @@ fn get_obj(obj: &JsValue, k: &str) -> JsValue {
 }
 
 fn get_str(obj: &JsValue, k: &str) -> String {
-	Reflect::get(obj, &k.into())
+	Reflect::get(obj, &k.into()).unwrap().as_string().unwrap()
-		.unwrap()
-		.as_string()
-		.unwrap()
 }
 
 fn get_config(scramjet: &Object) -> Config {
@@ -57,13 +54,33 @@ fn get_config(scramjet: &Object) -> Config {
 	}
 }
 
+#[cfg(feature = "drm")]
+#[inline(always)]
+fn drmcheck() -> bool {
+	use js_sys::global;
+	use obfstr::obfstr;
+
+	let true_origin = get_str(&get_obj(&global(), obfstr!("location")), obfstr!("origin"));
+	return vec![obfstr!("http://localhost:1337")].contains(&true_origin.as_str());
+}
+
 #[wasm_bindgen]
 pub fn rewrite_js(js: &str, url: &str, scramjet: &Object) -> Vec<u8> {
+	#[cfg(feature = "drm")]
+	if !drmcheck() {
+		return Vec::new();
+	}
+
 	rewrite(js, Url::from_str(url).unwrap(), get_config(scramjet))
 }
 
 #[wasm_bindgen]
 pub fn rewrite_js_from_arraybuffer(js: &[u8], url: &str, scramjet: &Object) -> Vec<u8> {
+	#[cfg(feature = "drm")]
+	if !drmcheck() {
+		return Vec::new();
+	}
+
 	// we know that this is a valid utf-8 string
 	let js = unsafe { std::str::from_utf8_unchecked(js) };
 

rewriter/src/main.rs

@@ -107,29 +107,138 @@ fn encode_string(s: String) -> String {
 	encode(&s).to_string()
 }
 
+fn dorewrite(source_text: &str) -> String {
+	from_utf8(
+		rewrite(
+			&source_text,
+			Url::from_str("https://google.com/glorngle/si.js").unwrap(),
+			Config {
+				prefix: "/scrammedjet/".to_string(),
+				encode: Box::new(encode_string),
+				wrapfn: "$wrap".to_string(),
+				importfn: "$import".to_string(),
+				rewritefn: "$rewrite".to_string(),
+			},
+		)
+		.as_slice(),
+	)
+	.unwrap()
+	.to_string()
+}
+
 fn main() -> std::io::Result<()> {
 	let name = env::args().nth(1).unwrap_or_else(|| "test.js".to_string());
 	let path = Path::new(&name);
 	let source_text = std::fs::read_to_string(path)?;
 
-	println!(
+	println!("{}", dorewrite(&source_text));
-		"{}",
-		from_utf8(
-			rewrite(
-				&source_text,
-				Url::from_str("https://google.com/glorngle/si.js").unwrap(),
-				Config {
-					prefix: "/scrammedjet/".to_string(),
-					encode: Box::new(encode_string),
-					wrapfn: "$wrap".to_string(),
-					importfn: "$import".to_string(),
-					rewritefn: "$rewrite".to_string(),
-				}
-			)
-			.as_slice()
-		)
-		.unwrap()
-	);
 
 	Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+	use std::{fs, path::Path};
+
+	use boa_engine::{
+		js_str, js_string,
+		object::ObjectInitializer,
+		property::{Attribute, PropertyDescriptorBuilder},
+		Context, NativeFunction, Source,
+	};
+
+	use crate::dorewrite;
+
+	#[test]
+	fn google() {
+		// sanity check- just making sure it won't crash
+		let source_text = include_str!("../sample/google.js");
+		dorewrite(source_text);
+	}
+
+	#[test]
+	fn test() {
+		let files = fs::read_dir("./tests").unwrap();
+
+		for file in files {
+			if !file
+				.as_ref()
+				.unwrap()
+				.file_name()
+				.to_str()
+				.unwrap()
+				.ends_with(".js")
+			{
+				continue;
+			}
+
+			let content = fs::read_to_string(file.unwrap().path()).unwrap();
+
+			let mut context = Context::default();
+
+			let window = ObjectInitializer::new(&mut context).build();
+			context
+				.register_global_property(js_str!("window"), window, Attribute::READONLY)
+				.unwrap();
+			context
+				.global_object()
+				.define_property_or_throw(
+					js_str!("location"),
+					PropertyDescriptorBuilder::new()
+						.get(
+							NativeFunction::from_copy_closure(|_, _, _| {
+								Ok(js_str!("location").into())
+							})
+							.to_js_function(context.realm()),
+						)
+						.set(
+							NativeFunction::from_copy_closure(|_, _, _| {
+								panic!("fail: window.location got set")
+							})
+							.to_js_function(context.realm()),
+						)
+						.build(),
+					&mut context,
+				)
+				.unwrap();
+
+			context
+				.register_global_callable(
+					js_string!("fail"),
+					0,
+					NativeFunction::from_copy_closure(|_, _, _| {
+						panic!("fail");
+					}),
+				)
+				.unwrap();
+
+			let result = context
+				.eval(Source::from_bytes(
+					br#"
+function $wrap(val) {
+	if (val === window || val === "location" || val === globalThis) return "";
+
+    return val;
+}
+
+function assert(val) {
+	if (!val) fail();
+}
+
+function check(val) {
+    if (val === window || val === "location") fail();
+}
+			    "#,
+				))
+				.unwrap();
+
+			let rewritten = dorewrite(&content);
+			println!("{}", rewritten);
+
+			context
+				.eval(Source::from_bytes(rewritten.as_bytes()))
+				.unwrap();
+			println!("PASS");
+		}
+	}
+}

rewriter/tests/0-sanity.js

@@ -0,0 +1,5 @@
+check(window);
+check(this);
+check(globalThis);
+check(location);
+check(globalThis["win" + "dow"])

rewriter/tests/1-destructure.js

@@ -0,0 +1,3 @@
+const { location: x } = globalThis;
+
+check(x);

rewriter/tests/2-function-default.js

@@ -0,0 +1,6 @@
+function f(g = globalThis, l = location) {
+  check(g);
+  check(l);
+}
+
+f();

rewriter/tests/3-shadow-globals.js

@@ -0,0 +1,7 @@
+
+function f(location, globalThis) {
+  assert(location === 1)
+  assert(globalThis === 2)
+}
+
+f(1, 2);

rewriter/tests/eval-overridden-locals.js.disabled

@@ -0,0 +1,8 @@
+
+let reached = false
+let eval = (t) => t === "location = 1" && (reached = true)
+
+// testing to make sure this doesn't get rewritten
+eval("location = 1")
+
+if (!reached) fail();

rewriter/tests/for-assignment.js.disabled

@@ -0,0 +1,3 @@
+for (location of ["https://google.com"]) {
+  //
+}