From a6f4156610b12d8d63e487599aec43e2fa533463 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Sat, 16 Feb 2019 20:46:24 +0200
Subject: [PATCH] Win/x64: MemoryLoader: Register exception handlers

---
 client/sources/MemoryModule.c | 43 +++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/client/sources/MemoryModule.c b/client/sources/MemoryModule.c
index d2b103c5..57eb8ddf 100644
--- a/client/sources/MemoryModule.c
+++ b/client/sources/MemoryModule.c
@@ -432,6 +432,42 @@ static void _FreeLibrary(HCUSTOMMODULE module, void *userdata)
     FreeLibrary((HMODULE) module);
 }
 
+//===============================================================================================//
+#if defined(_WIN64)
+BOOL WINAPI RegisterExceptionTable(PMEMORYMODULE pModule)
+{
+	UINT_PTR uiLibraryAddress = 0;
+    UINT_PTR uiAddressArray = 0;
+    UINT_PTR uiNameArray    = 0;
+    UINT_PTR uiNameOrdinals = 0;
+    PIMAGE_NT_HEADERS pNtHeaders             = NULL;
+    PIMAGE_DATA_DIRECTORY pDataDirectory     = NULL;
+    PIMAGE_RUNTIME_FUNCTION_ENTRY pExceptionDirectory = NULL;
+    DWORD dwCount;
+    BOOL bResult;
+    unsigned char *codeBase = pModule->codeBase;
+
+	if( pModule == NULL )
+		return FALSE;
+
+    pDataDirectory = GET_HEADER_DICTIONARY(pModule, IMAGE_DIRECTORY_ENTRY_EXCEPTION);
+
+    if (pDataDirectory->Size == 0 || pDataDirectory->VirtualAddress == 0)
+        return TRUE;
+
+    // get the VA of the export directory
+    pExceptionDirectory = (PIMAGE_RUNTIME_FUNCTION_ENTRY)(codeBase + pDataDirectory->VirtualAddress);
+
+    if (!pExceptionDirectory)
+            return TRUE;
+
+    dwCount = (pDataDirectory->Size / sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)) - 1;
+
+    return RtlAddFunctionTable((PRUNTIME_FUNCTION)pExceptionDirectory, dwCount, (UINT_PTR)codeBase);
+}
+#endif
+
+
 HMEMORYMODULE MemoryLoadLibrary(const void *data)
 {
     return MemoryLoadLibraryEx(data, _LoadLibrary, _GetProcAddress, _FreeLibrary, NULL);
@@ -551,6 +587,13 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data,
         goto error;
     }
 
+#if defined(_WIN64)
+    // Enable exceptions
+    if (!RegisterExceptionTable(result)) {
+        goto error;
+    }
+#endif
+
     // mark memory pages depending on section headers and release
     // sections that are marked as "discardable"
     if (!FinalizeSections(result)) {