3 files changed, 74 insertions, 19 deletions

diff --git a/NorthstarDedicatedTest/dedicated.cpp b/NorthstarDedicatedTest/dedicated.cpp
index c97ac91d..c00bed07 100644
--- a/NorthstarDedicatedTest/dedicated.cpp
+++ b/NorthstarDedicatedTest/dedicated.cpp
@@ -184,13 +184,6 @@ void InitialiseDedicated(HMODULE engineAddress)
 	//	NSMem::NOP(ea + 0x1C4E07, 5);
 	//}
 
-	// not sure if this should be done, not loading ui at least is good, but should everything be gone?
-	{
-		// Host_Init
-		// change the number of rpaks to load from 6 to 1, so we only load common.rpak
-		NSMem::BytePatch(ea + 0x15653B + 1, "01");
-	}
-
 	{
 		// Host_Init
 		// remove call to ui loading stuff
diff --git a/NorthstarDedicatedTest/rpakfilesystem.cpp b/NorthstarDedicatedTest/rpakfilesystem.cpp
index e9d5f0c6..b4ddcd1a 100644
--- a/NorthstarDedicatedTest/rpakfilesystem.cpp
+++ b/NorthstarDedicatedTest/rpakfilesystem.cpp
@@ -2,12 +2,14 @@
 #include "rpakfilesystem.h"
 #include "hookutils.h"
 #include "modmanager.h"
+#include "dedicated.h"
 
 typedef void* (*LoadCommonPaksForMapType)(char* map);
 LoadCommonPaksForMapType LoadCommonPaksForMap;
 
 typedef void* (*LoadPakSyncType)(const char* path, void* unknownSingleton, int flags);
 typedef int (*LoadPakAsyncType)(const char* path, void* unknownSingleton, int flags, void* callback0, void* callback1);
+typedef void* (*UnloadPakType)(int pakHandle, void* callback);
 typedef void* (*ReadFullFileFromDiskType)(const char* requestedPath, void* a2);
 
 // there are more i'm just too lazy to add
@@ -16,7 +18,9 @@ struct PakLoadFuncs
 	void* unk0[2];
 	LoadPakSyncType LoadPakSync;
 	LoadPakAsyncType LoadPakAsync;
-	void* unk1[20];
+	void* unk1[2];
+	UnloadPakType UnloadPak;
+	void* unk2[17];
 	ReadFullFileFromDiskType ReadFullFileFromDisk;
 };
 
@@ -25,7 +29,21 @@ void** pUnknownPakLoadSingleton;
 
 PakLoadManager* g_PakLoadManager;
 void PakLoadManager::LoadPakSync(const char* path) { g_pakLoadApi->LoadPakSync(path, *pUnknownPakLoadSingleton, 0); }
-void PakLoadManager::LoadPakAsync(const char* path) { g_pakLoadApi->LoadPakAsync(path, *pUnknownPakLoadSingleton, 2, nullptr, nullptr); }
+void PakLoadManager::LoadPakAsync(const char* path, bool bMarkForUnload)
+{
+	int handle = g_pakLoadApi->LoadPakAsync(path, *pUnknownPakLoadSingleton, 2, nullptr, nullptr);
+
+	if (bMarkForUnload)
+		m_pakHandlesToUnload.push_back(handle);
+}
+
+void PakLoadManager::UnloadPaks()
+{
+	for (int pakHandle : m_pakHandlesToUnload)
+		g_pakLoadApi->UnloadPak(pakHandle, nullptr);
+
+	m_pakHandlesToUnload.clear();
+}
 
 void HandlePakAliases(char** map)
 {
@@ -57,14 +75,14 @@ void LoadPreloadPaks()
 
 		for (ModRpakEntry& pak : mod.Rpaks)
 			if (pak.m_bAutoLoad)
-				g_PakLoadManager->LoadPakAsync((modPakPath / pak.m_sPakName).string().c_str());
+				g_PakLoadManager->LoadPakAsync((modPakPath / pak.m_sPakName).string().c_str(), false);
 	}
 }
 
-void LoadCustomMapPaks(char** pakName)
+void LoadCustomMapPaks(char** pakName, bool* bNeedToFreePakName)
 {
 	// whether the vanilla game has this rpak
-	bool bHasOriginalPak = fs::exists(fs::path("./r2/paks") / *pakName);
+	bool bHasOriginalPak = fs::exists(fs::path("./r2/paks/Win64/") / *pakName);
 
 	// note, loading from ./ is necessary otherwise paks will load from gamedir/r2/paks
 	for (Mod& mod : g_ModManager->m_loadedMods)
@@ -82,11 +100,17 @@ void LoadCustomMapPaks(char** pakName)
 				// if the game doesn't have the original pak, let it handle loading this one as if it was the one it was loading originally
 				if (!bHasOriginalPak)
 				{
-					*pakName = &pak.m_sPakName[0];
+					std::string path = (modPakPath / pak.m_sPakName).string();
+					*pakName = new char[path.size() + 1];
+					strcpy(*pakName, &path[0]);
+					(*pakName)[path.size()] = '\0';
+
 					bHasOriginalPak = true;
+					*bNeedToFreePakName =
+						true; // we can't free this memory until we're done with the pak, so let whatever's calling this deal with it
 				}
 				else
-					g_PakLoadManager->LoadPakAsync((modPakPath / pak.m_sPakName).string().c_str());
+					g_PakLoadManager->LoadPakAsync((modPakPath / pak.m_sPakName).string().c_str(), true);
 			}
 		}
 	}
