Fix exception handling registration.

fancycode · fancycode · commit f89837636a37 · 2019-02-24T15:35:31.000+01:00
- Only compile on 64bit, 32bit exception handling works differently.
- Unregister table when module is unloaded.
- Handle empty tables.
- Fix wrong number of entries being passed to "RtlAddFunctionTable".
diff --git a/MemoryModule.c b/MemoryModule.c
@@ -438,10 +438,39 @@ PerformBaseRelocation(PMEMORYMODULE module, ptrdiff_t delta)
 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);
+#ifdef _WIN64
+    PRUNTIME_FUNCTION pEntry;
+    DWORD count;
+    PIMAGE_DATA_DIRECTORY pDir = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_EXCEPTION);
+    if (pDir->Size == 0) {
+        return TRUE;
+    }
+
+    pEntry = (PRUNTIME_FUNCTION)(module->codeBase + pDir->VirtualAddress);
+    count = (pDir->Size / sizeof(RUNTIME_FUNCTION));
+    return RtlAddFunctionTable(pEntry, count, (DWORD64) module->codeBase);
+#else
+    // TODO(fancycode): Support 32bit exception handling.
+    return TRUE;
+#endif
+}
+
+static BOOL
+UnregisterExceptionHandling(PMEMORYMODULE module)
+{
+#ifdef _WIN64
+    PRUNTIME_FUNCTION pEntry;
+    PIMAGE_DATA_DIRECTORY pDir = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_EXCEPTION);
+    if (pDir->Size == 0) {
+        return TRUE;
+    }
+
+    pEntry = (PRUNTIME_FUNCTION)(module->codeBase + pDir->VirtualAddress);
+    return RtlDeleteFunctionTable(pEntry);
+#else
+    // TODO(fancycode): Support 32bit exception handling.
+    return TRUE;
+#endif
 }
 
 static BOOL
@@ -734,6 +763,7 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data, size_t size,
         goto error;
     }
 
+    // register exception handling table so "try { } catch ( ) { }"" works
     if (!RegisterExceptionHandling(result)) {
         goto error;
     }
@@ -880,6 +910,8 @@ void MemoryFreeLibrary(HMEMORYMODULE mod)
         (*DllEntry)((HINSTANCE)module->codeBase, DLL_PROCESS_DETACH, 0);
     }
 
+    UnregisterExceptionHandling(module);
+
     free(module->nameExportsTable);
     if (module->modules != NULL) {
         // free previously opened libraries
