Alexandre Julliard : kernel32: Fix GetLogicalProcessorInformationEx() usage.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 171de0a7f6510a9fa1e2de4f601e31c3208254fd
URL:    https://source.winehq.org/git/wine.git/?a=commit;h=171de0a7f6510a9fa1e2de4f601e31c3208254fd

Author: Alexandre Julliard <julliard at winehq.org>
Date:   Wed May 26 21:19:42 2021 +0200

kernel32: Fix GetLogicalProcessorInformationEx() usage.

Signed-off-by: Alexandre Julliard <julliard at winehq.org>

---

 dlls/kernel32/process.c          | 57 ++++++++++++++++------------------------
 dlls/ntoskrnl.exe/tests/driver.c |  2 +-
 2 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/dlls/kernel32/process.c b/dlls/kernel32/process.c
index aa0922f5164..17a23cca08d 100644
--- a/dlls/kernel32/process.c
+++ b/dlls/kernel32/process.c
@@ -604,24 +604,34 @@ HRESULT WINAPI RegisterApplicationRecoveryCallback(APPLICATION_RECOVERY_CALLBACK
     return S_OK;
 }
 
+static SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *get_logical_processor_info(void)
+{
+    DWORD size = 0;
+    SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *info;
+
+    GetLogicalProcessorInformationEx( RelationGroup, NULL, &size );
+    if (GetLastError() != ERROR_INSUFFICIENT_BUFFER) return NULL;
+    if (!(info = HeapAlloc( GetProcessHeap(), 0, size ))) return NULL;
+    if (!GetLogicalProcessorInformationEx( RelationGroup, info, &size ))
+    {
+        HeapFree( GetProcessHeap(), 0, info );
+        return NULL;
+    }
+    return info;
+}
+
+
 /***********************************************************************
  *           GetActiveProcessorGroupCount (KERNEL32.@)
  */
 WORD WINAPI GetActiveProcessorGroupCount(void)
 {
     WORD groups;
-    DWORD size = 0;
     SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *info;
 
     TRACE("()\n");
 
-    if (!GetLogicalProcessorInformationEx(RelationGroup, NULL, &size)) return 0;
-    if (!(info = HeapAlloc(GetProcessHeap(), 0, size))) return 0;
-    if (!GetLogicalProcessorInformationEx(RelationGroup, info, &size))
-    {
-        HeapFree(GetProcessHeap(), 0, info);
-        return 0;
-    }
+    if (!(info = get_logical_processor_info())) return 0;
 
     groups = info->Group.ActiveGroupCount;
 
@@ -635,18 +645,11 @@ WORD WINAPI GetActiveProcessorGroupCount(void)
 DWORD WINAPI GetActiveProcessorCount(WORD group)
 {
     DWORD cpus = 0;
-    DWORD size = 0;
     SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *info;
 
     TRACE("(0x%x)\n", group);
 
-    if (!GetLogicalProcessorInformationEx(RelationGroup, NULL, &size)) return 0;
-    if (!(info = HeapAlloc(GetProcessHeap(), 0, size))) return 0;
-    if (!GetLogicalProcessorInformationEx(RelationGroup, info, &size))
-    {
-        HeapFree(GetProcessHeap(), 0, info);
-        return 0;
-    }
+    if (!(info = get_logical_processor_info())) return 0;
 
     if (group == ALL_PROCESSOR_GROUPS)
     {
@@ -669,27 +672,20 @@ DWORD WINAPI GetActiveProcessorCount(WORD group)
 DWORD WINAPI GetMaximumProcessorCount(WORD group)
 {
     DWORD cpus = 0;
-    DWORD size = 0;
     SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *info;
 
     TRACE("(0x%x)\n", group);
 
-    if (!GetLogicalProcessorInformationEx(RelationGroup, NULL, &size)) return 0;
-    if (!(info = HeapAlloc(GetProcessHeap(), 0, size))) return 0;
-    if (!GetLogicalProcessorInformationEx(RelationGroup, info, &size))
-    {
-        HeapFree(GetProcessHeap(), 0, info);
-        return 0;
-    }
+    if (!(info = get_logical_processor_info())) return 0;
 
     if (group == ALL_PROCESSOR_GROUPS)
     {
-        for (group = 0; group < info->Group.ActiveGroupCount; group++)
+        for (group = 0; group < info->Group.MaximumGroupCount; group++)
             cpus += info->Group.GroupInfo[group].MaximumProcessorCount;
     }
     else
     {
-        if (group < info->Group.ActiveGroupCount)
+        if (group < info->Group.MaximumGroupCount)
             cpus = info->Group.GroupInfo[group].MaximumProcessorCount;
     }
 
@@ -703,18 +699,11 @@ DWORD WINAPI GetMaximumProcessorCount(WORD group)
 WORD WINAPI GetMaximumProcessorGroupCount(void)
 {
     WORD groups;
-    DWORD size = 0;
     SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *info;
 
     TRACE("()\n");
 
-    if (!GetLogicalProcessorInformationEx(RelationGroup, NULL, &size)) return 0;
-    if (!(info = HeapAlloc(GetProcessHeap(), 0, size))) return 0;
-    if (!GetLogicalProcessorInformationEx(RelationGroup, info, &size))
-    {
-        HeapFree(GetProcessHeap(), 0, info);
-        return 0;
-    }
+    if (!(info = get_logical_processor_info())) return 0;
 
     groups = info->Group.MaximumGroupCount;
 
diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index 0d991cd98d6..f1c72d8fdaf 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -1832,7 +1832,7 @@ static void test_affinity(void)
     ok(!!pKeRevertToUserAffinityThreadEx, "KeRevertToUserAffinityThreadEx is not available.\n");
 
     count = pKeQueryActiveProcessorCountEx(1);
-    todo_wine ok(!count, "Got unexpected count %u.\n", count);
+    ok(!count, "Got unexpected count %u.\n", count);
 
     cpu_count = pKeQueryActiveProcessorCountEx(0);
     ok(cpu_count, "Got unexpected cpu_count %u.\n", cpu_count);