[scramjet/core] create plugin api, add fetch and rewriter.html hooks

Author: velzie

packages/scramjet/packages/core/src/Tap.ts

@@ -0,0 +1,94 @@
+type Description = {
+	context?: object;
+	props?: object;
+};
+
+type Callback<T extends Description> = (
+	context: T["context"],
+	props: T["props"]
+) => void | Promise<void>;
+
+type Sorter = (other: Plugin) => number;
+
+type CallbackInfo<T extends Description> = {
+	callback: Callback<T>;
+	plugin: Plugin;
+	sorter: Sorter;
+};
+
+type InternalHookDescription = {
+	tap: TapInternal;
+	key: string;
+};
+
+type TapInternal = {
+	callbacks: Record<string, CallbackInfo<Description>[]>;
+};
+
+export type TapInstance<T extends Record<string, Description>> = {
+	[K in keyof T]: T[K] & InternalHookDescription;
+};
+
+export class Plugin {
+	constructor(public name: string) {}
+
+	tap<T extends Description>(
+		hook: T,
+		callback: Callback<T>,
+		sorter?: Sorter
+	): void {
+		sorter ??= () => 0;
+		Tap.tap(hook, callback, this, sorter);
+	}
+}
+
+export class Tap {
+	static dispatch<T extends Description>(
+		hook: T,
+		context: T["context"],
+		props: T["props"]
+	): Promise<void[]> {
+		let internal = hook as unknown as InternalHookDescription;
+		let callbacks = internal.tap.callbacks[internal.key];
+		if (!callbacks || callbacks.length === 0) return;
+
+		callbacks = [...callbacks];
+		callbacks.sort((a, b) => a.sorter(b.plugin));
+
+		const results = callbacks.map((cb) => cb.callback(context, props));
+		return Promise.all(results);
+	}
+
+	static tap<T extends Description>(
+		hook: T,
+		callback: Callback<T>,
+		plugin: Plugin,
+		sorter: Sorter
+	) {
+		let internal = hook as unknown as InternalHookDescription;
+		let callbacks = internal.tap.callbacks;
+		if (!callbacks[internal.key]) callbacks[internal.key] = [];
+		callbacks[internal.key]!.push({
+			callback,
+			plugin,
+			sorter,
+		});
+	}
+
+	static create<T extends Record<string, Description>>(): TapInstance<T> {
+		const internal: TapInternal = {
+			callbacks: {},
+		};
+		const hooks: Record<string, InternalHookDescription> = {};
+
+		return new Proxy(internal as unknown as TapInstance<T>, {
+			get(target, key: string) {
+				if (key === "callbacks") return internal.callbacks;
+				if (!hooks[key]) {
+					hooks[key] = { tap: internal, key };
+				}
+				return hooks[key];
+			},
+		});
+	}
+}

packages/scramjet/packages/core/src/client/client.ts

@@ -9,7 +9,12 @@ import { createLocationProxy } from "@client/location";
 import { createWrapFn } from "@client/shared/wrap";
 import { NavigateEvent } from "@client/events";
 import { rewriteUrl, unrewriteUrl, type URLMeta } from "@rewriters/url";
-import { flagEnabled, ScramjetContext, ScramjetInterface } from "@/shared";
+import {
+	flagEnabled,
+	HtmlRewriterTap,
+	ScramjetContext,
+	ScramjetInterface,
+} from "@/shared";
 import { CookieJar } from "@/shared/cookie";
 import { iswindow } from "./entry";
 import { SingletonBox } from "./singletonbox";
