fix(browser): bound screenshot comparison memory

This commit is contained in:
morluto 2026-10-11 11:03:39 +08:00
parent a99b2fd3b5
commit b3e977e6d0
5 changed files with 289 additions and 22 deletions

View File

@ -24,6 +24,29 @@ REA can attach to a user-owned Chrome-family browser through the Chrome DevTools
Target discovery returns the complete in-scope target array in one result.
## Screenshot comparison memory budget
`compare_web_screenshots` admits both PNGs against a 256 MiB estimated
working-memory budget before inflating either image. This is a processing
buffer budget for this operation, not a PNG dimension limit, CDP transport
limit, or process RSS guarantee. Header-valid dimensions that exceed the
budget return `resource_constraint` with `details.resource: "memory"` and the
estimated and maximum working-set bytes. Invalid PNG structure remains
malformed input; encodings outside supported pixel-comparison coverage remain
unsupported-target errors.
The conservative estimate includes the two decoder-owned decoded input
buffers, the current image's concatenated IDAT buffer, Node zlib's default
16 KiB output chunk backing rounded up for retained chunks, the concatenated
inflate result while its output chunks are retained, one possible extra output
chunk before zlib checks `maxOutputLength`, unfiltered scanline storage,
RGB-to-RGBA expansion, and the first retained RGBA image while the second is
decoded. RGBA output aliases its unfiltered buffer and is counted once.
Admission uses checked integer arithmetic before image-sized allocation.
Caller-owned JSON/base64 strings, input parsing and validation, transport
buffers, garbage-collector behavior, and other runtime memory are outside this
estimate; the budget does not promise a bound on total process memory.
Electron `file://` pages use a separate provider and target boundary; see [electron-observation.md](electron-observation.md).
## Provider authorities

View File

@ -3,7 +3,10 @@ import { crc32, deflateSync } from "node:zlib";
import { describe, expect, it } from "vitest";
import { comparePngScreenshots } from "./PngVisualDiff.js";
import { AnalysisUnsupportedTargetError } from "../domain/analysisErrorCore.js";
import {
AnalysisResourceConstraintError,
AnalysisUnsupportedTargetError,
} from "../domain/analysisErrorCore.js";
import {
compareWebScreenshotsInputSchema,
createWebScreenshotArtifact,
@ -118,7 +121,7 @@ describe("PNG visual diff", () => {
});
it("rejects malformed PNG dimensions after validating image data", () => {
const image = artifact(32_000_001, 1, [0, 0, 0, 255]);
const image = artifact(2, 1, [0, 0, 0, 255]);
expect(() =>
comparePngScreenshots(
compareWebScreenshotsInputSchema.parse({
@ -150,6 +153,32 @@ describe("PNG visual diff", () => {
});
});
describe("PNG comparison memory budget", () => {
it("classifies legal huge dimensions as a memory constraint before inflate", () => {
const image = createWebScreenshotArtifact(headerOnlyPng(0x7fffffff, 1));
const input = compareWebScreenshotsInputSchema.parse({
before: image,
after: image,
});
try {
comparePngScreenshots(input);
throw new Error("expected PNG working-memory rejection");
} catch (cause: unknown) {
expect(cause).toBeInstanceOf(AnalysisResourceConstraintError);
expect(cause).toMatchObject({
resource: "memory",
reportedLimits: {
maximum_working_memory_bytes: 256 * 1024 * 1024,
estimated_working_memory_bytes: "34359738499",
before_dimensions: { width: 0x7fffffff, height: 1 },
after_dimensions: { width: 0x7fffffff, height: 1 },
},
});
}
});
});
describe("PNG chunk integrity", () => {
it.each(["IHDR", "IDAT", "IEND"])(
"rejects a corrupt %s chunk checksum",
@ -398,6 +427,19 @@ const png = (
]);
};
const headerOnlyPng = (width: number, height: number): Buffer => {
const header = Buffer.alloc(13);
header.writeUInt32BE(width, 0);
header.writeUInt32BE(height, 4);
header.set([8, 6, 0, 0, 0], 8);
return Buffer.concat([
Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]),
chunk("IHDR", header),
chunk("IDAT", deflateSync(Buffer.from([0]))),
chunk("IEND", Buffer.alloc(0)),
]);
};
const chunk = (type: string, data: Buffer): Buffer => {
const result = Buffer.alloc(12 + data.byteLength);
result.writeUInt32BE(data.byteLength, 0);

View File

@ -1,11 +1,14 @@
import { crc32, inflateSync } from "node:zlib";
import { constants as zlibConstants, crc32, inflateSync } from "node:zlib";
import type {
CompareWebScreenshotsInput,
WebScreenshotArtifact,
WebScreenshotDiff,
} from "../domain/webScreenshot.js";
import { AnalysisUnsupportedTargetError } from "../domain/analysisErrorCore.js";
import {
AnalysisResourceConstraintError,
AnalysisUnsupportedTargetError,
} from "../domain/analysisErrorCore.js";
import { decodeCanonicalBase64 } from "../domain/webScreenshot.js";
interface DecodedPng {
@ -14,12 +17,32 @@ interface DecodedPng {
readonly rgba: Buffer;
}
interface ParsedPng {
readonly bytes: Buffer;
readonly width: number;
readonly height: number;
readonly channels: number;
readonly compressed: readonly Buffer[];
readonly transparentColor: readonly number[] | undefined;
readonly expectedRawBytes: bigint;
readonly unfilteredBytes: bigint;
readonly rgbaBytes: bigint;
}
/** Conservative allowance for the simultaneously live PNG comparison buffers. */
const PNG_COMPARISON_WORKING_MEMORY_BYTES = 256 * 1024 * 1024;
const PNG_INFLATE_CHUNK_BYTES = zlibConstants.Z_DEFAULT_CHUNK;
/** Compare two validated screenshots using deterministic RGBA metrics. */
export const comparePngScreenshots = (
input: CompareWebScreenshotsInput,
): WebScreenshotDiff => {
const before = decodePng(input.before, "before");
const after = decodePng(input.after, "after");
admitEncodedInputPair(input.before, input.after);
const beforePlan = parsePng(input.before, "before");
const afterPlan = parsePng(input.after, "after");
admitPngPair(beforePlan, afterPlan);
const before = decodePng(beforePlan);
const after = decodePng(afterPlan);
if (before.width !== after.width || before.height !== after.height)
return {
status: "dimension_mismatch",
@ -74,10 +97,10 @@ export const comparePngScreenshots = (
};
};
const decodePng = (
const parsePng = (
artifact: WebScreenshotArtifact,
field: "before" | "after",
): DecodedPng => {
): ParsedPng => {
const bytes = decodeCanonicalBase64(artifact.data_base64);
if (bytes === undefined || !bytes.subarray(0, 8).equals(PNG_SIGNATURE))
throw new TypeError("Invalid PNG signature");
@ -180,24 +203,136 @@ const decodePng = (
field,
"PNG encoding is not supported for pixel comparison.",
);
const rowBytes = header.width * header.channels;
const expected = (rowBytes + 1) * header.height;
const raw = inflatePng(Buffer.concat(compressed), expected + 1);
if (raw.byteLength !== expected) throw new TypeError("Unexpected PNG size");
const pixels = BigInt(header.width) * BigInt(header.height);
const expectedRawBytes =
(BigInt(header.width) * BigInt(header.channels) + 1n) *
BigInt(header.height);
return {
bytes,
width: header.width,
height: header.height,
channels: header.channels,
compressed,
transparentColor,
expectedRawBytes,
unfilteredBytes: pixels * BigInt(header.channels),
rgbaBytes: pixels * 4n,
};
};
const decodePng = (plan: ParsedPng): DecodedPng => {
const expected = Number(plan.expectedRawBytes);
const raw = inflatePng(Buffer.concat(plan.compressed), expected + 1);
if (raw.byteLength !== expected) throw new TypeError("Unexpected PNG size");
return {
width: plan.width,
height: plan.height,
rgba: expandRgba(
unfilter(raw, header.width, header.height, header.channels),
header.channels,
transparentColor,
unfilter(raw, plan.width, plan.height, plan.channels),
plan.channels,
plan.transparentColor,
),
};
};
const admitEncodedInputPair = (
before: WebScreenshotArtifact,
after: WebScreenshotArtifact,
): void => {
const sourceBytes = BigInt(before.bytes) + BigInt(after.bytes);
if (sourceBytes > BigInt(PNG_COMPARISON_WORKING_MEMORY_BYTES))
throw pngWorkingMemoryError(sourceBytes, null, null);
};
const admitPngPair = (before: ParsedPng, after: ParsedPng): void => {
// Input buffers remain live for both plans. For each decode, count the
// concatenated IDAT, rounded zlib output chunks, any concatenated output and
// unfiltered rows, plus RGB-to-RGBA output. The first completed RGBA image
// remains live while the second is decoded.
const inputBytes = BigInt(before.bytes.byteLength + after.bytes.byteLength);
const estimate = (current: ParsedPng, retained: ParsedPng | undefined) => {
const compressedBytes = current.compressed.reduce(
(total, chunk) => total + BigInt(chunk.byteLength),
0n,
);
const rgbaExpansion = current.channels === 3 ? current.rgbaBytes : 0n;
return (
inputBytes +
(retained?.rgbaBytes ?? 0n) +
compressedBytes +
inflateBufferAllowance(current.expectedRawBytes) +
current.unfilteredBytes +
rgbaExpansion
);
};
const estimateBytes = [
estimate(before, undefined),
estimate(after, before),
].reduce((maximum, value) => (value > maximum ? value : maximum), 0n);
if (estimateBytes > BigInt(PNG_COMPARISON_WORKING_MEMORY_BYTES))
throw pngWorkingMemoryError(estimateBytes, before, after);
};
const inflateBufferAllowance = (expectedRawBytes: bigint): bigint => {
const chunkBytes = BigInt(PNG_INFLATE_CHUNK_BYTES);
const maximumOutputBytes = expectedRawBytes + 1n;
const roundUpToChunk = (bytes: bigint): bigint =>
((bytes + chunkBytes - 1n) / chunkBytes) * chunkBytes;
// Node checks maxOutputLength after writing a chunk, so a rejected stream can
// retain one chunk's worth of output past the limit before it throws.
const rejectedOutputChunks = roundUpToChunk(
maximumOutputBytes + chunkBytes - 1n,
);
// Sync zlib retains each output chunk while Buffer.concat creates the final
// result. A full final chunk also allocates the next scratch chunk first.
const returnedOutputChunks = roundUpToChunk(maximumOutputBytes);
const scratchChunk =
expectedRawBytes % chunkBytes === 0n ||
maximumOutputBytes % chunkBytes === 0n
? chunkBytes
: 0n;
const concatenatedOutput =
maximumOutputBytes > chunkBytes ? maximumOutputBytes : 0n;
const returnedOutput =
returnedOutputChunks + scratchChunk + concatenatedOutput;
return rejectedOutputChunks > returnedOutput
? rejectedOutputChunks
: returnedOutput;
};
const pngWorkingMemoryError = (
estimateBytes: bigint,
before: Pick<ParsedPng, "width" | "height"> | null,
after: Pick<ParsedPng, "width" | "height"> | null,
) =>
new AnalysisResourceConstraintError(
"compare_web_screenshots",
"memory",
"Screenshot comparison exceeds REA's 256 MiB estimated working-memory budget.",
{
boundary: "png-screenshot-comparison",
estimated_working_memory_bytes: estimateBytes.toString(),
maximum_working_memory_bytes: PNG_COMPARISON_WORKING_MEMORY_BYTES,
before_dimensions:
before === null ? null : { width: before.width, height: before.height },
after_dimensions:
after === null ? null : { width: after.width, height: after.height },
},
{
remediationAction:
"Compare smaller screenshots or reduce their viewport dimensions before capture.",
},
);
const inflatePng = (compressed: Buffer, maxOutputLength: number): Buffer => {
try {
return inflateSync(compressed, { maxOutputLength });
return inflateSync(compressed, {
maxOutputLength,
chunkSize: PNG_INFLATE_CHUNK_BYTES,
});
} catch (cause: unknown) {
if (
cause instanceof Error &&
@ -217,8 +352,6 @@ const parseHeader = (data: Buffer) => {
const height = data.readUInt32BE(4);
const bitDepth = data[8] ?? 0;
const colorType = data[9] ?? -1;
const safePixelCount =
BigInt(width) * BigInt(height) <= BigInt(Number.MAX_SAFE_INTEGER);
const validDepth =
(colorType === 0 && [1, 2, 4, 8, 16].includes(bitDepth)) ||
(colorType === 2 && [8, 16].includes(bitDepth)) ||
@ -244,10 +377,7 @@ const parseHeader = (data: Buffer) => {
interlace: data[12],
channels: colorType === 6 ? 4 : 3,
supported:
safePixelCount &&
bitDepth === 8 &&
(colorType === 2 || colorType === 6) &&
data[12] === 0,
bitDepth === 8 && (colorType === 2 || colorType === 6) && data[12] === 0,
};
};

View File

@ -197,4 +197,39 @@ const assertPngInputErrorProjections = async (): Promise<void> => {
code: "unsupported_target",
details: { operation: "compare_web_screenshots" },
});
const oversizedHeader = Buffer.alloc(13);
oversizedHeader.writeUInt32BE(0x7fffffff, 0);
oversizedHeader.writeUInt32BE(1, 4);
oversizedHeader.set([8, 6, 0, 0, 0], 8);
const hugeDimensions = JSON.stringify(
createWebScreenshotArtifact(
Buffer.concat([
Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]),
pngChunk("IHDR", oversizedHeader),
pngChunk("IDAT", deflateSync(Buffer.from([0]))),
pngChunk("IEND", Buffer.alloc(0)),
]),
),
);
expect(
await runCli([
"compare-web-screenshots",
hugeDimensions,
hugeDimensions,
"--json",
]),
).toMatchObject({
code: "resource_constraint",
details: {
operation: "compare_web_screenshots",
resource: "memory",
reported_limits: {
maximum_working_memory_bytes: 256 * 1024 * 1024,
estimated_working_memory_bytes: "34359738499",
before_dimensions: { width: 0x7fffffff, height: 1 },
after_dimensions: { width: 0x7fffffff, height: 1 },
},
},
});
};

View File

@ -361,6 +361,43 @@ it("does not attach to a target outside the request's allowed origin scope", asy
expect(browser.commands).toHaveLength(0);
});
it("projects legal oversized PNG dimensions as a memory constraint", async () => {
const header = Buffer.alloc(13);
header.writeUInt32BE(0x7fffffff, 0);
header.writeUInt32BE(1, 4);
header.set([8, 6, 0, 0, 0], 8);
const artifact = createWebScreenshotArtifact(
Buffer.concat([
Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]),
pngChunk("IHDR", header),
pngChunk("IDAT", deflateSync(Buffer.from([0]))),
pngChunk("IEND", Buffer.alloc(0)),
]),
);
const connected = await connectBrowser();
const result = await connected.client.callTool({
name: "compare_web_screenshots",
arguments: { before: artifact, after: artifact },
});
expect(parseMcpToolError(result)).toMatchObject({
error: {
code: "resource_constraint",
category: "resource_constraint",
details: {
operation: "compare_web_screenshots",
resource: "memory",
reported_limits: {
maximum_working_memory_bytes: 256 * 1024 * 1024,
estimated_working_memory_bytes: "34359738499",
before_dimensions: { width: 0x7fffffff, height: 1 },
after_dimensions: { width: 0x7fffffff, height: 1 },
},
},
},
});
});
const connectBrowser = async () => {
const session = createTestBinarySession(() => ({
execute: () => Promise.resolve(observed(null)),