From fda39b207c5b000219dc15c6971a8a6e81310551 Mon Sep 17 00:00:00 2001
From: Erich Hoover <ehoover@mines.edu>
Date: Fri, 19 Feb 2010 19:07:41 -0700
Subject: [PATCH] ntdll: Allow 'all processors' flag used in Vista and newer.

---
 dlls/ntdll/tests/info.c | 17 +++++++++++++++++
 dlls/ntdll/thread.c     | 11 +++++++----
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/dlls/ntdll/tests/info.c b/dlls/ntdll/tests/info.c
index 6414cac01e8..c7d451c70b3 100644
--- a/dlls/ntdll/tests/info.c
+++ b/dlls/ntdll/tests/info.c
@@ -927,11 +927,28 @@ static void test_affinity(void)
     status = pNtQueryInformationThread( GetCurrentThread(), ThreadBasicInformation, &tbi, sizeof(tbi), NULL );
     ok( tbi.AffinityMask == 1, "Unexpected thread affinity\n" );
 
+    /* NOTE: Pre-Vista does not recognize the "all processors" flag (all bits set) */
+    thread_affinity = ~0UL;
+    status = pNtSetInformationThread( GetCurrentThread(), ThreadAffinityMask, &thread_affinity, sizeof(thread_affinity) );
+    ok( broken(status == STATUS_INVALID_PARAMETER) || status == STATUS_SUCCESS,
+        "Expected STATUS_SUCCESS, got %08x\n", status);
+
     if (si.dwNumberOfProcessors <= 1)
     {
         skip("only one processor, skipping affinity testing\n");
         return;
     }
+
+    /* Test thread affinity mask resulting from "all processors" flag */
+    if (status == STATUS_SUCCESS)
+    {
+        status = pNtQueryInformationThread( GetCurrentThread(), ThreadBasicInformation, &tbi, sizeof(tbi), NULL );
+        ok( broken(tbi.AffinityMask == 1) || tbi.AffinityMask == (1 << si.dwNumberOfProcessors) - 1,
+            "Unexpected thread affinity\n" );
+    }
+    else
+        skip("Cannot test thread affinity mask for 'all processors' flag\n");
+
     proc_affinity = 2;
     status = pNtSetInformationProcess( GetCurrentProcess(), ProcessAffinityMask, &proc_affinity, sizeof(proc_affinity) );
     ok( status == STATUS_SUCCESS, "Expected STATUS_SUCCESS, got %08x\n", status);
diff --git a/dlls/ntdll/thread.c b/dlls/ntdll/thread.c
index 65fc7374b9b..d632f3968fb 100644
--- a/dlls/ntdll/thread.c
+++ b/dlls/ntdll/thread.c
@@ -1134,14 +1134,17 @@ NTSTATUS WINAPI NtSetInformationThread( HANDLE handle, THREADINFOCLASS class,
     case ThreadAffinityMask:
         {
             const ULONG_PTR affinity_mask = ((ULONG_PTR)1 << NtCurrentTeb()->Peb->NumberOfProcessors) - 1;
-            const ULONG_PTR *paff = data;
+            ULONG_PTR req_aff;
+
             if (length != sizeof(ULONG_PTR)) return STATUS_INVALID_PARAMETER;
-            if (*paff & ~affinity_mask) return STATUS_INVALID_PARAMETER;
-            if (!*paff) return STATUS_INVALID_PARAMETER;
+            req_aff = *(const ULONG_PTR *)data;
+            if (req_aff == ~0UL) req_aff = affinity_mask;
+            else if (req_aff & ~affinity_mask) return STATUS_INVALID_PARAMETER;
+            else if (!req_aff) return STATUS_INVALID_PARAMETER;
             SERVER_START_REQ( set_thread_info )
             {
                 req->handle   = wine_server_obj_handle( handle );
-                req->affinity = *paff;
+                req->affinity = req_aff;
                 req->mask     = SET_THREAD_INFO_AFFINITY;
                 status = wine_server_call( req );
             }