From b70fec0d3fd0245660c3103c6d1466a2ffd36ee9 Mon Sep 17 00:00:00 2001
From: Elisha Riedlinger <elisha@novicemail.com>
Date: Wed, 25 Sep 2024 15:44:49 -0700
Subject: [PATCH] Pin modules to ensure they stay loaded on exit

---
 Common/Utils.cpp                    | 6 ++++++
 Common/Utils.h                      | 1 +
 Resources/BuildNo.rc                | 2 +-
 Wrappers/d3d8/d3d8wrapper.cpp       | 3 +++
 Wrappers/d3d9/d3d9wrapper.cpp       | 6 ++++++
 Wrappers/dinput8/dinput8wrapper.cpp | 1 +
 Wrappers/dsound/dsoundwrapper.cpp   | 2 ++
 dllmain.cpp                         | 3 +++
 8 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/Common/Utils.cpp b/Common/Utils.cpp
index c6b7fd2a..7cd6d42e 100644
--- a/Common/Utils.cpp
+++ b/Common/Utils.cpp
@@ -1177,6 +1177,12 @@ void CopyReplaceSlash(char* DestStr, size_t Size, LPCSTR SrcStr)
 	}
 }
 
+void PinModule(HMODULE dll)
+{
+	HMODULE dummy = nullptr;
+	GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_PIN | GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<LPCSTR>(dll), &dummy);
+}
+
 BOOL GetAppsLightMode()
 {
 	HKEY hKey;
diff --git a/Common/Utils.h b/Common/Utils.h
index baa8b663..91af3eba 100644
--- a/Common/Utils.h
+++ b/Common/Utils.h
@@ -44,6 +44,7 @@ bool GetConfigName(char* ConfigName, rsize_t size, char* ext);
 bool GetConfigName(wchar_t* ConfigName, rsize_t size, wchar_t* ext);
 bool CheckPathNameMatch(LPCSTR lpFileName1, LPCSTR lpFileName2);
 void CopyReplaceSlash(char* DestStr, size_t Size, LPCSTR SrcStr);
+void PinModule(HMODULE dll);
 BOOL GetAppsLightMode();
 void ClearGDISurface(HWND hWnd, COLORREF color);
 HMONITOR GetMonitorHandle();
diff --git a/Resources/BuildNo.rc b/Resources/BuildNo.rc
index f072b981..c2f3f101 100644
--- a/Resources/BuildNo.rc
+++ b/Resources/BuildNo.rc
@@ -1 +1 @@
-#define BUILD_NUMBER 2217 
+#define BUILD_NUMBER 2218 
diff --git a/Wrappers/d3d8/d3d8wrapper.cpp b/Wrappers/d3d8/d3d8wrapper.cpp
index b1507944..981b11fc 100644
--- a/Wrappers/d3d8/d3d8wrapper.cpp
+++ b/Wrappers/d3d8/d3d8wrapper.cpp
@@ -51,6 +51,7 @@ HMODULE GetD3d8ScriptDll()
 	if (script_d3d8_dll)
 	{
 		Logging::Log() << "Loaded d3d8.dll from: " << script_path_dll.c_str();
+		PinModule(script_d3d8_dll);
 	}
 
 	// Load d3d8.dll from 'plugins' folder
@@ -61,6 +62,7 @@ HMODULE GetD3d8ScriptDll()
 		if (script_d3d8_dll)
 		{
 			Logging::Log() << "Loaded d3d8.dll from: " << plugin_path_dll.c_str();
+			PinModule(script_d3d8_dll);
 		}
 	}
 
@@ -118,6 +120,7 @@ bool GetLocalDirect3DCreate8()
 		m_pDirect3DCreate8_local = (Direct3DCreate8Proc)GetProcAddress(h_d3d8, "Direct3DCreate8");
 		if (m_pDirect3DCreate8_local)
 		{
+			PinModule(h_d3d8);
 			return true;
 		}
 	}
diff --git a/Wrappers/d3d9/d3d9wrapper.cpp b/Wrappers/d3d9/d3d9wrapper.cpp
index 8ff512e4..d74cc410 100644
--- a/Wrappers/d3d9/d3d9wrapper.cpp
+++ b/Wrappers/d3d9/d3d9wrapper.cpp
@@ -38,6 +38,10 @@ HMODULE GetSystemD3d9()
 		GetSystemDirectoryA(Path, MAX_PATH);
 		strcat_s(Path, "\\d3d9.dll");
 		GetModuleHandleExA(NULL, Path, &h_d3d9);
+		if (h_d3d9)
+		{
+			PinModule(h_d3d9);
+		}
 	}
 
 	return h_d3d9;
@@ -85,6 +89,7 @@ bool GetDirect3DCreate9()
 		m_pDirect3DCreate9 = (Direct3DCreate9Proc)GetProcAddress(h_d3d9, "Direct3DCreate9");
 		if (m_pDirect3DCreate9)
 		{
+			PinModule(h_d3d9);
 			return true;
 		}
 	}
@@ -113,6 +118,7 @@ bool GetLocalDirect3DCreate9()
 		m_pDirect3DCreate9_local = (Direct3DCreate9Proc)GetProcAddress(h_d3d9, "Direct3DCreate9");
 		if (m_pDirect3DCreate9_local)
 		{
+			PinModule(h_d3d9);
 			return true;
 		}
 	}
diff --git a/Wrappers/dinput8/dinput8wrapper.cpp b/Wrappers/dinput8/dinput8wrapper.cpp
index 1a41d314..332d689c 100644
--- a/Wrappers/dinput8/dinput8wrapper.cpp
+++ b/Wrappers/dinput8/dinput8wrapper.cpp
@@ -69,6 +69,7 @@ bool GetLocalDirectInput8Create()
 		m_pDirectInput8Create_local = (DirectInput8CreateProc)GetProcAddress(h_dinput8, "DirectInput8Create");
 		if (m_pDirectInput8Create_local)
 		{
+			PinModule(h_dinput8);
 			return true;
 		}
 	}
diff --git a/Wrappers/dsound/dsoundwrapper.cpp b/Wrappers/dsound/dsoundwrapper.cpp
index 436f464d..db489bfc 100644
--- a/Wrappers/dsound/dsoundwrapper.cpp
+++ b/Wrappers/dsound/dsoundwrapper.cpp
@@ -66,6 +66,7 @@ bool GetLocalDirectSoundCreate8()
 		m_pDirectSoundCreate8_local = (DirectSoundCreate8Proc)GetProcAddress(h_dsound, "DirectSoundCreate8");
 		if (m_pDirectSoundCreate8_local)
 		{
+			PinModule(h_dsound);
 			return true;
 		}
 	}
@@ -93,6 +94,7 @@ bool GetSystem32DirectSoundCreate8()
 		m_pDirectSoundCreate8_system32 = (DirectSoundCreate8Proc)GetProcAddress(h_dsound, "DirectSoundCreate8");
 		if (m_pDirectSoundCreate8_system32)
 		{
+			PinModule(h_dsound);
 			return true;
 		}
 	}
diff --git a/dllmain.cpp b/dllmain.cpp
index 89862b38..237328ee 100644
--- a/dllmain.cpp
+++ b/dllmain.cpp
@@ -903,6 +903,9 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD fdwReason, LPVOID lpReserved)
 		{
 			DelayedStart();
 		}
+
+		// Pin current module
+		PinModule(m_hModule);
 	}
 	break;
 	case DLL_THREAD_ATTACH: