[scramjet/controller] fix controller

Author: velzie

packages/scramjet/packages/controller/src/index.ts

@@ -73,7 +73,7 @@ export class Controller {
 	frames: Frame[] = [];
 	cookieJar = new $scramjet.CookieJar();
 
-	private rpc: RpcHelper<Controllerbound, SWbound>;
+	rpc: RpcHelper<Controllerbound, SWbound>;
 	private ready: Promise<void>;
 	private readyResolve: () => void;
 
@@ -228,29 +228,58 @@ function yieldGetInjectScripts(
 				"data:text/javascript;base64," +
 					btoa(`
 					(()=>{
-						$scramjet.setWasm(Uint8Array.from(atob(self.WASM), (c) => c.charCodeAt(0)));
+						const { ScramjetClient, CookieJar, setWasm } = $scramjet;
+
+						setWasm(Uint8Array.from(atob(self.WASM), (c) => c.charCodeAt(0)));
 						delete self.WASM;
-						const cookieJar = new $scramjet.CookieJar();
+
+						const cookieJar = new CookieJar();
 						const config = ${JSON.stringify(config)};
 						const sjconfig = ${JSON.stringify(sjconfig)};
 					 	cookieJar.load(${cookieJar.dump()});
 
 						const prefix = new URL("${prefix.href}");
+						const sw = navigator.serviceWorker.controller;
 
-						$scramjet.loadAndHook({
-							context: {
-								interface: {
-									getInjectScripts: (${yieldGetInjectScripts.toString()})(cookieJar, config, sjconfig, prefix),
-									codecEncode: ${codecEncode.toString()},
-									codecDecode: ${codecDecode.toString()},
-								},
-								prefix,
-								cookieJar,
-								config: sjconfig
+						const context = {
+							interface: {
+								getInjectScripts: (${yieldGetInjectScripts.toString()})(cookieJar, config, sjconfig, prefix),
+								codecEncode: ${codecEncode.toString()},
+								codecDecode: ${codecDecode.toString()},
 							},
+							prefix,
+							cookieJar,
+							config: sjconfig
+						};
+						function createFrameId() {
+							return \`\${Array(8)
+								.fill(0)
+								.map(() => Math.floor(Math.random() * 36).toString(36))
+								.join("")}\`;
+						}
+
+						const frame = globalThis.frameElement;
+						if (frame && !frame.name) {
+							frame.name = createFrameId();
+						}
+
+						const client = new ScramjetClient(globalThis, {
+							context,
 							transport: null,
-						})
+							sendSetCookie: async (url, cookie) => {
+								// sw.postMessage({
+								// 	$controller$setCookie: {
+								// 		url,
+								// 		cookie
+								// 	}
+								// });
+							},
+							shouldPassthroughWebsocket: (url) => {
+								return url === "wss://anura.pro/";
+							}
+						});
 
+						client.hook();
 						document.currentScript.remove();
 					})();
 				`)

packages/scramjet/packages/controller/src/sw.ts

@@ -3,6 +3,25 @@ import { RpcHelper } from "@mercuryworkshop/rpc";
 import { Controllerbound, SWbound } from "./types";
 import type { BareHeaders } from "@mercuryworkshop/bare-mux-custom";
 
+function makeId(): string {
+	return Math.random().toString(36).substring(2, 10);
+}
+
+let cookieResolvers: Record<string, (value: void) => void> = {};
+addEventListener("message", (e) => {
+	if (!e.data) return;
+	if (typeof e.data != "object") return;
+	if (!e.data.$sw$setCookieDone) return;
+	if (typeof e.data.$sw$setCookieDone != "object") return;
+	const done = e.data.$sw$setCookieDone;
+
+	const resolver = cookieResolvers[done.id];
+	if (resolver) {
+		resolver();
+		delete cookieResolvers[done.id];
+	}
+});
+
 class Tab {
 	rpc: RpcHelper<SWbound, Controllerbound>;
 
@@ -11,9 +30,45 @@ class Tab {
 		public id: string,
 		port: MessagePort
 	) {
-		this.rpc = new RpcHelper({}, "tabchannel-" + id, (data, transfer) => {
-			port.postMessage(data, transfer);
-		});
+		this.rpc = new RpcHelper(
+			{
+				sendSetCookie: async ({ url, cookie }) => {
+					let clients = await self.clients.matchAll();
+					let promises = [];
+
+					for (const client of clients) {
+						let id = makeId();
+						client.postMessage({
+							$controller$setCookie: {
+								url,
+								cookie,
+								id,
+							},
+						});
+						promises.push(
+							new Promise<void>((resolve) => {
+								cookieResolvers[id] = resolve;
+							})
+						);
+					}
+					await Promise.race([
+						new Promise<void>((resolve) =>
+							setTimeout(() => {
+								console.error(
+									"timed out waiting for set cookie response (deadlock?)"
+								);
+								resolve();
+							}, 1000)
+						),
+						promises,
+					]);
+				},
+			},
+			"tabchannel-" + id,
+			(data, transfer) => {
+				port.postMessage(data, transfer);
+			}
+		);
 		port.addEventListener("message", (e) => {
 			this.rpc.recieve(e.data);
 		});

packages/scramjet/packages/controller/src/types.d.ts

@@ -29,6 +29,19 @@ export type TransferResponse = {
 export type Controllerbound = {
 	ready: [];
 	request: [TransferRequest, TransferResponse];
+	sendSetCookie: [
+		{
+			url: string;
+			cookie: string;
+		},
+	];
 };
 
-export type SWbound = {};
+export type SWbound = {
+	sendSetCookie: [
+		{
+			url: string;
+			cookie: string;
+		},
+	];
+};

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

@@ -10,7 +10,7 @@ export default function (client: ScramjetClient, self: Self) {
 	const handlers = {
 		message: {
 			_init() {
-				if (client.init.shouldBlockMessageEvent(this)) {
+				if (client.init.shouldBlockMessageEvent?.(this)) {
 					return false;
 				}
 

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

@@ -8,7 +8,7 @@ export default function (client: ScramjetClient, self: Self) {
 		"return import(url)"
 	);
 
-	Object.defineProperty(self, config.globals.importfn, {
+	Object.defineProperty(self, client.config.globals.importfn, {
 		value: function (base: string, url: string) {
 			const resolved = new URL(url, base).href;
 

packages/scramjet/packages/core/src/client/shared/requests/websocket.ts

@@ -30,7 +30,7 @@ export default function (client: ScramjetClient, self: typeof globalThis) {
 		new WeakMap();
 	client.Proxy("WebSocket", {
 		construct(ctx) {
-			if (client.init.shouldPassthroughWebsocket(ctx.args[0])) {
+			if (client.init.shouldPassthroughWebsocket?.(ctx.args[0])) {
 				return ctx.return(client.natives.construct("WebSocket", ...ctx.args));
 			}
 			const fakeWebSocket = new EventTarget() as WebSocket;