fix(screenshot): validate complete PNG IHDR (#6412)

Co-authored-by: eco-null <eco-null@users.noreply.github.com>

src/domain/pngHeader.test.ts

@@ -1,3 +1,4 @@
+import { crc32 } from "node:zlib";
 import { expect, it } from "vitest";
 import { readPngHeaderDimensions } from "./pngHeader.js";
 
@@ -10,6 +11,7 @@ it("reads encoded dimensions without changing screenshot bytes", () => {
   const bytes = Buffer.from(png);
   bytes.writeUInt32BE(257, 16);
   bytes.writeUInt32BE(129, 20);
+  bytes.writeUInt32BE(crc32(bytes.subarray(12, 29)), 29);
   const original = Buffer.from(bytes);
   expect(readPngHeaderDimensions(bytes)).toEqual({ width: 257, height: 129 });
   expect(bytes.equals(original)).toBe(true);
@@ -18,6 +20,7 @@ it("reads encoded dimensions without changing screenshot bytes", () => {
 it.each([12, 13, 14, 15])("rejects ASCII-folded IHDR byte at %s", (offset) => {
   const bytes = Buffer.from(png);
   bytes[offset] = (bytes[offset] ?? 0) | 0x80;
+  bytes.writeUInt32BE(crc32(bytes.subarray(12, 29)), 29);
   expect(bytes.subarray(12, 16).toString("ascii")).toBe("IHDR");
   expect(readPngHeaderDimensions(bytes)).toBeUndefined();
 });
@@ -30,6 +33,7 @@ it("rejects incomplete headers, invalid signatures and zero dimensions", () => {
   for (const offset of [16, 20]) {
     const zero = Buffer.from(png);
     zero.writeUInt32BE(0, offset);
+    zero.writeUInt32BE(crc32(zero.subarray(12, 29)), 29);
     expect(readPngHeaderDimensions(zero)).toBeUndefined();
   }
 });
@@ -43,14 +47,26 @@ it.each([0, 12, 14, 0xffff_ffff])("rejects IHDR data length %s", (length) => {
 it.each([16, 20])("rejects out-of-range PNG dimension at %s", (offset) => {
   const bytes = Buffer.from(png);
   bytes.writeUInt32BE(0x8000_0000, offset);
+  bytes.writeUInt32BE(crc32(bytes.subarray(12, 29)), 29);
   expect(readPngHeaderDimensions(bytes)).toBeUndefined();
 });
 
 it("accepts the largest dimension representable by the PNG specification", () => {
   const bytes = Buffer.from(png);
   bytes.writeUInt32BE(0x7fff_ffff, 16);
+  bytes.writeUInt32BE(crc32(bytes.subarray(12, 29)), 29);
   expect(readPngHeaderDimensions(bytes)).toEqual({
     width: 0x7fff_ffff,
     height: 1,
   });
 });
+
+it.each([24, 32, "crc"])("rejects incomplete or corrupt IHDR %s", (kind) => {
+  const bytes = Buffer.from(png);
+  if (kind === "crc") bytes[29] = (bytes[29] ?? 0) ^ 1;
+  expect(
+    readPngHeaderDimensions(
+      typeof kind === "number" ? bytes.subarray(0, kind) : bytes,
+    ),
+  ).toBeUndefined();
+});

src/domain/pngHeader.ts

@@ -1,3 +1,5 @@
+import { crc32 } from "node:zlib";
+
 const PNG_SIGNATURE = Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]);
 const IHDR_TYPE = Buffer.from("IHDR");
 
@@ -6,10 +8,11 @@ export const readPngHeaderDimensions = (
   bytes: Buffer,
 ): { readonly width: number; readonly height: number } | undefined => {
   if (
-    bytes.byteLength < 24 ||
+    bytes.byteLength < 33 ||
     !bytes.subarray(0, 8).equals(PNG_SIGNATURE) ||
     bytes.readUInt32BE(8) !== 13 ||
-    !bytes.subarray(12, 16).equals(IHDR_TYPE)
+    !bytes.subarray(12, 16).equals(IHDR_TYPE) ||
+    crc32(bytes.subarray(12, 29)) !== bytes.readUInt32BE(29)
   )
     return undefined;
   const width = bytes.readUInt32BE(16);

src/native/NativeUiObservation.test.ts

@@ -1,4 +1,5 @@
 import { createHash } from "node:crypto";
+import { crc32 } from "node:zlib";
 
 import { describe, expect, it } from "vitest";
 import { observeNativeUi } from "./NativeUiObservation.js";
@@ -73,11 +74,24 @@ it.each([
   async ({ offset, value }) => {
     const bytes = Buffer.from(validPng);
     bytes.writeUInt32BE(value, offset);
+    bytes.writeUInt32BE(crc32(bytes.subarray(12, 29)), 29);
     const result = await observeWithScreenshot(screenshotOf(bytes));
     expect(result.ok).toBe(false);
   },
 );
 describe("native UI screenshot validation", () => {
+  it.each([24, 32, "crc"])(
+    "rejects incomplete or corrupt IHDR despite matching native digest: %s",
+    async (kind) => {
+      const bytes = Buffer.from(validPng);
+      if (kind === "crc") bytes[29] = (bytes[29] ?? 0) ^ 1;
+      const payload =
+        typeof kind === "number" ? bytes.subarray(0, kind) : bytes;
+      expect((await observeWithScreenshot(screenshotOf(payload))).ok).toBe(
+        false,
+      );
+    },
+  );
   it("rejects malformed or mismatched screenshot bytes from the helper", async () => {
     const result = await observeWithScreenshot({
       ...screenshotOf(validPng),

tests/boundary/browser/cdpScreenshotDocumentIdentity.test.ts

@@ -1,3 +1,4 @@
+import { crc32 } from "node:zlib";
 import { describe, expect, it, onTestFinished } from "vitest";
 import { CdpBrowserProvider } from "../../../src/browser/CdpBrowserProvider.js";
 import { captureWebScreenshotInputSchema } from "../../../src/domain/webScreenshot.js";
@@ -28,6 +29,7 @@ const screenshotData = (width: number, height: number) => {
   );
   bytes.writeUInt32BE(width, 16);
   bytes.writeUInt32BE(height, 20);
+  bytes.writeUInt32BE(crc32(bytes.subarray(12, 29)), 29);
   return { data: bytes.toString("base64") };
 };
 
@@ -791,3 +793,33 @@ describe("screenshot navigation event applicability", () => {
     expect(result.ok).toBe(true);
   });
 });
+
+it.each([24, 32, "crc"])(
+  "rejects incomplete or corrupt PNG IHDR: %s",
+  async (kind) => {
+    const png = Buffer.from(
+      "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jv4sAAAAASUVORK5CYII=",
+      "base64",
+    );
+    if (kind === "crc") png[29] = (png[29] ?? 0) ^ 1;
+    const bytes = typeof kind === "number" ? png.subarray(0, kind) : png;
+    const browser = await startFakeCdpBrowser({
+      commandResult: ({ method }) =>
+        method === "Page.captureScreenshot"
+          ? { data: bytes.toString("base64") }
+          : undefined,
+    });
+    onTestFinished(async () => browser.close());
+    const result = await new CdpBrowserProvider().captureScreenshot(
+      captureWebScreenshotInputSchema.parse({
+        cdp_endpoint: browser.endpoint,
+        allowed_origins: [browser.allowedOrigin],
+        target_id: "allowed-page",
+      }),
+    );
+    expect(result).toMatchObject({
+      ok: false,
+      error: { reason: "protocol_error" },
+    });
+  },
+);

tests/boundary/mcp/browserMcp.test.ts

@@ -608,3 +608,37 @@ const artifactOf = (value: unknown): unknown => {
     throw new TypeError("Missing screenshot artifact");
   return normalized.artifact;
 };
+
+it.each([24, 32, "crc"])(
+  "rejects corrupt captured PNG through MCP: %s",
+  async (kind) => {
+    const bytes = Buffer.from(
+      "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jv4sAAAAASUVORK5CYII=",
+      "base64",
+    );
+    if (kind === "crc") bytes[29] = (bytes[29] ?? 0) ^ 1;
+    const data = (
+      typeof kind === "number" ? bytes.subarray(0, kind) : bytes
+    ).toString("base64");
+    const browser = await startFakeCdpBrowser({
+      commandResult: ({ method }) =>
+        method === "Page.captureScreenshot" ? { data } : undefined,
+    });
+    browsers.push(browser);
+    const { client } = await connectBrowser();
+    const result = await client.callTool({
+      name: "capture_web_screenshot",
+      arguments: {
+        cdp_endpoint: browser.endpoint,
+        allowed_origins: [browser.allowedOrigin],
+        target_id: "allowed-page",
+      },
+    });
+    expect(parseMcpToolError(result)).toMatchObject({
+      error: {
+        code: "unreadable_output",
+        details: { reason: "protocol_error" },
+      },
+    });
+  },
+);