Module: wine Branch: master Commit: c3bc7689b12a9d62841850d65c786c9cd56725a0 URL: https://source.winehq.org/git/wine.git/?a=commit;h=c3bc7689b12a9d62841850d65...
Author: Paul Gofman pgofman@codeweavers.com Date: Tue May 26 13:05:07 2020 +0300
ntoskrnl.exe: Add KeSetSystemAffinityThreadEx() function.
Signed-off-by: Paul Gofman pgofman@codeweavers.com Signed-off-by: Alexandre Julliard julliard@winehq.org
---
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 f80709a9a0..01a8e67376 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..c7b3c93f6d 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(~(KAFFINITY)0); + 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*);