From 07dfd12613175af467ac86312a9d802b7ac36bac Mon Sep 17 00:00:00 2001 From: hhang <17309584214@163.com> Date: Thu, 29 Oct 2026 07:45:32 +0800 Subject: [PATCH] fix(pe): reject non-MZ inputs before image inspection (#7709) Co-authored-by: hhang <169027046+xu91102@users.noreply.github.com> --- src/native/metadata/PeImage.test.ts | 29 ++++++++++++ src/native/metadata/PeImage.ts | 2 + tests/acceptance/peImportsCli.test.ts | 56 +++++++++++++++++++++++- tests/acceptance/peLoadConfigCli.test.ts | 50 ++++++++++++++++++++- 4 files changed, 135 insertions(+), 2 deletions(-) diff --git a/src/native/metadata/PeImage.test.ts b/src/native/metadata/PeImage.test.ts index 87ca005b5c..baaad17772 100644 --- a/src/native/metadata/PeImage.test.ts +++ b/src/native/metadata/PeImage.test.ts @@ -2,6 +2,7 @@ import { expect, it } from "vitest"; import { readGoBinaryImage } from "../../go/GoBinaryImage.js"; import { readPeImage } from "./PeImage.js"; import { pe } from "../../domain/binaryTarget.fixture.js"; +import { parseMzHeaderEvidence } from "../../domain/dosMz.js"; import { BinaryFormatFailure } from "./BinaryContainer.js"; import { readBinaryContainer } from "./readBinaryContainer.js"; import { @@ -9,6 +10,34 @@ import { setGoPeVirtualSize, } from "../../../tests/fixtures/go/image.js"; +for (const bits of [32, 64] as const) { + it.each([ + ["ZZ", Buffer.from("ZZ")], + ["high-bit MZ", Buffer.from([0xcd, 0xda])], + ] as const)( + `rejects %s before reading PE${String(bits)} facts`, + (_, prefix) => { + const { bytes } = createGoBinaryFixture({ format: "pe", bits }); + prefix.copy(bytes); + const declarations = parseMzHeaderEvidence(bytes).declarations; + if (declarations.ok) throw new Error("Expected invalid MZ declarations"); + expect(() => readPeImage(bytes)).toThrow( + expect.objectContaining({ + kind: "malformed", + message: declarations.error, + }), + ); + }, + ); + + it(`reads PE${String(bits)} without requiring a valid DOS load module`, () => { + const { bytes } = createGoBinaryFixture({ format: "pe", bits }); + bytes.writeUInt16LE(0, 4); + expect(parseMzHeaderEvidence(bytes).dos.ok).toBe(false); + expect(readPeImage(bytes)).toMatchObject({ format: "pe", bits }); + }); +} + it("classifies a confirmed overlapping PE as an unsupported layout", () => { try { readPeImage(pe(0x014c, 4)); diff --git a/src/native/metadata/PeImage.ts b/src/native/metadata/PeImage.ts index 3edf910cba..89a66873a6 100644 --- a/src/native/metadata/PeImage.ts +++ b/src/native/metadata/PeImage.ts @@ -17,6 +17,8 @@ import { const peHeaderOffset = (reader: BinaryReader): number => { const evidence = parseMzHeaderEvidence(reader.bytes); + if (!evidence.declarations.ok) + throw new BinaryFormatFailure("malformed", evidence.declarations.error); const { windowsHeader } = evidence; if (windowsHeader.state === "candidate") throw new BinaryFormatFailure( diff --git a/tests/acceptance/peImportsCli.test.ts b/tests/acceptance/peImportsCli.test.ts index 31fd580f79..b0213a8e6e 100644 --- a/tests/acceptance/peImportsCli.test.ts +++ b/tests/acceptance/peImportsCli.test.ts @@ -1,4 +1,4 @@ -import { writeFile } from "node:fs/promises"; +import { readFile, writeFile } from "node:fs/promises"; import { join } from "node:path"; import { expect } from "vitest"; import { toolContract } from "../../src/contracts/toolContracts.js"; @@ -12,8 +12,62 @@ import { peWithoutImportTableFixture, } from "../fixtures/binaryDiagnostics/peImports.js"; import { createTestTempDirectory } from "../fixtures/temporaryDirectory.js"; +import { parseMcpToolError } from "../fixtures/mcpToolError.js"; import { cliTest } from "../support/cli/cliFixture.js"; +for (const bits of [32, 64] as const) { + cliTest( + `rejects non-MZ PE${String(bits)} imports across CLI and MCP`, + async ({ cli }) => { + const root = await createTestTempDirectory("rea-pe-imports-signature-"); + const call = await cli.sharedMcpCall!(); + for (const [name, prefix] of [ + ["zz", Buffer.from("ZZ")], + ["high-bit", Buffer.from([0xcd, 0xda])], + ] as const) { + const bytes = peImportImage({ bits }); + prefix.copy(bytes); + const path = join(root, `${name}.exe`); + await writeFile(path, bytes); + const result = await cli.run({ + arguments: ["inspect-pe-imports", path, "--json"], + }); + expect(result.exitCode).toBe(1); + expect(result.json).toMatchObject({ + code: "invalid_request", + details: { operation: "inspect_pe_imports" }, + }); + const response = await call("inspect_pe_imports", { path }); + expect(response.isError).toBe(true); + expect(parseMcpToolError(response)).toEqual({ error: result.json }); + expect((await readFile(path)).equals(bytes)).toBe(true); + } + const path = join(root, "valid.exe"); + const bytes = peImportImage({ bits }); + await writeFile(path, bytes); + const result = await cli.run({ + arguments: ["inspect-pe-imports", path, "--json"], + }); + expect(result.exitCode, result.stderr).toBe(0); + const response = await call("inspect_pe_imports", { path }); + expect(response.isError).not.toBe(true); + for (const value of [result.json, response.structuredContent]) { + const report = + toolContract("inspect_pe_imports").outputSchema.parse( + value, + ).normalized_result; + expect(report.bits).toBe(bits); + expect(report.imports[0]?.functions.map((row) => row.name)).toEqual([ + "GetProcAddress", + "ExitProcess", + ]); + expect(report.coverage.status).toBe("complete"); + } + expect((await readFile(path)).equals(bytes)).toBe(true); + }, + ); +} + cliTest( "distinguishes pefile's bounded name read from an invalid import name", async ({ cli }) => { diff --git a/tests/acceptance/peLoadConfigCli.test.ts b/tests/acceptance/peLoadConfigCli.test.ts index 077d1d4734..72243d33ea 100644 --- a/tests/acceptance/peLoadConfigCli.test.ts +++ b/tests/acceptance/peLoadConfigCli.test.ts @@ -1,4 +1,4 @@ -import { writeFile } from "node:fs/promises"; +import { readFile, writeFile } from "node:fs/promises"; import { join } from "node:path"; import { expect } from "vitest"; import { toolContract } from "../../src/contracts/toolContracts.js"; @@ -7,9 +7,57 @@ import { peLoadConfigImage, } from "../../src/native/pe/PeLoadConfig.fixture.js"; import { createTestTempDirectory } from "../fixtures/temporaryDirectory.js"; +import { parseMcpToolError } from "../fixtures/mcpToolError.js"; import { cliTest } from "../support/cli/cliFixture.js"; for (const bits of [32, 64] as const) { + cliTest( + `rejects non-MZ PE${String(bits)} load config across CLI and MCP`, + async ({ cli }) => { + const root = await createTestTempDirectory("rea-pe-config-signature-"); + const call = await cli.sharedMcpCall!(); + for (const [name, prefix] of [ + ["zz", Buffer.from("ZZ")], + ["high-bit", Buffer.from([0xcd, 0xda])], + ] as const) { + const bytes = peLoadConfigImage({ bits }); + prefix.copy(bytes); + const path = join(root, `${name}.exe`); + await writeFile(path, bytes); + const result = await cli.run({ + arguments: ["inspect-pe-load-config", path, "--json"], + }); + expect(result.exitCode).toBe(1); + expect(result.json).toMatchObject({ + code: "invalid_request", + details: { operation: "inspect_pe_load_config" }, + }); + const response = await call("inspect_pe_load_config", { path }); + expect(response.isError).toBe(true); + expect(parseMcpToolError(response)).toEqual({ error: result.json }); + expect((await readFile(path)).equals(bytes)).toBe(true); + } + const path = join(root, "valid.exe"); + const bytes = peLoadConfigImage({ bits }); + await writeFile(path, bytes); + const result = await cli.run({ + arguments: ["inspect-pe-load-config", path, "--json"], + }); + expect(result.exitCode, result.stderr).toBe(0); + const response = await call("inspect_pe_load_config", { path }); + expect(response.isError).not.toBe(true); + for (const value of [result.json, response.structuredContent]) { + const report = toolContract( + "inspect_pe_load_config", + ).outputSchema.parse(value).normalized_result; + expect(report.bits).toBe(bits); + expect(report.tls?.callbacks).toHaveLength(3); + expect(report.load_config?.guard_flags?.raw).toBe(0x500); + } + expect((await readFile(path)).equals(bytes)).toBe(true); + }, + ); + cliTest(`reads PE${bits} TLS callbacks through the CLI`, async ({ cli }) => { const root = await createTestTempDirectory("rea-pe-config-cli-"); await writeFile(join(root, "app.exe"), peLoadConfigImage({ bits }));