From 599099f65670e4e4eb478970e2f7a845eaa53ebd Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Sun, 25 Oct 2026 05:28:16 +0800 Subject: [PATCH] refactor(pe): share optional-header structural facts Decode optional-header declarations once and pass detached directory words, locations and SizeOfHeaders to native consumers. Preserve executable machine admission, raw GPU/resource machine reporting, RVA mapping and certificate file-offset semantics. Validation: fresh build/catalog, full static checks, 466 core/provider fixture tests, 122 workflow tests and 14 selected file-backed PE admission tests passed. Public CLI and actual stdio MCP used authored artifacts; installed external providers were not qualified. Added 17 structural cases and two GPU machine-policy controls. Eager directory decoding costs remain unmeasured. --- src/domain/peMachine.ts | 52 ++--- src/domain/peOptionalHeader.test.ts | 181 ++++++++++++++++++ src/domain/peOptionalHeader.ts | 82 ++++++++ src/native/exports/NativeExportParser.ts | 10 +- src/native/metadata/BinaryContainer.ts | 7 +- src/native/metadata/NativeSectionDirectory.ts | 32 ++-- src/native/metadata/PeImage.ts | 4 +- src/native/pe/PeCertificateParser.ts | 29 +-- src/native/pe/PeImportParser.ts | 14 +- src/native/pe/PeLoadConfigParser.ts | 2 +- src/native/pe/PeResourceLayout.ts | 74 ++++--- tests/boundary/gpu/codeObjects.test.ts | 26 +++ 12 files changed, 394 insertions(+), 119 deletions(-) create mode 100644 src/domain/peOptionalHeader.test.ts create mode 100644 src/domain/peOptionalHeader.ts diff --git a/src/domain/peMachine.ts b/src/domain/peMachine.ts index a08432d2f5..85072455dd 100644 --- a/src/domain/peMachine.ts +++ b/src/domain/peMachine.ts @@ -1,4 +1,8 @@ import type { HeaderArchitecture } from "./binaryTargetTypes.js"; +import { + readPeOptionalHeaderStructure, + type PeOptionalHeaderStructure, +} from "./peOptionalHeader.js"; import { err, ok, type Result } from "./result.js"; export type PeArchitecture = Exclude; @@ -74,14 +78,9 @@ export const peSubjectArchitecture = ( }; /** Optional-header facts shared by admission and native metadata readers. */ -export interface PeOptionalHeaderFacts { - readonly bits: 32 | 64; +export interface PeOptionalHeaderFacts extends PeOptionalHeaderStructure { readonly machine: PeMachine | null; readonly managed: boolean; - readonly dataDirectories: readonly { - readonly virtual_address: number; - readonly size: number; - }[]; } export interface PeHeaderFailure { @@ -94,43 +93,24 @@ export const readPeOptionalHeader = ( optionalHeader: Buffer, machineType: number, ): Result => { - if (optionalHeader.length < 2) - return err({ kind: "malformed", message: "truncated PE optional header" }); - const magic = optionalHeader.readUInt16LE(0); - if (magic !== 0x10b && magic !== 0x20b) - return err({ - kind: "unsupported", - message: "Only PE32 and PE32+ image optional headers are supported", - }); - const bits = magic === 0x20b ? 64 : 32; - const directoryOffset = bits === 64 ? 112 : 96; - if (optionalHeader.length < directoryOffset) - return err({ kind: "malformed", message: "truncated PE optional header" }); - const directoryCount = optionalHeader.readUInt32LE(directoryOffset - 4); - if (directoryCount * 8 > optionalHeader.length - directoryOffset) - return err({ - kind: "malformed", - message: "truncated PE data-directory table", - }); - const dataDirectories = Array.from({ length: directoryCount }, (_, index) => { - const offset = directoryOffset + index * 8; - return { - virtual_address: optionalHeader.readUInt32LE(offset), - size: optionalHeader.readUInt32LE(offset + 4), - }; - }); - const cli = dataDirectories[14]; - const managed = - cli !== undefined && (cli.virtual_address !== 0 || cli.size !== 0); + const parsed = readPeOptionalHeaderStructure(optionalHeader); + if (!parsed.ok) + return err({ kind: parsed.error.kind, message: parsed.error.message }); + const facts = parsed.value; + const cli = facts.dataDirectories[14]; + const managed = cli !== undefined && (cli.address !== 0 || cli.size !== 0); const decoded = decodePeMachine(machineType); // An encoded machine without a CLI declaration remains an unknown CPU to // metadata readers; admission separately rejects unsupported identities. const machine = decoded?.nativeOsOverride && !managed ? null : (decoded ?? null); - if (machine !== null && magic !== peOptionalHeaderMagic(machine.architecture)) + if ( + machine !== null && + facts.magic !== peOptionalHeaderMagic(machine.architecture) + ) return err({ kind: "malformed", message: "PE optional-header magic does not match its architecture", }); - return ok({ bits, machine, managed, dataDirectories }); + return ok({ ...facts, machine, managed }); }; diff --git a/src/domain/peOptionalHeader.test.ts b/src/domain/peOptionalHeader.test.ts new file mode 100644 index 0000000000..c019995951 --- /dev/null +++ b/src/domain/peOptionalHeader.test.ts @@ -0,0 +1,181 @@ +import { expect, it } from "vitest"; +import { readPeOptionalHeader } from "./peMachine.js"; +import { readPeOptionalHeaderStructure } from "./peOptionalHeader.js"; + +const formats = [ + { format: "pe32", magic: 0x10b, bits: 32, directoryOffset: 96 }, + { format: "pe32-plus", magic: 0x20b, bits: 64, directoryOffset: 112 }, +] as const; + +const header = ( + format: (typeof formats)[number], + directoryCount: number, + byteLength = format.directoryOffset + directoryCount * 8, +): Buffer => { + const bytes = Buffer.alloc(byteLength); + bytes.writeUInt16LE(format.magic); + bytes.writeUInt32LE(0x400, 60); + bytes.writeUInt32LE(directoryCount, format.directoryOffset - 4); + return bytes; +}; + +it.each(formats)( + "copies $format declarations and raw directory addresses with their locations", + (format) => { + const bytes = header(format, 5); + bytes.writeUInt32LE(0x2000, format.directoryOffset); + bytes.writeUInt32LE(0x180, format.directoryOffset + 4); + const securityOffset = format.directoryOffset + 4 * 8; + bytes.writeUInt32LE(0x600, securityOffset); + bytes.writeUInt32LE(0x80, securityOffset + 4); + const result = readPeOptionalHeaderStructure(bytes); + if (!result.ok) throw new Error(result.error.message); + expect(result.value).toMatchObject({ + ...format, + sizeOfHeaders: 0x400, + directoryCount: 5, + }); + expect(result.value.dataDirectories).toEqual([ + { address: 0x2000, size: 0x180, offset: format.directoryOffset }, + { address: 0, size: 0, offset: format.directoryOffset + 8 }, + { address: 0, size: 0, offset: format.directoryOffset + 16 }, + { address: 0, size: 0, offset: format.directoryOffset + 24 }, + { address: 0x600, size: 0x80, offset: securityOffset }, + ]); + }, +); + +it.each(formats)( + "keeps $format zero directories and SizeOfHeaders as declarations", + (format) => { + const bytes = header(format, 0); + bytes.writeUInt32LE(0xffffffff, 60); + expect(readPeOptionalHeaderStructure(bytes)).toEqual({ + ok: true, + value: { + ...format, + sizeOfHeaders: 0xffffffff, + directoryCount: 0, + dataDirectories: [], + }, + }); + }, +); + +it.each([ + { + name: "missing magic", + bytes: new Uint8Array(), + kind: "malformed", + reason: "truncated-magic", + }, + { + name: "one magic byte", + bytes: new Uint8Array([0x0b]), + kind: "malformed", + reason: "truncated-magic", + }, + { + name: "unsupported magic before absent fixed fields", + bytes: new Uint8Array([0, 0]), + kind: "unsupported", + reason: "unsupported-magic", + }, + ...formats.map((format) => ({ + name: `truncated ${format.format} fixed fields`, + bytes: header(format, 0).subarray(0, format.directoryOffset - 1), + kind: "malformed", + reason: "truncated-fixed-fields", + })), +])("rejects $name with a structural failure tag", ({ bytes, kind, reason }) => { + expect(readPeOptionalHeaderStructure(bytes)).toMatchObject({ + ok: false, + error: { kind, reason, message: expect.any(String) }, + }); +}); + +it.each(formats)("rejects a partial $format directory entry", (format) => { + const bytes = header(format, 1).subarray(0, format.directoryOffset + 7); + expect(readPeOptionalHeaderStructure(bytes)).toEqual({ + ok: false, + error: { + kind: "malformed", + reason: "truncated-directory-table", + message: "truncated PE data-directory table", + }, + }); +}); + +it.each(formats)( + "refuses the maximum declared $format directory count without its bytes", + (format) => { + const bytes = header(format, 0); + bytes.writeUInt32LE(0xffffffff, format.directoryOffset - 4); + expect(readPeOptionalHeaderStructure(bytes)).toMatchObject({ + ok: false, + error: { kind: "malformed", reason: "truncated-directory-table" }, + }); + }, +); + +it.each(formats)( + "reads every $format directory that fits the maximum COFF optional-header extent", + (format) => { + const byteLength = 0xffff; + const directoryCount = Math.floor( + (byteLength - format.directoryOffset) / 8, + ); + const bytes = header(format, directoryCount, byteLength); + const offset = format.directoryOffset + (directoryCount - 1) * 8; + bytes.writeUInt32LE(0xffffffff, offset); + bytes.writeUInt32LE(0xfffffffe, offset + 4); + const result = readPeOptionalHeaderStructure(bytes); + if (!result.ok) throw new Error(result.error.message); + expect(result.value.directoryCount).toBe(directoryCount); + expect(result.value.dataDirectories).toHaveLength(directoryCount); + expect(result.value.dataDirectories.at(-1)).toEqual({ + address: 0xffffffff, + size: 0xfffffffe, + offset, + }); + }, +); + +it("honors a Uint8Array window and detaches facts from later source changes", () => { + const bytes = header(formats[1], 1); + bytes.writeUInt32LE(0x12345678, 112); + bytes.writeUInt32LE(0x90, 116); + const backing = Buffer.alloc(bytes.length + 12, 0xcc); + bytes.copy(backing, 5); + const window = new Uint8Array( + backing.buffer, + backing.byteOffset + 5, + bytes.length, + ); + const result = readPeOptionalHeaderStructure(window); + if (!result.ok) throw new Error(result.error.message); + backing.fill(0); + expect(result.value).toEqual({ + ...formats[1], + sizeOfHeaders: 0x400, + directoryCount: 1, + dataDirectories: [{ address: 0x12345678, size: 0x90, offset: 112 }], + }); +}); + +it("keeps structural refusal ahead of conflicting machine width", () => { + const bytes = header(formats[0], 0); + bytes.writeUInt32LE(1, 92); + expect(readPeOptionalHeader(bytes, 0x8664)).toEqual({ + ok: false, + error: { kind: "malformed", message: "truncated PE data-directory table" }, + }); + bytes.writeUInt32LE(0, 92); + expect(readPeOptionalHeader(bytes, 0x8664)).toEqual({ + ok: false, + error: { + kind: "malformed", + message: "PE optional-header magic does not match its architecture", + }, + }); +}); diff --git a/src/domain/peOptionalHeader.ts b/src/domain/peOptionalHeader.ts new file mode 100644 index 0000000000..9d27350468 --- /dev/null +++ b/src/domain/peOptionalHeader.ts @@ -0,0 +1,82 @@ +import { err, ok, type Result } from "./result.js"; + +/** Raw data-directory words; entry 4 names file bytes, while resource entry 2 names an RVA. */ +export interface PeDataDirectory { + readonly address: number; + readonly size: number; + /** Byte offset of this entry relative to the optional header. */ + readonly offset: number; +} + +/** Structural facts copied from one declared PE optional header. */ +export interface PeOptionalHeaderStructure { + readonly magic: 0x10b | 0x20b; + readonly bits: 32 | 64; + readonly format: "pe32" | "pe32-plus"; + readonly sizeOfHeaders: number; + readonly directoryOffset: number; + readonly directoryCount: number; + readonly dataDirectories: readonly PeDataDirectory[]; +} + +export interface PeOptionalHeaderFailure { + readonly kind: "malformed" | "unsupported"; + readonly reason: + | "truncated-magic" + | "unsupported-magic" + | "truncated-fixed-fields" + | "truncated-directory-table"; + readonly message: string; +} + +/** Decode structural declarations without admitting a machine, file extent, or directory address. */ +export const readPeOptionalHeaderStructure = ( + bytes: Uint8Array, +): Result => { + if (bytes.byteLength < 2) + return err({ + kind: "malformed", + reason: "truncated-magic", + message: "truncated PE optional header", + }); + const view = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength); + const magic = view.getUint16(0, true); + if (magic !== 0x10b && magic !== 0x20b) + return err({ + kind: "unsupported", + reason: "unsupported-magic", + message: "Only PE32 and PE32+ image optional headers are supported", + }); + const bits = magic === 0x20b ? 64 : 32; + const directoryOffset = bits === 64 ? 112 : 96; + if (bytes.byteLength < directoryOffset) + return err({ + kind: "malformed", + reason: "truncated-fixed-fields", + message: "truncated PE optional header", + }); + const directoryCount = view.getUint32(directoryOffset - 4, true); + if (directoryCount > Math.floor((bytes.byteLength - directoryOffset) / 8)) + return err({ + kind: "malformed", + reason: "truncated-directory-table", + message: "truncated PE data-directory table", + }); + const dataDirectories = Array.from({ length: directoryCount }, (_, index) => { + const offset = directoryOffset + index * 8; + return { + address: view.getUint32(offset, true), + size: view.getUint32(offset + 4, true), + offset, + }; + }); + return ok({ + magic, + bits, + format: bits === 64 ? "pe32-plus" : "pe32", + sizeOfHeaders: view.getUint32(60, true), + directoryOffset, + directoryCount, + dataDirectories, + }); +}; diff --git a/src/native/exports/NativeExportParser.ts b/src/native/exports/NativeExportParser.ts index dfe49d776b..20fcb18663 100644 --- a/src/native/exports/NativeExportParser.ts +++ b/src/native/exports/NativeExportParser.ts @@ -98,7 +98,7 @@ const parsePeExports = (bytes: Buffer, artifact: Artifact): NativeExports => { architecture: image.architecture, bits: image.bits, }; - if (directory === undefined || directory.virtual_address === 0) + if (directory === undefined || directory.address === 0) return { ...result, export_directory: null, @@ -129,7 +129,7 @@ const parsePeExports = (bytes: Buffer, artifact: Artifact): NativeExports => { return value.toString("latin1"); }; - const header = map(directory.virtual_address, PE_EXPORT_DIRECTORY_BYTES); + const header = map(directory.address, PE_EXPORT_DIRECTORY_BYTES); const ordinalBase = reader.u32(header + 16); const functionCount = reader.u32(header + 20); const nameCount = reader.u32(header + 24); @@ -180,7 +180,7 @@ const parsePeExports = (bytes: Buffer, artifact: Artifact): NativeExports => { } } - const forwarderEnd = directory.virtual_address + directory.size; + const forwarderEnd = directory.address + directory.size; const retain = createRetention(); const exports: PeExport[] = []; for (let index = 0; index < functionCount; index++) { @@ -189,7 +189,7 @@ const parsePeExports = (bytes: Buffer, artifact: Artifact): NativeExports => { if (rva === 0) continue; const ordinal = ordinalBase + index; let target: Pick; - if (rva >= directory.virtual_address && rva < forwarderEnd) { + if (rva >= directory.address && rva < forwarderEnd) { try { target = { address: null, forwarder: cString(rva) }; } catch (cause) { @@ -216,7 +216,7 @@ const parsePeExports = (bytes: Buffer, artifact: Artifact): NativeExports => { return { ...result, export_directory: { - rva: directory.virtual_address, + rva: directory.address, size: directory.size, dll_name: dllName, ordinal_base: ordinalBase, diff --git a/src/native/metadata/BinaryContainer.ts b/src/native/metadata/BinaryContainer.ts index c7c92e4818..dad27cfde8 100644 --- a/src/native/metadata/BinaryContainer.ts +++ b/src/native/metadata/BinaryContainer.ts @@ -1,3 +1,5 @@ +import type { PeDataDirectory } from "../../domain/peOptionalHeader.js"; + /** Distinguish corrupt recognized binaries from containers outside parser coverage. */ export class BinaryFormatFailure extends Error { constructor( @@ -120,10 +122,7 @@ interface BinaryContainerFacts { */ readonly pe_header_mapping?: FileMapping; /** PE optional-header data directories, indexed as declared. */ - readonly pe_data_directories?: readonly { - readonly virtual_address: number; - readonly size: number; - }[]; + readonly pe_data_directories?: readonly PeDataDirectory[]; /** Mach-O load-command observation, when available. */ readonly macho_chained_fixups?: boolean; /** Section-name lookup was incomplete; absence by name is not established. */ diff --git a/src/native/metadata/NativeSectionDirectory.ts b/src/native/metadata/NativeSectionDirectory.ts index 6bff184571..85c4f9f4cc 100644 --- a/src/native/metadata/NativeSectionDirectory.ts +++ b/src/native/metadata/NativeSectionDirectory.ts @@ -1,4 +1,5 @@ import { readElfTableDeclarations } from "../../domain/elfDeclaredRanges.js"; +import { readPeOptionalHeaderStructure } from "../../domain/peOptionalHeader.js"; import { BinaryFormatFailure, BinaryReader, @@ -286,23 +287,24 @@ export function* readPeSectionDirectory( const coff = new BinaryReader(yield request(fileSize, BigInt(pe), 24n), true); if (coff.u32(0) !== 0x4550) reject("PE signature is invalid"); const optionalSize = coff.u16(20); - const optional = new BinaryReader( + const optional = readPeOptionalHeaderStructure( yield request(fileSize, BigInt(pe + 24), BigInt(optionalSize)), - true, ); - const magic = optional.u16(0); - const bits = magic === 0x10b ? 32 : magic === 0x20b ? 64 : null; - if (bits === null) - throw new BinaryFormatFailure( - "unsupported", - "Only PE32/PE32+ images are supported", - ); - const minimum = bits === 64 ? 112 : 96; - if ( - optionalSize < minimum || - optional.u32(minimum - 4) * 8 > optionalSize - minimum - ) - reject("PE optional header or data directories are truncated"); + if (!optional.ok) { + switch (optional.error.reason) { + case "truncated-magic": + return reject("16-bit field range is truncated or outside the file"); + case "unsupported-magic": + throw new BinaryFormatFailure( + "unsupported", + "Only PE32/PE32+ images are supported", + ); + case "truncated-fixed-fields": + case "truncated-directory-table": + return reject("PE optional header or data directories are truncated"); + } + } + const { bits } = optional.value; const count = coff.u16(6); const start = pe + 24 + optionalSize; const table = yield request(fileSize, BigInt(start), BigInt(count) * 40n); diff --git a/src/native/metadata/PeImage.ts b/src/native/metadata/PeImage.ts index 285da0ec1f..6524b37068 100644 --- a/src/native/metadata/PeImage.ts +++ b/src/native/metadata/PeImage.ts @@ -82,13 +82,13 @@ export const readPeImage = (bytes: Buffer): PeImage => { ); if (!facts.ok) throw new BinaryFormatFailure(facts.error.kind, facts.error.message); - const { bits, dataDirectories } = facts.value; + const { bits, dataDirectories, sizeOfHeaders } = facts.value; const imageBase = reader.word(optional + (bits === 64 ? 24 : 28), bits); const fileAlignment = reader.u32(optional + 36); const headerMapping: FileMapping = { address: imageBase, offset: 0, - size: Math.min(reader.u32(optional + 60), bytes.length), + size: Math.min(sizeOfHeaders, bytes.length), }; const count = reader.u16(pe + 6); const start = reader.table( diff --git a/src/native/pe/PeCertificateParser.ts b/src/native/pe/PeCertificateParser.ts index 4b55e2f61e..9806626bd3 100644 --- a/src/native/pe/PeCertificateParser.ts +++ b/src/native/pe/PeCertificateParser.ts @@ -54,27 +54,13 @@ export const parsePeCertificates = ( artifact: PeCertificates["artifact"], input: InspectPeCertificatesInput, ): PeCertificates => { - const { pe, optional, magic, directoriesAt, directoryCount, headerSize } = + const { machine, format, dataDirectories, headerSize, optionalHeaderOffset } = readPeHeaders(bytes); - const machine = bytes.readUInt16LE(pe + 4); - const format = magic === 0x20b ? "pe32-plus" : "pe32"; - const securityIndex = 4; - if (directoryCount <= securityIndex) { - return { - artifact, - format, - machine, - table: null, - entries: [], - coverage: { status: "complete", examined_entries: 0, certificates: 0 }, - limitations: [...ABSENT_LIMITATIONS], - }; - } - const directoryAt = optional + directoriesAt + securityIndex * 8; - requireRange(bytes, directoryAt, 8); - const tableOffset = bytes.readUInt32LE(directoryAt); - const tableSize = bytes.readUInt32LE(directoryAt + 4); - if (tableOffset === 0 && tableSize === 0) { + const security = dataDirectories[4]; + if ( + security === undefined || + (security.address === 0 && security.size === 0) + ) { return { artifact, format, @@ -85,6 +71,9 @@ export const parsePeCertificates = ( limitations: [...ABSENT_LIMITATIONS], }; } + const tableOffset = security.address; + const tableSize = security.size; + const directoryAt = optionalHeaderOffset + security.offset; if (tableSize === 0 || tableOffset < headerSize) fail("PE certificate table offset overlaps headers or has zero size."); requireRange(bytes, tableOffset, tableSize); diff --git a/src/native/pe/PeImportParser.ts b/src/native/pe/PeImportParser.ts index fa801d135e..83fc3a76a3 100644 --- a/src/native/pe/PeImportParser.ts +++ b/src/native/pe/PeImportParser.ts @@ -389,12 +389,12 @@ export const parsePeImports = async ( const delayDirectory = image.pe_data_directories?.[PE_DELAY_IMPORT_DIRECTORY]; const imports: PeImportDll[] = []; let importDirectoryResult: PeImports["import_directory"] = null; - if (importDirectory !== undefined && importDirectory.virtual_address !== 0) { + if (importDirectory !== undefined && importDirectory.address !== 0) { const spans: { offset: number; bytes: number }[] = []; for (let index = 0; ; index++) { await checkpoint(); const descriptorRva = - importDirectory.virtual_address + index * IMPORT_DESCRIPTOR_BYTES; + importDirectory.address + index * IMPORT_DESCRIPTOR_BYTES; const descriptor = mapDescriptor( descriptorRva, IMPORT_DESCRIPTOR_BYTES, @@ -453,7 +453,7 @@ export const parsePeImports = async ( }); } importDirectoryResult = { - rva: importDirectory.virtual_address, + rva: importDirectory.address, size: importDirectory.size, locations: spans, }; @@ -463,12 +463,12 @@ export const parsePeImports = async ( const delayImports: PeDelayImportDll[] = []; let delayDirectoryResult: PeImports["delay_import_directory"] = null; - if (delayDirectory !== undefined && delayDirectory.virtual_address !== 0) { + if (delayDirectory !== undefined && delayDirectory.address !== 0) { const spans: { offset: number; bytes: number }[] = []; for (let index = 0; ; index++) { await checkpoint(); const descriptorRva = - delayDirectory.virtual_address + index * DELAY_DESCRIPTOR_BYTES; + delayDirectory.address + index * DELAY_DESCRIPTOR_BYTES; const descriptor = mapDescriptor( descriptorRva, DELAY_DESCRIPTOR_BYTES, @@ -534,7 +534,7 @@ export const parsePeImports = async ( }); } delayDirectoryResult = { - rva: delayDirectory.virtual_address, + rva: delayDirectory.address, size: delayDirectory.size, locations: spans, }; @@ -559,7 +559,7 @@ export const parsePeImports = async ( mappingCompatibility: standardMappingCompatibility, truncatedDescriptor: pefileTruncatedDescriptor, bits, - directoryRva: importDirectory?.virtual_address ?? 0, + directoryRva: importDirectory?.address ?? 0, fileBytes: bytes.length, }), coverage: coverage( diff --git a/src/native/pe/PeLoadConfigParser.ts b/src/native/pe/PeLoadConfigParser.ts index 3dd24b1903..18663530e1 100644 --- a/src/native/pe/PeLoadConfigParser.ts +++ b/src/native/pe/PeLoadConfigParser.ts @@ -164,7 +164,7 @@ export const parsePeLoadConfig = async ( const dataDirectory = (index: number): DirectorySpan | null => { const entry = image.pe_data_directories[index]; if (entry === undefined) return null; - const { virtual_address: rva, size } = entry; + const { address: rva, size } = entry; if (rva === 0 && size === 0) return null; if (rva === 0 || size === 0) throw new BinaryFormatFailure( diff --git a/src/native/pe/PeResourceLayout.ts b/src/native/pe/PeResourceLayout.ts index bf91866ebf..b8fd425031 100644 --- a/src/native/pe/PeResourceLayout.ts +++ b/src/native/pe/PeResourceLayout.ts @@ -1,4 +1,18 @@ import { ArtifactReaderFailure } from "../../artifacts/ArtifactReader.js"; +import { + readPeOptionalHeaderStructure, + type PeOptionalHeaderFailure, +} from "../../domain/peOptionalHeader.js"; + +const OPTIONAL_HEADER_FAILURES: Readonly< + Record +> = { + "truncated-magic": "Truncated PE optional header.", + "unsupported-magic": "Unsupported PE optional header magic.", + "truncated-fixed-fields": "Truncated PE data directory count.", + "truncated-directory-table": + "PE directory count exceeds its optional header.", +}; export const peFailure = (message: string): never => { throw new ArtifactReaderFailure("format", message); @@ -71,31 +85,31 @@ export const readPeHeaders = (bytes: Buffer, fileSize = bytes.length) => { if (bytes.readUInt32LE(pe) !== 0x4550) peFailure("Expected a PE signature."); const sectionCount = bytes.readUInt16LE(pe + 6); const optionalSize = bytes.readUInt16LE(pe + 20); - const optional = pe + 24; - requirePeRange(bytes, optional, optionalSize); + const optionalHeaderOffset = pe + 24; + requirePeRange(bytes, optionalHeaderOffset, optionalSize); if (optionalSize < 64) peFailure("Truncated PE optional header."); - const magic = bytes.readUInt16LE(optional); - if (magic !== 0x10b && magic !== 0x20b) - peFailure("Unsupported PE optional header magic."); - const countAt = magic === 0x10b ? 92 : 108; - const directoriesAt = countAt + 4; - if (optionalSize < directoriesAt) - peFailure("Truncated PE data directory count."); - const directoryCount = bytes.readUInt32LE(optional + countAt); - if (directoryCount > Math.floor((optionalSize - directoriesAt) / 8)) - peFailure("PE directory count exceeds its optional header."); - const headerSize = bytes.readUInt32LE(optional + 60); - const table = optional + optionalSize; + const parsed = readPeOptionalHeaderStructure( + bytes.subarray(optionalHeaderOffset, optionalHeaderOffset + optionalSize), + ); + if (!parsed.ok) + return peFailure(OPTIONAL_HEADER_FAILURES[parsed.error.reason]); + const { + sizeOfHeaders: headerSize, + format, + directoryOffset, + dataDirectories, + } = parsed.value; + const table = optionalHeaderOffset + optionalSize; requirePeRange(bytes, table, sectionCount * 40); if (headerSize < table + sectionCount * 40 || headerSize > fileSize) peFailure("Invalid PE SizeOfHeaders."); return { - pe, + machine: bytes.readUInt16LE(pe + 4), + format, sectionCount, - optional, - magic, - directoriesAt, - directoryCount, + optionalHeaderOffset, + directoryOffset, + dataDirectories, headerSize, table, }; @@ -107,12 +121,12 @@ export const readPeResourceLayout = ( fileSize = bytes.length, ) => { const { - pe, + machine, + format, sectionCount, - optional, - magic, - directoriesAt, - directoryCount, + optionalHeaderOffset, + directoryOffset, + dataDirectories, headerSize, table, } = readPeHeaders(bytes, fileSize); @@ -194,14 +208,16 @@ export const readPeResourceLayout = ( requirePeExtent(fileSize, offset, size); return offset; }; - const directoryAt = optional + directoriesAt + 16; - const rva = directoryCount > 2 ? bytes.readUInt32LE(directoryAt) : 0; - const size = directoryCount > 2 ? bytes.readUInt32LE(directoryAt + 4) : 0; + const resource = dataDirectories[2]; + const directoryAt = + optionalHeaderOffset + (resource?.offset ?? directoryOffset + 2 * 8); + const rva = resource?.address ?? 0; + const size = resource?.size ?? 0; if ((rva === 0) !== (size === 0)) peFailure("Inconsistent PE resource directory RVA/size."); return { - format: magic === 0x10b ? ("pe32" as const) : ("pe32-plus" as const), - machine: bytes.readUInt16LE(pe + 4), + format, + machine, directoryAt, rva, size, diff --git a/tests/boundary/gpu/codeObjects.test.ts b/tests/boundary/gpu/codeObjects.test.ts index fa51746f17..cf1f8c9071 100644 --- a/tests/boundary/gpu/codeObjects.test.ts +++ b/tests/boundary/gpu/codeObjects.test.ts @@ -1079,6 +1079,32 @@ it("retains an embedded CCOB with a corrupted signature as unsupported", async ( }), ); }); +it.each([ + ["unknown COFF machine", 0xffff], + ["x86 machine with a PE32+ optional header", 0x14c], +] as const)( + "discovers PE GPU sections with %s", + async (_description, machine) => { + const bytes = peGpuFixture(); + const pe = bytes.readUInt32LE(60); + bytes.writeUInt16LE(machine, pe + 4); + const result = await inspect(bytes); + if (!result.ok) throw result.error; + expect(result.value.format).toBe("pe"); + expect(result.value.artifact.sha256).toBe( + createHash("sha256").update(bytes).digest("hex"), + ); + expect(result.value.coverage).toMatchObject({ + discovery: "complete", + payloads: "complete", + }); + expect(result.value.kernels.map(({ name }) => name)).toEqual([ + "add_one", + "add_one", + ]); + expect(result.value.ptx_artifacts[0]!.text).toBe(ptxFixture); + }, +); it("does not dispatch a high-bit-corrupted DOS signature to the PE parser", async () => { const bytes = peGpuFixture(Buffer.from(ptxFixture)); bytes[0] = bytes[0]! | 0x80;