[PATCH v5 1/1] ntoskrnl.exe: add KeNumberProcessors export.

[Etaash Mathamsetty]

From: Etaash Mathamsetty

fix various issues with original implementation

use winxp+ behavior

add tests
---
 dlls/ntoskrnl.exe/ntoskrnl.c        | 3 +++
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec | 2 +-
 dlls/ntoskrnl.exe/tests/driver.c    | 9 +++++++++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index 9320e02dca9..ae508a1a432 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -2595,6 +2595,8 @@ KAFFINITY WINAPI KeQueryActiveProcessors( void )
     return affinity_mask;
 }

+volatile CCHAR KeNumberProcessors = 0;
+
 ULONG WINAPI KeQueryActiveProcessorCountEx(USHORT group_number)
 {
     TRACE("group_number %u.\n", group_number);
@@ -4544,6 +4546,7 @@ BOOL WINAPI DllMain( HINSTANCE inst, DWORD reason, LPVOID reserved )
         NtBuildNumber = NtCurrentTeb()->Peb->OSBuildNumber;
         ntoskrnl_heap = HeapCreate( HEAP_CREATE_ENABLE_EXECUTE, 0, 0 );
         dpc_call_tls_index = TlsAlloc();
+        KeNumberProcessors = (CCHAR)min( 64UL, KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS) );
         LdrRegisterDllNotification( 0, ldr_notify_callback, NULL, &ldr_notify_cookie );
         break;
     case DLL_PROCESS_DETACH:
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index 8f8edcac7a0..322ed5329de 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -593,7 +593,7 @@
 @ stdcall KeLeaveCriticalRegion()
 @ stub KeLoaderBlock
 @ stdcall -arch=x86_64 KeLowerIrql(long)
-@ stub KeNumberProcessors
+@ extern KeNumberProcessors
 @ stub KeProfileInterrupt
 @ stub KeProfileInterruptWithSource
 @ stub KePulseEvent
diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index d293889c823..a0a8741f744 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -2003,6 +2003,7 @@ static void test_affinity(void)
     void (WINAPI *pKeRevertToUserAffinityThreadEx)(KAFFINITY affinity);
     ULONG (WINAPI *pKeQueryActiveProcessorCountEx)(USHORT);
     KAFFINITY (WINAPI *pKeQueryActiveProcessors)(void);
+    CCHAR *pKeNumberProcessors;
     KAFFINITY mask, mask_all_cpus;
     ULONG cpu_count, count;

@@ -2022,6 +2023,9 @@ static void test_affinity(void)
     pKeRevertToUserAffinityThreadEx = get_proc_address("KeRevertToUserAffinityThreadEx");
     ok(!!pKeRevertToUserAffinityThreadEx, "KeRevertToUserAffinityThreadEx is not available.\n");

+    pKeNumberProcessors = get_proc_address("KeNumberProcessors");
+    ok(!!pKeNumberProcessors, "KeNumberProcessors is not available.\n");
+
     count = pKeQueryActiveProcessorCountEx(1);
     ok(!count, "Got unexpected count %lu.\n", count);

@@ -2031,6 +2035,11 @@ static void test_affinity(void)
     count = pKeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
     ok(count == cpu_count, "Got unexpected count %lu.\n", count);

+    if(count >= 64)
+        ok(*pKeNumberProcessors == 64, "Got unexpected count %lu.\n", count);
+    else
+        ok(*pKeNumberProcessors == cpu_count, "Got unexpected count %lu.\n", count);
+
     if (cpu_count >= 8 * sizeof(KAFFINITY))
         mask_all_cpus = ~(KAFFINITY)0;
     else