From: Derek Lesho <dereklesho52@gmail.com>
Subject: [PATCH 1/5] ntoskrnl.exe: Implement KeAreApcsDisabled using critical region functions.
Message-Id: <20190424142509.528-1-dereklesho52@Gmail.com>
Date: Wed, 24 Apr 2019 10:25:05 -0400

Signed-off-by: Derek Lesho <dereklesho52@Gmail.com>
---
 dlls/ntoskrnl.exe/ntoskrnl.c         | 16 ++++++++++++++--
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec  |  2 +-
 dlls/ntoskrnl.exe/ntoskrnl_private.h |  1 +
 dlls/ntoskrnl.exe/tests/driver.c     | 17 +++++++++++++++++
 include/ddk/ntddk.h                  |  1 +
 5 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index c1e6a9cf02..76046c8d90 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -2515,6 +2515,8 @@ static void *create_thread_object( HANDLE handle )
     if (!NtQueryInformationThread( handle, ThreadBasicInformation, &info, sizeof(info), NULL ))
         thread->id = info.ClientId;
 
+    thread->critical_region_count = 0;
+
     return thread;
 }
 
@@ -3417,7 +3419,8 @@ NTSTATUS WINAPI IoCsqInitialize(PIO_CSQ csq, PIO_CSQ_INSERT_IRP insert_irp, PIO_
  */
 void WINAPI KeEnterCriticalRegion(void)
 {
-    FIXME(": stub\n");
+    TRACE(": semi-stub\n");
+    KeGetCurrentThread()->critical_region_count++;
 }
 
 /***********************************************************************
@@ -3425,7 +3428,8 @@ void WINAPI KeEnterCriticalRegion(void)
  */
 void WINAPI KeLeaveCriticalRegion(void)
 {
-    FIXME(": stub\n");
+    TRACE(": semi-stub\n");
+    KeGetCurrentThread()->critical_region_count--;
 }
 
 /***********************************************************************
@@ -4357,3 +4361,11 @@ ULONG WINAPI ExSetTimerResolution(ULONG time, BOOLEAN set_resolution)
     FIXME("stub: %u %d\n", time, set_resolution);
     return KeQueryTimeIncrement();
 }
+
+/*********************************************************************
+ *           KeAreApcsDisabled    (NTOSKRNL.@)
+ */
+BOOLEAN WINAPI KeAreApcsDisabled(void)
+{
+    return !!KeGetCurrentThread()->critical_region_count;
+}
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index 633a8c4b6c..97fc4ec0db 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -521,7 +521,7 @@
 @ stdcall KeAcquireSpinLockAtDpcLevel(ptr)
 @ stdcall -arch=arm,arm64,x86_64 KeAcquireSpinLockRaiseToDpc(ptr)
 @ stub KeAddSystemServiceTable
-@ stub KeAreApcsDisabled
+@ stdcall KeAreApcsDisabled()
 @ stub KeAttachProcess
 @ stub KeBugCheck
 @ stub KeBugCheckEx
diff --git a/dlls/ntoskrnl.exe/ntoskrnl_private.h b/dlls/ntoskrnl.exe/ntoskrnl_private.h
index 5f309616d3..4f5ae14cfb 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl_private.h
+++ b/dlls/ntoskrnl.exe/ntoskrnl_private.h
@@ -36,6 +36,7 @@ struct _KTHREAD
 {
     DISPATCHER_HEADER header;
     CLIENT_ID id;
+    unsigned int critical_region_count;
 };
 
 void *alloc_kernel_object( POBJECT_TYPE type, HANDLE handle, SIZE_T size, LONG ref ) DECLSPEC_HIDDEN;
diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index 030c95507c..f70f865c40 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -1187,6 +1187,22 @@ static void test_lookup_thread(void)
        "PsLookupThreadByThreadId returned %#x\n", status);
 }
 
+static void test_critical_regions(void)
+{
+    BOOLEAN result;
+
+    result = KeAreApcsDisabled();
+    ok(!result, "got %u, expected 0\n", result);
+
+    KeEnterCriticalRegion();
+
+    result = KeAreApcsDisabled();
+
+    KeLeaveCriticalRegion();
+
+    ok(result, "got %u, expected 1\n", result);
+}
+
 static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *stack, ULONG_PTR *info)
 {
     ULONG length = stack->Parameters.DeviceIoControl.OutputBufferLength;
@@ -1231,6 +1247,7 @@ static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *st
     test_ob_reference(test_input->path);
     test_resource();
     test_lookup_thread();
+    test_critical_regions();
 
     /* print process report */
     if (winetest_debug)
diff --git a/include/ddk/ntddk.h b/include/ddk/ntddk.h
index 719ba67c6a..f09f879032 100644
--- a/include/ddk/ntddk.h
+++ b/include/ddk/ntddk.h
@@ -213,6 +213,7 @@ NTSTATUS  WINAPI IoQueryDeviceDescription(PINTERFACE_TYPE,PULONG,PCONFIGURATION_
                                   PCONFIGURATION_TYPE,PULONG,PIO_QUERY_DEVICE_ROUTINE,PVOID);
 void      WINAPI IoRegisterDriverReinitialization(PDRIVER_OBJECT,PDRIVER_REINITIALIZE,PVOID);
 NTSTATUS  WINAPI IoRegisterShutdownNotification(PDEVICE_OBJECT);
+BOOLEAN   WINAPI KeAreApcsDisabled(void);
 NTSTATUS  WINAPI KeExpandKernelStackAndCallout(PEXPAND_STACK_CALLOUT,void*,SIZE_T);
 void      WINAPI KeSetTargetProcessorDpc(PRKDPC,CCHAR);
 BOOLEAN   WINAPI MmIsAddressValid(void *);

-- 
2.20.1