[reactos] 01/01: [NTOS:PS] Rewrite NtSetInformationThread to match NtQueryInformationThread

[Eric Kohl]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=2e88e2b90448a8a320fa2…

commit 2e88e2b90448a8a320fa293e95d8766a11a63a53
Author:     Eric Kohl <eric.kohl(a)reactos.org>
AuthorDate: Sun May 2 13:46:22 2021 +0200
Commit:     Eric Kohl <eric.kohl(a)reactos.org>
CommitDate: Sun May 2 13:46:22 2021 +0200

    [NTOS:PS] Rewrite NtSetInformationThread to match NtQueryInformationThread
    
    The Information length must always be checked before referencing the thread object.
This fixes a test failure.
---
 ntoskrnl/ps/query.c | 149 ++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 127 insertions(+), 22 deletions(-)

diff --git a/ntoskrnl/ps/query.c b/ntoskrnl/ps/query.c
index fa8d201c249..b8598a0b8a0 100644
--- a/ntoskrnl/ps/query.c
+++ b/ntoskrnl/ps/query.c
@@ -2018,7 +2018,6 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                        IN ULONG ThreadInformationLength)
 {
     PETHREAD Thread;
-    ULONG Access;
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
     NTSTATUS Status;
     HANDLE TokenHandle = NULL;
@@ -2046,23 +2045,6 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
     if (!NT_SUCCESS(Status)) return Status;
 #endif
 
-    /* Check what class this is */
-    Access = THREAD_SET_INFORMATION;
-    if (ThreadInformationClass == ThreadImpersonationToken)
-    {
-        /* Setting the impersonation token needs a special mask */
-        Access = THREAD_SET_THREAD_TOKEN;
-    }
-
-    /* Reference the thread */
-    Status = ObReferenceObjectByHandle(ThreadHandle,
-                                       Access,
-                                       PsThreadType,
-                                       PreviousMode,
-                                       (PVOID*)&Thread,
-                                       NULL);
-    if (!NT_SUCCESS(Status)) return Status;
-
     /* Check what kind of information class this is */
     switch (ThreadInformationClass)
     {
@@ -2099,8 +2081,21 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 break;
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Set the priority */
             KeSetPriorityThread(&Thread->Tcb, Priority);
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadBasePriority:
@@ -2145,8 +2140,21 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 }
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Set the base priority */
             KeSetBasePriorityThread(&Thread->Tcb, Priority);
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadAffinityMask:
@@ -2180,6 +2188,16 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 break;
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Get the process */
             Process = Thread->ThreadsProcess;
 
@@ -2214,7 +2232,8 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 Status = STATUS_PROCESS_IS_TERMINATING;
             }
 
-            /* Return status */
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadImpersonationToken:
@@ -2240,8 +2259,21 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
             }
             _SEH2_END;
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_THREAD_TOKEN,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Assign the actual token */
             Status = PsAssignImpersonationToken(Thread, TokenHandle);
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadQuerySetWin32StartAddress:
@@ -2267,8 +2299,21 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
             }
             _SEH2_END;
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Set the address */
             Thread->Win32StartAddress = Address;
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadIdealProcessor:
@@ -2302,6 +2347,16 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 break;
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Set the ideal */
             Status = KeSetIdealProcessorThread(&Thread->Tcb,
                                                (CCHAR)IdealProcessor);
@@ -2317,6 +2372,8 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 ExReleaseRundownProtection(&Thread->RundownProtect);
             }
 
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadPriorityBoost:
@@ -2342,8 +2399,21 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
             }
             _SEH2_END;
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Call the kernel */
             KeSetDisableBoostThread(&Thread->Tcb, (BOOLEAN)DisableBoost);
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadZeroTlsCell:
@@ -2377,6 +2447,16 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 break;
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Get the process */
             Process = Thread->ThreadsProcess;
 
@@ -2421,7 +2501,8 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 ProcThread = PsGetNextProcessThread(Process, ProcThread);
             }
 
-            /* All done */
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadBreakOnTermination:
@@ -2455,6 +2536,16 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 break;
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Set or clear the flag */
             if (Break)
             {
@@ -2464,6 +2555,9 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
             {
                 PspClearCrossThreadFlag(Thread, CT_BREAK_ON_TERMINATION_BIT);
             }
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         case ThreadHideFromDebugger:
@@ -2475,8 +2569,21 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
                 break;
             }
 
+            /* Reference the thread */
+            Status = ObReferenceObjectByHandle(ThreadHandle,
+                                               THREAD_SET_INFORMATION,
+                                               PsThreadType,
+                                               PreviousMode,
+                                               (PVOID*)&Thread,
+                                               NULL);
+            if (!NT_SUCCESS(Status))
+                break;
+
             /* Set the flag */
             PspSetCrossThreadFlag(Thread, CT_HIDE_FROM_DEBUGGER_BIT);
+
+            /* Dereference the thread */
+            ObDereferenceObject(Thread);
             break;
 
         default:
@@ -2485,8 +2592,6 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
             Status = STATUS_NOT_IMPLEMENTED;
     }
 
-    /* Dereference and return status */
-    ObDereferenceObject(Thread);
     return Status;
 }