[MEMORY][TRANSACTION] Use CRT initializer

- [MEMORY] Fix private heap didn't initialized - [TRANSACTION] Load `LdrRegisterDllNotification` dynamically
KNSoft · Jun 22, 2024 · c7e6609 · c7e6609
1 parent a3daab8
commit c7e6609
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 28 deletions.
diff --git a/Source/SlimDetours/Memory.c b/Source/SlimDetours/Memory.c
@@ -60,10 +60,10 @@ static ULONG_PTR s_ulSystemRegionLowUpperBound = SYSTEM_RESERVED_REGION_HIGHEST;
 static ULONG_PTR s_ulSystemRegionLowLowerBound = SYSTEM_RESERVED_REGION_LOWEST;
 #endif
 
-VOID detour_memory_init()
+MSVC_INITIALIZER(detour_memory_init)
 {
+    /* Initialize memory management information */
     NtQuerySystemInformation(SystemBasicInformation, &g_sbi, sizeof(g_sbi), NULL);
-
 #if defined(_WIN64)
     PLDR_DATA_TABLE_ENTRY NtdllLdrEntry;
 
@@ -79,21 +79,22 @@ VOID detour_memory_init()
     }
 #endif
 
+    /* Initialize private heap */
+    g_hHeap = RtlCreateHeap(HEAP_NO_SERIALIZE | HEAP_GROWABLE, NULL, 0, 0, NULL, NULL);
     if (g_hHeap == NULL)
     {
-        g_hHeap = RtlCreateHeap(HEAP_NO_SERIALIZE | HEAP_GROWABLE, NULL, 0, 0, NULL, NULL);
+        DETOUR_TRACE("RtlCreateHeap failed, fallback to process default heap\n");
+        g_hHeap = NtGetProcessHeap();
     }
+
+    return 0;
 }
 
 _Must_inspect_result_
 _Ret_maybenull_
 _Post_writable_byte_size_(Size)
 PVOID detour_memory_alloc(_In_ SIZE_T Size)
 {
-    if (g_hHeap == NULL)
-    {
-        g_hHeap = NtGetProcessHeap();
-    }
     return RtlAllocateHeap(g_hHeap, 0, Size);
 }
 

diff --git a/Source/SlimDetours/SlimDetours.cpp b/Source/SlimDetours/SlimDetours.cpp
diff --git a/Source/SlimDetours/SlimDetours.inl b/Source/SlimDetours/SlimDetours.inl
@@ -111,8 +111,6 @@ struct _DETOUR_OPERATION
 
 /* Memory management */
 
-VOID detour_memory_init();
-
 _Must_inspect_result_
 _Ret_maybenull_
 _Post_writable_byte_size_(Size)

diff --git a/Source/SlimDetours/SlimDetours.vcxproj b/Source/SlimDetours/SlimDetours.vcxproj
@@ -29,7 +29,6 @@
   <ItemGroup>
     <ClCompile Include="Instruction.c" />
     <ClCompile Include="Memory.c" />
-    <ClCompile Include="SlimDetours.cpp" />
     <ClCompile Include="Disasm.c" />
     <ClCompile Include="Thread.c" />
     <ClCompile Include="Trampoline.c" />

diff --git a/Source/SlimDetours/SlimDetours.vcxproj.filters b/Source/SlimDetours/SlimDetours.vcxproj.filters
@@ -2,7 +2,6 @@
 <Project ToolsVersion="4.0" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
   <ItemGroup>
     <ClCompile Include="Disasm.c" />
-    <ClCompile Include="SlimDetours.cpp" />
     <ClCompile Include="Memory.c" />
     <ClCompile Include="Instruction.c" />
     <ClCompile Include="Trampoline.c" />

diff --git a/Source/SlimDetours/Transaction.c b/Source/SlimDetours/Transaction.c
@@ -12,7 +12,14 @@
 
 #include "SlimDetours.inl"
 
-#pragma comment(lib, "KNSoft.NDK.WinAPI.lib")
+typedef
+NTSTATUS
+NTAPI
+FN_LdrRegisterDllNotification(
+    _In_ ULONG Flags,
+    _In_ PLDR_DLL_NOTIFICATION_FUNCTION NotificationFunction,
+    _In_opt_ PVOID Context,
+    _Out_ PVOID* Cookie);
 
 typedef struct _DETOUR_DELAY_ATTACH DETOUR_DELAY_ATTACH, *PDETOUR_DELAY_ATTACH;
 typedef VOID(CALLBACK* DETOUR_DELAY_ATTACH_CALLBACK)(
@@ -33,6 +40,8 @@ struct _DETOUR_DELAY_ATTACH
     PVOID Context;
 };
 
+static const ANSI_STRING g_asLdrRegisterDllNotification = RTL_CONSTANT_STRING("LdrRegisterDllNotification");
+static FN_LdrRegisterDllNotification* g_pfnLdrRegisterDllNotification = NULL;
 static HANDLE s_nPendingThreadId = 0; // Thread owning pending transaction.
 static PHANDLE s_phSuspendedThreads = NULL;
 static ULONG s_ulSuspendedThreadCount = 0;
@@ -41,6 +50,22 @@ static RTL_SRWLOCK g_DelayedAttachesLock = RTL_SRWLOCK_INIT;
 static PVOID g_DllNotifyCookie = NULL;
 static PDETOUR_DELAY_ATTACH g_DelayedAttaches = NULL;
 
+MSVC_INITIALIZER(detour_transaction_init)
+{
+    NTSTATUS Status;
+
+    Status = LdrGetProcedureAddress(NtGetNtdllBase(),
+                                    (PANSI_STRING)&g_asLdrRegisterDllNotification,
+                                    0,
+                                    (PVOID*)&g_pfnLdrRegisterDllNotification);
+    if (!NT_SUCCESS(Status))
+    {
+        DETOUR_TRACE("LdrGetProcedureAddress failed to get LdrRegisterDllNotification with 0x%08lX\n", Status);
+        return Status;
+    }
+    return 0;
+}
+
 NTSTATUS NTAPI SlimDetoursTransactionBegin()
 {
     NTSTATUS Status;
@@ -658,7 +683,7 @@ NTSTATUS NTAPI SlimDetoursDelayAttach(
 
     if (g_DllNotifyCookie == NULL)
     {
-        Status = LdrRegisterDllNotification(0, detour_dll_notify_proc, NULL, &g_DllNotifyCookie);
+        Status = g_pfnLdrRegisterDllNotification(0, detour_dll_notify_proc, NULL, &g_DllNotifyCookie);
         if (!NT_SUCCESS(Status))
         {
             RtlReleaseSRWLockExclusive(&g_DelayedAttachesLock);