inject: Prevent the C/C++ runtime from destroying things when process terminates.

2 files changed, 48 insertions, 1 deletions

diff --git a/inject/CMakeLists.txt b/inject/CMakeLists.txt
index 0b556502..6c24370d 100644
--- a/inject/CMakeLists.txt
+++ b/inject/CMakeLists.txt
@@ -8,6 +8,16 @@ add_library (injectee MODULE
 set_target_properties (injectee PROPERTIES
     PREFIX ""
 )
+if (MSVC)
+    set_target_properties (injectee PROPERTIES LINK_FLAGS "/ENTRY:DllMainStartup")
+else ()
+    if (CMAKE_SIZEOF_VOID_P EQUAL 4)
+        set_target_properties (injectee PROPERTIES LINK_FLAGS "-Wl,--entry=_DllMainStartup@12")
+    else ()
+        set_target_properties (injectee PROPERTIES LINK_FLAGS "-Wl,--entry=DllMainStartup")
+    endif ()
+endif ()
+
 install (TARGETS injectee LIBRARY DESTINATION bin)
 install_pdb (injectee DESTINATION bin)
 
diff --git a/inject/injectee.cpp b/inject/injectee.cpp
index 05fc5ab7..95e194c8 100644
--- a/inject/injectee.cpp
+++ b/inject/injectee.cpp
@@ -991,7 +991,7 @@ dumpRegisteredHooks(void)
 
 
 EXTERN_C BOOL WINAPI
-DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved)
+DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
 {
     const char *szNewDllName = NULL;
     const char *szNewDllBaseName;
@@ -1112,6 +1112,8 @@ DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved)
             debugPrintf("inject: DLL_PROCESS_DETACH\n");
         }
 
+        assert(!lpvReserved);
+
         patchAllModules(ACTION_UNHOOK);
 
         if (g_hHookModule) {
@@ -1121,3 +1123,38 @@ DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved)
     }
     return TRUE;
 }
+
+
+/*
+ * Prevent the C/C++ runtime from destroying things when the program
+ * terminates.
+ *
+ * There is no effective way to control the order DLLs receive
+ * DLL_PROCESS_DETACH -- patched DLLs might get detacched after we are --, and
+ * unpatching our hooks doesn't always work.  So instead just do nothing (and
+ * prevent C/C++ runtime from doing anything too), so our hooks can still work
+ * after we are dettached.
+ */
+
+#ifdef _MSC_VER
+#  define DLLMAIN_CRT_STARTUP _DllMainCRTStartup
+#else
+#  define DLLMAIN_CRT_STARTUP DllMainCRTStartup
+#  pragma GCC optimize ("no-stack-protector")
+#endif
+
+EXTERN_C BOOL WINAPI
+DLLMAIN_CRT_STARTUP(HANDLE hDllHandle, DWORD dwReason, LPVOID lpvReserved);
+
+EXTERN_C BOOL WINAPI
+DllMainStartup(HANDLE hDllHandle, DWORD dwReason, LPVOID lpvReserved)
+{
+    if (dwReason == DLL_PROCESS_DETACH && lpvReserved) {
+        if (VERBOSITY > 0) {
+            debugPrintf("inject: DLL_PROCESS_DETACH\n");
+        }
+        return TRUE;
+    }
+
+    return DLLMAIN_CRT_STARTUP(hDllHandle, dwReason, lpvReserved);
+}