actually use custom allocation, override allocators of curl and rapidjson

8 files changed, 99 insertions, 25 deletions

diff --git a/NorthstarDedicatedTest/dllmain.cpp b/NorthstarDedicatedTest/dllmain.cpp
index dfc3afe1..691c9bc7 100644
--- a/NorthstarDedicatedTest/dllmain.cpp
+++ b/NorthstarDedicatedTest/dllmain.cpp
@@ -65,7 +65,8 @@ void WaitForDebugger(HMODULE baseAddress)
     }
 }
 
-// in the future this will be called from launcher instead of dllmain
+SourceAllocator* g_SourceAllocator;
+
 bool InitialiseNorthstar()
 {
     if (initialised)
@@ -81,6 +82,8 @@ bool InitialiseNorthstar()
     InstallInitialHooks();
     InitialiseInterfaceCreationHooks();
 
+    g_SourceAllocator = new SourceAllocator;
+
     AddDllLoadCallback("tier0.dll", InitialiseTier0GameUtilFunctions);
     AddDllLoadCallback("engine.dll", WaitForDebugger);
     AddDllLoadCallback("engine.dll", InitialiseEngineGameUtilFunctions);
@@ -129,5 +132,9 @@ bool InitialiseNorthstar()
     // mod manager after everything else
     AddDllLoadCallback("engine.dll", InitialiseModManager);
 
+    // TODO: If you wanna make it more flexible and for example injectable with old Icepick injector
+    // in this place you should iterate over all already loaded DLLs and execute their callbacks and mark them as executed
+    // (as they will never get called otherwise and stuff will fail)
+
     return true;
 }
\ No newline at end of file
diff --git a/NorthstarDedicatedTest/hooks.cpp b/NorthstarDedicatedTest/hooks.cpp
index e58b7afc..5723a8ab 100644
--- a/NorthstarDedicatedTest/hooks.cpp
+++ b/NorthstarDedicatedTest/hooks.cpp
@@ -77,7 +77,7 @@ LPSTR GetCommandLineAHook()
 		}
 
 		auto len = args.length();
