/** @file toybox end to end: a real workerd, an admin socket, a public port, a switchboard with two handles. */
import { spawn } from "node:child_process";
import { mkdir, mkdtemp, readFile, rm } from "node:fs/promises";
import { createServer as createHttpServer, request } from "node:http";
import { connect, createServer } from "node:net";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { after, before, test } from "node:test";
import assert from "node:assert/strict";
import { WebSocketServer } from "ws";

const here = new URL("..", import.meta.url).pathname;
// the package binary, not the .bin shim: the shim forks it, so a kill would orphan workerd
const workerd = join(here, "node_modules/workerd/bin/workerd");
const app = (name) => readFile(join(here, "apps", name), "utf8");

const freePort = () =>
	new Promise((resolve) => {
		const s = createServer().listen(0, "127.0.0.1", () => {
			const { port } = s.address();
			s.close(() => resolve(port));
		});
	});

/** Resolves once something accepts connections at `opts` (a port or a socket path). */
const accepting = (opts) =>
	new Promise((resolve) => {
		const s = connect(opts, () => (s.end(), resolve(true)));
		s.on("error", () => resolve(false));
	});

/**
 * The admin socket: holding it is the authority to create and upgrade vats.
 *
 * @param {string} socketPath
 */
const makeAdmin = (socketPath) => {
	const send = (method, path, src, { egress, mount } = {}) =>
		new Promise((resolve, reject) => {
			const q = new URLSearchParams();
			if (egress) q.set("egress", typeof egress === "string" ? egress : JSON.stringify(egress));
			if (mount) q.set("mount", JSON.stringify(mount));
			const req = request({ socketPath, path: `${path}?${q}`, method, agent: false }, (res) => {
				let body = "";
				res.on("data", (c) => (body += c));
				res.on("end", () => resolve(res.statusCode === 200 ? JSON.parse(body) : { status: res.statusCode, body }));
			});
			req.on("error", reject);
			req.end(src);
		});
	return {
		/** `createVat`: a new vat, reachable only by its id here, and by its mount */
		create: (src, grants) => send("POST", "/", src, grants),
		upgrade: (vat, src, grants) => send("PUT", `/${vat}`, src, grants),
		/** a vat from before ids */
		byName: (name, src, grants) => send("PUT", `/by-name/${name}`, src, grants),
	};
};

/** @param {string} base the public origin */
const makePublic = (base) => ({
	get: async (pathAndQuery, headers = {}) => {
		const resp = await fetch(`${base}${pathAndQuery}`, { headers });
		return { status: resp.status, text: await resp.text(), referrerPolicy: resp.headers.get("Referrer-Policy") };
	},
	text: async (pathAndQuery) => (await fetch(`${base}${pathAndQuery}`)).text(),
	put: async (pathAndQuery, src) => (await fetch(`${base}${pathAndQuery}`, { method: "PUT", body: src })).status,
});

/**
 * Run workerd on `config` until `stop()`; resolves once `listening` accepts.
 *
 * @param {string} config
 * @param {string[]} args socket, directory, and external overrides
 * @param {{ port: number } | { path: string }} listening
 * @param {string[]} stale socket files a crash leaves behind; bind() would fail on them
 */
const startWorkerd = async (config, args, listening, stale = []) => {
	for (const f of stale) await rm(f, { force: true });
	const child = spawn(workerd, ["serve", "--experimental", join(here, config), ...args], {
		stdio: ["ignore", "ignore", "pipe"],
	});
	let log = "";
	child.stderr.on("data", (c) => (log += c));
	const exited = new Promise((resolve) => child.on("exit", resolve));
	for (let i = 0; i < 100; i++) {
		// SIGKILL: a crash is the harder durability test
		if (await accepting(listening)) return { stop: () => (child.kill("SIGKILL"), exited) };
		await new Promise((r) => setTimeout(r, 50));
	}
	child.kill("SIGKILL");
	throw new Error(`workerd did not start: ${log}`);
};

/** An HTTP server that answers GET with "plain ok", echoes WebSocket messages, and counts requests. */
const startEcho = async () => {
	let hits = 0;
	const server = createHttpServer((_req, res) => (hits++, res.end("plain ok")));
	const wss = new WebSocketServer({ server });
	wss.on("connection", (ws) => ws.on("message", (m) => ws.send(String(m))));
	await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve));
	return {
		host: `127.0.0.1:${server.address().port}`,
		hits: () => hits,
		stop: () => new Promise((resolve) => (wss.close(), server.closeAllConnections(), server.close(resolve))),
	};
};

let dir, port, server, switchboard, echo, echoLabels, admin, pub;

