From afb082ee88017a2e0069c29416a5efec0a247b2f Mon Sep 17 00:00:00 2001
From: AeonLucid <aeonlucid@gmail.com>
Date: Sat, 30 Oct 2021 23:31:56 +0200
Subject: [PATCH] Improve hook reliability - fixes GMalloc nullptr and proper
 alloc size

---
 src/Prospect.Agent/SDK.h        | 18 +++++++++++++-----
 src/Prospect.Agent/SDK_Memory.h |  2 +-
 src/Prospect.Agent/main.cpp     |  5 +----
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/src/Prospect.Agent/SDK.h b/src/Prospect.Agent/SDK.h
index 86cea59..d79019e 100644
--- a/src/Prospect.Agent/SDK.h
+++ b/src/Prospect.Agent/SDK.h
@@ -48,13 +48,21 @@ namespace SDK
 	{
 		FString(const wchar_t* other)
 		{
-			Max = Count = *other ? std::wcslen(other) + 1 : 0;
+			const size_t charLen = std::wcslen(other);
 
-			if (Count)
+			if (charLen == 0)
 			{
-				Data = static_cast<wchar_t*>(GMalloc->Malloc(Count));
-				
-				wcscpy_s(Data, Count, other);
+				Count = 0;
+				Max = 0;
+				Data = nullptr;
+			}
+			else
+			{
+				Count = charLen + 1;
+				Max = Count;
+				Data = static_cast<wchar_t*>((*GMalloc)->Malloc(Count * 2));
+
+				memcpy(Data, other, charLen * 2);
 			}
 		}
 
diff --git a/src/Prospect.Agent/SDK_Memory.h b/src/Prospect.Agent/SDK_Memory.h
index af651a9..8927b14 100644
--- a/src/Prospect.Agent/SDK_Memory.h
+++ b/src/Prospect.Agent/SDK_Memory.h
@@ -75,5 +75,5 @@ namespace SDK
 		virtual const TCHAR* GetDescriptiveName() = 0;
 	}; 
 
-	FMalloc* GMalloc;
+	FMalloc** GMalloc;
 }
diff --git a/src/Prospect.Agent/main.cpp b/src/Prospect.Agent/main.cpp
index 1351fd4..d22c5c3 100644
--- a/src/Prospect.Agent/main.cpp
+++ b/src/Prospect.Agent/main.cpp
@@ -38,10 +38,7 @@ DWORD WINAPI OnDllAttach(LPVOID base)
 	const auto hModulePtr = reinterpret_cast<uintptr_t>(hModule);
 
 	// Initialize GMalloc.
-	do
-	{
-		SDK::GMalloc = *reinterpret_cast<SDK::FMalloc**>(hModulePtr + OffsetGMalloc);
-	} while (SDK::GMalloc == nullptr);
+	SDK::GMalloc = reinterpret_cast<SDK::FMalloc**>(hModulePtr + OffsetGMalloc);
 
 	// Initialize MinHook.
 	if (MH_Initialize() != MH_OK)