From d0632b0bca74e7aed85a58d65245e951b31aa63b 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: Mon, 23 Jul 2018 01:17:14 +0200
Subject: [PATCH 1/1] [WIN32K:NTUSER] In UserSetProcessWindowStation(), use a
 duplicated window station handle to be set in the EPROCESS:Win32WindowStation
 cache.

Fixes most of the user32:desktop window station handle reference count tests.
---
 win32ss/user/ntuser/winsta.c | 75 ++++++++++++++++++++++++++++--------
 1 file changed, 60 insertions(+), 15 deletions(-)

diff --git a/win32ss/user/ntuser/winsta.c b/win32ss/user/ntuser/winsta.c
index 13b968516c1..12a1b90d3ad 100644
--- a/win32ss/user/ntuser/winsta.c
+++ b/win32ss/user/ntuser/winsta.c
@@ -972,7 +972,6 @@ NtUserCloseWindowStation(
                                             0,
                                             &Object,
                                             NULL);
-
     if (!NT_SUCCESS(Status))
     {
         ERR("Validation of window station handle (%p) failed\n", hWinSta);
@@ -1319,11 +1318,11 @@ NtUserGetProcessWindowStation(VOID)
 BOOL FASTCALL
 UserSetProcessWindowStation(HWINSTA hWindowStation)
 {
-    PPROCESSINFO ppi;
     NTSTATUS Status;
-    HWINSTA hwinstaOld;
+    PPROCESSINFO ppi;
     OBJECT_HANDLE_INFORMATION ObjectHandleInfo;
     PWINSTATION_OBJECT NewWinSta = NULL, OldWinSta;
+    HWINSTA hCacheWinSta;
 
     ppi = PsGetCurrentProcessWin32Process();
 
@@ -1337,15 +1336,14 @@ UserSetProcessWindowStation(HWINSTA hWindowStation)
                                                 &ObjectHandleInfo);
         if (!NT_SUCCESS(Status))
         {
-            TRACE("Validation of window station handle (%p) failed\n",
-                  hWindowStation);
+            TRACE("Validation of window station handle 0x%p failed\n", hWindowStation);
             SetLastNtError(Status);
             return FALSE;
         }
     }
 
     OldWinSta = ppi->prpwinsta;
-    hwinstaOld = PsGetProcessWin32WindowStation(ppi->peProcess);
+    hCacheWinSta = PsGetProcessWin32WindowStation(ppi->peProcess);
 
     /* Dereference the previous window station */
     if (OldWinSta != NULL)
@@ -1353,17 +1351,64 @@ UserSetProcessWindowStation(HWINSTA hWindowStation)
         ObDereferenceObject(OldWinSta);
     }
 
-    /* Check if we have a stale handle (it should happen for console apps) */
-    if (hwinstaOld != ppi->hwinsta)
-    {
-        ObCloseHandle(hwinstaOld, UserMode);
-    }
-
     /*
-     * FIXME: Don't allow changing the window station if there are threads that are attached to desktops and own GUI objects.
+     * FIXME: Don't allow changing the window station if there are threads that are attached to desktops and own GUI objects?
      */
 
-    PsSetProcessWindowStation(ppi->peProcess, hWindowStation);
+    /* Close the cached EPROCESS window station handle if needed */
+    if (hCacheWinSta != NULL)
+    {
+        /* Reference the window station */
+        Status = ObReferenceObjectByHandle(hCacheWinSta,
+                                           0,
+                                           ExWindowStationObjectType,
+                                           UserMode,
+                                           (PVOID*)&OldWinSta,
+                                           NULL);
+        if (!NT_SUCCESS(Status))
+        {
+            ERR("Failed to reference the inherited window station, Status 0x%08lx\n", Status);
+            /* We failed, reset the cache */
+            hCacheWinSta = NULL;
+            PsSetProcessWindowStation(ppi->peProcess, hCacheWinSta);
+        }
+        else
+        {
+            /*
+             * Close the old handle and reset the cache only
+             * if we are setting a different window station.
+             */
+            if (NewWinSta != OldWinSta)
+            {
+                ObCloseHandle(hCacheWinSta, UserMode);
+                hCacheWinSta = NULL;
+                PsSetProcessWindowStation(ppi->peProcess, hCacheWinSta);
+            }
+
+            /* Dereference the window station */
+            ObDereferenceObject(OldWinSta);
+        }
+    }
+
+    /* Duplicate and save a new cached EPROCESS window station handle */
+    if ((hCacheWinSta == NULL) && (hWindowStation != NULL))
+    {
+        Status = ZwDuplicateObject(ZwCurrentProcess(),
+                                   hWindowStation,
+                                   ZwCurrentProcess(),
+                                   (PHANDLE)&hCacheWinSta,
+                                   0,
+                                   0,
+                                   DUPLICATE_SAME_ACCESS);
+        if (!NT_SUCCESS(Status))
+        {
+            ERR("UserSetProcessWindowStation: Failed to duplicate the window station handle, Status 0x%08lx\n", Status);
+        }
+        else
+        {
+            PsSetProcessWindowStation(ppi->peProcess, hCacheWinSta);
+        }
+    }
 
     ppi->prpwinsta = NewWinSta;
     ppi->hwinsta = hWindowStation;
@@ -1383,7 +1428,7 @@ UserSetProcessWindowStation(HWINSTA hWindowStation)
     {
         ppi->W32PF_flags |= W32PF_IOWINSTA;
     }
-    else // Might be closed if the handle is null.
+    else /* Might be closed if the handle is NULL */
     {
         ppi->W32PF_flags &= ~W32PF_IOWINSTA;
     }
-- 
2.17.1