/**
 * A toybox in `dir/<name>` that holds one switchboard handle.
 *
 * @param {string} name
 * @param {number} port
 * @param {"public" | "local"} handle which switchboard socket it gets
 */
const startToybox = async (name, port, handle) => {
	const home = join(dir, name);
	await mkdir(join(home, "state"), { recursive: true });
	return startWorkerd(
		"config.capnp",
		[
			"-s", `http=127.0.0.1:${port}`, "-s", `admin=unix:${join(home, "admin.sock")}`,
			"-d", `state=${join(home, "state")}`, "-e", `switchboard=unix:${join(dir, `sb-${handle}.sock`)}`,
		],
		{ port, host: "127.0.0.1" },
		[join(home, "admin.sock")],
	);
};

before(async () => {
	dir = await mkdtemp(join(tmpdir(), "toybox-"));
	port = await freePort();
	echo = await startEcho();
	echoLabels = { host: "127.0.0.1", port: echo.host.split(":")[1] };
	switchboard = await startWorkerd(
		"switchboard.capnp",
		["-s", `public=unix:${join(dir, "sb-public.sock")}`, "-s", `local=unix:${join(dir, "sb-local.sock")}`],
		{ path: join(dir, "sb-local.sock") },
		[join(dir, "sb-public.sock"), join(dir, "sb-local.sock")],
	);
	// the tests' echo server is on localhost, so this toybox holds the local handle
	server = await startToybox("main", port, "local");
	admin = makeAdmin(join(dir, "main", "admin.sock"));
	pub = makePublic(`http://127.0.0.1:${port}`);
});

after(async () => {
	await server?.stop();
	await switchboard?.stop();
	await echo?.stop();
	await rm(dir, { recursive: true, force: true });
});

const at = (name) => ({ mount: { vat: name } });

test("nothing mounted: 404", async () => {
	assert.equal((await pub.get("/?vat=nobody")).status, 404);
});

test("createVat returns an id and a code hash; unmounted, the vat is unreachable", async () => {
	const { vat, hash } = await admin.create(await app("counter-v1.js"));
	assert.match(vat, /^[0-9a-f]{64}$/);
	assert.match(hash, /^[0-9a-f]{64}$/);
	assert.equal((await pub.get("/")).status, 404);
});

test("install, run, upgrade: the app's state carries over; every start states its grants", async () => {
	const { vat } = await admin.create(await app("counter-v1.js"), at("alice"));
	assert.equal(await pub.text("/?vat=alice"), "v1 count 1\n");
	assert.equal(await pub.text("/?vat=alice"), "v1 count 2\n");
	await admin.upgrade(vat, await app("counter-v2.js"), at("alice"));
	assert.equal(await pub.text("/?vat=alice"), "v2 count 3\n");
	await admin.upgrade(vat, await app("counter-v2.js"));
	assert.equal((await pub.get("/?vat=alice")).status, 404);
	await admin.upgrade(vat, await app("counter-v2.js"), at("alice"));
	assert.equal(await pub.text("/?vat=alice"), "v2 count 4\n");
});

test("vats don't share state", async () => {
	await admin.create(await app("counter-v1.js"), at("bob"));
	await admin.create(await app("counter-v1.js"), at("carol"));
	assert.equal(await pub.text("/?vat=bob"), "v1 count 1\n");
	assert.equal(await pub.text("/?vat=carol"), "v1 count 1\n");
});

test("an app gets no fetch", async () => {
	await admin.create(await app("try-fetch.js"), at("bob-fetch"));
	assert.match(await pub.text("/?vat=bob-fetch"), /^fetch refused: This worker is not permitted to access the internet/);
});

test("the public port can't install code", async () => {
	assert.equal(await pub.put("/", "export class App {}"), 404);
});

test("the path is the app's; its labels are the mount's and the request's together", async () => {
	await admin.create(await app("show-labels.js"), at("lab"));
	const got = await pub.get("/a/b?vat=lab&page=2");
	assert.equal(got.text, `/a/b ${JSON.stringify({ vat: "lab", page: "2" })}\n`);
	assert.equal(got.referrerPolicy, "no-referrer");
});

test("a label given twice is a 400", async () => {
	assert.equal((await pub.get("/?vat=lab&vat=lab")).status, 400);
});

test("mounts that could both match a request are refused", async () => {
	const src = await app("show-labels.js");
	await admin.create(src, at("x"));
	const narrower = await admin.create(src, { mount: { vat: "x", n: "2" } });
	assert.equal(narrower.status, 409);
	assert.match(narrower.body, /overlaps another/);
	// no key in common with the {vat} mounts: ?vat=x&app=x would carry both
	assert.equal((await admin.create(src, { mount: { app: "x" } })).status, 409);
	assert.match((await admin.create(src, { mount: { vat: "y", n: "2" } })).vat, /^[0-9a-f]{64}$/);
});

