Module: wine Branch: master Commit: 7d4f4783a5d166b0d2987ca4c2b9188b405bdcbe URL: https://source.winehq.org/git/wine.git/?a=commit;h=7d4f4783a5d166b0d2987ca4c...
Author: Paul Gofman pgofman@codeweavers.com Date: Tue May 26 13:05:09 2020 +0300
ntoskrnl.exe: Implement KeRevertToUserAffinityThreadEx() function.
Signed-off-by: Paul Gofman pgofman@codeweavers.com Signed-off-by: Alexandre Julliard julliard@winehq.org
---
dlls/ntoskrnl.exe/ntoskrnl.c | 24 +++++++++++++++++++++++- dlls/ntoskrnl.exe/ntoskrnl.exe.spec | 1 + dlls/ntoskrnl.exe/ntoskrnl_private.h | 1 + dlls/ntoskrnl.exe/tests/driver.c | 16 +++++++++++++++- include/ddk/wdm.h | 1 + 5 files changed, 41 insertions(+), 2 deletions(-)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index a0657b52c5..fc7828cf15 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2267,6 +2267,7 @@ static void *create_thread_object( HANDLE handle )
thread->header.Type = 6; thread->header.WaitListHead.Blink = INVALID_HANDLE_VALUE; /* mark as kernel object */ + thread->user_affinity = 0;
if (!NtQueryInformationThread( handle, ThreadBasicInformation, &info, sizeof(info), NULL )) { @@ -2458,6 +2459,7 @@ VOID WINAPI KeSetSystemAffinityThread(KAFFINITY affinity) KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity) { DWORD_PTR system_affinity = KeQueryActiveProcessors(); + PKTHREAD thread = KeGetCurrentThread(); GROUP_AFFINITY old, new;
TRACE("affinity %#lx.\n", affinity); @@ -2467,11 +2469,14 @@ KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity) NtQueryInformationThread(GetCurrentThread(), ThreadGroupInformation, &old, sizeof(old), NULL);
+ if (old.Mask != system_affinity) + thread->user_affinity = old.Mask; + memset(&new, 0, sizeof(new)); new.Mask = affinity;
return NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new)) - ? 0 : old.Mask; + ? 0 : thread->user_affinity; }
@@ -2483,6 +2488,23 @@ void WINAPI KeRevertToUserAffinityThread(void) FIXME("() stub\n"); }
+void WINAPI KeRevertToUserAffinityThreadEx(KAFFINITY affinity) +{ + DWORD_PTR system_affinity = KeQueryActiveProcessors(); + PRKTHREAD thread = KeGetCurrentThread(); + GROUP_AFFINITY new; + + TRACE("affinity %#lx.\n", affinity); + + affinity &= system_affinity; + + memset(&new, 0, sizeof(new)); + new.Mask = affinity ? affinity + : (thread->user_affinity ? thread->user_affinity : system_affinity); + + NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new)); + thread->user_affinity = affinity; +}
/*********************************************************************** * IoRegisterFileSystem (NTOSKRNL.EXE.@) diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec index 01a8e67376..3c7a9d1950 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec +++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec @@ -619,6 +619,7 @@ @ stdcall KeResetEvent(ptr) @ stub KeRestoreFloatingPointState @ stdcall KeRevertToUserAffinityThread() +@ stdcall KeRevertToUserAffinityThreadEx(long) @ stub KeRundownQueue @ stub KeSaveFloatingPointState @ stub KeSaveStateForHibernate diff --git a/dlls/ntoskrnl.exe/ntoskrnl_private.h b/dlls/ntoskrnl.exe/ntoskrnl_private.h index 23f03c336d..a1e1b892e8 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl_private.h +++ b/dlls/ntoskrnl.exe/ntoskrnl_private.h @@ -49,6 +49,7 @@ struct _KTHREAD PEPROCESS process; CLIENT_ID id; unsigned int critical_region; + KAFFINITY user_affinity; };
struct _ETHREAD diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c index c7b3c93f6d..458815a0f2 100644 --- a/dlls/ntoskrnl.exe/tests/driver.c +++ b/dlls/ntoskrnl.exe/tests/driver.c @@ -1713,6 +1713,7 @@ static void test_executable_pool(void) static void test_affinity(void) { KAFFINITY (WINAPI *pKeSetSystemAffinityThreadEx)(KAFFINITY affinity); + void (WINAPI *pKeRevertToUserAffinityThreadEx)(KAFFINITY affinity); ULONG (WINAPI *pKeQueryActiveProcessorCountEx)(USHORT); KAFFINITY (WINAPI *pKeQueryActiveProcessors)(void); KAFFINITY mask, mask_all_cpus; @@ -1731,6 +1732,9 @@ static void test_affinity(void) pKeSetSystemAffinityThreadEx = get_proc_address("KeSetSystemAffinityThreadEx"); ok(!!pKeSetSystemAffinityThreadEx, "KeSetSystemAffinityThreadEx is not available.\n");
+ pKeRevertToUserAffinityThreadEx = get_proc_address("KeRevertToUserAffinityThreadEx"); + ok(!!pKeRevertToUserAffinityThreadEx, "KeRevertToUserAffinityThreadEx is not available.\n"); + count = pKeQueryActiveProcessorCountEx(1); todo_wine ok(!count, "Got unexpected count %u.\n", count);
@@ -1745,14 +1749,24 @@ static void test_affinity(void) mask = pKeQueryActiveProcessors(); ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask);
+ pKeRevertToUserAffinityThreadEx(0x2); + mask = pKeSetSystemAffinityThreadEx(0); ok(!mask, "Got unexpected mask %#lx.\n", mask);
+ pKeRevertToUserAffinityThreadEx(0x2); + mask = pKeSetSystemAffinityThreadEx(0x1); - ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask); + ok(mask == 0x2, "Got unexpected mask %#lx.\n", mask);
mask = pKeSetSystemAffinityThreadEx(~(KAFFINITY)0); ok(mask == 0x1, "Got unexpected mask %#lx.\n", mask); + + pKeRevertToUserAffinityThreadEx(~(KAFFINITY)0); + mask = pKeSetSystemAffinityThreadEx(0x1); + ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask); + + pKeRevertToUserAffinityThreadEx(mask_all_cpus); }
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 42ff0c72b1..b0ad0e70d3 100644 --- a/include/ddk/wdm.h +++ b/include/ddk/wdm.h @@ -1706,6 +1706,7 @@ void WINAPI KeReleaseSpinLock(KSPIN_LOCK*,KIRQL); void WINAPI KeReleaseSpinLockFromDpcLevel(KSPIN_LOCK*); LONG WINAPI KeResetEvent(PRKEVENT); void WINAPI KeRevertToUserAffinityThread(void); +void WINAPI KeRevertToUserAffinityThreadEx(KAFFINITY affinity); LONG WINAPI KeSetEvent(PRKEVENT,KPRIORITY,BOOLEAN); KPRIORITY WINAPI KeSetPriorityThread(PKTHREAD,KPRIORITY); void WINAPI KeSetSystemAffinityThread(KAFFINITY);