From 3957e9b9dd191e0c4f7fc41d15a865357c097d9e Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Thu, 5 Aug 2021 14:56:48 +0200
Subject: [PATCH] memmod: register exception handler tables

Otherwise recent WDK binaries fail on ARM64, where an exception handler
is used for trapping an illegal instruction when ARMv8.1 atomics are
being tested for functionality.

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>
---
 tun/wintun/memmod/memmod_windows.go  | 13 +++++++++++++
 tun/wintun/memmod/syscall_windows.go |  6 ++++++
 2 files changed, 19 insertions(+)

diff --git a/tun/wintun/memmod/memmod_windows.go b/tun/wintun/memmod/memmod_windows.go
index 59450e7..075c03a 100644
--- a/tun/wintun/memmod/memmod_windows.go
+++ b/tun/wintun/memmod/memmod_windows.go
@@ -159,6 +159,16 @@ func (module *Module) finalizeSection(sectionData *sectionFinalizeData) error {
 	return nil
 }
 
+var rtlAddFunctionTable = windows.NewLazySystemDLL("ntdll.dll").NewProc("RtlAddFunctionTable")
+
+func (module *Module) registerExceptionHandlers() {
+	directory := module.headerDirectory(IMAGE_DIRECTORY_ENTRY_EXCEPTION)
+	if directory.Size == 0 || directory.VirtualAddress == 0 {
+		return
+	}
+	rtlAddFunctionTable.Call(module.codeBase+uintptr(directory.VirtualAddress), uintptr(directory.Size)/unsafe.Sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY{}), module.codeBase)
+}
+
 func (module *Module) finalizeSections() error {
 	sections := module.headers.Sections()
 	imageOffset := module.headers.OptionalHeader.imageOffset()
@@ -500,6 +510,9 @@ func LoadLibrary(data []byte) (module *Module, err error) {
 		return
 	}
 
+	// Register exception tables, if they exist.
+	module.registerExceptionHandlers()
+
 	// TLS callbacks are executed BEFORE the main loading.
 	module.executeTLS()
 
diff --git a/tun/wintun/memmod/syscall_windows.go b/tun/wintun/memmod/syscall_windows.go
index b79be69..a111f92 100644
--- a/tun/wintun/memmod/syscall_windows.go
+++ b/tun/wintun/memmod/syscall_windows.go
@@ -370,6 +370,12 @@ const (
 	IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT       = 28
 )
 
+type IMAGE_RUNTIME_FUNCTION_ENTRY struct {
+	BeginAddress      uint32
+	EndAddress        uint32
+	UnwindInfoAddress uint32
+}
+
 const (
 	DLL_PROCESS_ATTACH = 1
 	DLL_THREAD_ATTACH  = 2