[WINLOGON] Display the shutdown message popup dialog on the current input desktop. (#3259)

CORE-17050

Display the shutdown message popup dialog on the current input desktop,
and periodically monitor for any change of the input desktop. When the
latter changes, close the dialog and recreate it on the new input desktop.

In addition, retain the current dialog position and restore it when the
dialog is recreated on the new desktop.

base/system/winlogon/shutdown.c

@@ -5,6 +5,7 @@
  * PURPOSE:         System shutdown dialog
  * PROGRAMMERS:     Edward Bronsten <tanki.alpha5056@gmail.com>
  *                  Eric Kohl
+ *                  Hermes Belusca-Maito
  */
 
 /* INCLUDES ******************************************************************/
@@ -26,11 +27,15 @@ typedef struct _SYS_SHUTDOWN_PARAMS
 {
     PWSTR pszMessage;
     ULONG dwTimeout;
+
+    HDESK hShutdownDesk;
+    WCHAR DesktopName[512];
+    WINDOWPLACEMENT wpPos;
+
+    BOOLEAN bShuttingDown;
     BOOLEAN bRebootAfterShutdown;
     BOOLEAN bForceAppsClosed;
     DWORD dwReason;
-
-    BOOLEAN bShuttingDown;
 } SYS_SHUTDOWN_PARAMS, *PSYS_SHUTDOWN_PARAMS;
 
 /* GLOBALS *******************************************************************/
@@ -46,12 +51,6 @@ DoSystemShutdown(
 {
     BOOL Success;
 
-    if (pShutdownParams->pszMessage)
-    {
-        HeapFree(GetProcessHeap(), 0, pShutdownParams->pszMessage);
-        pShutdownParams->pszMessage = NULL;
-    }
-
     /* If shutdown has been cancelled, bail out now */
     if (!pShutdownParams->bShuttingDown)
         return TRUE;
@@ -74,17 +73,67 @@ OnTimer(
     IN HWND hwndDlg,
     IN PSYS_SHUTDOWN_PARAMS pShutdownParams)
 {
+    HDESK hInputDesktop;
+    BOOL bSuccess;
+    DWORD dwSize;
+    INT iSeconds, iMinutes, iHours, iDays;
     WCHAR szFormatBuffer[32];
     WCHAR szBuffer[32];
-    INT iSeconds, iMinutes, iHours, iDays;
+    WCHAR DesktopName[512];
 
     if (!pShutdownParams->bShuttingDown)
     {
         /* Shutdown has been cancelled, close the dialog and bail out */
-        EndDialog(hwndDlg, 0);
+        EndDialog(hwndDlg, IDABORT);
         return;
     }
 
+    /*
+     * Check whether the input desktop has changed. If so, close the dialog,
+     * and let the shutdown thread recreate it on the new desktop.
+     */
+
+    // TODO: Investigate: It would be great if we could also compare with
+    // our internally maintained desktop handles, before calling that heavy
+    // comparison.
+    // (Note that we cannot compare handles with arbitrary input desktop,
+    // since OpenInputDesktop() creates new handle instances everytime.)
+
+    hInputDesktop = OpenInputDesktop(0, FALSE, GENERIC_ALL);
+    if (!hInputDesktop)
+    {
+        /* No input desktop but we have a dialog: kill it */
+        ERR("OpenInputDesktop() failed, error 0x%lx\n", GetLastError());
+        EndDialog(hwndDlg, 0);
+        return;
+    }
+    bSuccess = GetUserObjectInformationW(hInputDesktop,
+                                         UOI_NAME,
+                                         DesktopName,
+                                         sizeof(DesktopName),
+                                         &dwSize);
+    if (!bSuccess)
+    {
+        ERR("GetUserObjectInformationW(0x%p) failed, error 0x%lx\n",
+            hInputDesktop, GetLastError());
+    }
+    CloseDesktop(hInputDesktop);
+
+    if (bSuccess && (wcscmp(DesktopName, pShutdownParams->DesktopName) != 0))
+    {
+        TRACE("Input desktop has changed: '%S' --> '%S'\n",
+              pShutdownParams->DesktopName, DesktopName);
+
+        /* Save the original dialog position to be restored later */
+        pShutdownParams->wpPos.length = sizeof(pShutdownParams->wpPos);
+        GetWindowPlacement(hwndDlg, &pShutdownParams->wpPos);
+
+        /* Close the dialog */
+        EndDialog(hwndDlg, IDCANCEL);
+        return;
+    }
+
+    /* Update the shutdown timeout */
     if (pShutdownParams->dwTimeout < SECONDS_PER_DAY)
     {
         iSeconds = (INT)pShutdownParams->dwTimeout;
@@ -108,9 +157,8 @@ OnTimer(
 
     if (pShutdownParams->dwTimeout == 0)
     {
-        /* Close the dialog and perform the system shutdown */
+        /* Close the dialog and let the shutdown thread perform the system shutdown */
         EndDialog(hwndDlg, 0);
-        DoSystemShutdown(pShutdownParams);
         return;
     }
 
@@ -128,15 +176,16 @@ ShutdownDialogProc(
 {
     PSYS_SHUTDOWN_PARAMS pShutdownParams;
 
-    pShutdownParams = (PSYS_SHUTDOWN_PARAMS)GetWindowLongPtr(hwndDlg, DWLP_USER);
+    pShutdownParams = (PSYS_SHUTDOWN_PARAMS)GetWindowLongPtrW(hwndDlg, DWLP_USER);
 
     switch (uMsg)
     {
         case WM_INITDIALOG:
         {
             pShutdownParams = (PSYS_SHUTDOWN_PARAMS)lParam;
-            SetWindowLongPtr(hwndDlg, DWLP_USER, (LONG_PTR)pShutdownParams);
+            SetWindowLongPtrW(hwndDlg, DWLP_USER, (LONG_PTR)pShutdownParams);
 
+            /* Display the shutdown message */
             if (pShutdownParams->pszMessage)
             {
                 SetDlgItemTextW(hwndDlg,
@@ -144,10 +193,18 @@ ShutdownDialogProc(
                                 pShutdownParams->pszMessage);
             }
 
+            /* Remove the Close menu item */
             DeleteMenu(GetSystemMenu(hwndDlg, FALSE), SC_CLOSE, MF_BYCOMMAND);
-            SetWindowPos(hwndDlg, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOSIZE | SWP_NOMOVE);
 
-            PostMessage(hwndDlg, WM_TIMER, 0, 0);
+            /* Position the window (initial position, or restore from old) */
+            if (pShutdownParams->wpPos.length == sizeof(pShutdownParams->wpPos))
+                SetWindowPlacement(hwndDlg, &pShutdownParams->wpPos);
+
+            SetWindowPos(hwndDlg, HWND_TOPMOST, 0, 0, 0, 0,
+                         SWP_SHOWWINDOW | SWP_NOMOVE | SWP_NOSIZE);
+
+            /* Initialize the timer */
+            PostMessageW(hwndDlg, WM_TIMER, 0, 0);
             SetTimer(hwndDlg, SHUTDOWN_TIMER_ID, 1000, NULL);
             break;
         }
@@ -175,27 +232,95 @@ InitiateSystemShutdownThread(
     IN LPVOID lpParameter)
 {
     PSYS_SHUTDOWN_PARAMS pShutdownParams;
-    INT_PTR status;
+    HDESK hInputDesktop;
+    DWORD dwSize;
+    INT_PTR res;
 
     pShutdownParams = (PSYS_SHUTDOWN_PARAMS)lpParameter;
 
-    status = DialogBoxParamW(hAppInstance,
+    /* Default to initial dialog position */
+    pShutdownParams->wpPos.length = 0;
+
+    /* Continuously display the shutdown dialog on the current input desktop */
+    while (TRUE)
+    {
+        /* Retrieve the current input desktop */
+        hInputDesktop = OpenInputDesktop(0, FALSE, GENERIC_ALL);
+        if (!hInputDesktop)
+        {
+            /* No input desktop on the current WinSta0, just shut down */
+            ERR("OpenInputDesktop() failed, error 0x%lx\n", GetLastError());
+            break;
+        }
+
+        /* Remember it for checking desktop changes later */
+        pShutdownParams->hShutdownDesk = hInputDesktop;
+        if (!GetUserObjectInformationW(pShutdownParams->hShutdownDesk,
+                                       UOI_NAME,
+                                       pShutdownParams->DesktopName,
+                                       sizeof(pShutdownParams->DesktopName),
+                                       &dwSize))
+        {
+            ERR("GetUserObjectInformationW(0x%p) failed, error 0x%lx\n",
+                pShutdownParams->hShutdownDesk, GetLastError());
+        }
+
+        /* Assign the desktop to the current thread */
+        SetThreadDesktop(hInputDesktop);
+
+        /* Display the shutdown dialog on the current input desktop */
+        res = DialogBoxParamW(hAppInstance,
                               MAKEINTRESOURCEW(IDD_SYSSHUTDOWN),
                               NULL,
                               ShutdownDialogProc,
                               (LPARAM)pShutdownParams);
 
+        /* Close the desktop */
+        CloseDesktop(hInputDesktop);
+
+        /*
+         * Check why the dialog has been closed.
+         *
+         * - If it failed to be created (returned -1), don't care about
+         *   re-creating it, and proceed directly to shutdown.
+         *
+         * - If it closed unexpectedly (returned != 1), check whether a
+         *   shutdown is in progress. If the shutdown has been cancelled,
+         *   just bail out; if a shutdown is in progress and the timeout
+         *   is 0, bail out and proceed to shutdown.
+         *
+         * - If the dialog has closed because the input desktop changed,
+         *   loop again and recreate it on the new desktop.
+         */
+        if ((res == -1) || (res != IDCANCEL) ||
+            !(pShutdownParams->bShuttingDown && (pShutdownParams->dwTimeout > 0)))
+        {
+            break;
+        }
+    }
+
+    /* Reset dialog information */
+    pShutdownParams->hShutdownDesk = NULL;
+    ZeroMemory(&pShutdownParams->DesktopName, sizeof(pShutdownParams->DesktopName));
+    ZeroMemory(&pShutdownParams->wpPos, sizeof(pShutdownParams->wpPos));
+
     if (pShutdownParams->pszMessage)
     {
         HeapFree(GetProcessHeap(), 0, pShutdownParams->pszMessage);
         pShutdownParams->pszMessage = NULL;
     }
 
-    if (status >= 0)
+    if (pShutdownParams->bShuttingDown)
+    {
+        /* Perform the system shutdown */
+        if (DoSystemShutdown(pShutdownParams))
             return ERROR_SUCCESS;
+        else
+            return GetLastError();
+    }
 
     pShutdownParams->bShuttingDown = FALSE;
-    return GetLastError();
+    return ERROR_SUCCESS;
 }
 
 
@@ -225,7 +350,7 @@ StartSystemShutdown(
     if (_InterlockedCompareExchange8((volatile char*)&g_ShutdownParams.bShuttingDown, TRUE, FALSE) == TRUE)
         return ERROR_SHUTDOWN_IN_PROGRESS;
 
-    if (pMessage && pMessage->Length && pMessage->Buffer)
+    if ((dwTimeout != 0) && pMessage && pMessage->Length && pMessage->Buffer)
     {
         g_ShutdownParams.pszMessage = HeapAlloc(GetProcessHeap(),
                                                 HEAP_ZERO_MEMORY,