Delay closing Burn's UI thread so that it can log the bundle's restart.

5 files changed, 159 insertions, 40 deletions

diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index 25124c8a..3c1ed117 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -1963,6 +1963,21 @@ LExit:
     return hr;
 }
 
+extern "C" void CoreUpdateRestartState(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in BURN_RESTART_STATE restartState
+    )
+{
+    ::EnterCriticalSection(&pEngineState->csRestartState);
+
+    if (pEngineState->fRestarting && restartState > pEngineState->restartState)
+    {
+        pEngineState->restartState = restartState;
+    }
+
+    ::LeaveCriticalSection(&pEngineState->csRestartState);
+}
+
 extern "C" void CoreFunctionOverride(
     __in_opt PFN_CREATEPROCESSW pfnCreateProcessW,
     __in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 7e594b52..28b5ba5d 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -82,6 +82,16 @@ enum BURN_AU_PAUSE_ACTION
     BURN_AU_PAUSE_ACTION_IFELEVATED_NORESUME,
 };
 
+enum BURN_RESTART_STATE
+{
+    BURN_RESTART_STATE_NONE,
+    BURN_RESTART_STATE_REQUESTING,
+    BURN_RESTART_STATE_REQUESTED,
+    BURN_RESTART_STATE_INITIATING,
+    BURN_RESTART_STATE_INITIATED,
+    BURN_RESTART_STATE_BLOCKED,
+};
+
 
 // structs
 
@@ -160,6 +170,10 @@ typedef struct _BURN_ENGINE_STATE
     BURN_PIPE_CONNECTION companionConnection;
     BURN_PIPE_CONNECTION embeddedConnection;
 
+    CRITICAL_SECTION csRestartState;
+    BOOL fRestarting;
+    BURN_RESTART_STATE restartState;
+
     BOOL fCriticalShutdownInitiated;
     BURN_RESUME_MODE resumeMode;
     LPCWSTR wzRestartInitiatedPackageId;
@@ -300,6 +314,10 @@ HRESULT CoreParseCommandLine(
     __inout HANDLE* phSectionFile,
     __inout HANDLE* phSourceEngineFile
     );
