Hans Leidekker : ntdll: Fix calculation of process and thread affinity masks on systems with a large number of processors.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: faf70e2fbddc5299afd21d8298a81f257127c230
URL:    http://source.winehq.org/git/wine.git/?a=commit;h=faf70e2fbddc5299afd21d8298a81f257127c230

Author: Hans Leidekker <hans at codeweavers.com>
Date:   Wed Jan 21 13:24:14 2015 +0100

ntdll: Fix calculation of process and thread affinity masks on systems with a large number of processors.

---

 dlls/ntdll/ntdll_misc.h |  1 +
 dlls/ntdll/process.c    | 17 ++++++++++++++---
 dlls/ntdll/thread.c     |  6 +++---
 dlls/ntdll/virtual.c    |  2 +-
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/dlls/ntdll/ntdll_misc.h b/dlls/ntdll/ntdll_misc.h
index b7ea6dc..92e5baf 100644
--- a/dlls/ntdll/ntdll_misc.h
+++ b/dlls/ntdll/ntdll_misc.h
@@ -46,6 +46,7 @@ struct drive_info
 };
 
 extern NTSTATUS close_handle( HANDLE ) DECLSPEC_HIDDEN;
+extern ULONG_PTR get_system_affinity_mask(void) DECLSPEC_HIDDEN;
 
 /* exceptions */
 extern void wait_suspend( CONTEXT *context ) DECLSPEC_HIDDEN;
diff --git a/dlls/ntdll/process.c b/dlls/ntdll/process.c
index 3defb6e..3ac8d52 100644
--- a/dlls/ntdll/process.c
+++ b/dlls/ntdll/process.c
@@ -98,6 +98,13 @@ static UINT process_error_mode;
         ret = STATUS_INVALID_INFO_CLASS; \
         break
 
+ULONG_PTR get_system_affinity_mask(void)
+{
+    ULONG num_cpus = NtCurrentTeb()->Peb->NumberOfProcessors;
+    if (num_cpus >= sizeof(ULONG_PTR) * 8) return ~(ULONG_PTR)0;
+    return ((ULONG_PTR)1 << num_cpus) - 1;
+}
+
 /******************************************************************************
 *  NtQueryInformationProcess		[NTDLL.@]
 *  ZwQueryInformationProcess		[NTDLL.@]
@@ -145,7 +152,7 @@ NTSTATUS WINAPI NtQueryInformationProcess(
     case ProcessBasicInformation:
         {
             PROCESS_BASIC_INFORMATION pbi;
-            const ULONG_PTR affinity_mask = ((ULONG_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
+            const ULONG_PTR affinity_mask = get_system_affinity_mask();
 
             if (ProcessInformationLength >= sizeof(PROCESS_BASIC_INFORMATION))
             {
@@ -389,7 +396,7 @@ NTSTATUS WINAPI NtQueryInformationProcess(
         len = sizeof(ULONG_PTR);
         if (ProcessInformationLength == len)
         {
-            const ULONG_PTR system_mask = ((ULONG_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
+            const ULONG_PTR system_mask = get_system_affinity_mask();
 
             SERVER_START_REQ(get_process_info)
             {
@@ -487,8 +494,11 @@ NTSTATUS WINAPI NtSetInformationProcess(
         process_error_mode = *(UINT *)ProcessInformation;
         break;
     case ProcessAffinityMask:
+    {
+        const ULONG_PTR system_mask = get_system_affinity_mask();
+
         if (ProcessInformationLength != sizeof(DWORD_PTR)) return STATUS_INVALID_PARAMETER;
-        if (*(PDWORD_PTR)ProcessInformation & ~(((DWORD_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1))
+        if (*(PDWORD_PTR)ProcessInformation & ~system_mask)
             return STATUS_INVALID_PARAMETER;
         if (!*(PDWORD_PTR)ProcessInformation)
             return STATUS_INVALID_PARAMETER;
@@ -501,6 +511,7 @@ NTSTATUS WINAPI NtSetInformationProcess(
         }
         SERVER_END_REQ;
         break;
+    }
     case ProcessPriorityClass:
         if (ProcessInformationLength != sizeof(PROCESS_PRIORITY_CLASS))
             return STATUS_INVALID_PARAMETER;
diff --git a/dlls/ntdll/thread.c b/dlls/ntdll/thread.c
index c8461b0..7dd13f2 100644
--- a/dlls/ntdll/thread.c
+++ b/dlls/ntdll/thread.c
@@ -900,7 +900,7 @@ NTSTATUS WINAPI NtQueryInformationThread( HANDLE handle, THREADINFOCLASS class,
     case ThreadBasicInformation:
         {
             THREAD_BASIC_INFORMATION info;
-            const ULONG_PTR affinity_mask = ((ULONG_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
+            const ULONG_PTR affinity_mask = get_system_affinity_mask();
 
             SERVER_START_REQ( get_thread_info )
             {
@@ -927,7 +927,7 @@ NTSTATUS WINAPI NtQueryInformationThread( HANDLE handle, THREADINFOCLASS class,
         return status;
     case ThreadAffinityMask:
         {
-            const ULONG_PTR affinity_mask = ((ULONG_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
+            const ULONG_PTR affinity_mask = get_system_affinity_mask();
             ULONG_PTR affinity = 0;
 
             SERVER_START_REQ( get_thread_info )
@@ -1170,7 +1170,7 @@ NTSTATUS WINAPI NtSetInformationThread( HANDLE handle, THREADINFOCLASS class,
         return status;
     case ThreadAffinityMask:
         {
-            const ULONG_PTR affinity_mask = ((ULONG_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
+            const ULONG_PTR affinity_mask = get_system_affinity_mask();
             ULONG_PTR req_aff;
 
             if (length != sizeof(ULONG_PTR)) return STATUS_INVALID_PARAMETER;
diff --git a/dlls/ntdll/virtual.c b/dlls/ntdll/virtual.c
index 4c4c05d..9f82204 100644
--- a/dlls/ntdll/virtual.c
+++ b/dlls/ntdll/virtual.c
@@ -1409,7 +1409,7 @@ void virtual_get_system_info( SYSTEM_BASIC_INFORMATION *info )
     info->AllocationGranularity   = get_mask(0) + 1;
     info->LowestUserAddress       = (void *)0x10000;
     info->HighestUserAddress      = (char *)user_space_limit - 1;
-    info->ActiveProcessorsAffinityMask = (1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
+    info->ActiveProcessorsAffinityMask = get_system_affinity_mask();
     info->NumberOfProcessors      = NtCurrentTeb()->Peb->NumberOfProcessors;
 }