elevate: consider ImageBaseAddress to be sufficiently stable

GetModuleHandle technically returns an opaque value, so comparing to the
PEB might in some theoretical sense be cleaner.

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

3 files changed, 33 insertions, 25 deletions

diff --git a/elevate/loader.go b/elevate/loader.go
new file mode 100644
index 00000000..0bb275da
--- /dev/null
+++ b/elevate/loader.go
@@ -0,0 +1,31 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2019 WireGuard LLC. All Rights Reserved.
+ */
+
+package elevate
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+/* We could use the undocumented LdrFindEntryForAddress function instead, but that's undocumented, and we're trying
+ * to be as rock-solid as possible here. */
+func findCurrentDataTableEntry() (entry *cLDR_DATA_TABLE_ENTRY, err error) {
+	peb := rtlGetCurrentPeb()
+	if peb == nil || peb.Ldr == nil {
+		err = windows.ERROR_INVALID_ADDRESS
+		return
+	}
+	for cur := peb.Ldr.InMemoryOrderModuleList.Flink; cur != &peb.Ldr.InMemoryOrderModuleList; cur = cur.Flink {
+		entry = (*cLDR_DATA_TABLE_ENTRY)(unsafe.Pointer(uintptr(unsafe.Pointer(cur)) - unsafe.Offsetof(cLDR_DATA_TABLE_ENTRY{}.InMemoryOrderLinks)))
+		if entry.DllBase == peb.ImageBaseAddress {
+			return
+		}
+	}
+	entry = nil
+	err = windows.ERROR_OBJECT_NOT_FOUND
+	return
+}
diff --git a/elevate/shellexecute.go b/elevate/shellexecute.go
index 00f2d915..149a4c57 100644
--- a/elevate/shellexecute.go
+++ b/elevate/shellexecute.go
@@ -22,29 +22,6 @@ const (
 	cSEE_MASK_DEFAULT = 0
 )
 
-/* We could use the undocumented LdrFindEntryForAddress function instead, but that's undocumented, and we're trying
- * to be as rock-solid as possible here. */
-func findCurrentDataTableEntry() (entry *cLDR_DATA_TABLE_ENTRY, err error) {
-	ourBase, err := getModuleHandle(nil) /* This is the same as peb->ImageBaseAddress, but that member is undocumented. */
-	if err != nil {
-		return
-	}
-	peb := rtlGetCurrentPeb()
-	if peb == nil || peb.Ldr == nil {
-		err = windows.ERROR_INVALID_ADDRESS
-		return
-	}
-	for cur := peb.Ldr.InMemoryOrderModuleList.Flink; cur != &peb.Ldr.InMemoryOrderModuleList; cur = cur.Flink {
-		entry = (*cLDR_DATA_TABLE_ENTRY)(unsafe.Pointer(uintptr(unsafe.Pointer(cur)) - unsafe.Offsetof(cLDR_DATA_TABLE_ENTRY{}.InMemoryOrderLinks)))
-		if entry.DllBase == ourBase {
-			return
-		}
-	}
-	entry = nil
-	err = windows.ERROR_OBJECT_NOT_FOUND
-	return
-}
-
 func ShellExecute(program string, arguments string, directory string, show int32) (err error) {
 	var (
 		program16   *uint16
diff --git a/elevate/syscall_windows.go b/elevate/syscall_windows.go
index c7def8fa..d50e9c87 100644
--- a/elevate/syscall_windows.go
+++ b/elevate/syscall_windows.go
@@ -54,7 +54,8 @@ type cPEB struct {
 	Reserved1              [2]byte
 	BeingDebugged          byte
 	Reserved2              [1]byte
-	Reserved3              [2]uintptr
+	Reserved3              uintptr
+	ImageBaseAddress       uintptr
 	Ldr                    *cPEB_LDR_DATA
 	ProcessParameters      uintptr
 	Reserved4              [3]uintptr
@@ -77,7 +78,6 @@ const (
 	cCOINIT_APARTMENTTHREADED = 2
 )
 
-//sys	getModuleHandle(moduleName *uint16) (moduleHandle uintptr, err error) [failretval==0] = kernel32.GetModuleHandleW
 //sys	getWindowsDirectory(windowsDirectory *uint16, inLen uint32) (outLen uint32, err error) [failretval==0] = kernel32.GetWindowsDirectoryW
 
 //sys	rtlInitUnicodeString(destinationString *cUNICODE_STRING, sourceString *uint16) = ntdll.RtlInitUnicodeString