Only block shutdown during Apply.

8 files changed, 31 insertions, 50 deletions

diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index ec996c4c..19b739f9 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -681,6 +681,8 @@ extern "C" HRESULT CoreApply(
     hr = ApplyLock(FALSE, &hLock);
     ExitOnFailure(hr, "Another per-user setup is already executing.");
 
+    pEngineState->plan.fApplying = TRUE;
+
     // Initialize only after getting a lock.
     fApplyInitialize = TRUE;
     ApplyInitialize();
@@ -814,6 +816,8 @@ LExit:
         ApplyUninitialize();
     }
 
+    pEngineState->plan.fApplying = FALSE;
+
     if (hLock)
     {
         ::ReleaseMutex(hLock);
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index d2c5c8c6..75b24ec3 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -94,6 +94,7 @@ typedef struct _BURN_ELEVATION_CHILD_MESSAGE_CONTEXT
     HANDLE hPipe;
     HANDLE* phLock;
     BOOL* pfDisabledAutomaticUpdates;
+    BOOL* pfApplying;
     BURN_APPROVED_EXES* pApprovedExes;
     BURN_CACHE* pCache;
     BURN_CONTAINERS* pContainers;
@@ -176,6 +177,7 @@ static HRESULT OnApplyInitialize(
     __in BURN_PACKAGES* pPackages,
     __in HANDLE* phLock,
     __in BOOL* pfDisabledWindowsUpdate,
+    __in BOOL* pfApplying,
     __in BYTE* pbData,
     __in SIZE_T cbData
     );
@@ -185,7 +187,8 @@ static HRESULT ElevatedProcessDetect(
     __in BURN_PACKAGES* pPackages
     );
 static HRESULT OnApplyUninitialize(
-    __in HANDLE* phLock
+    __in HANDLE* phLock,
+    __in BOOL* pfApplying
     );
 static HRESULT OnSessionBegin(
     __in BURN_CACHE* pCache,
@@ -1538,7 +1541,8 @@ extern "C" HRESULT ElevationChildPumpMessages(
     __out HANDLE* phLock,
     __out BOOL* pfDisabledAutomaticUpdates,
     __out DWORD* pdwChildExitCode,
-    __out BOOL* pfRestart
+    __out BOOL* pfRestart,
+    __out BOOL* pfApplying
     )
 {
     HRESULT hr = S_OK;
@@ -1561,6 +1565,7 @@ extern "C" HRESULT ElevationChildPumpMessages(
     context.hPipe = hPipe;
     context.phLock = phLock;
     context.pfDisabledAutomaticUpdates = pfDisabledAutomaticUpdates;
+    context.pfApplying = pfApplying;
     context.pApprovedExes = pApprovedExes;
     context.pCache = pCache;
     context.pContainers = pContainers;
@@ -2114,11 +2119,11 @@ static HRESULT ProcessElevatedChildMessage(
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE:
-        hrResult = OnApplyInitialize(pContext->hPipe, pContext->pVariables, pContext->pRegistration, pContext->pPackages, pContext->phLock, pContext->pfDisabledAutomaticUpdates, (BYTE*)pMsg->pvData, pMsg->cbData);
+        hrResult = OnApplyInitialize(pContext->hPipe, pContext->pVariables, pContext->pRegistration, pContext->pPackages, pContext->phLock, pContext->pfDisabledAutomaticUpdates, pContext->pfApplying, (BYTE*)pMsg->pvData, pMsg->cbData);
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_APPLY_UNINITIALIZE:
-        hrResult = OnApplyUninitialize(pContext->phLock);
+        hrResult = OnApplyUninitialize(pContext->phLock, pContext->pfApplying);
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_SESSION_BEGIN:
@@ -2267,6 +2272,7 @@ static HRESULT OnApplyInitialize(
     __in BURN_PACKAGES* pPackages,
     __in HANDLE* phLock,
     __in BOOL* pfDisabledWindowsUpdate,
+    __in BOOL* pfApplying,
     __in BYTE* pbData,
     __in SIZE_T cbData
     )
@@ -2296,6 +2302,8 @@ static HRESULT OnApplyInitialize(
     hr = ApplyLock(TRUE, phLock);
     ExitOnFailure(hr, "Failed to acquire lock due to setup in other session.");
 
+    *pfApplying = TRUE;
+
     // Detect.
     hr = ElevatedProcessDetect(pRegistration, pVariables, pPackages);
     ExitOnFailure(hr, "Failed to run detection in elevated process.");
@@ -2419,13 +2427,16 @@ LExit:
 }
 
 static HRESULT OnApplyUninitialize(
-    __in HANDLE* phLock
+    __in HANDLE* phLock,
+    __in BOOL* pfApplying
     )
 {
     Assert(phLock);
 
     // TODO: end system restore point.
 
+    *pfApplying = FALSE;
+
     if (*phLock)
     {
         ::ReleaseMutex(*phLock);
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index 0dc189c5..d08d2edc 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -180,7 +180,8 @@ HRESULT ElevationChildPumpMessages(
     __out HANDLE* phLock,
     __out BOOL* pfDisabledAutomaticUpdates,
     __out DWORD* pdwChildExitCode,
-    __out BOOL* pfRestart
+    __out BOOL* pfRestart,
+    __out BOOL* pfApplying
     );
 HRESULT ElevationChildResumeAutomaticUpdates();
 
diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
index 821626f8..ee848acf 100644
--- a/src/burn/engine/engine.cpp
+++ b/src/burn/engine/engine.cpp
@@ -653,7 +653,7 @@ static HRESULT RunElevated(
     SrpInitialize(TRUE);
 
     // Pump messages from parent process.
-    hr = ElevationChildPumpMessages(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe, pEngineState->companionConnection.hCachePipe, &pEngineState->approvedExes, &pEngineState->cache, &pEngineState->containers, &pEngineState->packages, &pEngineState->payloads, &pEngineState->variables, &pEngineState->registration, &pEngineState->userExperience, &hLock, &fDisabledAutomaticUpdates, &pEngineState->userExperience.dwExitCode, &pEngineState->fRestart);
+    hr = ElevationChildPumpMessages(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe, pEngineState->companionConnection.hCachePipe, &pEngineState->approvedExes, &pEngineState->cache, &pEngineState->containers, &pEngineState->packages, &pEngineState->payloads, &pEngineState->variables, &pEngineState->registration, &pEngineState->userExperience, &hLock, &fDisabledAutomaticUpdates, &pEngineState->userExperience.dwExitCode, &pEngineState->fRestart, &pEngineState->plan.fApplying);
     LogRedirect(NULL, NULL); // reset logging so the next failure gets written to "log buffer" for the failure log.
     ExitOnFailure(hr, "Failed to pump messages from parent process.");
 
diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h
index 2b04e097..46a9a363 100644
--- a/src/burn/engine/plan.h
+++ b/src/burn/engine/plan.h
@@ -264,6 +264,7 @@ typedef struct _BURN_PLAN
     LPWSTR sczLayoutDirectory;
     BOOL fPlanPackageCacheRollback;
     BOOL fDowngrade;
+    BOOL fApplying;
 
     DWORD64 qwCacheSizeTotal;
 
diff --git a/src/burn/engine/uithread.cpp b/src/burn/engine/uithread.cpp
index bb1e5972..cdd8613c 100644
--- a/src/burn/engine/uithread.cpp
+++ b/src/burn/engine/uithread.cpp
@@ -17,8 +17,7 @@ struct UITHREAD_CONTEXT
 struct UITHREAD_INFO
 {
     BOOL fElevatedEngine;
-    BURN_USER_EXPERIENCE* pUserExperience;
-    BOOL* pfCriticalShutdownInitiated;
+    BURN_ENGINE_STATE* pEngineState;
 };
 
 
@@ -132,8 +131,7 @@ static DWORD WINAPI ThreadProc(
     fRegistered = TRUE;
 
     info.fElevatedEngine = fElevatedEngine;
-    info.pUserExperience = &pEngineState->userExperience;
-    info.pfCriticalShutdownInitiated = &pEngineState->fCriticalShutdownInitiated;
+    info.pEngineState = pEngineState;
 
     // Create the window to handle reboots without activating it.
     hWnd = ::CreateWindowExW(WS_EX_NOACTIVATE, wc.lpszClassName, NULL, WS_POPUP, 0, 0, 0, 0, HWND_DESKTOP, NULL, pContext->hInstance, &info);
@@ -196,21 +194,17 @@ static LRESULT CALLBACK WndProc(
         {
         DWORD dwEndSession = static_cast<DWORD>(lParam);
         BOOL fCritical = ENDSESSION_CRITICAL & dwEndSession;
-        BOOL fCancel = TRUE;
+        BOOL fCancel = FALSE;
         BOOL fRet = FALSE;
 
-        // Always block shutdown in the elevated process, but ask the BA in the non-elevated.
+        // Always block shutdown during apply.
         UITHREAD_INFO* pInfo = reinterpret_cast<UITHREAD_INFO*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
-        if (!pInfo->fElevatedEngine)
+        if (pInfo->pEngineState->plan.fApplying)
         {
-            // TODO: instead of recommending canceling all non-critical shutdowns, maybe we should only recommend cancel
-            //       when the engine is doing work?
-            fCancel = !fCritical;
-            // TODO: There's a race condition here where the BA may not have been loaded, or already was unloaded.
-            UserExperienceOnSystemShutdown(pInfo->pUserExperience, dwEndSession, &fCancel);
+            fCancel = TRUE;
         }
 
-        *pInfo->pfCriticalShutdownInitiated |= fCritical;
+        pInfo->pEngineState->fCriticalShutdownInitiated |= fCritical;
 
         fRet = !fCancel;
         LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN, LoggingBoolToString(fCritical), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fRet));
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index a3cbcd4a..4325a6ee 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -2631,31 +2631,6 @@ LExit:
     return hr;
 }
 
-EXTERN_C BAAPI UserExperienceOnSystemShutdown(
-    __in BURN_USER_EXPERIENCE* pUserExperience,
-    __in DWORD dwEndSession,
-    __inout BOOL* pfCancel
-    )
-{
-    HRESULT hr = S_OK;
-    BA_ONSYSTEMSHUTDOWN_ARGS args = { };
-    BA_ONSYSTEMSHUTDOWN_RESULTS results = { };
-
-    args.cbSize = sizeof(args);
-    args.dwEndSession = dwEndSession;
-
-    results.cbSize = sizeof(results);
-    results.fCancel = *pfCancel;
-
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMSHUTDOWN, &args, &results);
-    ExitOnFailure(hr, "BA OnSystemShutdown failed.");
-    
-    *pfCancel = results.fCancel;
-
-LExit:
-    return hr;
-}
-
 EXTERN_C BAAPI UserExperienceOnUnregisterBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index 2059b521..e7489710 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -583,11 +583,6 @@ BAAPI UserExperienceOnSystemRestorePointComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in HRESULT hrStatus
     );
-BAAPI UserExperienceOnSystemShutdown(
-    __in BURN_USER_EXPERIENCE* pUserExperience,
-    __in DWORD dwEndSession,
-    __inout BOOL* pfCancel
-    );
 BAAPI UserExperienceOnUnregisterBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType