From 24fec2301c2f8b08c346edb06aa8e08a664f2ae1 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Fri, 22 Dec 2023 05:01:04 +0000
Subject: [PATCH] [+] NtWaitForMultipleObjects awareness

---
 Source/AuProcAddresses.NT.cpp   |  1 +
 Source/AuProcAddresses.NT.hpp   |  9 +++++++++
 Source/IO/Loop/LoopQueue.NT.cpp | 16 ++++++++++++++--
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/Source/AuProcAddresses.NT.cpp b/Source/AuProcAddresses.NT.cpp
index 48f24d51..e1034954 100644
--- a/Source/AuProcAddresses.NT.cpp
+++ b/Source/AuProcAddresses.NT.cpp
@@ -241,6 +241,7 @@ namespace Aurora
         ADD_GET_PROC(Nt, NtTerminateProcess)
         ADD_GET_PROC(Nt, NtQuerySymbolicLinkObject)
         ADD_GET_PROC(Nt, NtOpenSymbolicLinkObject)
+        ADD_GET_PROC(Nt, NtWaitForMultipleObjects)
             
         ADD_GET_PROC_BI(Kernel32, KernelBase, VirtualAlloc2)
         ADD_GET_PROC_BI(Kernel32, KernelBase, MapViewOfFile3)
diff --git a/Source/AuProcAddresses.NT.hpp b/Source/AuProcAddresses.NT.hpp
index 823adf9f..c229b6dc 100644
--- a/Source/AuProcAddresses.NT.hpp
+++ b/Source/AuProcAddresses.NT.hpp
@@ -25,6 +25,7 @@ struct _SP_DEVICE_INTERFACE_DETAIL_DATA_W;
 struct _NETRESOURCEW;
 enum _SE_OBJECT_TYPE;
 enum _MINIDUMP_TYPE;
+enum _OBJECT_WAIT_TYPE;
 
 //#if defined(AURORA_COMPILER_MSVC)
     struct _IP_ADAPTER_ADDRESSES_LH;
@@ -94,6 +95,14 @@ namespace Aurora
         PVOID                       Address
     );
 
+    inline NTSTATUS(__stdcall *pNtWaitForMultipleObjects)(
+        ULONG                       ObjectCount,
+        PHANDLE                     ObjectsArray,
+        _OBJECT_WAIT_TYPE           WaitType,
+        BOOLEAN                     Alertable,
+        AuUInt64 *                  pTimeOut
+    );
+
     inline NTSTATUS(__stdcall *pNtQuerySymbolicLinkObject)(
         HANDLE                      LinkHandle,
         PUNICODE_STRING             LinkTarget,
diff --git a/Source/IO/Loop/LoopQueue.NT.cpp b/Source/IO/Loop/LoopQueue.NT.cpp
index 8f5bae38..bd0cd015 100644
--- a/Source/IO/Loop/LoopQueue.NT.cpp
+++ b/Source/IO/Loop/LoopQueue.NT.cpp
@@ -914,7 +914,7 @@ namespace Aurora::IO::Loop
         bool active = this->hEvent_ == INVALID_HANDLE_VALUE;
 
 
-        while (count != index )
+        while (count != index)
         {
             auto next = AuMin(count - index, AuUInt32(MAXIMUM_WAIT_OBJECTS));
 
@@ -960,7 +960,19 @@ namespace Aurora::IO::Loop
             {
                 do
                 {
-                    status = ::WaitForMultipleObjectsEx(next, this->handleArrayOr_.data() + index, false, sleepMS, true);
+                    if (pNtWaitForMultipleObjects)
+                    {
+                        AuUInt64 uTimeout = -(AuMSToNS<AuUInt64>(1) / 2ull / 100ull);
+                        if (AuSwInfo::IsWindows10OrGreater())
+                        {
+                            uTimeout /= 2ull;
+                        }
+                        status = pNtWaitForMultipleObjects(next, this->handleArrayOr_.data() + index, (_OBJECT_WAIT_TYPE)1, TRUE, &uTimeout);
+                    }
+                    else
+                    {
+                        status = ::WaitForMultipleObjectsEx(next, this->handleArrayOr_.data() + index, false, sleepMS, true);
+                    }
                 }
                 while (status == WAIT_IO_COMPLETION);
             }