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

@@ -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.@)

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

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

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)

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);