[NTUSER] Fix desktop and window station assignment for csrss

- NtUserSetInformationThread: Stub UserThreadUseActiveDesktop and UserThreadRestoreDesktop
- Properly mark the first thread that enters win32k belonging to csrss. At this point we assume that since gpepCSRSS isn't initialized yet, it probably is the first thread.

[WINSRV] Use NtUserSetInformationThread to set the current desktop when needed
-When csrss needs to use user32 or enter win32k, it first needs to assign the current thread to a desktop.

win32ss/user/ntuser/main.c

@@ -535,7 +535,7 @@ InitThreadCallback(PETHREAD Thread)
     // FIXME: Flag SYSTEM threads with... TIF_SYSTEMTHREAD !!
 
     /* CSRSS threads have some special features */
-    if (Process == gpepCSRSS)
+    if (Process == gpepCSRSS || !gpepCSRSS)
         ptiCurrent->TIF_flags = TIF_CSRSSTHREAD | TIF_DONTATTACHQUEUE;
 
     ptiCurrent->pcti = &ptiCurrent->cti;
@@ -575,11 +575,8 @@ InitThreadCallback(PETHREAD Thread)
      * Do not try to open a desktop or window station before the very first
      * (interactive) window station has been created by Winlogon.
      */
-    // if (ptiCurrent->ppi->hdeskStartup == NULL && InputWindowStation != NULL)
+    if (!(ptiCurrent->TIF_flags & (TIF_SYSTEMTHREAD | TIF_CSRSSTHREAD)) &&
-    /* Last things to do only if we are not a SYSTEM or CSRSS thread */
+        ptiCurrent->ppi->hdeskStartup == NULL &&
-    // HACK Part #1: Temporarily disabled to have our current USERSRV running, but normally this is its duty to connect itself to the required desktop!
-    if (// !(ptiCurrent->TIF_flags & (TIF_SYSTEMTHREAD | TIF_CSRSSTHREAD)) &&
-        /**/ptiCurrent->ppi->hdeskStartup == NULL &&/**/
         InputWindowStation != NULL)
     {
         HWINSTA hWinSta = NULL;
@@ -587,15 +584,10 @@ InitThreadCallback(PETHREAD Thread)
         UNICODE_STRING DesktopPath;
         PDESKTOP pdesk;
 
-        // HACK Part #2: We force USERSRV to connect to WinSta0 by setting the STARTF_INHERITDESKTOP flag.
-        if (ptiCurrent->TIF_flags & (TIF_SYSTEMTHREAD | TIF_CSRSSTHREAD))
-            ProcessParams->WindowFlags |= STARTF_INHERITDESKTOP;
-
         /*
          * Inherit the thread desktop and process window station (if not yet inherited)
          * from the process startup info structure. See documentation of CreateProcess().
          */
-
         Status = STATUS_UNSUCCESSFUL;
         if (ProcessParams && ProcessParams->DesktopInfo.Length > 0)
         {

win32ss/user/ntuser/ntstubs.c

@@ -691,6 +691,9 @@ NtUserSetInformationProcess(
     return 0;
 }
 
+HDESK FASTCALL
+IntGetDesktopObjectHandle(PDESKTOP DesktopObject);
+
 NTSTATUS
 APIENTRY
 NtUserSetInformationThread(IN HANDLE ThreadHandle,
@@ -820,6 +823,32 @@ NtUserSetInformationThread(IN HANDLE ThreadHandle,
             break;
         }
 
+        case UserThreadUseActiveDesktop:
+        {
+            HDESK hdesk;
+
+            if (Thread != PsGetCurrentThread())
+            {
+                Status = STATUS_NOT_IMPLEMENTED;
+                break;
+            }
+
+            hdesk = IntGetDesktopObjectHandle(gpdeskInputDesktop);
+            IntSetThreadDesktop(hdesk, FALSE);
+
+            break;
+        }
+        case UserThreadRestoreDesktop:
+        {
+            if (Thread != PsGetCurrentThread())
+            {
+                Status = STATUS_NOT_IMPLEMENTED;
+                break;
+            }
+
+            IntSetThreadDesktop(NULL, FALSE);
+            break;
+        }
         default:
         {
             STUB;

win32ss/user/winsrv/usersrv/harderror.c

@@ -1012,34 +1012,16 @@ UserpMessageBox(
     IN ULONG Timeout)
 {
     ULONG MessageBoxResponse;
-    HDESK hDesk, hOldDesk;
 
     DPRINT("Text = '%S', Caption = '%S', Type = 0x%lx\n",
            TextStringU->Buffer, CaptionStringU->Buffer, Type);
 
-    // TEMPORARY HACK to fix desktop assignment for harderror message boxes.
-    hDesk = OpenInputDesktop(0, FALSE, GENERIC_WRITE);
-    if (!hDesk)
-        return ResponseNotHandled;
-
-    /* Assign the desktop to this thread */
-    hOldDesk = GetThreadDesktop(GetCurrentThreadId());
-    if (!SetThreadDesktop(hDesk))
-    {
-        CloseDesktop(hDesk);
-        return ResponseNotHandled;
-    }
-
     /* Display a message box */
     MessageBoxResponse = MessageBoxTimeoutW(NULL,
                                             TextStringU->Buffer,
                                             CaptionStringU->Buffer,
                                             Type, 0, Timeout);
 
-    /* Restore the original desktop */
-    SetThreadDesktop(hOldDesk);
-    CloseDesktop(hDesk);
-
     /* Return response value */
     switch (MessageBoxResponse)
     {
@@ -1107,6 +1089,7 @@ UserServerHardError(
     UNICODE_STRING TextU, CaptionU;
     WCHAR LocalTextBuffer[256];
     WCHAR LocalCaptionBuffer[256];
+    NTSTATUS Status;
 
     ASSERT(ThreadData->Process != NULL);
 
@@ -1136,6 +1119,16 @@ UserServerHardError(
         // (Message->UnicodeStringParameterMask & 0x3)
     }
 
+    Status = NtUserSetInformationThread(NtCurrentThread(),
+                                        UserThreadUseActiveDesktop,
+                                        NULL,
+                                        0);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("Failed to set thread desktop!\n");
+        return;
+    }
+
     /* Re-initialize the hard errors cache */
     UserInitHardErrorsCache();
 
@@ -1187,6 +1180,8 @@ Quit:
     if (CaptionU.Buffer != LocalCaptionBuffer)
         RtlFreeUnicodeString(&CaptionU);
 
+    NtUserSetInformationThread(NtCurrentThread(), UserThreadRestoreDesktop, NULL, 0);
+
     return;
 }
 

win32ss/user/winsrv/usersrv/shutdown.c

@@ -191,8 +191,17 @@ EndNowThreadProc(LPVOID Parameter)
     PNOTIFY_CONTEXT NotifyContext = (PNOTIFY_CONTEXT)Parameter;
     MSG Msg;
 
-    // SetThreadDesktop(NotifyContext->Desktop);
+#if 0
-    // SwitchDesktop(NotifyContext->Desktop);
+    SetThreadDesktop(NotifyContext->Desktop);
+    SwitchDesktop(NotifyContext->Desktop);
+#else
+    /* For now show the end task dialog in the active desktop */
+    NtUserSetInformationThread(NtCurrentThread(),
+                               UserThreadUseActiveDesktop,
+                               NULL,
+                               0);
+#endif
+
     CallInitCommonControls();
     NotifyContext->Dlg = CreateDialogParam(UserServerDllInstance,
                                            MAKEINTRESOURCE(IDD_END_NOW), NULL,
@@ -811,20 +820,43 @@ CSR_API(SrvExitWindowsEx)
     NTSTATUS Status;
     PUSER_EXIT_REACTOS ExitReactOSRequest = &((PUSER_API_MESSAGE)ApiMessage)->Data.ExitReactOSRequest;
 
+    Status = NtUserSetInformationThread(NtCurrentThread(),
+                                        UserThreadUseActiveDesktop,
+                                        NULL,
+                                        0);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("Failed to set thread desktop!\n");
+        return Status;
+    }
+
     Status = UserExitReactOS(CsrGetClientThread(), ExitReactOSRequest->Flags);
     ExitReactOSRequest->Success   = NT_SUCCESS(Status);
     ExitReactOSRequest->LastError = GetLastError();
 
+    NtUserSetInformationThread(NtCurrentThread(), UserThreadRestoreDesktop, NULL, 0);
+
     return Status;
 }
 
 CSR_API(SrvEndTask)
 {
     PUSER_END_TASK EndTaskRequest = &((PUSER_API_MESSAGE)ApiMessage)->Data.EndTaskRequest;
+    NTSTATUS Status;
 
     // FIXME: This is HACK-plemented!!
     DPRINT1("SrvEndTask is HACKPLEMENTED!!\n");
 
+    Status = NtUserSetInformationThread(NtCurrentThread(),
+                                        UserThreadUseActiveDesktop,
+                                        NULL,
+                                        0);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("Failed to set thread desktop!\n");
+        return Status;
+    }
+
     SendMessageW(EndTaskRequest->WndHandle, WM_CLOSE, 0, 0);
     // PostMessageW(EndTaskRequest->WndHandle, WM_CLOSE, 0, 0);
 
@@ -846,6 +878,8 @@ CSR_API(SrvEndTask)
         EndTaskRequest->LastError = ERROR_SUCCESS;
     }
 
+    NtUserSetInformationThread(NtCurrentThread(), UserThreadRestoreDesktop, NULL, 0);
+
     return STATUS_SUCCESS;
 }