debug/pe: enable parsing of variable length optional header in PE file

The debug/pe package assumes there are always 16 entries in
DataDirectory in OptionalHeader32/64
ref pe.go:
...
NumberOfRvaAndSizes uint32
DataDirectory [16]DataDirectory
}
...

But that is not always the case, there could be less no of
entries (PE signed linux kernel for example):
$ sudo pev /boot/vmlinuz-4.15.0-47-generic
....
Data-dictionary entries:	6
....

In such case, the parsing gives incorrect results.
This changes aims to fix that by:
1. Determining type of optional header by looking at header
   magic instead of size
2. Parsing optional header in 2 steps:
   a. Fixed part
   b. Variable data directories part

Testing:
1. Fixed existing test cases to reflect the change
2. Added new file (modified linux kernel image)
   which has smaller number of data directories

Fixes #32126

Change-Id: Iee56ecc4369a0e75a4be805e7cb8555c7d81ae2f
Reviewed-on: https://go-review.googlesource.com/c/go/+/177959
Run-TryBot: Alex Brainman <alex.brainman@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Alex Brainman <alex.brainman@gmail.com>

1 files changed, 173 insertions, 27 deletions

diff --git a/src/debug/pe/file.go b/src/debug/pe/file.go
index 58814162bc..14ad245224 100644
--- a/src/debug/pe/file.go
+++ b/src/debug/pe/file.go
@@ -58,11 +58,6 @@ func (f *File) Close() error {
 	return err
 }
 
-var (
-	sizeofOptionalHeader32 = uint16(binary.Size(OptionalHeader32{}))
-	sizeofOptionalHeader64 = uint16(binary.Size(OptionalHeader64{}))
-)
-
 // TODO(brainman): add Load function, as a replacement for NewFile, that does not call removeAuxSymbols (for performance)
 
 // NewFile creates a new File for accessing a PE binary in an underlying reader.
@@ -114,31 +109,17 @@ func NewFile(r io.ReaderAt) (*File, error) {
 		return nil, err
 	}
 
+	// Seek past file header.
+	_, err = sr.Seek(base+int64(binary.Size(f.FileHeader)), seekStart)
+	if err != nil {
+		return nil, fmt.Errorf("failure to seek past the file header: %v", err)
+	}
+
 	// Read optional header.
