1 files changed, 46 insertions, 22 deletions

diff --git a/NorthstarDedicatedTest/ExploitFixes.cpp b/NorthstarDedicatedTest/ExploitFixes.cpp
index 0f8569b5..2f4e2b5c 100644
--- a/NorthstarDedicatedTest/ExploitFixes.cpp
+++ b/NorthstarDedicatedTest/ExploitFixes.cpp
@@ -272,13 +272,11 @@ INVALID_CMD:
 // basically: by default r2 isn't set as a valve mod, meaning that m_bRestrictServerCommands is false
 // this is HORRIBLE for security, because it means servers can run arbitrary concommands on clients
 // especially since we have script commands this could theoretically be awful
-
-typedef void (*IsValveModType)();
-IsValveModType IsValveMod;
-
-bool IsValveModHook()
+KHOOK(IsValveMod, ("engine.dll", "48 83 EC 28 48 8B 0D ? ? ? ? 48 8D 15 ? ? ? ? E8 ? ? ? ? 85 C0 74 63"), bool, __fastcall, ())
 {
-	return !CommandLine()->CheckParm("-norestrictservercommands");
+	bool result = !CommandLine()->CheckParm("-norestrictservercommands");
+	spdlog::info("ExploitFixes: Overriding IsValveMod to {}...", result);
+	return result;
 }
 
 // Fix respawn's crappy UTF8 parser so it doesn't crash -_-
@@ -490,36 +488,62 @@ KHOOK(
 	return oSpecialClientCommand(player, command);
 }
 
-void ExploitFixes::LoadCallback(HMODULE baseAddress)
+void SetupKHook(KHook* hook)
 {
-	spdlog::info("ExploitFixes::LoadCallback ...");
-
-	spdlog::info("\tByte patching...");
-	DoBytePatches();
-
-	if (KHook::InitAllHooks())
+	if (hook->Setup())
 	{
-		spdlog::info("\tInitialized " + std::to_string(KHook::_allHooks.size()) + " exploit-patch hooks.");
+		spdlog::debug("KHook::Setup(): Hooked at {}", hook->targetFuncAddr);
 	}
 	else
 	{
 		spdlog::critical("\tFAILED to initialize all exploit patches.");
 
-		// Force exit?
+		// Force exit
 		MessageBoxA(0, "FAILED to initialize all exploit patches.", "Northstar", MB_ICONERROR);
 		exit(0);
 	}
+}
 
-	ns_exploitfixes_log =
-		new ConVar("ns_exploitfixes_log", "1", FCVAR_GAMEDLL, "Whether to log whenever ExploitFixes.cpp blocks/corrects something");
+void ExploitFixes::LoadCallback_MultiModule(HMODULE baseAddress)
+{
 
-	HookEnabler hook;
-	ENABLER_CREATEHOOK(hook, (char*)baseAddress + 0x2a8a50, &GetEntByIndexHook, reinterpret_cast<LPVOID*>(&GetEntByIndex));
+	spdlog::info("ExploitFixes::LoadCallback_MultiModule({}) ...", (void*)baseAddress);
+
+	int hooksEnabled = 0;
+	for (auto itr = KHook::_allHooks.begin(); itr != KHook::_allHooks.end(); itr++)
+	{
+		auto curHook = *itr;
+		if (GetModuleHandleA(curHook->targetFunc.moduleName) == baseAddress)
+		{
+			SetupKHook(curHook);
+			itr = KHook::_allHooks.erase(itr); // Prevent repeated initialization
+
+			hooksEnabled++;
+
+			if (itr == KHook::_allHooks.end())
+				break;
+		}
+	}
+
+	spdlog::info("\tEnabled {} hooks.", hooksEnabled);
 }
 
-void ExploitFixes::LoadCallbackEngine(HMODULE baseAddress)
+void ExploitFixes::LoadCallback_Full(HMODULE baseAddress)
 {
-	spdlog::info("ExploitFixes::LoadCallbackEngine ...");
+	spdlog::info("ExploitFixes::LoadCallback_Full ...");
+
+	spdlog::info("\tByte patching...");
+	DoBytePatches();
+
+	for (KHook* hook : KHook::_allHooks)
+		SetupKHook(hook);
+
+	spdlog::info("\tInitialized " + std::to_string(KHook::_allHooks.size()) + " late exploit-patch hooks.");
+	KHook::_allHooks.clear();
+
+	ns_exploitfixes_log =
+		new ConVar("ns_exploitfixes_log", "1", FCVAR_GAMEDLL, "Whether to log whenever ExploitFixes.cpp blocks/corrects something");
+
 	HookEnabler hook;
-	ENABLER_CREATEHOOK(hook, (char*)baseAddress + 0x1C6360, &IsValveModHook, reinterpret_cast<LPVOID*>(&IsValveMod));
+	ENABLER_CREATEHOOK(hook, (char*)baseAddress + 0x2a8a50, &GetEntByIndexHook, reinterpret_cast<LPVOID*>(&GetEntByIndex));
 }
\ No newline at end of file