From ca2f8fea7100a1eec346b57c47d7a6dd1eec4a70 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Sat, 9 Sep 2023 15:29:12 +0100
Subject: [PATCH] [*] Mitigate Kernel32 and Rtl mixing

---
 Source/AuProcAddresses.NT.cpp        |  2 ++
 Source/AuProcAddresses.NT.hpp        | 11 +++++++-
 Source/Threading/AuWakeOnAddress.cpp | 39 +++++++++++++++++++++++++---
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/Source/AuProcAddresses.NT.cpp b/Source/AuProcAddresses.NT.cpp
index 15af09a9..0e2c5512 100644
--- a/Source/AuProcAddresses.NT.cpp
+++ b/Source/AuProcAddresses.NT.cpp
@@ -88,6 +88,8 @@ namespace Aurora
         ADD_GET_PROC(Nt, NtOpenKeyedEvent)
         ADD_GET_PROC(Nt, NtCreateKeyedEvent)
         ADD_GET_PROC(Nt, RtlWaitOnAddress)
+        ADD_GET_PROC(Nt, RtlWakeByAddressAll)
+        ADD_GET_PROC(Nt, RtlWakeAddressSingle)
         ADD_GET_PROC(Nt, ZwSetTimerResolution)
         ADD_GET_PROC(Nt, NtQueryInformationProcess)
         ADD_GET_PROC(Nt, NtNotifyChangeDirectoryFile)
diff --git a/Source/AuProcAddresses.NT.hpp b/Source/AuProcAddresses.NT.hpp
index d3b3db66..b45eab26 100644
--- a/Source/AuProcAddresses.NT.hpp
+++ b/Source/AuProcAddresses.NT.hpp
@@ -146,7 +146,16 @@ namespace Aurora
         const void *                addr,
         const void *                cmp,
         SIZE_T                      size,
-        const LARGE_INTEGER *       timeout);
+        const LARGE_INTEGER *       timeout
+    );
+
+    inline void(__stdcall *pRtlWakeByAddressAll)(
+        const void *                addr
+    );
+
+    inline void(__stdcall *pRtlWakeAddressSingle)(
+        const void *                addr
+    );
 
 #if defined(AURORA_PLATFORM_WIN32)
     inline NTSTATUS(__stdcall *pRtlGetVersion)(
diff --git a/Source/Threading/AuWakeOnAddress.cpp b/Source/Threading/AuWakeOnAddress.cpp
index b44516ba..d4f9b197 100644
--- a/Source/Threading/AuWakeOnAddress.cpp
+++ b/Source/Threading/AuWakeOnAddress.cpp
@@ -462,7 +462,14 @@ namespace Aurora::Threading
                                           AuUInt8 uWordSize)
     {
     #if defined(AURORA_IS_MODERNNT_DERIVED)
-        return pWaitOnAddress((void *)pTargetAddress, (void *)pCompareAddress, uWordSize, INFINITE);
+        if (pRtlWaitOnAddress)
+        {
+            return pRtlWaitOnAddress((void *)pTargetAddress, (void *)pCompareAddress, uWordSize, nullptr);
+        }
+        else
+        {
+            return pWaitOnAddress((void *)pTargetAddress, (void *)pCompareAddress, uWordSize, INFINITE);
+        }
     #endif
 
     #if defined(AURORA_IS_LINUX_DERIVED)
@@ -746,9 +753,26 @@ namespace Aurora::Threading
     #endif
 
     #if defined(AURORA_IS_MODERNNT_DERIVED)
-        for (AuUInt i = 0; i < dwCount; i++)
+        if (pRtlWakeAddressSingle)
         {
-            pWakeByAddressSingle((void *)pAddress);
+            if (dwCount < 6)
+            {
+                for (AuUInt i = 0; i < dwCount; i++)
+                {
+                    pRtlWakeAddressSingle((void *)pAddress);
+                }
+            }
+            else
+            {
+                pRtlWakeByAddressAll((void *)pAddress);
+            }
+        }
+        else
+        {
+            for (AuUInt i = 0; i < dwCount; i++)
+            {
+                pWakeByAddressSingle((void *)pAddress);
+            }
         }
     #endif
     }
@@ -760,7 +784,14 @@ namespace Aurora::Threading
     #endif
 
     #if defined(AURORA_IS_MODERNNT_DERIVED)
-        pWakeByAddressAll((void *)pAddress);
+        if (pRtlWakeByAddressAll)
+        {
+            pRtlWakeByAddressAll((void *)pAddress);
+        }
+        else
+        {
+            pWakeByAddressAll((void *)pAddress);
+        }
     #endif
     }