From b003d68ca58ce94766141c779332e4c42b7fc430 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Herm=C3=A8s=20B=C3=A9lusca-Ma=C3=AFto?=
 <hermes.belusca-maito@reactos.org>
Date: Sun, 8 Jul 2018 20:39:19 +0200
Subject: [PATCH] [WIN32SS:NTUSER] Update SetWindowStationUser() and
 NtUserSetWindowStationUser() prototypes. Also, improve
 NtUserSetWindowStationUser() capture order, make psid optional as it should
 (and avoid a user-mode triggered BSOD), and initialize luidUser only when
 everything succeeded.

---
 win32ss/include/ntuser.h          |  8 ++--
 win32ss/user/ntuser/winsta.c      | 80 +++++++++++++++++++++----------
 win32ss/user/user32/misc/winsta.c |  6 +--
 3 files changed, 61 insertions(+), 33 deletions(-)

diff --git a/win32ss/include/ntuser.h b/win32ss/include/ntuser.h
index 724e3e612cd..9e4b70938be 100644
--- a/win32ss/include/ntuser.h
+++ b/win32ss/include/ntuser.h
@@ -3289,10 +3289,10 @@ NtUserSetWindowsHookEx(
 BOOL
 NTAPI
 NtUserSetWindowStationUser(
-    HWINSTA hWindowStation,
-    PLUID pluid,
-    PSID psid,
-    DWORD size);
+    IN HWINSTA hWindowStation,
+    IN PLUID pluid,
+    IN PSID psid OPTIONAL,
+    IN DWORD size);
 
 WORD
 NTAPI
diff --git a/win32ss/user/ntuser/winsta.c b/win32ss/user/ntuser/winsta.c
index e999ba8a2c0..44a8f405051 100644
--- a/win32ss/user/ntuser/winsta.c
+++ b/win32ss/user/ntuser/winsta.c
@@ -1482,16 +1482,18 @@ NtUserLockWorkStation(VOID)
     return ret;
 }
 
-BOOL APIENTRY
+BOOL
+NTAPI
 NtUserSetWindowStationUser(
-    HWINSTA hWindowStation,
-    PLUID pluid,
-    PSID psid,
-    DWORD size)
+    IN HWINSTA hWindowStation,
+    IN PLUID pluid,
+    IN PSID psid OPTIONAL,
+    IN DWORD size)
 {
+    BOOL Ret = FALSE;
     NTSTATUS Status;
     PWINSTATION_OBJECT WindowStation = NULL;
-    BOOL Ret = FALSE;
+    LUID luidUser;
 
     UserEnterExclusive();
 
@@ -1501,53 +1503,79 @@ NtUserSetWindowStationUser(
         goto Leave;
     }
 
+    /* Validate the window station */
     Status = IntValidateWindowStationHandle(hWindowStation,
                                             UserMode,
                                             0,
                                             &WindowStation,
-                                            0);
+                                            NULL);
     if (!NT_SUCCESS(Status))
     {
         goto Leave;
     }
 
-    if (WindowStation->psidUser)
-    {
-        ExFreePoolWithTag(WindowStation->psidUser, USERTAG_SECURITY);
-    }
-
-    WindowStation->psidUser = ExAllocatePoolWithTag(PagedPool, size, USERTAG_SECURITY);
-    if (WindowStation->psidUser == NULL)
-    {
-        EngSetLastError(ERROR_OUTOFMEMORY);
-        goto Leave;
-    }
-
+    /* Capture the user LUID */
     _SEH2_TRY
     {
-        ProbeForRead(psid, size, 1);
         ProbeForRead(pluid, sizeof(LUID), 1);
-
-        RtlCopyMemory(WindowStation->psidUser, psid, size);
-        WindowStation->luidUser = *pluid;
+        luidUser = *pluid;
     }
     _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
     {
         Status = _SEH2_GetExceptionCode();
+        _SEH2_YIELD(goto Leave);
     }
     _SEH2_END;
 
-    if (!NT_SUCCESS(Status))
+    /* Reset the window station user LUID */
+    RtlZeroMemory(&WindowStation->luidUser, sizeof(LUID));
+
+    /* Reset the window station user SID */
+    if (WindowStation->psidUser)
     {
         ExFreePoolWithTag(WindowStation->psidUser, USERTAG_SECURITY);
         WindowStation->psidUser = NULL;
-        goto Leave;
     }
 
+    /* Copy the new user SID if one has been provided */
+    if (psid)
+    {
+        WindowStation->psidUser = ExAllocatePoolWithTag(PagedPool, size, USERTAG_SECURITY);
+        if (WindowStation->psidUser == NULL)
+        {
+            EngSetLastError(ERROR_OUTOFMEMORY);
+            goto Leave;
+        }
+
+        Status = STATUS_SUCCESS;
+        _SEH2_TRY
+        {
+            ProbeForRead(psid, size, 1);
+            RtlCopyMemory(WindowStation->psidUser, psid, size);
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            Status = _SEH2_GetExceptionCode();
+        }
+        _SEH2_END;
+
+        if (!NT_SUCCESS(Status))
+        {
+            ExFreePoolWithTag(WindowStation->psidUser, USERTAG_SECURITY);
+            WindowStation->psidUser = NULL;
+            goto Leave;
+        }
+    }
+
+    /* Copy the new user LUID */
+    WindowStation->luidUser = luidUser;
+
     Ret = TRUE;
 
 Leave:
-    if (WindowStation) ObDereferenceObject(WindowStation);
+    if (WindowStation)
+        ObDereferenceObject(WindowStation);
+
     UserLeave();
     return Ret;
 }
diff --git a/win32ss/user/user32/misc/winsta.c b/win32ss/user/user32/misc/winsta.c
index 11d4bb5a007..24f6d893b26 100644
--- a/win32ss/user/user32/misc/winsta.c
+++ b/win32ss/user/user32/misc/winsta.c
@@ -399,8 +399,8 @@ BOOL
 WINAPI
 SetWindowStationUser(
     IN HWINSTA hWindowStation,
-    IN PLUID pluid OPTIONAL,
-    IN PSID psid,
+    IN PLUID pluid,
+    IN PSID psid OPTIONAL,
     IN DWORD size)
 {
     BOOL Success;
@@ -410,7 +410,7 @@ SetWindowStationUser(
     {
         /* Signal log-on/off to WINSRV */
 
-        /* User is logging on if pluid != LuidNone, otherwise it is a log-off */
+        /* User is logging on if *pluid != LuidNone, otherwise it is a log-off */
         LUID LuidNone = {0, 0};
         BOOL IsLogon = (pluid && !RtlEqualLuid(pluid, &LuidNone));
 
-- 
2.17.1