ntdll: Fix calculation of process and thread affinity masks on systems with a…

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;

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;

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;

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;
 }