diff --git a/MemoryModule.c b/MemoryModule.c index 2fb311b..bc5124a 100644 --- 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