WIP: Overlapping headers

file.go

@@ -13,8 +13,9 @@ type PE struct {
 	DOSStub          *Segment
 	PESignature      *Header[PESignature]
 	COFFHeader       *Header[pe.FileHeader]
-	OptionalHeader32 *Header[pe.OptionalHeader32]
-	OptionalHeader64 *Header[pe.OptionalHeader64]
+	OptionalHeader32 *Header[OptionalHeader32]
+	OptionalHeader64 *Header[OptionalHeader64]
+	DataDirectories  []Header[pe.DataDirectory]
 }
 
 // NewPE parses a Portable Executable or plain COFF file from reader.
@@ -68,6 +69,8 @@ func NewPE(reader io.ReaderAt) (*PE, error) {
 	optionalHeaderSize := p.COFFHeader.Data.SizeOfOptionalHeader
 	if optionalHeaderSize > 0 {
 		var magicBytes [2]byte
+		var numDirs uint32
+		var baseSize int64
 
 		// Read the magic number to determine if it's PE32 or PE32+.
 		_, err = reader.ReadAt(magicBytes[:], offset)
@@ -78,9 +81,17 @@ func NewPE(reader io.ReaderAt) (*PE, error) {
 		magic := binary.LittleEndian.Uint16(magicBytes[:])
 		switch magic {
 		case PE32Magic:
-			p.OptionalHeader32, err = NewHeader[pe.OptionalHeader32](reader, &offset)
+			p.OptionalHeader32, err = NewHeader[OptionalHeader32](reader, &offset)
+			if err == nil {
+				numDirs = p.OptionalHeader32.Data.NumberOfRvaAndSizes
+				baseSize = p.OptionalHeader32.Size()
+			}
 		case PE32PlusMagic:
-			p.OptionalHeader64, err = NewHeader[pe.OptionalHeader64](reader, &offset)
+			p.OptionalHeader64, err = NewHeader[OptionalHeader64](reader, &offset)
+			if err == nil {
+				numDirs = p.OptionalHeader64.Data.NumberOfRvaAndSizes
+				baseSize = p.OptionalHeader64.Size()
+			}
 		default:
 			err = fmt.Errorf("invalid optional header magic: %#x", magic)
 		}
@@ -89,15 +100,9 @@ func NewPE(reader io.ReaderAt) (*PE, error) {
 			return nil, err
 		}
 
-		var expectedSize uint16
-		if p.OptionalHeader32 != nil {
-			expectedSize = uint16(p.OptionalHeader32.Size())
-		} else {
-			expectedSize = uint16(p.OptionalHeader64.Size())
-		}
-
-		if optionalHeaderSize != expectedSize {
-			return nil, fmt.Errorf("optional header size does not match the expected size: %#x != %#x", optionalHeaderSize, expectedSize)
+		p.DataDirectories, err = parseDataDirectories(reader, &offset, numDirs, int64(optionalHeaderSize), baseSize)
+		if err != nil {
+			return nil, err
 		}
 	}
 
@@ -141,3 +146,22 @@ func isValidMachine(machine uint16) bool {
 	}
 	return false
 }
+
+func parseDataDirectories(reader io.ReaderAt, offset *int64, numDirs uint32, size int64, baseSize int64) ([]Header[pe.DataDirectory], error) {
+	if numDirs > 16 {
+		return nil, fmt.Errorf("NumberOfRvaAndSizes exceeds maximum: %d > %d", numDirs, 16)
+	}
+	if size < baseSize+int64(numDirs)*8 {
+		return nil, fmt.Errorf("optional header too small for NumberOfRvaAndSizes: %#x < %#x", size, baseSize+int64(numDirs)*8)
+	}
+	dirs := make([]Header[pe.DataDirectory], numDirs)
+	for i := range dirs {
+		dir, err := NewHeader[pe.DataDirectory](reader, offset)
+		if err != nil {
+			return nil, err
+		}
+		dirs[i] = *dir
+	}
+	*offset += size - baseSize - int64(numDirs)*8
+	return dirs, nil
+}

file_test.go

@@ -39,7 +39,31 @@ func TestFilePianoExe(t *testing.T) {
 	assert.Assert(t, p.OptionalHeader64 == nil)
 }
 
-// TestFileReadTruncated verifies that NewPE returns an error when the input is truncated at various points.
+func TestFile268Exe(t *testing.T) {
+	// 268.exe is a minimal 268-byte EXE where the PE structure starts inside the DOS header
+	// area (Lfanew = 4). This is a degenerate but valid technique used by size-optimized executables.
+	p, err := NewPE(openTestFile(t, "268.exe"))
+	assert.NilError(t, err)
+
+	// DOS Header.
+	assert.Assert(t, p.DOSHeader != nil)
+	assert.Equal(t, p.DOSHeader.Data.Lfanew, uint32(0x04))
+
+	// No DOS stub since Lfanew points inside the DOS header.
+	assert.Assert(t, p.DOSStub == nil)
+
+	// PE Signature.
+	assert.Assert(t, p.PESignature != nil)
+
+	// COFF Header.
+	assert.Equal(t, p.COFFHeader.Data.Machine, uint16(pe.IMAGE_FILE_MACHINE_AMD64))
+	assert.Equal(t, p.COFFHeader.Data.NumberOfSections, uint16(1))
+
+	// Optional Header (PE32+).
+	assert.Assert(t, p.OptionalHeader64 != nil)
+	assert.Assert(t, p.OptionalHeader32 == nil)
+}
+
 func TestFileReadTruncated(t *testing.T) {
 	tests := []struct {
 		name string

headers.go

@@ -59,6 +59,73 @@ type DOSHeader struct {
 // PESignature is the PE file signature type.
 type PESignature uint32
 
+// OptionalHeader32 contains the PE32 optional header fields, excluding data directories.
+type OptionalHeader32 struct {
+	Magic                       uint16
+	MajorLinkerVersion          uint8
+	MinorLinkerVersion          uint8
+	SizeOfCode                  uint32
+	SizeOfInitializedData       uint32
+	SizeOfUninitializedData     uint32
+	AddressOfEntryPoint         uint32
+	BaseOfCode                  uint32
+	BaseOfData                  uint32
+	ImageBase                   uint32
+	SectionAlignment            uint32
+	FileAlignment               uint32
+	MajorOperatingSystemVersion uint16
+	MinorOperatingSystemVersion uint16
+	MajorImageVersion           uint16
+	MinorImageVersion           uint16
+	MajorSubsystemVersion       uint16
+	MinorSubsystemVersion       uint16
+	Win32VersionValue           uint32
+	SizeOfImage                 uint32
+	SizeOfHeaders               uint32
+	CheckSum                    uint32
+	Subsystem                   uint16
+	DllCharacteristics          uint16
+	SizeOfStackReserve          uint32
+	SizeOfStackCommit           uint32
+	SizeOfHeapReserve           uint32
+	SizeOfHeapCommit            uint32
+	LoaderFlags                 uint32
+	NumberOfRvaAndSizes         uint32
+}
+
+// OptionalHeader64 contains the PE32+ optional header fields, excluding data directories.
+type OptionalHeader64 struct {
+	Magic                       uint16
+	MajorLinkerVersion          uint8
+	MinorLinkerVersion          uint8
+	SizeOfCode                  uint32
+	SizeOfInitializedData       uint32
+	SizeOfUninitializedData     uint32
+	AddressOfEntryPoint         uint32
+	BaseOfCode                  uint32
+	ImageBase                   uint64
+	SectionAlignment            uint32
+	FileAlignment               uint32
+	MajorOperatingSystemVersion uint16
+	MinorOperatingSystemVersion uint16
+	MajorImageVersion           uint16
+	MinorImageVersion           uint16
+	MajorSubsystemVersion       uint16
+	MinorSubsystemVersion       uint16
+	Win32VersionValue           uint32
+	SizeOfImage                 uint32
+	SizeOfHeaders               uint32
+	CheckSum                    uint32
+	Subsystem                   uint16
+	DllCharacteristics          uint16
+	SizeOfStackReserve          uint64
+	SizeOfStackCommit           uint64
+	SizeOfHeapReserve           uint64
+	SizeOfHeapCommit            uint64
+	LoaderFlags                 uint32
+	NumberOfRvaAndSizes         uint32
+}
+
 const (
 	DOSHeaderMagic   uint16      = 0x5a4d       // 'M', 'Z'
 	PESignatureMagic PESignature = 0x00004550   // 'P', 'E', 0, 0