test("a secret label: in a header, not the URL; omitting it doesn't match", async () => {
	const secret = "k".repeat(43);
	await admin.create(await app("show-labels.js"), { mount: { vat: "safe", key: secret } });
	assert.equal((await pub.get("/?vat=safe")).status, 404);
	assert.equal((await pub.get("/?vat=safe&key=wrong")).status, 404);
	const got = await pub.get("/?vat=safe", { "Toybox-Labels": JSON.stringify({ key: secret }) });
	assert.equal(got.status, 200);
});

test("code, state, and mount survive a restart", async () => {
	await admin.create(await app("counter-v1.js"), at("dave"));
	assert.equal(await pub.text("/?vat=dave"), "v1 count 1\n");
	await server.stop();
	server = await startToybox("main", port, "local");
	assert.equal(await pub.text("/?vat=dave"), "v1 count 2\n");
});

test("a vat from before ids is unreachable until the admin mounts it by its old name", async () => {
	await admin.byName("old", await app("counter-v1.js"));
	assert.equal((await pub.get("/?vat=old")).status, 404);
	await admin.byName("old", await app("counter-v1.js"), at("old"));
	assert.equal(await pub.text("/?vat=old"), "v1 count 1\n");
});

const probe = (vat, target, { ws, constrain } = {}) =>
	pub.text(
		`/?vat=${vat}&url=${encodeURIComponent(`http://${target}/`)}` +
			(ws ? "&ws=1" : "") +
			(constrain ? `&constrain=${encodeURIComponent(JSON.stringify(constrain))}` : ""),
	);

test("a vat granted a host reaches it through the switchboard: fetch and WebSocket", async () => {
	await admin.create(await app("egress-probe.js"), { ...at("erin"), egress: echoLabels });
	assert.equal(await probe("erin", echo.host), "got: plain ok\n");
	assert.equal(await probe("erin", echo.host, { ws: true }), "echo: hello\n");
});

test("same code, no grant: no egress", async () => {
	await admin.create(await app("egress-probe.js"), at("frank"));
	assert.equal(await probe("frank", echo.host), "no egress granted\n");
});

test("a grant is an object, not ambient: the global fetch still throws", async () => {
	await admin.create(await app("try-fetch.js"), { ...at("heidi"), egress: echoLabels });
	assert.match(await pub.text("/?vat=heidi"), /^fetch refused: This worker is not permitted to access the internet/);
});

test("the switchboard refuses what a grant doesn't cover", async () => {
	await admin.create(await app("egress-probe.js"), { ...at("grace"), egress: { host: "127.0.0.1", port: "1" } });
	assert.match(await probe("grace", echo.host), /^got: egress refused: .*egress to 127\.0\.0\.1 not granted/);
	assert.match(await probe("grace", echo.host, { ws: true }), /^no websocket: 403 egress refused: .*not granted/);
});

test("a vat can narrow its grant, but not contradict it", async () => {
	await admin.create(await app("egress-probe.js"), { ...at("ivan"), egress: { scheme: "http" } });
	assert.equal(await probe("ivan", echo.host, { constrain: echoLabels }), "got: plain ok\n");
	assert.match(
		await probe("ivan", echo.host, { constrain: { scheme: "https" } }),
		/^refused: label scheme is already set to a different value/,
	);
});

test("a grant stored before labels, a bare host:port, still works", async () => {
	await admin.create(await app("egress-probe.js"), { ...at("kim"), egress: echo.host });
	assert.equal(await probe("kim", echo.host), "got: plain ok\n");
});

test("a grant survives a restart", async () => {
	await server.stop();
	server = await startToybox("main", port, "local");
	assert.equal(await probe("erin", echo.host, { ws: true }), "echo: hello\n");
});

test("a toybox holding only the public handle can't reach localhost, whatever the grant", async () => {
	const port2 = await freePort();
	const other = await startToybox("public-only", port2, "public");
	try {
		const admin2 = makeAdmin(join(dir, "public-only", "admin.sock"));
		await admin2.create(await app("egress-probe.js"), { ...at("judy"), egress: echoLabels });
		const before = echo.hits();
		const got = await (
			await fetch(`http://127.0.0.1:${port2}/?vat=judy&url=${encodeURIComponent(`http://${echo.host}/`)}`)
		).text();
		// workerd reports the public network's refusal (restrictPeers) to JS only as "internal error"
		assert.match(got, /^got: egress refused: /);
		assert.equal(echo.hits(), before);
	} finally {
		await other.stop();
	}
});