From 54214872121cc00b01eff46855c91cb2731456b0 Mon Sep 17 00:00:00 2001
From: mkrsym1 <mkrsym1@gmail.com>
Date: Wed, 21 Jun 2023 14:25:22 +0300
Subject: [PATCH] Counter-based unloading logic

---
 game_payload/include/main.h |  4 ++++
 game_payload/include/tp6.h  |  2 +-
 game_payload/src/main.c     | 21 ++++++++++++++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)
 create mode 100644 game_payload/include/main.h

diff --git a/game_payload/include/main.h b/game_payload/include/main.h
new file mode 100644
index 0000000..84ac1d2
--- /dev/null
+++ b/game_payload/include/main.h
@@ -0,0 +1,4 @@
+#pragma once
+
+void unload_ctr_inc();
+void unload_ctr_dec();
diff --git a/game_payload/include/tp6.h b/game_payload/include/tp6.h
index 76bec34..c35b120 100644
--- a/game_payload/include/tp6.h
+++ b/game_payload/include/tp6.h
@@ -4,4 +4,4 @@
 
 #include <game.h>
 
-void tp6_setup_patcher(struct game_data *game, HMODULE thisModule, HMODULE baseModule);
+void tp6_setup_patcher(struct game_data *game, HMODULE baseModule);
diff --git a/game_payload/src/main.c b/game_payload/src/main.c
index b1e945d..fbc8d7c 100644
--- a/game_payload/src/main.c
+++ b/game_payload/src/main.c
@@ -6,12 +6,31 @@
 #include <tp6.h>
 #include <utils.h>
 
+#include <main.h>
+
+HMODULE this_module;
+size_t unload_ctr = 0;
+
+void unload_ctr_inc() {
+    unload_ctr++;
+}
+
+void unload_ctr_dec() {
+    unload_ctr--;
+    if (unload_ctr == 0) {
+        void *pFreeLibrary = GetProcAddress(GetModuleHandleA("kernel32.dll"), "FreeLibrary");
+        CreateThread(NULL, 0, pFreeLibrary, this_module, 0, NULL);
+    }
+}
+
 BOOL WINAPI DllMain(HINSTANCE instance, DWORD reason, LPVOID reserved) {
     // Only listen to attach
     if (reason != DLL_PROCESS_ATTACH) {
         return TRUE;
     }
 
+    this_module = instance;
+
     // Dynamically link functions from ntdll
     ntdll_link();
 
@@ -27,7 +46,7 @@ BOOL WINAPI DllMain(HINSTANCE instance, DWORD reason, LPVOID reserved) {
     ace_load_driver_module();
 
     // ...magic
-    tp6_setup_patcher(&game, instance, baseModule);
+    tp6_setup_patcher(&game, baseModule);
 
     // Load the UnityPlayer module and invoke the callback
     HMODULE unityModule = LoadLibraryA("UnityPlayer.dll");