From 771b993dd8b1c9e4118bf9fd3dd5df31963d1005 Mon Sep 17 00:00:00 2001
From: Eduard Urbach <admin@akyoto.dev>
Date: Sat, 17 Aug 2024 15:34:42 +0200
Subject: [PATCH] Implemented loading of multiple DLLs

---
 src/compiler/Result.go | 19 ++++++++-
 src/exe/pe/DLL.go      |  6 +++
 src/exe/pe/EXE.go      | 87 ++++++++++++++++--------------------------
 src/exe/pe/EXE_test.go |  2 +-
 4 files changed, 57 insertions(+), 57 deletions(-)
 create mode 100644 src/exe/pe/DLL.go

diff --git a/src/compiler/Result.go b/src/compiler/Result.go
index 3cb48d1..b4b0979 100644
--- a/src/compiler/Result.go
+++ b/src/compiler/Result.go
@@ -154,7 +154,24 @@ func write(writer io.Writer, code []byte, data []byte) error {
 	case "mac":
 		macho.Write(buffer, code, data)
 	case "windows":
-		pe.Write(buffer, code, data)
+		dlls := []pe.DLL{
+			{
+				Name: "kernel32.dll",
+				Functions: []string{
+					"ExitProcess",
+					"GetStdHandle",
+					"WriteFile",
+				},
+			},
+			{
+				Name: "user32.dll",
+				Functions: []string{
+					"MessageBoxA",
+				},
+			},
+		}
+
+		pe.Write(buffer, code, data, dlls)
 	default:
 		return fmt.Errorf("unsupported platform '%s'", config.TargetOS)
 	}
diff --git a/src/exe/pe/DLL.go b/src/exe/pe/DLL.go
new file mode 100644
index 0000000..78a2b48
--- /dev/null
+++ b/src/exe/pe/DLL.go
@@ -0,0 +1,6 @@
+package pe
+
+type DLL struct {
+	Name      string
+	Functions []string
+}
diff --git a/src/exe/pe/EXE.go b/src/exe/pe/EXE.go
index 69034d1..3e20183 100644
--- a/src/exe/pe/EXE.go
+++ b/src/exe/pe/EXE.go
@@ -17,80 +17,57 @@ type EXE struct {
 	Sections []SectionHeader
 }
 
-type DLL struct {
-	Name      string
-	Functions []string
-}
-
 // Write writes the EXE file to the given writer.
-func Write(writer io.Writer, code []byte, data []byte) {
+func Write(writer io.Writer, code []byte, data []byte, dlls []DLL) {
 	NumSections := 2
 	HeaderEnd := DOSHeaderSize + NTHeaderSize + OptionalHeader64Size + SectionHeaderSize*NumSections
 	codeStart, codePadding := exe.Align(HeaderEnd, config.Align)
 	dataStart, dataPadding := exe.Align(codeStart+len(code), config.Align)
 
-	dlls := []DLL{
-		{
-			Name: "kernel32.dll",
-			Functions: []string{
-				"ExitProcess",
-				"GetStdHandle",
-				"WriteFile",
-			},
-		},
-	}
-
-	dllAddresses := []uint64{}
 	dllImports := []DLLImport{}
 
-	dllName := len(data)
-	data = append(data, dlls[0].Name...)
-	data = append(data, 0x00)
-
-	for _, f := range dlls[0].Functions {
-		pos := len(data)
-		data = append(data, 0x00, 0x00)
-		data = append(data, f...)
+	for _, dll := range dlls {
+		dllAddresses := []uint64{}
+		dllNamePos := len(data)
+		data = append(data, dll.Name...)
 		data = append(data, 0x00)
 
-		if len(data)&1 != 0 {
-			data = append(data, 0x00) // align the next entry on an even boundary
+		for _, f := range dll.Functions {
+			pos := len(data)
+			data = append(data, 0x00, 0x00)
+			data = append(data, f...)
+			data = append(data, 0x00)
+
+			if len(data)&1 != 0 {
+				data = append(data, 0x00) // align the next entry on an even boundary
+			}
+
+			dllAddresses = append(dllAddresses, uint64(dataStart+pos))
 		}
 
-		dllAddresses = append(dllAddresses, uint64(dataStart+pos))
+		dllAddresses = append(dllAddresses, 0)
+
+		// Add the address table to the data section
+		functionAddressesStart := dataStart + len(data)
+		data, _ = binary.Append(data, binary.LittleEndian, &dllAddresses)
+
+		dllImports = append(dllImports, DLLImport{
+			RvaFunctionNameList:    uint32(functionAddressesStart),
+			TimeDateStamp:          0,
+			ForwarderChain:         0,
+			RvaModuleName:          uint32(dataStart + dllNamePos),
+			RvaFunctionAddressList: uint32(functionAddressesStart),
+		})
 	}
 
-	dllAddresses = append(dllAddresses, 0)
-
-	// Add the address table to the data section
-	functionAddressesStart := dataStart + len(data)
-	functionAddressesSize := 8 * len(dllAddresses)
-	data, err := binary.Append(data, binary.LittleEndian, &dllAddresses)
-
-	if err != nil {
-		panic(err)
-	}
-
-	dllImports = append(dllImports, DLLImport{
-		RvaFunctionNameList:    uint32(functionAddressesStart),
-		TimeDateStamp:          0,
-		ForwarderChain:         0,
-		RvaModuleName:          uint32(dataStart + dllName),
-		RvaFunctionAddressList: uint32(functionAddressesStart),
-	})
-
 	dllImports = append(dllImports, DLLImport{}) // a zeroed structure marks the end of the list
 
 	// Add imports to the data section
 	importsStart := dataStart + len(data)
 	importsSize := DLLImportSize * len(dllImports)
-	data, err = binary.Append(data, binary.LittleEndian, &dllImports)
+	data, _ = binary.Append(data, binary.LittleEndian, &dllImports)
 
-	if err != nil {
-		panic(err)
-	}
-
-	imageSize := functionAddressesStart + functionAddressesSize
+	imageSize := dataStart + len(data)
 	imageSize, _ = exe.Align(imageSize, config.Align)
 
 	pe := &EXE{
@@ -151,7 +128,7 @@ func Write(writer io.Writer, code []byte, data []byte) {
 				{VirtualAddress: 0, Size: 0},
 				{VirtualAddress: 0, Size: 0},
 				{VirtualAddress: 0, Size: 0},
-				{VirtualAddress: uint32(functionAddressesStart), Size: uint32(functionAddressesSize)}, // RVA of the import address table
+				{VirtualAddress: 0, Size: 0},
 				{VirtualAddress: 0, Size: 0},
 				{VirtualAddress: 0, Size: 0},
 				{VirtualAddress: 0, Size: 0},
diff --git a/src/exe/pe/EXE_test.go b/src/exe/pe/EXE_test.go
index 71eede8..5f8d07e 100644
--- a/src/exe/pe/EXE_test.go
+++ b/src/exe/pe/EXE_test.go
@@ -8,5 +8,5 @@ import (
 )
 
 func TestWrite(t *testing.T) {
-	pe.Write(io.Discard, nil, nil)
+	pe.Write(io.Discard, nil, nil, nil)
 }