From 35cd117ffc78228a8688ed118330294ff330e39b Mon Sep 17 00:00:00 2001
From: mkrsym1 <mkrsym1@gmail.com>
Date: Sun, 7 Jan 2024 20:51:55 +0200
Subject: [PATCH] Refactored dynamically linking ntdll

---
 game_payload/include/ntdll.h | 32 +++++++++++++++++++++++++++-----
 game_payload/meson.build     |  1 -
 game_payload/src/main.c      |  6 ++++--
 game_payload/src/ntdll.c     | 11 -----------
 4 files changed, 31 insertions(+), 19 deletions(-)
 delete mode 100644 game_payload/src/ntdll.c

diff --git a/game_payload/include/ntdll.h b/game_payload/include/ntdll.h
index 058ddc5..942bfdc 100644
--- a/game_payload/include/ntdll.h
+++ b/game_payload/include/ntdll.h
@@ -27,10 +27,32 @@ typedef union _LDR_DLL_NOTIFICATION_DATA {
 
 typedef void (*LdrDllNotification_t)(ULONG reason, const PLDR_DLL_NOTIFICATION_DATA data, void *context);
 
-typedef NTSTATUS (*LdrRegisterDllNotification_t)(ULONG flags, LdrDllNotification_t notificationFunction, void *context, void **cookie);
-typedef NTSTATUS (*LdrUnregisterDllNotification_t)(void *cookie);
+#define DYNAMIC_FN_TYPE(ret, name, args) typedef ret (*name##_t)args
 
-extern LdrRegisterDllNotification_t LdrRegisterDllNotification;
-extern LdrUnregisterDllNotification_t LdrUnregisterDllNotification;
+#ifdef NTDLL_DYNAMIC_LINK_IMPL
+    #define DYNAMIC_FN_VAR(name) extern name##_t name; name##_t name
+#else
+    #define DYNAMIC_FN_VAR(name) extern name##_t name
+#endif
 
-void ntdll_link();
+#define DYNAMIC_FN_DEF(ret, name, args) DYNAMIC_FN_TYPE(ret, name, args); DYNAMIC_FN_VAR(name)
+
+DYNAMIC_FN_DEF(NTSTATUS, LdrRegisterDllNotification, (ULONG flags, LdrDllNotification_t notification, void *context, void **cookie));
+DYNAMIC_FN_DEF(NTSTATUS, LdrUnregisterDllNotification, (void *cookie));
+
+#ifdef NTDLL_DYNAMIC_LINK_IMPL
+    #define DYNAMIC_FN_LINK(module, name) name = (name##_t)GetProcAddress(module, #name)
+
+    static void _ntdll_link() {
+        HMODULE ntdll = GetModuleHandleA("ntdll.dll");
+
+        DYNAMIC_FN_LINK(ntdll, LdrRegisterDllNotification);
+        DYNAMIC_FN_LINK(ntdll, LdrUnregisterDllNotification);
+    }
+
+    #undef DYNAMIC_FN_LINK
+#endif
+
+#undef DYNAMIC_FN_TYPE
+#undef DYNAMIC_FN_VAR
+#undef DYNAMIC_FN_DEF
diff --git a/game_payload/meson.build b/game_payload/meson.build
index cf71f6f..2ad1474 100644
--- a/game_payload/meson.build
+++ b/game_payload/meson.build
@@ -5,7 +5,6 @@ include_dir = include_directories('include')
 # Input files
 sources = [
     'src/main.c',
-    'src/ntdll.c',
     'src/ace.c',
     'src/pe.c',
     'src/game.c',
diff --git a/game_payload/src/main.c b/game_payload/src/main.c
index eeedfbb..6253b5c 100644
--- a/game_payload/src/main.c
+++ b/game_payload/src/main.c
@@ -1,6 +1,5 @@
 #include <windows.h>
 
-#include <ntdll.h>
 #include <ace.h>
 #include <game.h>
 #include <core.h>
@@ -8,6 +7,9 @@
 #include <msg.h>
 #include <tx.h>
 
+#define NTDLL_DYNAMIC_LINK_IMPL
+#include <ntdll.h>
+
 #include <main.h>
 
 HMODULE this_module;
@@ -76,7 +78,7 @@ BOOL WINAPI DllMain(HINSTANCE instance, DWORD reason, LPVOID reserved) {
     this_module = instance;
 
     // Dynamically link functions from ntdll
-    ntdll_link();
+    _ntdll_link();
 
     // Detect which game the user is trying to run
     struct game_data game;
diff --git a/game_payload/src/ntdll.c b/game_payload/src/ntdll.c
deleted file mode 100644
index 5ddef95..0000000
--- a/game_payload/src/ntdll.c
+++ /dev/null
@@ -1,11 +0,0 @@
-#include <ntdll.h>
-
-LdrRegisterDllNotification_t LdrRegisterDllNotification;
-LdrUnregisterDllNotification_t LdrUnregisterDllNotification;
-
-void ntdll_link() {
-    HMODULE ntdll = GetModuleHandleA("ntdll.dll");
-
-    LdrRegisterDllNotification = (LdrRegisterDllNotification_t)GetProcAddress(ntdll, "LdrRegisterDllNotification");
-    LdrUnregisterDllNotification = (LdrUnregisterDllNotification_t)GetProcAddress(ntdll, "LdrUnregisterDllNotification");
-}