[PATCH 2/5] ntoskrnl.exe: Add KeSetSystemAffinityThreadEx() function.

[Paul Gofman]

Signed-off-by: Paul Gofman <pgofman at codeweavers.com>
---
 dlls/ntoskrnl.exe/ntoskrnl.c        | 19 +++++++++++++++++++
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec |  1 +
 dlls/ntoskrnl.exe/tests/driver.c    | 19 +++++++++++++++++--
 include/ddk/wdm.h                   |  1 +
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index c10c830738..dc6bbe866c 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -2455,6 +2455,25 @@ VOID WINAPI KeSetSystemAffinityThread(KAFFINITY Affinity)
     FIXME("(%lx) stub\n", Affinity);
 }
 
+KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity)
+{
+    DWORD_PTR system_affinity = KeQueryActiveProcessors();
+    GROUP_AFFINITY old, new;
+
+    TRACE("affinity %#lx.\n", affinity);
+
+    affinity &= system_affinity;
+
+    NtQueryInformationThread(GetCurrentThread(), ThreadGroupInformation,
+            &old, sizeof(old), NULL);
+
+    memset(&new, 0, sizeof(new));
+    new.Mask = affinity;
+
+    return NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new))
+            ? 0 : old.Mask;
+}
+
 
 /***********************************************************************
  *           KeRevertToUserAffinityThread   (NTOSKRNL.EXE.@)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index f25ab9c5e0..e6e641d16c 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -634,6 +634,7 @@
 @ stdcall KeSetPriorityThread(ptr long)
 @ stub KeSetProfileIrql
 @ stdcall KeSetSystemAffinityThread(long)
+@ stdcall KeSetSystemAffinityThreadEx(long)
 @ stdcall KeSetTargetProcessorDpc(ptr long)
 @ stub KeSetTimeIncrement
 @ stub KeSetTimer
diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index 65be4a8d35..da6b9c1d63 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -1712,10 +1712,11 @@ static void test_executable_pool(void)
 
 static void test_affinity(void)
 {
+    KAFFINITY (WINAPI *pKeSetSystemAffinityThreadEx)(KAFFINITY affinity);
     ULONG (WINAPI *pKeQueryActiveProcessorCountEx)(USHORT);
     KAFFINITY (WINAPI *pKeQueryActiveProcessors)(void);
+    KAFFINITY mask, mask_all_cpus;
     ULONG cpu_count, count;
-    KAFFINITY mask;
 
     pKeQueryActiveProcessorCountEx = get_proc_address("KeQueryActiveProcessorCountEx");
     if (!pKeQueryActiveProcessorCountEx)
@@ -1727,6 +1728,9 @@ static void test_affinity(void)
     pKeQueryActiveProcessors = get_proc_address("KeQueryActiveProcessors");
     ok(!!pKeQueryActiveProcessors, "KeQueryActiveProcessors is not available.\n");
 
+    pKeSetSystemAffinityThreadEx = get_proc_address("KeSetSystemAffinityThreadEx");
+    ok(!!pKeSetSystemAffinityThreadEx, "KeSetSystemAffinityThreadEx is not available.\n");
+
     count = pKeQueryActiveProcessorCountEx(1);
     todo_wine ok(!count, "Got unexpected count %u.\n", count);
 
@@ -1736,8 +1740,19 @@ static void test_affinity(void)
     count = pKeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
     ok(count == cpu_count, "Got unexpected count %u.\n", count);
 
+    mask_all_cpus = ~((~0u) << cpu_count);
+
     mask = pKeQueryActiveProcessors();
-    ok(mask == ~((~0u) << cpu_count), "Got unexpected mask %#lx.\n", mask);
+    ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask);
+
+    mask = pKeSetSystemAffinityThreadEx(0);
+    ok(!mask, "Got unexpected mask %#lx.\n", mask);
+
+    mask = pKeSetSystemAffinityThreadEx(0x1);
+    ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask);
+
+    mask = pKeSetSystemAffinityThreadEx(~0ull);
+    ok(mask == 0x1, "Got unexpected mask %#lx.\n", mask);
 }
 
 static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *stack)
diff --git a/include/ddk/wdm.h b/include/ddk/wdm.h
index 1f80a5af29..42ff0c72b1 100644
--- a/include/ddk/wdm.h
+++ b/include/ddk/wdm.h
@@ -1709,6 +1709,7 @@ void      WINAPI KeRevertToUserAffinityThread(void);
 LONG      WINAPI KeSetEvent(PRKEVENT,KPRIORITY,BOOLEAN);
 KPRIORITY WINAPI KeSetPriorityThread(PKTHREAD,KPRIORITY);
 void      WINAPI KeSetSystemAffinityThread(KAFFINITY);
+KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity);
 BOOLEAN   WINAPI KeSetTimerEx(KTIMER*,LARGE_INTEGER,LONG,KDPC*);
 NTSTATUS  WINAPI KeWaitForMultipleObjects(ULONG,void*[],WAIT_TYPE,KWAIT_REASON,KPROCESSOR_MODE,BOOLEAN,LARGE_INTEGER*,KWAIT_BLOCK*);
 NTSTATUS  WINAPI KeWaitForSingleObject(void*,KWAIT_REASON,KPROCESSOR_MODE,BOOLEAN,LARGE_INTEGER*);
-- 
2.26.2