fix: close the newest capture when memory nears the container limit

README.md

@@ -47,6 +47,7 @@ curl -X POST http://localhost:3000/v1/screenshots \
 | --- | --- | --- |
 | `PORT` | `3000` | Port the API listens on |
 | `ALLOWED_HOSTS` | _(empty)_ | Comma-separated hostnames the browser may load even though they resolve to a private address, e.g. `appwrite` |
+| `MEMORY_BUDGET` | `0.85` | Fraction of the container memory limit captures may use. Past it the newest screenshot or report is closed and fails with `507`, instead of the kernel killing Chromium and every capture in it. No effect without a cgroup v2 memory limit |
 
 Screenshots and reports refuse any URL, frame or subresource whose host is not publicly routable, which includes private, loopback and link-local ranges such as the `169.254.169.254` metadata address. A self-hosted setup that captures its own internal services lists them in `ALLOWED_HOSTS`. Matching is exact and case-insensitive. Every other private host stays blocked, including on redirect hops and script navigations; see [Network Access](#network-access).
 

src/routes/reports.ts

@@ -8,11 +8,17 @@ import {
 	startEgressGuard,
 	watchNavigations,
 } from "../utils/egress.js";
+import {
+	type Capture,
+	MEMORY_BUDGET_EXCEEDED,
+	memoryGuard,
+} from "../utils/memory.js";
 import { assertAllowedUrl } from "../utils/ssrf.js";
 
 export async function handleReportsRequest(req: Request): Promise<Response> {
 	let context: BrowserContext | undefined;
 	let guard: EgressGuard | undefined;
+	let capture: Capture | undefined;
 
 	try {
 		const json = await req.json();
@@ -38,6 +44,7 @@ export async function handleReportsRequest(req: Request): Promise<Response> {
 
 		const browser = await getBrowser();
 		context = await browser.newContext(contextOptions);
+		capture = memoryGuard.track(context);
 
 		// Grant permissions if specified
 		if (body.permissions && body.permissions.length > 0) {
@@ -107,12 +114,20 @@ export async function handleReportsRequest(req: Request): Promise<Response> {
 			headers: { "Content-Type": "application/json" },
 		});
 	} catch (error) {
+		if (capture?.exceeded) {
+			return new Response(JSON.stringify({ error: MEMORY_BUDGET_EXCEEDED }), {
+				status: 507,
+				headers: { "Content-Type": "application/json" },
+			});
+		}
+
 		const message = error instanceof Error ? error.message : "Unknown error";
 		return new Response(JSON.stringify({ error: message }), {
 			status: error instanceof BlockedNavigationError ? 403 : 400,
 			headers: { "Content-Type": "application/json" },
 		});
 	} finally {
+		capture?.release();
 		await context?.close();
 		guard?.close();
 	}

src/routes/screenshots.ts

@@ -12,13 +12,19 @@ import {
 	startEgressGuard,
 	watchNavigations,
 } from "../utils/egress.js";
+import {
+	type Capture,
+	MEMORY_BUDGET_EXCEEDED,
+	memoryGuard,
+} from "../utils/memory.js";
 import { assertAllowedUrl } from "../utils/ssrf.js";
 
 export async function handleScreenshotsRequest(
 	req: Request,
 ): Promise<Response> {
 	let context: BrowserContext | undefined;
 	let guard: EgressGuard | undefined;
+	let capture: Capture | undefined;
 
 	try {
 		const json = await req.json();
@@ -49,6 +55,7 @@ export async function handleScreenshotsRequest(
 
 		const browser = await getBrowser();
 		context = await browser.newContext(contextOptions);
+		capture = memoryGuard.track(context);
 
 		// Grant permissions if specified
 		if (body.permissions && body.permissions.length > 0) {
@@ -103,12 +110,20 @@ export async function handleScreenshotsRequest(
 			},
 		});
 	} catch (error) {
+		if (capture?.exceeded) {
+			return new Response(JSON.stringify({ error: MEMORY_BUDGET_EXCEEDED }), {
+				status: 507,
+				headers: { "Content-Type": "application/json" },
+			});
+		}
+
 		const message = error instanceof Error ? error.message : "Unknown error";
 		return new Response(JSON.stringify({ error: message }), {
 			status: error instanceof BlockedNavigationError ? 403 : 400,
 			headers: { "Content-Type": "application/json" },
 		});
 	} finally {
+		capture?.release();
 		await context?.close();
 		guard?.close();
 	}

src/utils/memory.ts

@@ -0,0 +1,148 @@
+import { readFileSync } from "node:fs";
+
+/**
+ * Chromium shares one GPU process between pages, so no page's memory can be
+ * measured on its own. The guard watches the container's working set instead,
+ * the figure the kernel OOM-kills on, and closes the newest capture whenever
+ * it crosses the budget. One capture then fails instead of Chromium and every
+ * capture running in it.
+ */
+export interface Capture {
+	/** Whether the guard closed this capture to stay within the budget. */
+	readonly exceeded: boolean;
+	release(): void;
+}
+
+interface Closable {
+	close(): Promise<void>;
+}
+
+interface Tracked {
+	readonly target: Closable;
+	exceeded: boolean;
+}
+
+export interface MemoryUsage {
+	readonly used: number;
+	readonly limit: number;
+}
+
+export const MEMORY_BUDGET_EXCEEDED =
+	"Page exceeded the memory budget and was closed.";
+
+const CGROUP_PATH = "/sys/fs/cgroup";
+const DEFAULT_BUDGET = 0.85;
+const POLL_INTERVAL_MS = 100;
+const SETTLE_MS = 1_000;
+const MEBIBYTE = 1024 * 1024;
+
+export function readCgroupMemory(): MemoryUsage | null {
+	try {
+		const limit = readFileSync(`${CGROUP_PATH}/memory.max`, "utf8").trim();
+		if (limit === "max") {
+			return null;
+		}
+
+		const current = Number(
+			readFileSync(`${CGROUP_PATH}/memory.current`, "utf8"),
+		);
+		const stat = readFileSync(`${CGROUP_PATH}/memory.stat`, "utf8");
+		const inactiveFile = Number(/^inactive_file (\d+)$/m.exec(stat)?.[1] ?? 0);
+
+		return { used: current - inactiveFile, limit: Number(limit) };
+	} catch {
+		return null;
+	}
+}
+
+export function parseMemoryBudget(value: string | undefined): number {
+	if (value === undefined || value === "") {
+		return DEFAULT_BUDGET;
+	}
+
+	const budget = Number(value);
+	if (!(budget > 0 && budget <= 1)) {
+		throw new Error(
+			`MEMORY_BUDGET must be a fraction of the memory limit between 0 and 1, got "${value}".`,
+		);
+	}
+
+	return budget;
+}
+
+export class MemoryGuard {
+	private readonly tracked: Tracked[] = [];
+	private readonly enabled: boolean;
+	private timer: ReturnType<typeof setInterval> | null = null;
+	private settledAt = 0;
+
+	constructor(
+		private readonly read: () => MemoryUsage | null,
+		private readonly budget: number,
+	) {
+		this.enabled = read() !== null;
+	}
+
+	track(target: Closable): Capture {
+		const entry: Tracked = { target, exceeded: false };
+		this.tracked.push(entry);
+
+		if (this.enabled && this.timer === null) {
+			this.timer = setInterval(() => this.check(), POLL_INTERVAL_MS);
+		}
+
+		return {
+			get exceeded() {
+				return entry.exceeded;
+			},
+			release: () => {
+				const index = this.tracked.indexOf(entry);
+				if (index !== -1) {
+					this.tracked.splice(index, 1);
+				}
+
+				if (this.tracked.length === 0 && this.timer !== null) {
+					clearInterval(this.timer);
+					this.timer = null;
+				}
+			},
+		};
+	}
+
+	private check(): void {
+		if (Date.now() < this.settledAt) {
+			return;
+		}
+
+		const usage = this.read();
+		if (usage === null || usage.used < usage.limit * this.budget) {
+			return;
+		}
+
+		let newest: Tracked | undefined;
+		for (let index = this.tracked.length - 1; index >= 0; index--) {
+			if (!this.tracked[index].exceeded) {
+				newest = this.tracked[index];
+				break;
+			}
+		}
+
+		if (newest === undefined) {
+			return;
+		}
+
+		newest.exceeded = true;
+		this.settledAt = Date.now() + SETTLE_MS;
+
+		console.log(
+			`Memory budget exceeded: ${Math.round(usage.used / MEBIBYTE)} of ${Math.round(usage.limit / MEBIBYTE)} MiB, closing the newest capture`,
+		);
+
+		newest.target.close().catch(() => {});
+	}
+}
+
+export const memoryGuard = new MemoryGuard(
+	readCgroupMemory,
+	parseMemoryBudget(process.env.MEMORY_BUDGET),
+);

tests/e2e/ssrf.test.ts

@@ -4,8 +4,9 @@ const BASE_URL = process.env.BASE_URL || "http://localhost:3000";
 
 // A public site the first-hop URL check accepts, whose redirect or script
 // then sends the browser somewhere private.
-const OPEN_REDIRECT = "https://httpbin.org/redirect-to?url=";
-const RAW_HTML = "https://httpbin.org/base64/";
+const OPEN_REDIRECT = "https://httpbin.dev/redirect-to?url=";
+const RAW_HTML =
+	"https://httpbin.dev/response-headers?Content-Type=text/html&page=";
 
 // The browser service itself, reachable on loopback inside its container.
 const LOOPBACK = "http://127.0.0.1:3000/v1/health";
@@ -24,15 +25,15 @@ beforeAll(async () => {
 		httpbinUp = false;
 	}
 	if (!httpbinUp) {
-		console.warn("httpbin.org is unreachable; skipping its SSRF checks.");
+		console.warn("httpbin.dev is unreachable; skipping its SSRF checks.");
 	}
 });
 
 function scriptPage(target: string): string {
-	const html = `<script>location.href = ${JSON.stringify(target)};</script>`;
-	// httpbin decodes URL-safe base64 but insists on the padding.
-	const encoded = Buffer.from(html).toString("base64url");
-	return `${RAW_HTML}${encoded.padEnd(Math.ceil(encoded.length / 4) * 4, "=")}`;
+	// httpbin echoes this into a JSON body served as HTML, and JSON escapes
+	// double quotes, so the script quotes its target with backticks.
+	const html = `<script>location.href = \`${target}\`;</script>`;
+	return `${RAW_HTML}${encodeURIComponent(html)}`;
 }
 
 async function screenshot(url: string): Promise<Response> {

tests/recovery/memory.test.ts

@@ -0,0 +1,59 @@
+import { afterAll, beforeAll, describe, expect, test } from "bun:test";
+
+const IMAGE = "appwrite/browser:local";
+const NAME = "appwrite-browser-memory-budget";
+const PORT = 3001;
+const BASE_URL = `http://localhost:${PORT}`;
+
+const imageBuilt = Bun.spawnSync(["docker", "image", "inspect", IMAGE]).success;
+
+describe.skipIf(!imageBuilt)("E2E Tests - memory budget", () => {
+	beforeAll(async () => {
+		Bun.spawnSync(["docker", "rm", "-f", NAME]);
+		// Idle Chromium alone is over 5% of 1 GiB, so every capture trips the budget.
+		Bun.spawnSync([
+			"docker",
+			"run",
+			"-d",
+			"--name",
+			NAME,
+			"--memory",
+			"1g",
+			"-e",
+			"MEMORY_BUDGET=0.05",
+			"-p",
+			`${PORT}:3000`,
+			IMAGE,
+		]);
+
+		const deadline = Date.now() + 30_000;
+		while (Date.now() < deadline) {
+			const health = await fetch(`${BASE_URL}/v1/health`).catch(() => null);
+			if (health?.ok) return;
+			await Bun.sleep(500);
+		}
+		throw new Error("Memory budget container did not become healthy.");
+	}, 60_000);
+
+	afterAll(() => {
+		Bun.spawnSync(["docker", "rm", "-f", NAME]);
+	});
+
+	test("should fail a capture over the budget with 507 and keep Chromium up", async () => {
+		const capture = await fetch(`${BASE_URL}/v1/screenshots`, {
+			method: "POST",
+			headers: { "Content-Type": "application/json" },
+			body: JSON.stringify({ url: "https://example.com", sleep: 2000 }),
+		});
+		expect(capture.status).toBe(507);
+		expect(capture.headers.get("Content-Type")).toBe("application/json");
+		const data = await capture.json();
+		expect(data.error).toBe("Page exceeded the memory budget and was closed.");
+
+		const health = await fetch(`${BASE_URL}/v1/health`);
+		expect(health.status).toBe(200);
+
+		const logs = Bun.spawnSync(["docker", "logs", NAME]).stdout.toString();
+		expect(logs).toContain("Memory budget exceeded");
+	}, 60_000);
+});