+void CoreUpdateRestartState(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in BURN_RESTART_STATE restartState
+    );
 void CoreFunctionOverride(
     __in_opt PFN_CREATEPROCESSW pfnCreateProcessW,
     __in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
index 13d23ecd..9c8b6c1d 100644
--- a/src/burn/engine/engine.cpp
+++ b/src/burn/engine/engine.cpp
@@ -72,7 +72,9 @@ static HRESULT WaitForElevatedLoggingThread(
 static HRESULT WaitForUnelevatedLoggingThread(
     __in HANDLE hUnelevatedLoggingThread
     );
-static HRESULT Restart();
+static HRESULT Restart(
+    __in BURN_ENGINE_STATE* pEngineState
+    );
 static void CALLBACK BurnTraceError(
     __in_z LPCSTR szFile,
     __in int iLine,
@@ -131,7 +133,6 @@ extern "C" HRESULT EngineRun(
     BOOL fRunNormal = FALSE;
     BOOL fRunElevated = FALSE;
     BOOL fRunRunOnce = FALSE;
-    BOOL fRestart = FALSE;
 
     BURN_ENGINE_STATE engineState = { };
     engineState.command.cbSize = sizeof(BOOTSTRAPPER_COMMAND);
@@ -269,9 +270,7 @@ extern "C" HRESULT EngineRun(
         ExitOnFailure(hr, "Invalid run mode.");
     }
 
-    // set exit code and remember if we are supposed to restart.
     *pdwExitCode = engineState.userExperience.dwExitCode;
-    fRestart = engineState.fRestart;
 
 LExit:
     ReleaseStr(sczExePath);
@@ -289,17 +288,17 @@ LExit:
     CacheUninitialize(&engineState.cache);
 
     // If this is a related bundle (but not an update) suppress restart and return the standard restart error code.
-    if (fRestart && BOOTSTRAPPER_RELATION_NONE != engineState.command.relationType && BOOTSTRAPPER_RELATION_UPDATE != engineState.command.relationType)
+    if (engineState.fRestart && BOOTSTRAPPER_RELATION_NONE != engineState.command.relationType && BOOTSTRAPPER_RELATION_UPDATE != engineState.command.relationType)
     {
         LogId(REPORT_STANDARD, MSG_RESTART_ABORTED, LoggingRelationTypeToString(engineState.command.relationType));
 
-        fRestart = FALSE;
+        engineState.fRestart = FALSE;
         hr = SUCCEEDED(hr) ? HRESULT_FROM_WIN32(ERROR_SUCCESS_REBOOT_REQUIRED) : HRESULT_FROM_WIN32(ERROR_FAIL_REBOOT_REQUIRED);
     }
 
     if (fRunNormal)
     {
-        LogId(REPORT_STANDARD, MSG_EXITING, FAILED(hr) ? (int)hr : *pdwExitCode, LoggingBoolToString(fRestart));
+        LogId(REPORT_STANDARD, MSG_EXITING, FAILED(hr) ? (int)hr : *pdwExitCode, LoggingBoolToString(engineState.fRestart));
     }
     else if (fRunUntrusted)
     {
@@ -321,17 +320,20 @@ LExit:
         LogFlush();
     }
 
-    if (fRestart)
+    if (engineState.fRestart)
     {
         LogId(REPORT_STANDARD, MSG_RESTARTING);
 
-        HRESULT hrRestart = Restart();
+        HRESULT hrRestart = Restart(&engineState);
         if (FAILED(hrRestart))
         {
             LogErrorId(hrRestart, MSG_RESTART_FAILED);
         }
     }
 
+    // If the message window is still around, close it.
+    UiCloseMessageWindow(&engineState);
+
     UninitializeEngineState(&engineState);
 
     if (fXmlInitialized)
@@ -385,6 +387,8 @@ static HRESULT InitializeEngineState(
     HANDLE hSectionFile = hEngineFile;
     HANDLE hSourceEngineFile = INVALID_HANDLE_VALUE;
 
+    ::InitializeCriticalSection(&pEngineState->csRestartState);
+
     pEngineState->internalCommand.automaticUpdates = BURN_AU_PAUSE_ACTION_IFELEVATED;
     ::InitializeCriticalSection(&pEngineState->userExperience.csEngineActive);
     PipeConnectionInitialize(&pEngineState->companionConnection);
@@ -459,6 +463,8 @@ static void UninitializeEngineState(
     ReleaseStr(pEngineState->log.sczPath);
     ReleaseStr(pEngineState->log.sczPathVariable);
 
+    ::DeleteCriticalSection(&pEngineState->csRestartState);
+
     // clear struct
     memset(pEngineState, 0, sizeof(BURN_ENGINE_STATE));
 }
@@ -635,9 +641,6 @@ LExit:
 
     BurnExtensionUnload(&pEngineState->extensions);
 
-    // If the message window is still around, close it.
-    UiCloseMessageWindow(pEngineState);
-
     VariablesDump(&pEngineState->variables);
 
     // end per-machine process if running
@@ -720,9 +723,6 @@ LExit:
 
     LogRedirect(NULL, NULL); // we're done talking to the child so always reset logging now.
 
-    // If the message window is still around, close it.
-    UiCloseMessageWindow(pEngineState);
-
     if (fDeleteLoggingCs)
     {
         ::DeleteCriticalSection(&loggingContext.csBuffer);
@@ -1112,7 +1112,9 @@ LExit:
     return hr;
 }
 
-static HRESULT Restart()
+static HRESULT Restart(
+    __in BURN_ENGINE_STATE* pEngineState
+    )
 {
     HRESULT hr = S_OK;
     HANDLE hProcessToken = NULL;
@@ -1136,17 +1138,19 @@ static HRESULT Restart()
         ExitWithLastError(hr, "Failed to adjust token to add shutdown privileges.");
     }
 
+    pEngineState->fRestarting = TRUE;
+    CoreUpdateRestartState(pEngineState, BURN_RESTART_STATE_REQUESTING);
+
     do
     {
         hr = S_OK;
 
-        // Wait a second to let the companion process (assuming we did an elevated install) to get to the
-        // point where it too is thinking about restarting the computer. Only one will schedule the restart
-        // but both will have their log files closed and otherwise be ready to exit.
-        //
-        // On retry, we'll also wait a second to let the OS try to get to a place where the restart can
-        // be initiated.
-        ::Sleep(1000);
+        if (dwRetries)
+        {
+            // On retry, wait a second to let the OS try to get to a place where the restart can
+            // be initiated.
+            ::Sleep(1000);
+        }
 
         if (!vpfnInitiateSystemShutdownExW(NULL, NULL, 0, FALSE, TRUE, SHTDN_REASON_MAJOR_APPLICATION | SHTDN_REASON_MINOR_INSTALLATION | SHTDN_REASON_FLAG_PLANNED))
         {
@@ -1155,6 +1159,41 @@ static HRESULT Restart()
     } while (dwRetries++ < RESTART_RETRIES && (HRESULT_FROM_WIN32(ERROR_MACHINE_LOCKED) == hr || HRESULT_FROM_WIN32(ERROR_NOT_READY) == hr));
     ExitOnRootFailure(hr, "Failed to schedule restart.");
 
+    CoreUpdateRestartState(pEngineState, BURN_RESTART_STATE_REQUESTED);
+
+    // Give the UI thread approximately 15 seconds to get the WM_QUERYENDSESSION message.
+    for (DWORD i = 0; i < 60; ++i)
+    {
+        if (!::IsWindow(pEngineState->hMessageWindow))
+        {
+            ExitFunction();
+        }
+
+        if (BURN_RESTART_STATE_REQUESTED < pEngineState->restartState)
+        {
+            break;
+        }
+
+        ::Sleep(250);
+    }
+
+    if (BURN_RESTART_STATE_INITIATING > pEngineState->restartState)
+    {
+        LogId(REPORT_WARNING, MSG_RESTART_BLOCKED);
+        ExitFunction();
+    }
+
+    // Give the UI thread the chance to process the WM_ENDSESSION message.
+    for (;;)
+    {
+        if (!::IsWindow(pEngineState->hMessageWindow) || BURN_RESTART_STATE_INITIATING < pEngineState->restartState)
+        {
+            break;
+        }
+
+        ::Sleep(250);
+    }
+
 LExit:
     ReleaseHandle(hProcessToken);
     return hr;
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index dab1a504..d965d4ad 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -62,7 +62,6 @@ Severity=Warning
 SymbolicName=MSG_RESTARTING
 Language=English
 Restarting computer...
-=======================================
 .
 
 MessageId=6
@@ -191,6 +190,13 @@ Language=English
 The restart request failed, error: %1!ls!. The machine will need to be manually restarted.
 .
 
+MessageId=24
+Severity=Error
+SymbolicName=MSG_RESTART_BLOCKED
+Language=English
+The restart request was successful, but no system restart messages have been received. This may be caused by another application blocking the restart or taking too long to respond. The machine might need to be manually restarted.
+.
+
 MessageId=51
 Severity=Error
 SymbolicName=MSG_FAILED_PARSE_CONDITION
@@ -1138,9 +1144,16 @@ Apply complete, result: 0x%1!x!, restart: %2!hs!, ba requested restart:  %3!hs!
 
 MessageId=400
 Severity=Success
-SymbolicName=MSG_SYSTEM_SHUTDOWN
+SymbolicName=MSG_SYSTEM_SHUTDOWN_REQUEST
+Language=English
+Received system request to shut down the process: allowed: %1!hs!, elevated: %2!hs!, critical: %3!hs!, logoff: %4!hs!, close app: %5!hs!
+.
+
+MessageId=401
+Severity=Success
+SymbolicName=MSG_SYSTEM_SHUTDOWN_RESULT
 Language=English
-Received system request to shut down the process: critical: %1!hs!, elevated: %2!hs!, allowed: %3!hs!
+Received result of system request to shut down the process: closing: %1!hs!, elevated: %2!hs!, critical: %3!hs!, logoff: %4!hs!, close app: %5!hs!
 .
 
 MessageId=410
diff --git a/src/burn/engine/uithread.cpp b/src/burn/engine/uithread.cpp
index 26a9b723..673111f8 100644
--- a/src/burn/engine/uithread.cpp
+++ b/src/burn/engine/uithread.cpp
@@ -46,6 +46,11 @@ HRESULT UiCreateMessageWindow(
     HANDLE rgWaitHandles[2] = { };
     UITHREAD_CONTEXT context = { };
 
+    // Try to make this process the first one to receive WM_QUERYENDSESSION.
+    // When blocking shutdown during Apply, this prevents other applications from being closed even though the restart will be blocked.
+    // When initiating a restart, this makes it reasonable to assume WM_QUERYENDSESSION will be received quickly because otherwise other applications could delay indefinitely.
+    ::SetProcessShutdownParameters(0x3FF, 0);
+
     // Create event to signal after the UI thread / window is initialized.
     rgWaitHandles[0] = ::CreateEventW(NULL, TRUE, FALSE, NULL);
     ExitOnNullWithLastError(rgWaitHandles[0], hr, "Failed to create initialization event.");
@@ -79,10 +84,6 @@ void UiCloseMessageWindow(
     if (::IsWindow(pEngineState->hMessageWindow))
     {
         ::PostMessageW(pEngineState->hMessageWindow, WM_CLOSE, 0, 0);
-
-        // Give the window 15 seconds to close because if it stays open it can prevent
-        // the engine from starting a reboot (should a reboot actually be necessary).
-        ::WaitForSingleObject(pEngineState->hMessageWindowThread, 15 * 1000);
     }
 }
 
@@ -180,23 +181,56 @@ static LRESULT CALLBACK WndProc(
 
     case WM_QUERYENDSESSION:
         {
-        DWORD dwEndSession = static_cast<DWORD>(lParam);
-        BOOL fCritical = ENDSESSION_CRITICAL & dwEndSession;
-        BOOL fCancel = FALSE;
-        BOOL fRet = FALSE;
+        BOOL fCritical = ENDSESSION_CRITICAL & lParam;
+        BOOL fAllowed = FALSE;
 
-        // Always block shutdown during apply.
         UITHREAD_INFO* pInfo = reinterpret_cast<UITHREAD_INFO*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
-        if (pInfo->pEngineState->plan.fApplying)
+        if (!pInfo->pEngineState->plan.fApplying && // always block shutdown during apply.
+            !fCritical)                             // always block critical shutdowns to receive the WM_ENDSESSION message.
         {
-            fCancel = TRUE;
+            fAllowed = TRUE;
         }
 
+        CoreUpdateRestartState(pInfo->pEngineState, BURN_RESTART_STATE_INITIATING);
         pInfo->pEngineState->fCriticalShutdownInitiated |= fCritical;
 
-        fRet = !fCancel;
-        LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN, LoggingBoolToString(fCritical), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fRet));
-        return fRet;
+        LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN_REQUEST, LoggingBoolToString(fAllowed), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fCritical), LoggingBoolToString(lParam & ENDSESSION_LOGOFF), LoggingBoolToString(lParam & ENDSESSION_CLOSEAPP));
+        LogFlush();
+        return fAllowed;
+        }
+
+    case WM_ENDSESSION:
+        {
+        UITHREAD_INFO* pInfo = reinterpret_cast<UITHREAD_INFO*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
+        BOOL fAllowed = 0 != wParam;
+
+        LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN_RESULT, LoggingBoolToString(fAllowed), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(lParam & ENDSESSION_CRITICAL), LoggingBoolToString(lParam & ENDSESSION_LOGOFF), LoggingBoolToString(lParam & ENDSESSION_CLOSEAPP));
+
+        if (fAllowed)
+        {
+            // Windows will shutdown the process as soon as we return from this message.
+            // https://docs.microsoft.com/en-us/previous-versions/windows/desktop/ms700677(v=vs.85)
+
+            // Give Apply approximately 20 seconds to complete.
+            for (DWORD i = 0; i < 80; ++i)
+            {
+                if (!pInfo->pEngineState->plan.fApplying)
+                {
+                    break;
+                }
+
+                ::Sleep(250);
+            }
+
+            LogStringWorkRaw("=======================================\r\n");
+
+            // Close the log to try to make sure everything is flushed to disk.
+            LogClose(FALSE);
+        }
+
+        CoreUpdateRestartState(pInfo->pEngineState, fAllowed ? BURN_RESTART_STATE_INITIATED : BURN_RESTART_STATE_BLOCKED);
+
+        return 0;
         }
 
     case WM_DESTROY: