** WIP ** WINLOGON: Study check_setup handling

base/system/winlogon/sas.c

@@ -19,7 +19,6 @@
 #include <mmsystem.h>
 #include <userenv.h>
 #include <ndk/setypes.h>
-#include <ndk/sefuncs.h>
 
 /* GLOBALS ******************************************************************/
 
@@ -1313,18 +1312,17 @@ SASWindowProc(
         }
         case WM_CREATE:
         {
-            /* Get the session pointer from the create data */
+            /* Get the session pointer from the create data and save it */
             Session = (PWLSESSION)((LPCREATESTRUCT)lParam)->lpCreateParams;
-
-            /* Save the Session pointer */
             SetWindowLongPtrW(hwndDlg, GWLP_USERDATA, (LONG_PTR)Session);
-            if (GetSetupType())
+
+            if (IsNonOOBESetup(g_setupType))
                 return TRUE;
             return RegisterHotKeys(Session, hwndDlg);
         }
         case WM_DESTROY:
         {
-            if (!GetSetupType())
+            if (!IsNonOOBESetup(g_setupType))
                 UnregisterHotKeys(Session, hwndDlg);
             return TRUE;
         }

base/system/winlogon/setup.c

@@ -9,31 +9,32 @@
 /* INCLUDES *****************************************************************/
 
 #include "winlogon.h"
+#include <ndk/setypes.h>
+
+SETUP_TYPE g_setupType = SetupType_None;
 
 /* FUNCTIONS ****************************************************************/
 
-DWORD
+SETUP_TYPE
 GetSetupType(VOID)
 {
-    DWORD dwError;
     HKEY hKey;
+    DWORD dwError;
     DWORD dwType;
     DWORD dwSize;
     DWORD dwSetupType;
 
-    TRACE("GetSetupType()\n");
-
     /* Open key */
     dwError = RegOpenKeyExW(HKEY_LOCAL_MACHINE,
-                            L"SYSTEM\\Setup",
+                            L"SYSTEM\\Setup", // REGSTR_KEY_SYSTEM REGSTR_KEY_SETUP
                             0,
                             KEY_QUERY_VALUE,
                             &hKey);
     if (dwError != ERROR_SUCCESS)
         return 0;
 
-    /* Read key */
+    /* Read value */
-    dwSize = sizeof(DWORD);
+    dwSize = sizeof(dwSetupType);
     dwError = RegQueryValueExW(hKey,
                                L"SetupType",
                                NULL,
@@ -41,18 +42,108 @@ GetSetupType(VOID)
                                (LPBYTE)&dwSetupType,
                                &dwSize);
 
-    /* Close key, and check if returned values are correct */
+    /* Close key and check if returned values are correct */
     RegCloseKey(hKey);
     if (dwError != ERROR_SUCCESS || dwType != REG_DWORD || dwSize != sizeof(DWORD))
         return 0;
 
     TRACE("GetSetupType() returns %lu\n", dwSetupType);
-    return dwSetupType;
+    return (SETUP_TYPE)dwSetupType;
+}
+
+// Partially reverts commit c88c4b4ad (r28194)
+static BOOL
+SetSetupType(
+    _In_ DWORD dwSetupType)
+{
+    HKEY hKey;
+    DWORD dwError;
+
+    /* Open key */
+    dwError = RegOpenKeyExW(HKEY_LOCAL_MACHINE,
+                            L"SYSTEM\\Setup",
+                            0,
+                            KEY_SET_VALUE,
+                            &hKey);
+    if (dwError != ERROR_SUCCESS)
+        return FALSE;
+
+    /* Write value */
+    dwError = RegSetValueExW(hKey,
+                             L"SetupType",
+                             0,
+                             REG_DWORD,
+                             (LPBYTE)&dwSetupType,
+                             sizeof(dwSetupType));
+
+    /* Close key and check for success */
+    RegCloseKey(hKey);
+    return (dwError == ERROR_SUCCESS);
+}
+
+static BOOL
+IsSetupShutdownRequired(
+    _Out_ SHUTDOWN_ACTION* pShutdownAction)
+{
+    HKEY hKey;
+    DWORD dwError;
+    DWORD dwType;
+    DWORD dwSize;
+    DWORD dwValue;
+
+    /* Open key */
+    dwError = RegOpenKeyExW(HKEY_LOCAL_MACHINE,
+                            L"SYSTEM\\Setup",  // REGSTR_KEY_SYSTEM REGSTR_KEY_SETUP
+                            0,
+                            KEY_QUERY_VALUE | KEY_SET_VALUE,
+                            &hKey);
+    if (dwError != ERROR_SUCCESS)
+        return FALSE;
+
+    /* Read value */
+    dwSize = sizeof(dwValue);
+    dwError = RegQueryValueExW(hKey,
+                               L"SetupShutdownRequired",
+                               NULL,
+                               &dwType,
+                               (LPBYTE)&dwValue,
+                               &dwSize);
+
+    /* Close key and check if returned values are correct */
+    RegCloseKey(hKey);
+    if (dwError != ERROR_SUCCESS || dwType != REG_DWORD || dwSize != sizeof(DWORD))
+        return FALSE;
+
+    /* Delete the value */
+    RegDeleteValueW(hKey, L"SetupShutdownRequired");
+
+    TRACE("IsSetupShutdownRequired() returns %lu\n", dwValue);
+    *pShutdownAction = dwValue;
+    return TRUE;
+}
+
+// FIXME: Make the function generic -- see sas.c!HandleShutdown()
+static VOID
+DoSetupShutdown(
+    _In_ SHUTDOWN_ACTION Action)
+{
+    BOOLEAN Old;
+
+#if DBG
+    static const PCSTR s_pszShutdownAction[] =
+        { "Shutting down", "Restarting", "Powering off"};
+
+    ERR("WL: %s NT...\n", s_pszShutdownAction[Action]);
+#endif
+
+    RtlAdjustPrivilege(SE_SHUTDOWN_PRIVILEGE, TRUE, FALSE, &Old);
+    NtShutdownSystem(Action);
+    RtlAdjustPrivilege(SE_SHUTDOWN_PRIVILEGE, Old, FALSE, &Old);
+    ExitProcess(0);
 }
 
 
