From 4d7062abb6bcdfefd1c25d391020d98cb6538576 Mon Sep 17 00:00:00 2001
From: Eric Kohl <eric.kohl@reactos.org>
Date: Sun, 2 May 2021 13:55:29 +0200
Subject: [PATCH] [NTOS:PS] Add ThreadInformation probing to
 NtSetInformationThread

Also get rid of unused buffer check code.
This fixes two test failures.
---
 ntoskrnl/ps/query.c | 49 ++++++++++++++++++++++++++++++++++++---------
 1 file changed, 39 insertions(+), 10 deletions(-)

diff --git a/ntoskrnl/ps/query.c b/ntoskrnl/ps/query.c
index b8598a0b8a0..ea343c51eab 100644
--- a/ntoskrnl/ps/query.c
+++ b/ntoskrnl/ps/query.c
@@ -2032,18 +2032,47 @@ NtSetInformationThread(IN HANDLE ThreadHandle,
     ULONG_PTR TlsIndex = 0;
     PVOID *ExpansionSlots;
     PETHREAD ProcThread;
+    ULONG Alignment;
     PAGED_CODE();
 
-    /* Verify Information Class validity */
-#if 0
-    Status = DefaultSetInfoBufferCheck(ThreadInformationClass,
-                                       PsThreadInfoClass,
-                                       RTL_NUMBER_OF(PsThreadInfoClass),
-                                       ThreadInformation,
-                                       ThreadInformationLength,
-                                       PreviousMode);
-    if (!NT_SUCCESS(Status)) return Status;
-#endif
+    /* Check if we were called from user mode */
+    if (PreviousMode != KernelMode)
+    {
+        /* Enter SEH */
+        _SEH2_TRY
+        {
+            switch (ThreadInformationClass)
+            {
+                case ThreadPriority:
+                    Alignment = sizeof(KPRIORITY);
+                    break;
+
+                case ThreadAffinityMask:
+                case ThreadQuerySetWin32StartAddress:
+                    Alignment = sizeof(ULONG_PTR);
+                    break;
+
+                case ThreadEnableAlignmentFaultFixup:
+                    Alignment = sizeof(BOOLEAN);
+                    break;
+
+                default:
+                    Alignment = sizeof(ULONG);
+                    break;
+            }
+
+            /* Probe the buffer */
+            ProbeForRead(ThreadInformation,
+                         ThreadInformationLength,
+                         Alignment);
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            /* Return the exception code */
+            _SEH2_YIELD(return _SEH2_GetExceptionCode());
+        }
+        _SEH2_END;
+    }
 
     /* Check what kind of information class this is */
     switch (ThreadInformationClass)