From cf298734aee4a2bb335b09b21c0ff9a3de67406a Mon Sep 17 00:00:00 2001
From: Weilu Jia <optix2000@gmail.com>
Date: Sun, 19 Sep 2021 19:11:32 -0700
Subject: [PATCH] windows: Add tests for process modules

---
 windows/syscall_windows_test.go | 84 +++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/windows/syscall_windows_test.go b/windows/syscall_windows_test.go
index fd75c09c..2563bf13 100644
--- a/windows/syscall_windows_test.go
+++ b/windows/syscall_windows_test.go
@@ -701,3 +701,87 @@ func TestWinVerifyTrust(t *testing.T) {
 	}
 
 }
+
+func TestProcessModules(t *testing.T) {
+	process, err := windows.GetCurrentProcess()
+	if err != nil {
+		t.Fatalf("unable to get current process: %v", err)
+	}
+	// NB: Assume that we're always the first module. This technically isn't documented anywhere (that I could find), but seems to always hold.
+	var module windows.Handle
+	var cbNeeded uint32
+	err = windows.EnumProcessModules(process, &module, uint32(unsafe.Sizeof(module)), &cbNeeded)
+	if err != nil {
+		t.Errorf("unable to call EnumProcessModules: %v", err)
+	}
+
+	var moduleEx windows.Handle
+	err = windows.EnumProcessModulesEx(process, &moduleEx, uint32(unsafe.Sizeof(moduleEx)), &cbNeeded, windows.LIST_MODULES_DEFAULT)
+	if err != nil {
+		t.Errorf("unable to call EnumProcessModulesEx: %v", err)
+	}
+	if module != moduleEx {
+		t.Errorf("module from EnumProcessModules does not match EnumProcessModulesEx: %v != %v", module, moduleEx)
+	}
+
+	exePath, err := os.Executable()
+	if err != nil {
+		t.Errorf("unable to get current executable path: %v", err)
+	}
+
+	modulePathUTF := make([]uint16, len(exePath)+1)
+	err = windows.GetModuleFileNameEx(process, module, &modulePathUTF[0], uint32(len(modulePathUTF)))
+	if err != nil {
+		t.Errorf("unable to call GetModuleFileNameEx: %v", err)
+	}
+
+	modulePath := windows.UTF16ToString(modulePathUTF)
+	if modulePath != exePath {
+		t.Errorf("module does not match executable for GetModuleFileNameEx: %s != %s", modulePath, exePath)
+	}
+
+	err = windows.GetModuleBaseName(process, module, &modulePathUTF[0], uint32(len(modulePathUTF)))
+	if err != nil {
+		t.Errorf("unable to call GetModuleBaseName: %v", err)
+	}
+
+	modulePath = windows.UTF16ToString(modulePathUTF)
+	baseExePath := filepath.Base(exePath)
+	if modulePath != baseExePath {
+		t.Errorf("module does not match executable for GetModuleBaseName: %s != %s", modulePath, baseExePath)
+	}
+
+	var moduleInfo windows.ModuleInfo
+	err = windows.GetModuleInformation(process, module, &moduleInfo, uint32(unsafe.Sizeof(moduleInfo)))
+	if err != nil {
+		t.Errorf("unable to call GetModuleInformation: %v", err)
+	}
+
+	var arch int
+	var size uint32
+	peFile, err := pe.Open(exePath)
+	if err != nil {
+		t.Errorf("unable to open current executable: %v", err)
+	} else {
+		switch runtime.GOARCH {
+		case "amd64":
+		case "arm64":
+			arch = 64
+		case "386":
+		case "arm":
+			arch = 32
+		}
+
+		if arch == 32 {
+			size = peFile.OptionalHeader.(pe.OptionalHeader32).SizeOfImage
+		} else if arch == 64 {
+			size = peFile.OptionalHeader.(pe.OptionalHeader64).SizeOfImage
+		} else {
+			t.Logf("unable to test GetModuleInformation on arch %v", runtime.GOARCH)
+		}
+	}
+
+	if arch != 0 && moduleInfo.SizeOfImage != size {
+		t.Errorf("module size does not match executable: %v != %v", moduleInfo.SizeOfImage, size)
+	}
+}