Paul Gofman : ntoskrnl.exe: Implement KeRevertToUserAffinityThreadEx() function.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 7d4f4783a5d166b0d2987ca4c2b9188b405bdcbe
URL:    https://source.winehq.org/git/wine.git/?a=commit;h=7d4f4783a5d166b0d2987ca4c2b9188b405bdcbe

Author: Paul Gofman <pgofman at codeweavers.com>
Date:   Tue May 26 13:05:09 2020 +0300

ntoskrnl.exe: Implement KeRevertToUserAffinityThreadEx() function.

Signed-off-by: Paul Gofman <pgofman at codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard at 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);