From 5dd854df4835d8db8b35224fcb90fcd171b5a973 Mon Sep 17 00:00:00 2001
From: DaLynX <d@l.ynx.fr>
Date: Sun, 17 Feb 2019 22:47:19 +0100
Subject: [PATCH] Added x64 exception handling

---
 MemoryModule.c | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/MemoryModule.c b/MemoryModule.c
index 9f95a70..2fb311b 100644
--- a/MemoryModule.c
+++ b/MemoryModule.c
@@ -435,6 +435,15 @@ PerformBaseRelocation(PMEMORYMODULE module, ptrdiff_t delta)
     return TRUE;
 }
 
+static BOOL
+RegisterExceptionHandling(PMEMORYMODULE module)
+{
+	PIMAGE_DATA_DIRECTORY pDir = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_EXCEPTION);
+	PIMAGE_RUNTIME_FUNCTION_ENTRY pEntry = (PIMAGE_RUNTIME_FUNCTION_ENTRY)(module->codeBase + pDir->VirtualAddress);
+	UINT count = (pDir->Size / sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)) - 1;
+	return RtlAddFunctionTable(pEntry, count, (DWORD64)module->codeBase);
+}
+
 static BOOL
 BuildImportTable(PMEMORYMODULE module)
 {
@@ -725,6 +734,10 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data, size_t size,
         goto error;
     }
 
+    if (!RegisterExceptionHandling(result)) {
+        goto error;
+    }
+
     // mark memory pages depending on section headers and release
     // sections that are marked as "discardable"
     if (!FinalizeSections(result)) {