From cf3c659140fb93577ea1541eafb4965e0503e835 Mon Sep 17 00:00:00 2001
From: Elisha Riedlinger <elisha@novicemail.com>
Date: Tue, 10 Sep 2024 21:10:36 -0700
Subject: [PATCH] Pin real dll module and current module

---
 Dllmain/BuildNo.rc  |  2 +-
 Dllmain/Dllmain.cpp | 12 ++++++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/Dllmain/BuildNo.rc b/Dllmain/BuildNo.rc
index 234b1d24..f27c6ccd 100644
--- a/Dllmain/BuildNo.rc
+++ b/Dllmain/BuildNo.rc
@@ -1 +1 @@
-#define BUILD_NUMBER 7137 
+#define BUILD_NUMBER 7138 
diff --git a/Dllmain/Dllmain.cpp b/Dllmain/Dllmain.cpp
index 21eaa88e..e791b698 100644
--- a/Dllmain/Dllmain.cpp
+++ b/Dllmain/Dllmain.cpp
@@ -123,10 +123,6 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD fdwReason, LPVOID lpReserved)
 		// Get handle
 		hModule_dll = hModule;
 
-		// Pin current module
-		HMODULE dummy = nullptr;
-		GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_PIN | GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<LPCSTR>(DllMain), &dummy);
-
 		// Initialize config
 		Config.Init();
 
@@ -235,6 +231,10 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD fdwReason, LPVOID lpReserved)
 			HMODULE dll = Wrapper::CreateWrapper((Config.RealDllPath.size()) ? Config.RealDllPath.c_str() : nullptr, (Config.WrapperMode.size()) ? Config.WrapperMode.c_str() : nullptr, Config.WrapperName.c_str());
 			if (dll)
 			{
+				// Pin real dll module
+				HMODULE dummy = nullptr;
+				GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_PIN | GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<LPCSTR>(dll), &dummy);
+
 				Utils::AddHandleToVector(dll, Config.WrapperName.c_str());
 
 				// Hook GetProcAddress to handle wrapped functions that are missing or not available in the OS
@@ -247,6 +247,10 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD fdwReason, LPVOID lpReserved)
 			}
 		}
 
+		// Pin current module
+		HMODULE dummy = nullptr;
+		GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_PIN | GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<LPCSTR>(DllMain), &dummy);
+
 		// Launch processes
 		if (!Config.RunProcess.empty())
 		{