diff --git a/README.md b/README.md index 5d340fb..1c86096 100644 --- a/README.md +++ b/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). diff --git a/src/routes/reports.ts b/src/routes/reports.ts index 043205a..d8cc5e5 100755 --- a/src/routes/reports.ts +++ b/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 { 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 { 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 { 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(); } diff --git a/src/routes/screenshots.ts b/src/routes/screenshots.ts index ce6766b..6b3432e 100755 --- a/src/routes/screenshots.ts +++ b/src/routes/screenshots.ts @@ -12,6 +12,11 @@ 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( @@ -19,6 +24,7 @@ export async function handleScreenshotsRequest( ): Promise { 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(); } diff --git a/src/utils/memory.ts b/src/utils/memory.ts new file mode 100644 index 0000000..8efea00 --- /dev/null +++ b/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; +} + +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 | 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), +); diff --git a/tests/e2e/ssrf.test.ts b/tests/e2e/ssrf.test.ts index 90d3706..d8838e7 100644 --- a/tests/e2e/ssrf.test.ts +++ b/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 = ``; - // 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 = ``; + return `${RAW_HTML}${encodeURIComponent(html)}`; } async function screenshot(url: string): Promise { diff --git a/tests/recovery/memory.test.ts b/tests/recovery/memory.test.ts new file mode 100644 index 0000000..df5951d --- /dev/null +++ b/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); +});