/** @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));
	});

/** @param {string} socketPath the admin socket; holding it is the authority to install */
const makeAdmin = (socketPath) => ({
	install: (vat, src, labels) =>
		new Promise((resolve, reject) => {
			const grant = typeof labels === "string" ? labels : JSON.stringify(labels);
			const path = labels ? `/${vat}?egress=${encodeURIComponent(grant)}` : `/${vat}`;
			const req = request({ socketPath, path, method: "PUT", agent: false }, (res) => {
				let body = "";
				res.on("data", (c) => (body += c));
				res.on("end", () => resolve(body.trim()));
			});
			req.on("error", reject);
			req.end(src);
		}),
});

/** @param {string} base the public origin */
const makePublic = (base) => ({
	get: async (path) => (await fetch(`${base}/${path}`)).text(),
	put: async (vat, src) => (await fetch(`${base}/${vat}`, { 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 });
});

test("a vat with no code is 404", async () => {
	assert.equal(await pub.get("nobody"), "no code installed\n");
});

test("install, run, upgrade: the app's state carries over", async () => {
	assert.match(await admin.install("alice", await app("counter-v1.js")), /^[0-9a-f]{64}$/);
	assert.equal(await pub.get("alice"), "v1 count 1\n");
	assert.equal(await pub.get("alice"), "v1 count 2\n");
	await admin.install("alice", await app("counter-v2.js"));
	assert.equal(await pub.get("alice"), "v2 count 3\n");
});

test("vats don't share state", async () => {
	await admin.install("carol", await app("counter-v1.js"));
	assert.equal(await pub.get("carol"), "v1 count 1\n");
});

test("an app gets no fetch", async () => {
	await admin.install("bob", await app("try-fetch.js"));
	assert.match(await pub.get("bob"), /^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("mallory", "export class App {}"), 404);
	assert.equal(await pub.get("mallory"), "no code installed\n");
});

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

const probe = (vat, target, { ws, constrain } = {}) =>
	pub.get(
		`${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.install("erin", await app("egress-probe.js"), 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.install("frank", await app("egress-probe.js"));
	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.install("heidi", await app("try-fetch.js"), echoLabels);
	assert.match(await pub.get("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.install("grace", await app("egress-probe.js"), { 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.install("ivan", await app("egress-probe.js"), { 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.install("kim", await app("egress-probe.js"), 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.install("judy", await app("egress-probe.js"), echoLabels);
		const before = echo.hits();
		const got = await (await fetch(`http://127.0.0.1:${port2}/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();
	}
});