Improved Windows support

src/os/windows/pe/EXE.go

@@ -1,6 +1,7 @@
 package pe
 
 import (
+	"bytes"
 	"encoding/binary"
 	"io"
 
@@ -8,53 +9,58 @@ import (
 	"git.akyoto.dev/cli/q/src/os/common"
 )
 
+const NumSections = 2
+
 // EXE is the portable executable format used on Windows.
 type EXE struct {
 	DOSHeader
-	NTHeader
+	PEHeader
 	OptionalHeader64
-	CodeHeader SectionHeader
-	Code       []byte
-	Data       []byte
+	Sections    [NumSections]SectionHeader
+	CodePadding []byte
+	Code        []byte
+	DataPadding []byte
+	Data        []byte
 }
 
 // New creates a new EXE file.
 func New(code []byte, data []byte) *EXE {
-	const codeStart = 0x170
-	const optHeaderSize = 0xF0
+	codeStart := uint32(DOSHeaderSize + PEHeaderSize + OptionalHeader64Size + SectionHeaderSize*NumSections)
+	codePadding := common.Padding(codeStart, Align)
+	codeStart += codePadding
 
-	codeSize := uint32(len(code))
-	headerSize := uint32(codeStart)
-	sectionAlign := uint32(0x10)
-	fileAlign := uint32(0x10)
-	imageSize := uint32(codeStart + len(code))
-	imageSize += common.Padding(imageSize, sectionAlign)
+	dataStart := codeStart + uint32(len(code))
+	dataPadding := common.Padding(dataStart, Align)
+	dataStart += dataPadding
+
+	imageSize := uint32(dataStart + uint32(len(data)))
+	imageSize += common.Padding(imageSize, Align)
 
 	return &EXE{
 		DOSHeader: DOSHeader{
 			Magic:          [4]byte{'M', 'Z', 0, 0},
-			NTHeaderOffset: 0x40,
+			PEHeaderOffset: 0x40,
 		},
-		NTHeader: NTHeader{
+		PEHeader: PEHeader{
 			Signature:            [4]byte{'P', 'E', 0, 0},
 			Machine:              IMAGE_FILE_MACHINE_AMD64,
-			NumberOfSections:     1,
-			SizeOfOptionalHeader: optHeaderSize,
+			NumberOfSections:     NumSections,
+			SizeOfOptionalHeader: OptionalHeader64Size,
 			Characteristics:      IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE,
 		},
 		OptionalHeader64: OptionalHeader64{
 			Magic:                       0x020B, // PE32+ executable
 			MajorLinkerVersion:          0x0E,
 			MinorLinkerVersion:          0x16,
-			SizeOfCode:                  codeSize,
+			SizeOfCode:                  uint32(len(code)),
 			AddressOfEntryPoint:         codeStart,
 			ImageBase:                   config.BaseAddress,
-			SectionAlignment:            sectionAlign, // power of 2, must be greater than or equal to FileAlignment
-			FileAlignment:               fileAlign,    // power of 2
+			SectionAlignment:            Align, // power of 2, must be greater than or equal to FileAlignment
+			FileAlignment:               Align, // power of 2
 			MajorOperatingSystemVersion: 0x06,
 			MajorSubsystemVersion:       0x06,
 			SizeOfImage:                 imageSize,
-			SizeOfHeaders:               headerSize,
+			SizeOfHeaders:               codeStart, // section bodies begin here
 			Subsystem:                   IMAGE_SUBSYSTEM_WINDOWS_GUI,
 			DllCharacteristics:          IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA | IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE | IMAGE_DLLCHARACTERISTICS_NX_COMPAT | IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE,
 			SizeOfStackReserve:          0x100000,
@@ -81,25 +87,39 @@ func New(code []byte, data []byte) *EXE {
 				{VirtualAddress: 0, Size: 0},
 			},
 		},
-		CodeHeader: SectionHeader{
-			Name:            [8]byte{'.', 't', 'e', 'x', 't'},
-			VirtualSize:     uint32(len(code)),
-			VirtualAddress:  codeStart,
-			RawSize:         uint32(len(code)), // must be a multiple of FileAlignment
-			RawAddress:      codeStart,         // must be a multiple of FileAlignment
-			Characteristics: IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
+		Sections: [NumSections]SectionHeader{
+			{
+				Name:            [8]byte{'.', 'c', 'o', 'd', 'e'},
+				VirtualSize:     uint32(len(code)),
+				VirtualAddress:  codeStart,
+				RawSize:         uint32(len(code)), // must be a multiple of FileAlignment
+				RawAddress:      codeStart,         // must be a multiple of FileAlignment
+				Characteristics: IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
+			},
+			{
+				Name:            [8]byte{'.', 'd', 'a', 't', 'a'},
+				VirtualSize:     uint32(len(data)),
+				VirtualAddress:  dataStart,
+				RawSize:         uint32(len(data)), // must be a multiple of FileAlignment
+				RawAddress:      dataStart,         // must be a multiple of FileAlignment
+				Characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ,
+			},
 		},
-		Code: code,
-		// Data: data,
+		CodePadding: bytes.Repeat([]byte{0}, int(codePadding)),
+		Code:        code,
+		DataPadding: bytes.Repeat([]byte{0}, int(dataPadding)),
+		Data:        data,
 	}
 }
 
 // Write writes the EXE file to the given writer.
 func (pe *EXE) Write(writer io.Writer) {
 	binary.Write(writer, binary.LittleEndian, &pe.DOSHeader)
-	binary.Write(writer, binary.LittleEndian, &pe.NTHeader)
+	binary.Write(writer, binary.LittleEndian, &pe.PEHeader)
 	binary.Write(writer, binary.LittleEndian, &pe.OptionalHeader64)
-	binary.Write(writer, binary.LittleEndian, &pe.CodeHeader)
+	binary.Write(writer, binary.LittleEndian, &pe.Sections)
+	binary.Write(writer, binary.LittleEndian, &pe.CodePadding)
 	binary.Write(writer, binary.LittleEndian, &pe.Code)
-	// binary.Write(writer, binary.LittleEndian, &pe.Data)
+	binary.Write(writer, binary.LittleEndian, &pe.DataPadding)
+	binary.Write(writer, binary.LittleEndian, &pe.Data)
 }