pi-ext

import { readFile } from "node:fs/promises";
import { runInNewContext } from "node:vm";
import { describe, expect, it, vi } from "vitest";
import { validateDiagram } from "../diagram";

const valid = "flowchart TD\n  A[Input] --> B[Output]";

describe("validateDiagram", () => {
	it("accepts supported diagram types", () => {
		for (const source of [
			valid,
			"sequenceDiagram\n  A->>B: request",
			"stateDiagram-v2\n  [*] --> Active",
			"flowchart LR\n A <--> B",
			"sequenceDiagram\n A->>B: request\n B-->>A: response",
			"classDiagram\n  class Animal {\n    +name: string\n  }",
			"erDiagram\n  USER ||--o{ ORDER : places",
		]) {
			expect(validateDiagram(source)).toBe(source);
		}
	});

	it("rejects executable, linked, styled, malformed, or oversized input", () => {
		for (const source of [
			"%%{init: {}}%%\nflowchart TD\n A-->B",
			'flowchart TD\n click A "https://example.test"',
			"flowchart TD\n style A fill:#fff",
			"flowchart TD\n A[<img src=x>] --> B",
			"graph TD\n A-->B",
			"---\ntitle: injected\n---\nflowchart TD\n A-->B",
			"flowchart TD; click A callback",
			"flowchart TD\n A-->B; click A callback",
			"flowchart TD\n A-->B\n linkStyle 0 stroke:red",
			"flowchart TD\n A:::unsafe --> B",
			"flowchart TD\n A@{ img: 'remote' }",
			"flowchart TD\n A[&lt;script&gt;]",
			"classDiagram\n cssClass Animal unsafe",
			"flowchart TD\n A-->B %%{init: {}}%%",
			`flowchart TD\n${"A-->B\n".repeat(3000)}`,
		]) {
			expect(() => validateDiagram(source)).toThrow();
		}
	});
});

async function child() {
	const script = await readFile(
		new URL("../assets/diagram.js", import.meta.url),
		"utf8",
	);
	const parent = { postMessage: vi.fn() };
	const root = { textContent: "", innerHTML: "" };
	const svg = {
		localName: "svg",
		namespaceURI: "http://www.w3.org/2000/svg",
		attributes: [{ name: "marker-end", value: "url(#arrow)" }],
		querySelectorAll: () => [],
		outerHTML: '<svg marker-end="url(#arrow)"></svg>',
	};
	const mermaid = {
		initialize: vi.fn(),
		render: vi.fn(async () => ({ svg: "<svg/>" })),
	};
	let listener: (event: {
		source: unknown;
		origin: string;
		data: unknown;
	}) => void = () => {};
	runInNewContext(
		`const validateDiagram = ${validateDiagram.toString()};\n${script}`,
		{
			document: {
				documentElement: {
					dataset: { parentOrigin: "http://127.0.0.1:1234" },
					scrollHeight: 9000,
				},
				querySelector: () => root,
			},
			window: {
				parent,
				addEventListener: (_: string, fn: typeof listener) => {
					listener = fn;
				},
			},
			mermaid,
			DOMParser: class {
				parseFromString() {
					return { documentElement: svg, querySelector: () => null };
				}
			},
		},
	);
	const send = (
		source = valid,
		origin = "http://127.0.0.1:1234",
		sender: unknown = parent,
	) =>
		listener({
			source: sender,
			origin,
			data: { type: "strata.diagram.render", source, dark: true },
		});
	return { parent, root, svg, mermaid, send };
}

describe("opaque diagram child protocol", () => {
	it("accepts only its parent and origin, retaining internal SVG markers", async () => {
		const c = await child();
		expect(c.parent.postMessage).toHaveBeenCalledWith(
			{ type: "strata.diagram.ready" },
			"http://127.0.0.1:1234",
		);
		c.send(valid, "null");
		c.send(valid, "http://127.0.0.1:1234", {});
		expect(c.mermaid.render).not.toHaveBeenCalled();
		c.send();
		await vi.waitFor(() => expect(c.root.innerHTML).toContain("url(#arrow)"));
		expect(c.mermaid.initialize).toHaveBeenCalledWith(
			expect.objectContaining({
				securityLevel: "strict",
				htmlLabels: false,
				theme: "dark",
				maxEdges: 500,
			}),
		);
		expect(c.parent.postMessage).toHaveBeenCalledWith(
			{ type: "strata.diagram.size", height: 4000 },
			"http://127.0.0.1:1234",
		);
	});

	it("blocks external SVG references and root event handlers", async () => {
		for (const attribute of [
			{ name: "onload", value: "alert(1)" },
			{ name: "marker-end", value: "url(https://example.test/arrow)" },
			{ name: "href", value: "https://example.test/" },
		]) {
			const c = await child();
			c.svg.attributes = [attribute];
			c.send();
			await vi.waitFor(() =>
				expect(c.root.textContent).toBe("Diagram unavailable."),
			);
			expect(c.root.innerHTML).toBe("");
		}
	});

	it("rejects unsafe source before rendering and discards stale completions", async () => {
		const c = await child();
		let finish: ((value: { svg: string }) => void) | undefined;
		c.mermaid.render.mockImplementationOnce(
			() =>
				new Promise((resolve) => {
					finish = resolve;
				}),
		);
		c.send();
		c.send("flowchart TD\n A-->B; click A callback");
		expect(c.root.textContent).toBe("Diagram unavailable.");
		finish?.({ svg: "<svg/>" });
		await Promise.resolve();
		await Promise.resolve();
		expect(c.root.innerHTML).toBe("");
		expect(c.mermaid.render).toHaveBeenCalledTimes(1);
	});
});