@@ -97,6 +121,8 @@ void* LoadPakSyncHook(char* path, void* unknownSingleton, int flags)
 {
 	HandlePakAliases(&path);
 
+	bool bNeedToFreePakName = false;
+
 	// note: we don't handle loading any preloaded custom paks synchronously since LoadPakSync is never actually called in retail, just load
 	// them async instead
 	static bool bShouldLoadPaks = true;
@@ -106,13 +132,18 @@ void* LoadPakSyncHook(char* path, void* unknownSingleton, int flags)
 		bShouldLoadPaks = false;
 
 		LoadPreloadPaks();
-		LoadCustomMapPaks(&path);
+		LoadCustomMapPaks(&path, &bNeedToFreePakName);
 
 		bShouldLoadPaks = true;
 	}
 
 	spdlog::info("LoadPakSync {}", path);
-	return LoadPakSyncOriginal(path, unknownSingleton, flags);
+	void* ret = LoadPakSyncOriginal(path, unknownSingleton, flags);
+
+	if (bNeedToFreePakName)
+		delete[] path;
+
+	return ret;
 }
 
 LoadPakAsyncType LoadPakAsyncOriginal;
@@ -120,6 +151,8 @@ int LoadPakAsyncHook(char* path, void* unknownSingleton, int flags, void* callba
 {
 	HandlePakAliases(&path);
 
+	bool bNeedToFreePakName = false;
+
 	static bool bShouldLoadPaks = true;
 	if (bShouldLoadPaks)
 	{
@@ -127,16 +160,40 @@ int LoadPakAsyncHook(char* path, void* unknownSingleton, int flags, void* callba
 		bShouldLoadPaks = false;
 
 		LoadPreloadPaks();
-		LoadCustomMapPaks(&path);
+		LoadCustomMapPaks(&path, &bNeedToFreePakName);
 
 		bShouldLoadPaks = true;
+
+		// do this after custom paks load and in bShouldLoadPaks so we only ever call this on the root pakload call
+		// todo: could probably add some way to flag custom paks to not be loaded on dedicated servers in rpak.json
+		if (IsDedicated() && strncmp(path, "common", 6)) // dedicated only needs common and common_mp
+			return -1;
 	}
 
 	int ret = LoadPakAsyncOriginal(path, unknownSingleton, flags, callback0, callback1);
 	spdlog::info("LoadPakAsync {} {}", path, ret);
+
+	if (bNeedToFreePakName)
+		delete[] path;
+
 	return ret;
 }
 
+UnloadPakType UnloadPakOriginal;
+void* UnloadPakHook(int pakHandle, void* callback)
+{
+	static bool bShouldUnloadPaks = true;
+	if (bShouldUnloadPaks)
+	{
+		bShouldUnloadPaks = false;
+		g_PakLoadManager->UnloadPaks();
+		bShouldUnloadPaks = true;
+	}
+
+	spdlog::info("UnloadPak {}", pakHandle);
+	return UnloadPakOriginal(pakHandle, callback);
+}
+
 // we hook this exclusively for resolving stbsp paths, but seemingly it's also used for other stuff like vpk and rpak loads
 // possibly just async loading all together?
 ReadFullFileFromDiskType ReadFullFileFromDiskOriginal;
@@ -158,7 +215,7 @@ void* ReadFullFileFromDiskHook(const char* requestedPath, void* a2)
 			std::string newPath = (modFile->second.owningMod->ModDirectory / "mod" / modFile->second.path).string();
 			allocatedNewPath = new char[newPath.size() + 1];
 			strncpy(allocatedNewPath, newPath.c_str(), newPath.size());
-			allocatedNewPath[newPath.size() + 1] = '\0';
+			allocatedNewPath[newPath.size()] = '\0';
 			requestedPath = allocatedNewPath;
 		}
 	}
@@ -180,6 +237,7 @@ void InitialiseEngineRpakFilesystem(HMODULE baseAddress)
 	HookEnabler hook;
 	ENABLER_CREATEHOOK(hook, g_pakLoadApi->LoadPakSync, &LoadPakSyncHook, reinterpret_cast<LPVOID*>(&LoadPakSyncOriginal));
 	ENABLER_CREATEHOOK(hook, g_pakLoadApi->LoadPakAsync, &LoadPakAsyncHook, reinterpret_cast<LPVOID*>(&LoadPakAsyncOriginal));
+	ENABLER_CREATEHOOK(hook, g_pakLoadApi->UnloadPak, &UnloadPakHook, reinterpret_cast<LPVOID*>(&UnloadPakOriginal));
 	ENABLER_CREATEHOOK(
 		hook, g_pakLoadApi->ReadFullFileFromDisk, &ReadFullFileFromDiskHook, reinterpret_cast<LPVOID*>(&ReadFullFileFromDiskOriginal));
 }
\ No newline at end of file
diff --git a/NorthstarDedicatedTest/rpakfilesystem.h b/NorthstarDedicatedTest/rpakfilesystem.h
index 334104a9..3c104822 100644
--- a/NorthstarDedicatedTest/rpakfilesystem.h
+++ b/NorthstarDedicatedTest/rpakfilesystem.h
@@ -6,7 +6,11 @@ class PakLoadManager
 {
   public:
 	void LoadPakSync(const char* path);
-	void LoadPakAsync(const char* path);
+	void LoadPakAsync(const char* path, bool bMarkForUnload);
+	void UnloadPaks();
+
+  private:
+	std::vector<int> m_pakHandlesToUnload;
 };
 
 extern PakLoadManager* g_PakLoadManager;
\ No newline at end of file