-static
+static DWORD
-DWORD
 WINAPI
 RunSetupThreadProc(
     IN LPVOID lpParameter)
@@ -104,7 +195,7 @@ RunSetupThreadProc(
 
     TRACE("Should run '%s' now\n", debugstr_w(CommandLine));
 
-    SwitchDesktop(WLSession->ApplicationDesktop);
+    // SwitchDesktop(WLSession->ApplicationDesktop);
 
     /* Start process */
     StartupInfo.cb = sizeof(StartupInfo);
@@ -127,8 +218,8 @@ RunSetupThreadProc(
                             &ProcessInformation);
     if (!Result)
     {
-        TRACE("Failed to run setup process\n");
+        ERR("Failed to run setup process '%s'\n", debugstr_w(CommandLine));
-        SwitchDesktop(WLSession->WinlogonDesktop);
+        // SwitchDesktop(WLSession->WinlogonDesktop);
         return FALSE;
     }
 
@@ -141,29 +232,111 @@ RunSetupThreadProc(
     CloseHandle(ProcessInformation.hThread);
     CloseHandle(ProcessInformation.hProcess);
 
-    // SwitchDesktop(WLSession->WinlogonDesktop);
+    /* Reset the current setup type */
+    if (dwExitCode == 0)
+        SetSetupType(SetupType_None);
+
+    // // SwitchDesktop(WLSession->WinlogonDesktop);
 
     TRACE("RunSetup() done\n");
-
     return TRUE;
 }
 
-
+static BOOL
-BOOL
 RunSetup(VOID)
 {
-    HANDLE hThread;
+    // HANDLE hThread;
 
+    SwitchDesktop(WLSession->ApplicationDesktop);
+
+#if 0
     hThread = CreateThread(NULL,
                            0,
                            RunSetupThreadProc,
                            NULL,
                            0,
                            NULL);
-    if (hThread != NULL)
+    if (hThread)
-        CloseHandle(hThread);
+#else
+    if (RunSetupThreadProc(NULL))
+#endif
+    {
+        /* Message loop for the current thread */
+        /* Pump and dispatch any input events */
+        MSG msg;
+        // while (MsgWaitForMultipleObjects(1, &hThread, FALSE, INFINITE, QS_ALLINPUT))
+        // // if (dwRet == WAIT_OBJECT_0 + 1)
+        {
+            while (PeekMessageW(&msg, NULL, 0, 0, PM_REMOVE))
+            {
+                TranslateMessage(&msg);
+                DispatchMessageW(&msg);
+            }
+        }
 
-    return hThread != NULL;
+        // CloseHandle(hThread);
+    }
+
+    SwitchDesktop(WLSession->WinlogonDesktop);
+
+    // return (hThread != NULL);
+    return TRUE;
+}
+
+
+VOID
+CheckForSetup(VOID)
+{
+    TRACE("CheckForSetup() called\n");
+
+    /* Check for pending setup */
+    switch (g_setupType)
+    {
+    case SetupType_None:
+        /* Nothing to do */
+        break;
+
+    case SetupType_Full: case SetupType_Upgrade:
+    case SetupType_OOBE:
+    {
+        SHUTDOWN_ACTION shutdownAction = ShutdownReboot;
+
+#if DBG
+        static const PCSTR pszSetupType[] =
+            { "None", "", "OOBE", "Reserved", "Upgrade" };
+
+        TRACE("WL: %s%sSetup mode detected\n",
+              pszSetupType[g_setupType], pszSetupType[g_setupType][0] ? " " : "");
+#endif
+
+        /*
+         * We currently support three types of setup actions:
+         *
+         * - Full setup or upgrade: run it and reboot when done
+         *   or if the setup program crashed for whatever reason;
+         *
+         * - OOBE setup: run it, then check whether it requested
+         *   a power action and initiate it if so, otherwise
+         *   continue with regular logon.
+         */
+        RunSetup();
+
+        /* Reboot only if needed */
+        if ((g_setupType == SetupType_Full) || (g_setupType == SetupType_Upgrade) ||
+            /*(g_setupType == SetupType_OOBE) &&*/IsSetupShutdownRequired(&shutdownAction))
+        {
+            DoSetupShutdown(shutdownAction);
+        }
+
+        break;
+    }
+
+    case SetupType_Reserved:
+    default:
+        /* Unknown setup type: ignore */
+        WARN("WL: Unknown Setup type %lu, ignored\n", g_setupType);
+        break;
+    }
 }
 
 /* EOF */

base/system/winlogon/winlogon.c

@@ -523,6 +523,11 @@ WinMain(
     /* Init Notifications */
     InitNotifications();
 
+    /* Check for pending setup: if so, run it and reboot when done */
+__debugbreak();
+    g_setupType = GetSetupType();
+    // CheckForSetup();
+
     /* Load and initialize gina */
     if (!GinaInit(WLSession))
     {
@@ -586,17 +591,20 @@ WinMain(
     WLSession->LogonState = STATE_INIT;
     RemoveStatusMessage(WLSession);
 
-    /* Check for pending setup */
+
-    if (GetSetupType() != 0)
+    /* Check for pending setup: if so, run it and reboot when done */
-    {
+__debugbreak();
-        /* Run setup and reboot when done */
+    // g_setupType = GetSetupType();
-        TRACE("WL: Setup mode detected\n");
+    /*****CheckForSetup();******/
-        RunSetup();
+
-    }
+
-    else
+    /*
-    {
+     * Simulate an initial Ctrl-Alt-Del key press.
-        PostMessageW(WLSession->SASWindow, WLX_WM_SAS, WLX_SAS_TYPE_CTRL_ALT_DEL, 0);
+     * FIXME: This should quite probably be done only if some registry
-    }
+     * setting is set. However, since our current Ctrl-Alt-Del hotkey
+     * handling in Win32k is broken, we need to use this hack!
+     */
+    // PostMessageW(WLSession->SASWindow, WLX_WM_SAS, WLX_SAS_TYPE_CTRL_ALT_DEL, 0);
 
     (void)LoadLibraryW(L"sfc_os.dll");
 

base/system/winlogon/winlogon.h

@@ -361,11 +361,25 @@ AllowAccessOnSession(
     _In_ PWLSESSION Session);
 
 /* setup.c */
-DWORD
+typedef enum _SETUP_TYPE
+{
+    SetupType_None = 0,
+    SetupType_Full,
+    SetupType_OOBE,
+    SetupType_Reserved,
+    SetupType_Upgrade
+} SETUP_TYPE, *PSETUP_TYPE;
+
+extern SETUP_TYPE g_setupType;
+
+#define IsNonOOBESetup(setupType) \
+    (((setupType) == SetupType_Full) || ((setupType) == SetupType_Upgrade))
+
+SETUP_TYPE
 GetSetupType(VOID);
 
-BOOL
+VOID
-RunSetup(VOID);
+CheckForSetup(VOID);
 
 /* shutdown.h */
 DWORD