From da5045580e066878bcae45eae40cb34299c49218 Mon Sep 17 00:00:00 2001
From: pongo1231 <pongo1999712@gmail.com>
Date: Thu, 6 Feb 2025 21:30:51 +0000
Subject: [PATCH] ChaosMod: Mitigate crashes for direct mod hot reloading hack

---
 ChaosMod/Main.cpp    | 16 +++++++---------
 ChaosMod/Main.h      |  1 -
 ChaosMod/dllmain.cpp |  1 -
 3 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/ChaosMod/Main.cpp b/ChaosMod/Main.cpp
index dca2455df..fb7a66c4f 100644
--- a/ChaosMod/Main.cpp
+++ b/ChaosMod/Main.cpp
@@ -342,21 +342,16 @@ namespace Main
 	static HMODULE ms_ModuleHandle = NULL;
 	EXTERN_C IMAGE_DOS_HEADER __ImageBase;
 
-	void OnInit()
+	void OnRun()
 	{
-		LOG("Running mod init");
+		SetUnhandledExceptionFilter(CrashHandler);
 
-		if (DoesFileExist("ScriptHookV.dev"))
+		if (!ms_ModuleHandle && DoesFileExist("ScriptHookV.dev"))
 		{
 			WCHAR fileName[MAX_PATH] = {};
 			GetModuleFileName(reinterpret_cast<HINSTANCE>(&__ImageBase), fileName, MAX_PATH);
 			ms_ModuleHandle = LoadLibrary(fileName);
 		}
-	}
-
-	void OnRun()
-	{
-		SetUnhandledExceptionFilter(CrashHandler);
 
 		MainRun();
 	}
@@ -421,8 +416,11 @@ namespace Main
 			}
 			else if (key == 0x52 && ms_ModuleHandle) // R
 			{
+				// Prevention for (somehow) double unloading
+				auto handle     = ms_ModuleHandle;
+				ms_ModuleHandle = NULL;
 				OnCleanup();
-				FreeModule(ms_ModuleHandle);
+				FreeModule(handle);
 			}
 		}
 
diff --git a/ChaosMod/Main.h b/ChaosMod/Main.h
index 236d21174..10196d9b2 100644
--- a/ChaosMod/Main.h
+++ b/ChaosMod/Main.h
@@ -9,7 +9,6 @@ struct IDXGISwapChain;
 
 namespace Main
 {
-	void OnInit();
 	void OnRun();
 	void OnCleanup();
 	void OnKeyboardInput(DWORD key, WORD repeats, BYTE scanCode, BOOL isExtended, BOOL isWithAlt, BOOL wasDownBefore,
diff --git a/ChaosMod/dllmain.cpp b/ChaosMod/dllmain.cpp
index f79f38f62..b78ae76eb 100644
--- a/ChaosMod/dllmain.cpp
+++ b/ChaosMod/dllmain.cpp
@@ -17,7 +17,6 @@ BOOL APIENTRY DllMain(HMODULE instance, DWORD reason, LPVOID reserved)
 
 		RAW_LOG("Chaos Mod v" MOD_VERSION "\n\n");
 
-		Main::OnInit();
 		Memory::Init();
 
 		scriptRegister(instance, Main::OnRun);