[core] add ScramjetHTMLPreRewrite/ScramjetHTMLPostRewrite events exposing domhandler

Author: velzie

src/shared/rewriters/html.ts

@@ -1,4 +1,4 @@
-import { ElementType, Parser } from "htmlparser2";
+import { ElementType, Handler, Parser } from "htmlparser2";
 import { ChildNode, DomHandler, Element, Comment } from "domhandler";
 import render from "dom-serializer";
 import { URLMeta, rewriteUrl } from "@rewriters/url";
@@ -36,13 +36,16 @@ function rewriteHtmlInner(
 	html: string,
 	cookieStore: CookieStore,
 	meta: URLMeta,
-	fromTop: boolean = false
+	fromTop: boolean = false,
+	preRewrite?: (handler: DomHandler) => void,
+	postRewrite?: (handler: DomHandler) => void
 ) {
 	const handler = new DomHandler((err, dom) => dom);
 	const parser = new Parser(handler);
 
 	parser.write(html);
 	parser.end();
+	if (preRewrite) preRewrite(handler);
 	traverseParsedHtml(handler.root, cookieStore, meta);
 
 	function findhead(node) {
@@ -69,6 +72,8 @@ function rewriteHtmlInner(
 		head.children.unshift(...getInjectScripts(cookieStore, script));
 	}
 
+	if (postRewrite) postRewrite(handler);
+
 	return render(handler.root, {
 		encodeEntities: "utf8",
 		decodeEntities: false,
@@ -79,10 +84,19 @@ export function rewriteHtml(
 	html: string,
 	cookieStore: CookieStore,
 	meta: URLMeta,
-	fromTop: boolean = false
+	fromTop: boolean = false,
+	preRewrite?: (handler: DomHandler) => void,
+	postRewrite?: (handler: DomHandler) => void
 ) {
 	const before = performance.now();
-	const ret = rewriteHtmlInner(html, cookieStore, meta, fromTop);
+	const ret = rewriteHtmlInner(
+		html,
+		cookieStore,
+		meta,
+		fromTop,
+		preRewrite,
+		postRewrite
+	);
 	dbg.time(meta, before, "html rewrite");
 
 	return ret;

src/worker/fetch.ts

@@ -29,6 +29,7 @@ import { rewriteCss } from "@rewriters/css";
 import { rewriteWorkers } from "@rewriters/worker";
 import { ScramjetDownload } from "@client/events";
 import { ScramjetConfig } from "@/types";
+import DomHandler from "domhandler";
 
 export interface ScramjetFetchContext {
 	rawUrl: URL;
@@ -150,7 +151,7 @@ export async function handleFetch(
 	// }
 
 	if (response.body && !isRedirect(response)) {
-		responseBody = await rewriteBody(context, parsed, response);
+		responseBody = await rewriteBody.call(this, context, parsed, response);
 	}
 
 	// Clean up tracker if not a redirect
@@ -729,7 +730,23 @@ async function rewriteBody(
 					await response.text(),
 					context.cookieStore,
 					parsed.meta,
-					true
+					true,
+					(handler) => {
+						const evt = new ScramjetHTMLPreRewriteEvent(
+							handler,
+							context,
+							parsed
+						);
+						this.dispatchEvent(evt);
+					},
+					(handler) => {
+						const evt = new ScramjetHTMLPostRewriteEvent(
+							handler,
+							context,
+							parsed
+						);
+						this.dispatchEvent(evt);
+					}
 				);
 			} else {
 				return response.body;
@@ -760,6 +777,26 @@ async function rewriteBody(
 
 type BodyType = string | ArrayBuffer | Blob | ReadableStream<any>;
 
+export class ScramjetHTMLPreRewriteEvent extends Event {
+	constructor(
+		public handler: DomHandler,
+		public context: ScramjetFetchContext,
+		public parsed: ScramjetFetchParsed
+	) {
+		super("htmlPreRewrite");
+	}
+}
+
+export class ScramjetHTMLPostRewriteEvent extends Event {
+	constructor(
+		public handler: DomHandler,
+		public context: ScramjetFetchContext,
+		public parsed: ScramjetFetchParsed
+	) {
+		super("htmlPostRewrite");
+	}
+}
+
 export class ScramjetResponseEvent extends Event {
 	_response?: ScramjetFetchResponse | Promise<ScramjetFetchResponse>;
 	constructor(