From 5b7311ba580a1c10796a1f8f30f6c5f77a671a36 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Sun, 19 Apr 2026 21:52:51 +0800
Subject: [PATCH] memmod: be more resilient toward weird PE files

pick from https://github.com/WireGuard/wireguard-windows/commit/d99192528696c8c5617cf512f66cca0003a13392
---
 internal/wintun/memmod/memmod_windows.go | 42 +++++++++++++++++++-----
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/internal/wintun/memmod/memmod_windows.go b/internal/wintun/memmod/memmod_windows.go
index 05d71f1..985d48a 100644
--- a/internal/wintun/memmod/memmod_windows.go
+++ b/internal/wintun/memmod/memmod_windows.go
@@ -75,7 +75,7 @@ func (module *Module) copySections(address, size uintptr, oldHeaders *IMAGE_NT_H
 			continue
 		}
 
-		if size < uintptr(sections[i].PointerToRawData+sections[i].SizeOfRawData) {
+		if size < uintptr(sections[i].PointerToRawData)+uintptr(sections[i].SizeOfRawData) {
 			return errors.New("Incomplete section")
 		}
 
@@ -245,8 +245,16 @@ func (module *Module) performBaseRelocation(delta uintptr) (relocated bool, err
 		return delta == 0, nil
 	}
 
-	relocationHdr := (*IMAGE_BASE_RELOCATION)(a2p(module.codeBase + uintptr(directory.VirtualAddress)))
-	for relocationHdr.VirtualAddress > 0 {
+	relocBase := module.codeBase + uintptr(directory.VirtualAddress)
+	relocEnd := relocBase + uintptr(directory.Size)
+	relocationHdr := (*IMAGE_BASE_RELOCATION)(a2p(relocBase))
+	for uintptr(unsafe.Pointer(relocationHdr))+unsafe.Sizeof(*relocationHdr) <= relocEnd && relocationHdr.VirtualAddress > 0 {
+		if uintptr(relocationHdr.SizeOfBlock) < unsafe.Sizeof(*relocationHdr) {
+			return false, errors.New("Invalid relocation block size")
+		}
+		if uintptr(unsafe.Pointer(relocationHdr))+uintptr(relocationHdr.SizeOfBlock) > relocEnd {
+			return false, errors.New("Relocation block exceeds directory bounds")
+		}
 		dest := module.codeBase + uintptr(relocationHdr.VirtualAddress)
 
 		relInfos := unsafe.Slice(
@@ -456,16 +464,16 @@ func hookRtlPcToFileHeader() error {
 
 // LoadLibrary loads module image to memory.
 func LoadLibrary(data []byte) (module *Module, err error) {
-	addr := uintptr(unsafe.Pointer(&data[0]))
 	size := uintptr(len(data))
 	if size < unsafe.Sizeof(IMAGE_DOS_HEADER{}) {
 		return nil, errors.New("Incomplete IMAGE_DOS_HEADER")
 	}
+	addr := uintptr(unsafe.Pointer(&data[0]))
 	dosHeader := (*IMAGE_DOS_HEADER)(a2p(addr))
 	if dosHeader.E_magic != IMAGE_DOS_SIGNATURE {
 		return nil, fmt.Errorf("Not an MS-DOS binary (provided: %x, expected: %x)", dosHeader.E_magic, IMAGE_DOS_SIGNATURE)
 	}
-	if (size < uintptr(dosHeader.E_lfanew)+unsafe.Sizeof(IMAGE_NT_HEADERS{})) {
+	if dosHeader.E_lfanew < 0 || (size < uintptr(dosHeader.E_lfanew)+unsafe.Sizeof(IMAGE_NT_HEADERS{})) {
 		return nil, errors.New("Incomplete IMAGE_NT_HEADERS")
 	}
 	oldHeader := (*IMAGE_NT_HEADERS)(a2p(addr + uintptr(dosHeader.E_lfanew)))
@@ -475,9 +483,23 @@ func LoadLibrary(data []byte) (module *Module, err error) {
 	if oldHeader.FileHeader.Machine != imageFileProcess {
 		return nil, fmt.Errorf("Foreign platform (provided: %x, expected: %x)", oldHeader.FileHeader.Machine, imageFileProcess)
 	}
-	if (oldHeader.OptionalHeader.SectionAlignment & 1) != 0 {
+	if oldHeader.OptionalHeader.SectionAlignment == 0 || (oldHeader.OptionalHeader.SectionAlignment&(oldHeader.OptionalHeader.SectionAlignment-1)) != 0 {
 		return nil, errors.New("Unaligned section")
 	}
+	if oldHeader.FileHeader.NumberOfSections == 0 {
+		return nil, errors.New("No sections")
+	}
+	if uintptr(oldHeader.FileHeader.SizeOfOptionalHeader) < unsafe.Sizeof(oldHeader.OptionalHeader) {
+		return nil, errors.New("Incomplete optional header")
+	}
+	if oldHeader.OptionalHeader.NumberOfRvaAndSizes < IMAGE_NUMBEROF_DIRECTORY_ENTRIES {
+		return nil, errors.New("Incomplete data directory")
+	}
+	sectionHeadersEnd := uintptr(dosHeader.E_lfanew) + unsafe.Offsetof(oldHeader.OptionalHeader) + uintptr(oldHeader.FileHeader.SizeOfOptionalHeader) +
+		uintptr(oldHeader.FileHeader.NumberOfSections)*unsafe.Sizeof(IMAGE_SECTION_HEADER{})
+	if size < sectionHeadersEnd {
+		return nil, errors.New("Incomplete section headers")
+	}
 	lastSectionEnd := uintptr(0)
 	sections := oldHeader.Sections()
 	optionalSectionSize := oldHeader.OptionalHeader.SectionAlignment
@@ -564,6 +586,10 @@ func LoadLibrary(data []byte) (module *Module, err error) {
 			err = fmt.Errorf("Error relocating module: %w", err)
 			return
 		}
+		if !module.isRelocated {
+			err = errors.New("Module could not be relocated")
+			return
+		}
 	} else {
 		module.isRelocated = true
 	}
@@ -654,7 +680,7 @@ func (module *Module) ProcAddressByName(name string) (uintptr, error) {
 		return 0, errors.New("No functions exported by name")
 	}
 	if idx, ok := module.nameExports[name]; ok {
-		if uint32(idx) > exports.NumberOfFunctions {
+		if uint32(idx) >= exports.NumberOfFunctions {
 			return 0, errors.New("Ordinal number too high")
 		}
 		// AddressOfFunctions contains the RVAs to the "real" functions.
@@ -674,7 +700,7 @@ func (module *Module) ProcAddressByOrdinal(ordinal uint16) (uintptr, error) {
 		return 0, errors.New("Ordinal number too low")
 	}
 	idx := ordinal - uint16(exports.Base)
-	if uint32(idx) > exports.NumberOfFunctions {
+	if uint32(idx) >= exports.NumberOfFunctions {
 		return 0, errors.New("Ordinal number too high")
 	}
 	// AddressOfFunctions contains the RVAs to the "real" functions.