From a3ee040a2033789f5a716397d4663d0c4bbe7c11 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sun, 24 Jan 2016 13:06:52 +0100
Subject: [PATCH] Fix release builds.

---
 src/Components/Modules/Theatre.cpp | 21 ++++++++------
 src/Components/Modules/Theatre.hpp |  2 +-
 src/Main.cpp                       | 45 ++++++++++++++++++++----------
 3 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/src/Components/Modules/Theatre.cpp b/src/Components/Modules/Theatre.cpp
index df88c4d0..c265a690 100644
--- a/src/Components/Modules/Theatre.cpp
+++ b/src/Components/Modules/Theatre.cpp
@@ -5,7 +5,6 @@ namespace Components
 	Theatre::Container Theatre::DemoContainer;
 
 	char Theatre::BaselineSnapshot[131072] = { 0 };
-	PBYTE Theatre::BaselineSnapshotMsg = 0;
 	int Theatre::BaselineSnapshotMsgLen;
 	int Theatre::BaselineSnapshotMsgOff;
 
@@ -21,20 +20,24 @@ namespace Components
 		Game::FS_Write(&sequence, 4, *Game::demoFile);
 	}
 
-	void __declspec(naked) Theatre::BaselineStoreStub()
+	void Theatre::StoreBaseline(PBYTE snapshotMsg)
 	{
-		// Store snapshot message
-		__asm mov Theatre::BaselineSnapshotMsg, edi
-
 		// Store offset and length
-		Theatre::BaselineSnapshotMsgLen = *reinterpret_cast<int*>(Theatre::BaselineSnapshotMsg + 20);
-		Theatre::BaselineSnapshotMsgOff = *reinterpret_cast<int*>(Theatre::BaselineSnapshotMsg + 28) - 7;
+		Theatre::BaselineSnapshotMsgLen = *reinterpret_cast<int*>(snapshotMsg + 20);
+		Theatre::BaselineSnapshotMsgOff = *reinterpret_cast<int*>(snapshotMsg + 28) - 7;
 
 		// Copy to our snapshot buffer
-		memcpy(Theatre::BaselineSnapshot, *reinterpret_cast<DWORD**>(Theatre::BaselineSnapshotMsg + 8), *reinterpret_cast<DWORD*>(Theatre::BaselineSnapshotMsg + 20));
+		memcpy(Theatre::BaselineSnapshot, *reinterpret_cast<DWORD**>(snapshotMsg + 8), *reinterpret_cast<DWORD*>(snapshotMsg + 20));
+	}
 
-		__asm
+	void __declspec(naked) Theatre::BaselineStoreStub()
+	{
+		_asm
 		{
+			push edi
+			call Theatre::StoreBaseline
+			pop edi
+
 			mov edx, 5ABEF5h
 			jmp edx
 		}
diff --git a/src/Components/Modules/Theatre.hpp b/src/Components/Modules/Theatre.hpp
index 01b8ad6a..82b0239b 100644
--- a/src/Components/Modules/Theatre.hpp
+++ b/src/Components/Modules/Theatre.hpp
@@ -39,11 +39,11 @@ namespace Components
 		static Container DemoContainer;
 
 		static char BaselineSnapshot[131072];
-		static PBYTE BaselineSnapshotMsg;
 		static int BaselineSnapshotMsgLen;
 		static int BaselineSnapshotMsgOff;
 
 		static void WriteBaseline();
+		static void StoreBaseline(PBYTE snapshotMsg);
 
 		static void LoadDemos();
 		static void DeleteDemo();
diff --git a/src/Main.cpp b/src/Main.cpp
index ec1a8c0a..447179be 100644
--- a/src/Main.cpp
+++ b/src/Main.cpp
@@ -1,26 +1,43 @@
 #include "STDInclude.hpp"
 
-BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved)
+namespace Main
 {
 	static Utils::Hook EntryPointHook;
 
-	if (ul_reason_for_call == DLL_PROCESS_ATTACH)
+	void Initialize()
 	{
-		DWORD oldProtect;
-		VirtualProtect(GetModuleHandle(NULL), 0x6C73000, PAGE_EXECUTE_READWRITE, &oldProtect);
-
-		EntryPointHook.Initialize(0x6BAC0F, [] ()
-		{
-			EntryPointHook.Uninstall();
-			Components::Loader::Initialize();
-			__asm jmp EntryPointHook.Place
-
-		})->Install();
+		EntryPointHook.Uninstall();
+		Components::Loader::Initialize();
 	}
-	else if (ul_reason_for_call == DLL_PROCESS_DETACH)
+
+	void Uninitialize()
 	{
 		Components::Loader::Uninitialize();
 	}
 
-	return TRUE;
+	BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved)
+	{
+		if (ul_reason_for_call == DLL_PROCESS_ATTACH)
+		{
+			DWORD oldProtect;
+			VirtualProtect(GetModuleHandle(NULL), 0x6C73000, PAGE_EXECUTE_READWRITE, &oldProtect);
+
+			EntryPointHook.Initialize(0x6BAC0F, [] ()
+			{
+				__asm
+				{
+					call Main::Initialize
+					mov eax, 6BAC0Fh
+					jmp eax
+				}
+
+			})->Install();
+		}
+		else if (ul_reason_for_call == DLL_PROCESS_DETACH)
+		{
+			Main::Uninitialize();
+		}
+
+		return TRUE;
+	}
 }