fix(pe): reject non-MZ inputs before image inspection (#7709)

Co-authored-by: hhang <[email protected]>

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));

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(

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 }) => {

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 }));