-		cmdlineResult = reinterpret_cast<char*>(malloc(len + 1));
+		cmdlineResult = reinterpret_cast<char*>(_malloc_base(len + 1));
 		if (!cmdlineResult)
 		{
 			spdlog::error("malloc failed for command line");
diff --git a/NorthstarDedicatedTest/masterserver.cpp b/NorthstarDedicatedTest/masterserver.cpp
index fa3854d3..2fec6c82 100644
--- a/NorthstarDedicatedTest/masterserver.cpp
+++ b/NorthstarDedicatedTest/masterserver.cpp
@@ -3,7 +3,6 @@
 #include "concommand.h"
 #include "gameutils.h"
 #include "hookutils.h"
-#include "libcurl/include/curl/curl.h"
 #include "serverauthentication.h"
 #include "gameutils.h"
 #include "rapidjson/document.h"
@@ -137,7 +136,7 @@ void MasterServerManager::AuthenticateOriginWithMasterServer(char* uid, char* or
 			{
 				m_successfullyConnected = true;
 
-				rapidjson::Document originAuthInfo;
+				rapidjson_document originAuthInfo;
 				originAuthInfo.Parse(readBuffer.c_str());
 
 				if (originAuthInfo.HasParseError())
@@ -209,7 +208,7 @@ void MasterServerManager::RequestServerList()
 			{
 				m_successfullyConnected = true;
 
-				rapidjson::Document serverInfoJson;
+				rapidjson_document serverInfoJson;
 				serverInfoJson.Parse(readBuffer.c_str());
 
 				if (serverInfoJson.HasParseError())
@@ -231,7 +230,7 @@ void MasterServerManager::RequestServerList()
 					goto REQUEST_END_CLEANUP;
 				}
 
-				rapidjson::GenericArray<false, rapidjson::Value> serverArray = serverInfoJson.GetArray();
+				rapidjson::GenericArray<false, rapidjson_document::GenericValue> serverArray = serverInfoJson.GetArray();
 
 				spdlog::info("Got {} servers", serverArray.Size());
 
@@ -346,7 +345,7 @@ void MasterServerManager::RequestMainMenuPromos()
 			{
 				m_successfullyConnected = true;
 
-				rapidjson::Document mainMenuPromoJson;
+				rapidjson_document mainMenuPromoJson;
 				mainMenuPromoJson.Parse(readBuffer.c_str());
 
 				if (mainMenuPromoJson.HasParseError())
@@ -456,7 +455,7 @@ void MasterServerManager::AuthenticateWithOwnServer(char* uid, char* playerToken
 			{
 				m_successfullyConnected = true;
 
-				rapidjson::Document authInfoJson;
+				rapidjson_document authInfoJson;
 				authInfoJson.Parse(readBuffer.c_str());
 
 				if (authInfoJson.HasParseError())
@@ -588,7 +587,7 @@ void MasterServerManager::AuthenticateWithServer(char* uid, char* playerToken, c
 			{
 				m_successfullyConnected = true;
 
-				rapidjson::Document connectionInfoJson;
+				rapidjson_document connectionInfoJson;
 				connectionInfoJson.Parse(readBuffer.c_str());
 
 				if (connectionInfoJson.HasParseError())
@@ -672,9 +671,9 @@ void MasterServerManager::AddSelfToServerList(int port, int authPort, char* name
 			m_ownServerId[0] = 0;
 
 			// build modinfo obj
-			rapidjson::Document modinfoDoc;
+			rapidjson_document modinfoDoc;
 			modinfoDoc.SetObject();
-			modinfoDoc.AddMember("Mods", rapidjson::Value(rapidjson::kArrayType), modinfoDoc.GetAllocator());
+			modinfoDoc.AddMember("Mods", rapidjson_document::GenericValue(rapidjson::kArrayType), modinfoDoc.GetAllocator());
 
 		int currentModIndex = 0;
 		for (Mod& mod : g_ModManager->m_loadedMods)
@@ -682,7 +681,7 @@ void MasterServerManager::AddSelfToServerList(int port, int authPort, char* name
 			if (!mod.Enabled || (!mod.RequiredOnClient && !mod.Pdiff.size()))
 				continue;
 
-			modinfoDoc["Mods"].PushBack(rapidjson::Value(rapidjson::kObjectType), modinfoDoc.GetAllocator());
+			modinfoDoc["Mods"].PushBack(rapidjson_document::GenericValue(rapidjson::kObjectType), modinfoDoc.GetAllocator());
 			modinfoDoc["Mods"][currentModIndex].AddMember("Name", rapidjson::StringRef(&mod.Name[0]), modinfoDoc.GetAllocator());
 			modinfoDoc["Mods"][currentModIndex].AddMember("Version", rapidjson::StringRef(&mod.Version[0]), modinfoDoc.GetAllocator());
 			modinfoDoc["Mods"][currentModIndex].AddMember("RequiredOnClient", mod.RequiredOnClient, modinfoDoc.GetAllocator());
@@ -738,7 +737,7 @@ void MasterServerManager::AddSelfToServerList(int port, int authPort, char* name
 			{
 				m_successfullyConnected = true;
 
-				rapidjson::Document serverAddedJson;
+				rapidjson_document serverAddedJson;
 				serverAddedJson.Parse(readBuffer.c_str());
 
 				if (serverAddedJson.HasParseError())
@@ -1025,7 +1024,7 @@ void CHostState__State_GameShutdownHook(CHostState* hostState)
 
 MasterServerManager::MasterServerManager()
 {
-	curl_global_init(CURL_GLOBAL_DEFAULT);
+	curl_global_init_mem(CURL_GLOBAL_DEFAULT, _malloc_base, _free_base, _realloc_base, _strdup_base, _calloc_base);
 }
 
 void InitialiseSharedMasterServer(HMODULE baseAddress)
diff --git a/NorthstarDedicatedTest/memalloc.cpp b/NorthstarDedicatedTest/memalloc.cpp
index d301f1fa..c1fb70e7 100644
--- a/NorthstarDedicatedTest/memalloc.cpp
+++ b/NorthstarDedicatedTest/memalloc.cpp
@@ -20,7 +20,7 @@ char pStaticAllocBuf[STATIC_ALLOC_SIZE];
 
 // TODO: rename to malloc and free after removing statically compiled .libs
 
-void* malloc_(size_t n)
+extern "C" void* _malloc_base(size_t n)
 {
 	// allocate into static buffer if g_pMemAllocSingleton isn't initialised
 	if (g_pMemAllocSingleton)
@@ -41,7 +41,12 @@ void* malloc_(size_t n)
 	}
 }
 
-void free_(void* p)
+/*extern "C" void* malloc(size_t n)
+{
+	return _malloc_base(n);
+}*/
+
+extern "C" void _free_base(void* p)
 {
 	// if it was allocated into the static buffer, just do nothing, safest way to deal with it
 	if (p >= pStaticAllocBuf && p <= pStaticAllocBuf + STATIC_ALLOC_SIZE)
@@ -54,7 +59,7 @@ void free_(void* p)
 	g_pMemAllocSingleton->m_vtable->Free(g_pMemAllocSingleton, p);
 }
 
-void* realloc_(void* old_ptr, size_t size) {
+extern "C" void* _realloc_base(void* old_ptr, size_t size) {
 	// it was allocated into the static buffer
 	if (old_ptr >= pStaticAllocBuf && old_ptr <= pStaticAllocBuf + STATIC_ALLOC_SIZE)
 	{
@@ -69,20 +74,42 @@ void* realloc_(void* old_ptr, size_t size) {
 		}
 		else
 		{
-			return malloc_(size);
+			return _malloc_base(size);
 		}
 	}
 
 	if (g_pMemAllocSingleton)
 		return g_pMemAllocSingleton->m_vtable->Realloc(g_pMemAllocSingleton, old_ptr, size);
+	return nullptr;
+}
+
+extern "C" void* _calloc_base(size_t n, size_t size)
+{
+	return _malloc_base(n * size);
+}
+
+extern "C" char* _strdup_base(const char* src)
+{
+	char* str;
+	char* p;
+	int len = 0;
+
+	while (src[len])
+		len++;
+	str = reinterpret_cast<char*>(_malloc_base(len + 1));
+	p = str;
+	while (*src)
+		*p++ = *src++;
+	*p = '\0';
+	return str;
 }
 
 void* operator new(size_t n)
 {
-	return malloc_(n);
+	return _malloc_base(n);
 }
 
 void operator delete(void* p)
 {
-	free_(p);
-}
\ No newline at end of file
+	_free_base(p);
+}// /FORCE:MULTIPLE 
\ No newline at end of file
diff --git a/NorthstarDedicatedTest/memalloc.h b/NorthstarDedicatedTest/memalloc.h
index fe3c5255..d9277694 100644
--- a/NorthstarDedicatedTest/memalloc.h
+++ b/NorthstarDedicatedTest/memalloc.h
@@ -1,6 +1,44 @@
 #pragma once
 
+#include "include/rapidjson/document.h"
+//#include "include/rapidjson/allocators.h"
+
 extern size_t g_iStaticAllocated;
 
+extern "C" {
+	char* _strdup_base(const char* src);
+}
+
 void* operator new(size_t n);
-void operator delete(void* p);
\ No newline at end of file
+void operator delete(void* p);
+
+void* _malloc_base(size_t n);
+//void* malloc(size_t n);
+
+class SourceAllocator {
+public:
+    static const bool kNeedFree = true;
+    void* Malloc(size_t size) {
+        if (size) //  behavior of malloc(0) is implementation defined.
+            return _malloc_base(size);
+        else
+            return NULL; // standardize to returning NULL.
+    }
+    void* Realloc(void* originalPtr, size_t originalSize, size_t newSize) {
+        (void)originalSize;
+        if (newSize == 0) {
+            _free_base(originalPtr);
+            return NULL;
+        }
+        return _realloc_base(originalPtr, newSize);
+    }
+    static void Free(void* ptr) { _free_base(ptr); }
+};
+
+extern SourceAllocator* g_SourceAllocator;
+
+typedef rapidjson::GenericDocument<rapidjson::UTF8<>, rapidjson::MemoryPoolAllocator<SourceAllocator>, SourceAllocator> rapidjson_document;
+//typedef rapidjson::GenericDocument<rapidjson::UTF8<>, SourceAllocator, SourceAllocator> rapidjson_document;
+//typedef rapidjson::Document rapidjson_document;
+//using MyDocument = rapidjson::GenericDocument<rapidjson::UTF8<>, MemoryAllocator>;
+//using rapidjson_document = rapidjson::GenericDocument<rapidjson::UTF8<>, SourceAllocator, SourceAllocator>;
diff --git a/NorthstarDedicatedTest/modmanager.cpp b/NorthstarDedicatedTest/modmanager.cpp
index 49c69c78..a9119075 100644
--- a/NorthstarDedicatedTest/modmanager.cpp
+++ b/NorthstarDedicatedTest/modmanager.cpp
@@ -22,7 +22,7 @@ Mod::Mod(fs::path modDir, char* jsonBuf)
 
 	ModDirectory = modDir;
 
-	rapidjson::Document modJson;
+	rapidjson_document modJson;
 	modJson.Parse<rapidjson::ParseFlag::kParseCommentsFlag | rapidjson::ParseFlag::kParseTrailingCommasFlag>(jsonBuf);
 
 	// fail if parse error
@@ -379,7 +379,7 @@ void ModManager::UnloadMods()
 		// should we be doing this here or should scripts be doing this manually?
 		// main issue with doing this here is when we reload mods for connecting to a server, we write enabled mods, which isn't necessarily what we wanna do
 		if (!m_enabledModsCfg.HasMember(mod.Name.c_str()))
-			m_enabledModsCfg.AddMember(rapidjson::StringRef(mod.Name.c_str()), rapidjson::Value(false), m_enabledModsCfg.GetAllocator());
+			m_enabledModsCfg.AddMember(rapidjson_document::StringRefType(mod.Name.c_str()), rapidjson_document::GenericValue(false), m_enabledModsCfg.GetAllocator());
 
 		m_enabledModsCfg[mod.Name.c_str()].SetBool(mod.Enabled);
 	}
diff --git a/NorthstarDedicatedTest/modmanager.h b/NorthstarDedicatedTest/modmanager.h
index 5f2f6441..20cb0a42 100644
--- a/NorthstarDedicatedTest/modmanager.h
+++ b/NorthstarDedicatedTest/modmanager.h
@@ -4,6 +4,7 @@
 #include <vector>
 #include <filesystem>
 #include "rapidjson/document.h"
+#include "memalloc.h"
 
 namespace fs = std::filesystem;
 
@@ -100,7 +101,7 @@ class ModManager
 private:
 	bool m_hasLoadedMods = false;
 	bool m_hasEnabledModsCfg;
-	rapidjson::Document m_enabledModsCfg;
+	rapidjson_document m_enabledModsCfg;
 
 	// precalculated hashes
 	size_t m_hScriptsRsonHash;
diff --git a/NorthstarDedicatedTest/pch.h b/NorthstarDedicatedTest/pch.h
index 9ac5b8a9..a07d1401 100644
--- a/NorthstarDedicatedTest/pch.h
+++ b/NorthstarDedicatedTest/pch.h
@@ -11,10 +11,12 @@
 // httplib ssl
 
 // add headers that you want to pre-compile here
+#include "memalloc.h"
 #include <Windows.h>
 #include "logging.h"
 #include "include/MinHook.h"
 #include "spdlog/spdlog.h"
+#include "libcurl/include/curl/curl.h"
 #include "hookutils.h"
 
 #endif
\ No newline at end of file