@@ -135,6 +140,12 @@ export class ScramjetClient {
 
 	context: ScramjetContext;
 
+	hooks = {
+		rewriter: {
+			html: HtmlRewriterTap,
+		},
+	};
+
 	constructor(
 		public global: typeof globalThis,
 		public init: ScramjetClientInit
@@ -160,6 +171,10 @@ export class ScramjetClient {
 		this.box.registerClient(this, global as Self);
 
 		this.context = init.context;
+		this.context.hooks = {
+			rewriter: this.hooks.rewriter,
+		};
+
 		this.bare = new BareCompatibleClient(init.transport);
 
 		this.serviceWorker = this.global.navigator.serviceWorker;

packages/scramjet/packages/core/src/fetch/index.ts

@@ -15,12 +15,13 @@ import {
 } from "@rewriters/url";
 import { rewriteJs } from "@rewriters/js";
 import { ScramjetHeaders } from "@/shared/headers";
-import { flagEnabled, ScramjetContext } from "@/shared";
+import { flagEnabled, HtmlRewriterHooks, ScramjetContext } from "@/shared";
 import { rewriteHtml } from "@rewriters/html";
 import { rewriteCss } from "@rewriters/css";
 import { rewriteWorkers } from "@rewriters/worker";
 import { ScramjetConfig } from "@/types";
 import DomHandler from "domhandler";
+import { Tap, TapInstance } from "@/Tap";
 
 export interface ScramjetFetchRequest {
 	rawUrl: URL;
@@ -66,6 +67,13 @@ export class ScramjetFetchHandler extends EventTarget {
 	public crossOriginIsolated: boolean = false;
 	public context: ScramjetContext;
 
+	public hooks: {
+		rewriter: {
+			html: TapInstance<HtmlRewriterHooks>;
+		};
+		fetch: TapInstance<FetchHooks>;
+	};
+
 	public fetchDataUrl: (dataUrl: string) => Promise<Response>;
 	public fetchBlobUrl: (blobUrl: string) => Promise<Response>;
 	public sendSetCookie: (url: URL, cookie: string) => Promise<void>;
@@ -78,6 +86,15 @@ export class ScramjetFetchHandler extends EventTarget {
 		this.sendSetCookie = init.sendSetCookie;
 		this.fetchDataUrl = init.fetchDataUrl;
 		this.fetchBlobUrl = init.fetchBlobUrl;
+		this.hooks = {
+			rewriter: {
+				html: Tap.create<HtmlRewriterHooks>(),
+			},
+			fetch: Tap.create<FetchHooks>(),
+		};
+		this.context.hooks = {
+			rewriter: this.hooks.rewriter,
+		};
 	}
 
 	async handleFetch(
@@ -113,22 +130,20 @@ async function doHandleFetch(
 		redirect: "manual",
 	} as BareRequestInit;
 
-	const req = new ScramjetRequestEvent(
+	let reqcontext: typeof handler.hooks.fetch.request.context = {
+		client: handler.client,
 		request,
-		parsed.url,
 		parsed,
+	};
+	let reqprops: typeof handler.hooks.fetch.request.props = {
 		init,
-		handler.client
-	);
-	handler.dispatchEvent(req);
-
+		url: parsed.url,
+	};
+	await Tap.dispatch(handler.hooks.fetch.request, reqcontext, reqprops);
 	let response: BareResponse;
 
-	if (req._response) {
-		let resp = req._response;
-		if ("then" in resp) {
-			resp = await resp;
-		}
+	if (reqprops.earlyResponse) {
+		let resp = reqprops.earlyResponse;
 		if ("rawHeaders" in resp) {
 			// it's a bare response
 			response = resp;
@@ -137,7 +152,7 @@ async function doHandleFetch(
 			response = BareResponse.fromNativeResponse(resp);
 		}
 	} else {
-		response = await handler.client.fetch(req.url, req.init);
+		response = await handler.client.fetch(reqprops.url, reqprops.init);
 	}
 
 	let responseBody: BodyType;
@@ -191,18 +206,22 @@ async function doHandleFetch(
 	// await cleanTracker(parsed.url.toString());
 	// }
 
-	const resp = new ScramjetResponseEvent(request, parsed, {
-		body: responseBody,
-		headers: responseHeaders,
-		status: response.status,
-		statusText: response.statusText,
-	});
-	handler.dispatchEvent(resp);
+	let respcontext: typeof handler.hooks.fetch.response.context = {
+		request,
+		parsed,
+	};
+	let respprops: typeof handler.hooks.fetch.response.props = {
+		response: {
+			body: responseBody,
+			headers: responseHeaders,
+			status: response.status,
+			statusText: response.statusText,
+		},
+	};
 
-	let r = resp.response;
-	if (resp._response) r = await resp._response;
+	await Tap.dispatch(handler.hooks.fetch.response, respcontext, respprops);
 
-	return r;
+	return respprops.response;
 }
 
 function isRedirect(response: BareResponse) {
@@ -682,23 +701,7 @@ async function rewriteBody(
 					await response.text(),
 					handler.context,
 					parsed.meta,
-					true,
-					(domhandler) => {
-						const evt = new ScramjetHTMLPreRewriteEvent(
-							domhandler,
-							request,
-							parsed
-						);
-						handler.dispatchEvent(evt);
-					},
-					(domhandler) => {
-						const evt = new ScramjetHTMLPostRewriteEvent(
-							domhandler,
-							request,
-							parsed
-						);
-						handler.dispatchEvent(evt);
-					}
+					true
 				);
 			} else {
 				return response.body;
@@ -729,60 +732,28 @@ async function rewriteBody(
 	}
 }
 
-type BodyType = string | ArrayBuffer | Blob | ReadableStream<any>;
-
-export class ScramjetHTMLPreRewriteEvent extends Event {
-	constructor(
-		public handler: DomHandler,
-		public context: ScramjetFetchRequest,
-		public parsed: ScramjetFetchParsed
-	) {
-		super("htmlPreRewrite");
-	}
-}
-
-export class ScramjetHTMLPostRewriteEvent extends Event {
-	constructor(
-		public handler: DomHandler,
-		public context: ScramjetFetchRequest,
-		public parsed: ScramjetFetchParsed
-	) {
-		super("htmlPostRewrite");
-	}
-}
-
-export class ScramjetResponseEvent extends Event {
-	_response?: ScramjetFetchResponse | Promise<ScramjetFetchResponse>;
-	constructor(
-		public context: ScramjetFetchRequest,
-		public parsed: ScramjetFetchParsed,
-		public response: ScramjetFetchResponse
-	) {
-		super("handleResponse");
-	}
-	respondWith(
-		response: ScramjetFetchResponse | Promise<ScramjetFetchResponse>
-	) {
-		this._response = response;
-	}
-}
+export type FetchHooks = {
+	request: {
+		context: {
+			request: ScramjetFetchRequest;
+			parsed: ScramjetFetchParsed;
+			client: BareCompatibleClient;
+		};
+		props: {
+			init: BareRequestInit;
+			url: URL;
+			earlyResponse?: BareResponse;
+		};
+	};
+	response: {
+		context: {
+			request: ScramjetFetchRequest;
+			parsed: ScramjetFetchParsed;
+		};
+		props: {
+			response: ScramjetFetchResponse;
+		};
+	};
+};
 
-export class ScramjetRequestEvent extends Event {
-	_response?:
-		| BareResponse
-		| Promise<BareResponse>
-		| Response
-		| Promise<Response>;
-	constructor(
-		public context: ScramjetFetchRequest,
-		public url: URL,
-		public parsed: ScramjetFetchParsed,
-		public init: BareRequestInit,
-		public client: BareCompatibleClient
-	) {
-		super("request");
-	}
-	respondWith(response: BareResponse | Promise<BareResponse>) {
-		this._response = response;
-	}
-}
+type BodyType = string | ArrayBuffer | Blob | ReadableStream<any>;

packages/scramjet/packages/core/src/index.ts

@@ -10,6 +10,7 @@ export * from "./shared";
 export * from "./symbols";
 export * from "./types";
 export * from "./fetch";
+export * from "./Tap";
 
 declare const REWRITERWASM: string | undefined;