feat: charset parsing & misc fixes

MercuryWorkshop · Dec 6, 2024 · 34bd4e5 · 34bd4e5
1 parent d821554
commit 34bd4e5
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 97 deletions.
diff --git a/src/client/dom/element.ts b/src/client/dom/element.ts
@@ -240,75 +240,58 @@ export default function (client: ScramjetClient, self: typeof window) {
 		},
 	});
 
-	client.Trap("HTMLIFrameElement.prototype.contentWindow", {
-		get(ctx) {
-			const realwin = ctx.get() as Window;
-			if (!realwin) return realwin;
-
-			if (SCRAMJETCLIENT in realwin) {
-				return realwin[SCRAMJETCLIENT].globalProxy;
-			} else {
-				// hook the iframe
-				const newclient = new ScramjetClient(realwin);
-				newclient.hook();
-
-				return newclient.globalProxy;
-			}
-		},
-	});
-
-	client.Trap("HTMLIFrameElement.prototype.contentDocument", {
-		get(ctx) {
-			const contentwindow =
-				client.descriptors["HTMLIFrameElement.prototype.contentWindow"].get;
-			const realwin = contentwindow.apply(ctx.this);
-			if (!realwin) return realwin;
-
-			if (SCRAMJETCLIENT in realwin) {
-				return realwin[SCRAMJETCLIENT].documentProxy;
-			} else {
-				const newclient = new ScramjetClient(realwin);
-				newclient.hook();
-
-				return newclient.documentProxy;
-			}
-		},
-	});
-
-	client.Trap("HTMLFrameElement.prototype.contentWindow", {
-		get(ctx) {
-			const realwin = ctx.get() as Window;
-			if (!realwin) return realwin;
-
-			if (SCRAMJETCLIENT in realwin) {
-				return realwin[SCRAMJETCLIENT].globalProxy;
-			} else {
-				// hook the iframe
-				const newclient = new ScramjetClient(realwin);
-				newclient.hook();
-
-				return newclient.globalProxy;
-			}
-		},
-	});
+	client.Trap(
+		[
+			"HTMLIFrameElement.prototype.contentWindow",
+			"HTMLFrameElement.prototype.contentWindow",
+			"HTMLObjectElement.prototype.contentWindow",
+			"HTMLEmbedElement.prototype.contentWindow",
+		],
+		{
+			get(ctx) {
+				const realwin = ctx.get() as Window;
+				if (!realwin) return realwin;
 
-	client.Trap("HTMLFrameElement.prototype.contentDocument", {
-		get(ctx) {
-			const contentwindow =
-				client.descriptors["HTMLFrameElement.prototype.contentWindow"].get;
-			const realwin = contentwindow.apply(ctx.this);
-			if (!realwin) return realwin;
+				if (SCRAMJETCLIENT in realwin) {
+					return realwin[SCRAMJETCLIENT].globalProxy;
+				} else {
+					// hook the iframe
+					const newclient = new ScramjetClient(realwin);
+					newclient.hook();
 
-			if (SCRAMJETCLIENT in realwin) {
-				return realwin[SCRAMJETCLIENT].documentProxy;
-			} else {
-				const newclient = new ScramjetClient(realwin);
-				newclient.hook();
+					return newclient.globalProxy;
+				}
+			},
+		}
+	);
 
-				return newclient.documentProxy;
-			}
-		},
-	});
+	client.Trap(
+		[
+			"HTMLIFrameElement.prototype.contentDocument",
+			"HTMLFrameElement.prototype.contentDocument",
+			"HTMLObjectElement.prototype.contentDocument",
+			"HTMLEmbedElement.prototype.contentDocument",
+		],
+		{
+			get(ctx) {
+				const contentwindow =
+					client.descriptors[
+						`${ctx.this.constructor.name}.prototype.contentWindow`
+					].get;
+				const realwin = contentwindow.apply(ctx.this);
+				if (!realwin) return realwin;
+
+				if (SCRAMJETCLIENT in realwin) {
+					return realwin[SCRAMJETCLIENT].documentProxy;
+				} else {
+					const newclient = new ScramjetClient(realwin);
+					newclient.hook();
+
+					return newclient.documentProxy;
+				}
+			},
+		}
+	);
 
 	client.Proxy(
 		[
@@ -320,7 +303,7 @@ export default function (client: ScramjetClient, self: typeof window) {
 			apply(ctx) {
 				const doc = ctx.call();
 				if (doc) {
-					ctx.return(ctx.this.contentDocument);
+					return ctx.return(ctx.this.contentDocument);
 				}
 			},
 		}
@@ -331,17 +314,27 @@ export default function (client: ScramjetClient, self: typeof window) {
 			return ctx.get();
 		},
 		set(ctx, value) {
-			if (value == client.documentProxy) {
+			if (value === client.documentProxy) {
 				return ctx.set(self.document);
 			}
 
 			return ctx.set(value);
 		},
 	});
 
+	client.Proxy("Document.prototype.open", {
+		apply(ctx) {
+			const doc = ctx.call() as Document;
+
+			const scram: ScramjetClient = doc[SCRAMJETCLIENT];
+			if (!scram) return ctx.return(doc); // ??
+
+			return ctx.return(scram.documentProxy);
+		},
+	});
+
 	client.Trap("Node.prototype.ownerDocument", {
 		get(ctx) {
-			// return client.documentProxy;
 			const doc = ctx.get() as Document | null;
 			if (!doc) return null;
 

diff --git a/src/client/helpers.ts b/src/client/helpers.ts
@@ -2,26 +2,4 @@ export function getOwnPropertyDescriptorHandler(target, prop) {
 	const realDescriptor = Reflect.getOwnPropertyDescriptor(target, prop);
 
 	return realDescriptor;
-	// const d: PropertyDescriptor = {};
-
-	// if (realDescriptor.enumerable !== undefined)
-	// 	d.enumerable = realDescriptor.enumerable;
-	// if (realDescriptor.configurable !== undefined)
-	// 	d.configurable = realDescriptor.configurable;
-	// if (realDescriptor.writable !== undefined)
-	// 	d.writable = realDescriptor.writable;
-
-	// if (realDescriptor.get) {
-	// 	d.get = () => this.get(target, prop);
-	// }
-
-	// if (realDescriptor.set) {
-	// 	d.set = (value) => this.set(target, prop, value);
-	// }
-
-	// if (realDescriptor.value) {
-	// 	d.value = this.get(target, prop);
-	// }
-
-	// return d;
 }
diff --git a/src/client/shared/function.ts b/src/client/shared/function.ts
@@ -11,11 +11,6 @@ function rewriteFunction(ctx: ProxyCtx, client: ScramjetClient) {
 export default function (client: ScramjetClient, _self: Self) {
 	const handler: Proxy = {
 		apply(ctx: ProxyCtx) {
-			if ((ctx.fn as any).alreadyProxied) {
-				debugger;
-				throw new Error("blah slop");
-			}
-			(ctx.fn as any).alreadyProxied = true;
 			rewriteFunction(ctx, client);
 		},
 		construct(ctx) {

diff --git a/src/worker/fetch.ts b/src/worker/fetch.ts
@@ -319,7 +319,19 @@ async function rewriteBody(
 		case "iframe":
 		case "document":
 			if (response.headers.get("content-type")?.startsWith("text/html")) {
-				return rewriteHtml(await response.text(), cookieStore, meta, true);
+				// note from percs: i think this has the potential to be slow asf, but for right now its fine (we should probably look for a better solution)
+				const buf = await response.arrayBuffer();
+				const decode = new TextDecoder("utf-8").decode(buf);
+				const charsetHeader = response.headers.get("content-type");
+				const charset =
+					charsetHeader?.split("charset=")[1] ||
+					decode.match(/charset=([^"]+)/)?.[1] ||
+					"utf-8";
+				const htmlContent = charset
+					? new TextDecoder(charset).decode(buf)
+					: decode;
+
+				return rewriteHtml(htmlContent, cookieStore, meta, true);
 			} else {
 				return response.body;
 			}