-	sr.Seek(base, seekStart)
-	if err := binary.Read(sr, binary.LittleEndian, &f.FileHeader); err != nil {
+	f.OptionalHeader, err = readOptionalHeader(sr, f.FileHeader.SizeOfOptionalHeader)
+	if err != nil {
 		return nil, err
 	}
-	var oh32 OptionalHeader32
-	var oh64 OptionalHeader64
-	switch f.FileHeader.SizeOfOptionalHeader {
-	case sizeofOptionalHeader32:
-		if err := binary.Read(sr, binary.LittleEndian, &oh32); err != nil {
-			return nil, err
-		}
-		if oh32.Magic != 0x10b { // PE32
-			return nil, fmt.Errorf("pe32 optional header has unexpected Magic of 0x%x", oh32.Magic)
-		}
-		f.OptionalHeader = &oh32
-	case sizeofOptionalHeader64:
-		if err := binary.Read(sr, binary.LittleEndian, &oh64); err != nil {
-			return nil, err
-		}
-		if oh64.Magic != 0x20b { // PE32+
-			return nil, fmt.Errorf("pe32+ optional header has unexpected Magic of 0x%x", oh64.Magic)
-		}
-		f.OptionalHeader = &oh64
-	}
 
 	// Process sections.
 	f.Sections = make([]*Section, f.FileHeader.NumberOfSections)
@@ -453,3 +434,168 @@ type FormatError struct {
 func (e *FormatError) Error() string {
 	return "unknown error"
 }
+
+// readOptionalHeader accepts a io.ReadSeeker pointing to optional header in the PE file
+// and its size as seen in the file header.
+// It parses the given size of bytes and returns optional header. It infers whether the
+// bytes being parsed refer to 32 bit or 64 bit version of optional header.
+func readOptionalHeader(r io.ReadSeeker, sz uint16) (interface{}, error) {
+	// If optional header size is 0, return empty optional header.
+	if sz == 0 {
+		return nil, nil
+	}
+
+	var (
+		// First couple of bytes in option header state its type.
+		// We need to read them first to determine the type and
+		// validity of optional header.
+		ohMagic   uint16
+		ohMagicSz = binary.Size(ohMagic)
+	)
+
+	// If optional header size is greater than 0 but less than its magic size, return error.
+	if sz < uint16(ohMagicSz) {
+		return nil, fmt.Errorf("optional header size is less than optional header magic size")
+	}
+
+	// read reads from io.ReadSeeke, r, into data.
+	var err error
+	read := func(data interface{}) bool {
+		err = binary.Read(r, binary.LittleEndian, data)
+		return err == nil
+	}
+
+	if !read(&ohMagic) {
+		return nil, fmt.Errorf("failure to read optional header magic: %v", err)
+
+	}
+
+	switch ohMagic {
+	case 0x10b: // PE32
+		var (
+			oh32 OptionalHeader32
+			// There can be 0 or more data directories. So the minimum size of optional
+			// header is calculated by substracting oh32.DataDirectory size from oh32 size.
+			oh32MinSz = binary.Size(oh32) - binary.Size(oh32.DataDirectory)
+		)
+
+		if sz < uint16(oh32MinSz) {
+			return nil, fmt.Errorf("optional header size(%d) is less minimum size (%d) of PE32 optional header", sz, oh32MinSz)
+		}
+
+		// Init oh32 fields
+		oh32.Magic = ohMagic
+		if !read(&oh32.MajorLinkerVersion) ||
+			!read(&oh32.MinorLinkerVersion) ||
+			!read(&oh32.SizeOfCode) ||
+			!read(&oh32.SizeOfInitializedData) ||
+			!read(&oh32.SizeOfUninitializedData) ||
+			!read(&oh32.AddressOfEntryPoint) ||
+			!read(&oh32.BaseOfCode) ||
+			!read(&oh32.BaseOfData) ||
+			!read(&oh32.ImageBase) ||
+			!read(&oh32.SectionAlignment) ||
+			!read(&oh32.FileAlignment) ||
+			!read(&oh32.MajorOperatingSystemVersion) ||
+			!read(&oh32.MinorOperatingSystemVersion) ||
+			!read(&oh32.MajorImageVersion) ||
+			!read(&oh32.MinorImageVersion) ||
+			!read(&oh32.MajorSubsystemVersion) ||
+			!read(&oh32.MinorSubsystemVersion) ||
+			!read(&oh32.Win32VersionValue) ||
+			!read(&oh32.SizeOfImage) ||
+			!read(&oh32.SizeOfHeaders) ||
+			!read(&oh32.CheckSum) ||
+			!read(&oh32.Subsystem) ||
+			!read(&oh32.DllCharacteristics) ||
+			!read(&oh32.SizeOfStackReserve) ||
+			!read(&oh32.SizeOfStackCommit) ||
+			!read(&oh32.SizeOfHeapReserve) ||
+			!read(&oh32.SizeOfHeapCommit) ||
+			!read(&oh32.LoaderFlags) ||
+			!read(&oh32.NumberOfRvaAndSizes) {
+			return nil, fmt.Errorf("failure to read PE32 optional header: %v", err)
+		}
+
+		dd, err := readDataDirectories(r, sz-uint16(oh32MinSz), oh32.NumberOfRvaAndSizes)
+		if err != nil {
+			return nil, err
+		}
+
+		copy(oh32.DataDirectory[:], dd)
+
+		return &oh32, nil
+	case 0x20b: // PE32+
+		var (
+			oh64 OptionalHeader64
+			// There can be 0 or more data directories. So the minimum size of optional
+			// header is calculated by substracting oh64.DataDirectory size from oh64 size.
+			oh64MinSz = binary.Size(oh64) - binary.Size(oh64.DataDirectory)
+		)
+
+		if sz < uint16(oh64MinSz) {
+			return nil, fmt.Errorf("optional header size(%d) is less minimum size (%d) for PE32+ optional header", sz, oh64MinSz)
+		}
+
+		// Init oh64 fields
+		oh64.Magic = ohMagic
+		if !read(&oh64.MajorLinkerVersion) ||
+			!read(&oh64.MinorLinkerVersion) ||
+			!read(&oh64.SizeOfCode) ||
+			!read(&oh64.SizeOfInitializedData) ||
+			!read(&oh64.SizeOfUninitializedData) ||
+			!read(&oh64.AddressOfEntryPoint) ||
+			!read(&oh64.BaseOfCode) ||
+			!read(&oh64.ImageBase) ||
+			!read(&oh64.SectionAlignment) ||
+			!read(&oh64.FileAlignment) ||
+			!read(&oh64.MajorOperatingSystemVersion) ||
+			!read(&oh64.MinorOperatingSystemVersion) ||
+			!read(&oh64.MajorImageVersion) ||
+			!read(&oh64.MinorImageVersion) ||
+			!read(&oh64.MajorSubsystemVersion) ||
+			!read(&oh64.MinorSubsystemVersion) ||
+			!read(&oh64.Win32VersionValue) ||
+			!read(&oh64.SizeOfImage) ||
+			!read(&oh64.SizeOfHeaders) ||
+			!read(&oh64.CheckSum) ||
+			!read(&oh64.Subsystem) ||
+			!read(&oh64.DllCharacteristics) ||
+			!read(&oh64.SizeOfStackReserve) ||
+			!read(&oh64.SizeOfStackCommit) ||
+			!read(&oh64.SizeOfHeapReserve) ||
+			!read(&oh64.SizeOfHeapCommit) ||
+			!read(&oh64.LoaderFlags) ||
+			!read(&oh64.NumberOfRvaAndSizes) {
+			return nil, fmt.Errorf("failure to read PE32+ optional header: %v", err)
+		}
+
+		dd, err := readDataDirectories(r, sz-uint16(oh64MinSz), oh64.NumberOfRvaAndSizes)
+		if err != nil {
+			return nil, err
+		}
+
+		copy(oh64.DataDirectory[:], dd)
+
+		return &oh64, nil
+	default:
+		return nil, fmt.Errorf("optional header has unexpected Magic of 0x%x", ohMagic)
+	}
+}
+
+// readDataDirectories accepts a io.ReadSeeker pointing to data directories in the PE file,
+// its size and number of data directories as seen in optional header.
+// It parses the given size of bytes and returns given number of data directories.
+func readDataDirectories(r io.ReadSeeker, sz uint16, n uint32) ([]DataDirectory, error) {
+	ddSz := binary.Size(DataDirectory{})
+	if uint32(sz) != n*uint32(ddSz) {
+		return nil, fmt.Errorf("size of data directories(%d) is inconsistent with number of data directories(%d)", sz, n)
+	}
+
+	dd := make([]DataDirectory, n)
+	if err := binary.Read(r, binary.LittleEndian, dd); err != nil {
+		return nil, fmt.Errorf("failure to read data directories: %v", err)
+	}
+
+	return dd, nil
+}