Only block shutdown during Apply.

25 files changed, 107 insertions, 231 deletions

diff --git a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index f0b5dad4..228cb6ff 100644
--- a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -149,7 +149,6 @@ enum BOOTSTRAPPER_APPLICATION_MESSAGE
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPLETE,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONSTARTUP,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONSHUTDOWN,
-    BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMSHUTDOWN,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATEBEGIN,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATE,
@@ -1463,18 +1462,6 @@ struct BA_ONSYSTEMRESTOREPOINTCOMPLETE_RESULTS
     DWORD cbSize;
 };
 
-struct BA_ONSYSTEMSHUTDOWN_ARGS
-{
-    DWORD cbSize;
-    DWORD dwEndSession;
-};
-
-struct BA_ONSYSTEMSHUTDOWN_RESULTS
-{
-    DWORD cbSize;
-    BOOL fCancel;
-};
-
 struct BA_ONUNREGISTERBEGIN_ARGS
 {
     DWORD cbSize;
diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
index 6738a4a6..bb34a33e 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
@@ -22,15 +22,12 @@ namespace WixToolset.Mba.Core
         /// </summary>
         protected readonly IEngine engine;
 
-        private bool applying;
-
         /// <summary>
         /// Creates a new instance of the <see cref="BootstrapperApplication"/> class.
         /// </summary>
         protected BootstrapperApplication(IEngine engine)
         {
             this.engine = engine;
-            this.applying = false;
             this.asyncExecution = true;
         }
 
@@ -41,9 +38,6 @@ namespace WixToolset.Mba.Core
         public event EventHandler<ShutdownEventArgs> Shutdown;
 
         /// <inheritdoc/>
-        public event EventHandler<SystemShutdownEventArgs> SystemShutdown;
-
-        /// <inheritdoc/>
         public event EventHandler<DetectBeginEventArgs> DetectBegin;
 
         /// <inheritdoc/>
@@ -332,25 +326,6 @@ namespace WixToolset.Mba.Core
         }
 
         /// <summary>
-        /// Called by the engine, raises the <see cref="SystemShutdown"/> event.
-        /// </summary>
-        /// <param name="args">Additional arguments for this event.</param>
-        protected virtual void OnSystemShutdown(SystemShutdownEventArgs args)
-        {
-            EventHandler<SystemShutdownEventArgs> handler = this.SystemShutdown;
-            if (null != handler)
-            {
-                handler(this, args);
-            }
-            else if (null != args)
-            {
-                // Allow requests to shut down when critical or not applying.
-                bool critical = EndSessionReasons.Critical == (EndSessionReasons.Critical & args.Reasons);
-                args.Cancel = !critical && this.applying;
-            }
-        }
-
-        /// <summary>
         /// Called by the engine, raises the <see cref="DetectBegin"/> event.
         /// </summary>
         /// <param name="args">Additional arguments for this event.</param>
@@ -1433,15 +1408,6 @@ namespace WixToolset.Mba.Core
             return args.HResult;
         }
 
-        int IBootstrapperApplication.OnSystemShutdown(EndSessionReasons dwEndSession, ref bool fCancel)
-        {
-            SystemShutdownEventArgs args = new SystemShutdownEventArgs(dwEndSession, fCancel);
-            this.OnSystemShutdown(args);
-
-            fCancel = args.Cancel;
-            return args.HResult;
-        }
-
         int IBootstrapperApplication.OnDetectBegin(bool fCached, RegistrationType registrationType, int cPackages, ref bool fCancel)
         {
             DetectBeginEventArgs args = new DetectBeginEventArgs(fCached, registrationType, cPackages, fCancel);
@@ -1693,8 +1659,6 @@ namespace WixToolset.Mba.Core
 
         int IBootstrapperApplication.OnApplyBegin(int dwPhaseCount, ref bool fCancel)
         {
-            this.applying = true;
-
             ApplyBeginEventArgs args = new ApplyBeginEventArgs(dwPhaseCount, fCancel);
             this.OnApplyBegin(args);
 
@@ -1949,8 +1913,6 @@ namespace WixToolset.Mba.Core
             ApplyCompleteEventArgs args = new ApplyCompleteEventArgs(hrStatus, restart, recommendation, pAction);
             this.OnApplyComplete(args);
 
-            this.applying = false;
-
             pAction = args.Action;
             return args.HResult;
         }
diff --git a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
index d79ac402..be113700 100644
--- a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
+++ b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
@@ -217,32 +217,6 @@ namespace WixToolset.Mba.Core
     }
 
     /// <summary>
-    /// Event arguments for <see cref="IDefaultBootstrapperApplication.SystemShutdown"/>
-    /// </summary>
-    [Serializable]
-    public class SystemShutdownEventArgs : CancellableHResultEventArgs
-    {
-        /// <summary />
-        public SystemShutdownEventArgs(EndSessionReasons reasons, bool cancelRecommendation)
-            : base(cancelRecommendation)
-        {
-            this.Reasons = reasons;
-        }
-
-        /// <summary>
-        /// Gets the reason the application is requested to close or being closed.
-        /// </summary>
-        /// <remarks>
-        /// <para>To prevent shutting down or logging off, set <see cref="CancellableHResultEventArgs.Cancel"/> to
-        /// true; otherwise, set it to false.</para>
-        /// <para>If <see cref="SystemShutdownEventArgs.Reasons"/> contains <see cref="EndSessionReasons.Critical"/>
-        /// the bootstrapper cannot prevent the shutdown and only has a few seconds to save state or perform any other
-        /// critical operations before being closed by the operating system.</para>
-        /// </remarks>
-        public EndSessionReasons Reasons { get; private set; }
-    }
-
-    /// <summary>
     /// Event arguments for <see cref="IDefaultBootstrapperApplication.DetectBegin"/>
     /// </summary>
     [Serializable]
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
index 2877d4de..8ce99808 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
@@ -51,19 +51,6 @@ namespace WixToolset.Mba.Core
         int OnShutdown(ref BOOTSTRAPPER_SHUTDOWN_ACTION action);
 
         /// <summary>
-        /// See <see cref="IDefaultBootstrapperApplication.SystemShutdown"/>.
-        /// </summary>
-        /// <param name="dwEndSession"></param>
-        /// <param name="fCancel"></param>
-        /// <returns></returns>
-        [PreserveSig]
-        [return: MarshalAs(UnmanagedType.I4)]
-        int OnSystemShutdown(
-            [MarshalAs(UnmanagedType.U4)] EndSessionReasons dwEndSession,
-            [MarshalAs(UnmanagedType.Bool)] ref bool fCancel
-            );
-
-        /// <summary>
         /// See <see cref="IDefaultBootstrapperApplication.DetectBegin"/>.
         /// </summary>
         [PreserveSig]
diff --git a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
index 77089e83..2535f756 100644
--- a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
@@ -419,22 +419,6 @@ namespace WixToolset.Mba.Core
         event EventHandler<SystemRestorePointCompleteEventArgs> SystemRestorePointComplete;
 
         /// <summary>
-        /// Fired when the system is shutting down or user is logging off.
-        /// </summary>
-        /// <remarks>
-        /// <para>To prevent shutting down or logging off, set <see cref="CancellableHResultEventArgs.Cancel"/> to
-        /// true; otherwise, set it to false.</para>
-        /// <para>By default setup will prevent shutting down or logging off between
-        /// <see cref="IDefaultBootstrapperApplication.ApplyBegin"/> and <see cref="IDefaultBootstrapperApplication.ApplyComplete"/>.
-        /// Derivatives can change this behavior by handling <see cref="IDefaultBootstrapperApplication.SystemShutdown"/>.</para>
-        /// <para>If <see cref="SystemShutdownEventArgs.Reasons"/> contains <see cref="EndSessionReasons.Critical"/>
-        /// the bootstrapper cannot prevent the shutdown and only has a few seconds to save state or perform any other
-        /// critical operations before being closed by the operating system.</para>
-        /// <para>This event may be fired on a different thread.</para>
-        /// </remarks>
-        event EventHandler<SystemShutdownEventArgs> SystemShutdown;
-
-        /// <summary>
         /// Fired when the engine unregisters the bundle.
         /// </summary>
         event EventHandler<UnregisterBeginEventArgs> UnregisterBegin;
diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 9be3f62f..c1057ab6 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -15,7 +15,6 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONSTARTUP = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSTARTUP,
     BA_FUNCTIONS_MESSAGE_ONSHUTDOWN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSHUTDOWN,
-    BA_FUNCTIONS_MESSAGE_ONSYSTEMSHUTDOWN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMSHUTDOWN,
     BA_FUNCTIONS_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE,
     BA_FUNCTIONS_MESSAGE_ONDETECTUPDATEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATEBEGIN,
     BA_FUNCTIONS_MESSAGE_ONDETECTUPDATE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATE,
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctions.h b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
index 9ff58d2b..ca070553 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -98,14 +98,6 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
-    virtual STDMETHODIMP OnSystemShutdown(
-        __in DWORD /*dwEndSession*/,
-        __inout BOOL* /*pfCancel*/
-        )
-    {
-        return S_OK;
-    }
-
     virtual STDMETHODIMP OnDetectBegin(
         __in BOOL /*fCached*/,
         __in BOOTSTRAPPER_REGISTRATION_TYPE /*registrationType*/,
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
index 4564ad0c..ff92717d 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -86,7 +86,6 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_ONPLANCOMPLETE:
         case BA_FUNCTIONS_MESSAGE_ONSTARTUP:
         case BA_FUNCTIONS_MESSAGE_ONSHUTDOWN:
-        case BA_FUNCTIONS_MESSAGE_ONSYSTEMSHUTDOWN:
         case BA_FUNCTIONS_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE:
         case BA_FUNCTIONS_MESSAGE_ONDETECTUPDATEBEGIN:
         case BA_FUNCTIONS_MESSAGE_ONDETECTUPDATE:
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
index 2e848df7..fc9c4dd7 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -94,19 +94,6 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
-    virtual STDMETHODIMP OnSystemShutdown(
-        __in DWORD dwEndSession,
-        __inout BOOL* pfCancel
-        )
-    {
-        HRESULT hr = S_OK;
-
-        // Allow requests to shut down when critical or not applying.
-        *pfCancel = !(ENDSESSION_CRITICAL & dwEndSession || !m_fApplying);
-
-        return hr;
-    }
-
     virtual STDMETHODIMP OnDetectBegin(
         __in BOOL /*fCached*/,
         __in BOOTSTRAPPER_REGISTRATION_TYPE /*registrationType*/,
@@ -406,8 +393,6 @@ public: // IBootstrapperApplication
         __inout BOOL* pfCancel
         )
     {
-        m_fApplying = TRUE;
-
         m_dwProgressPercentage = 0;
         m_dwOverallProgressPercentage = 0;
 
@@ -859,8 +844,6 @@ public: // IBootstrapperApplication
             *pAction = BOOTSTRAPPER_APPLYCOMPLETE_ACTION_RESTART;
         }
 
-        m_fApplying = FALSE;
-
         return hr;
     }
 
@@ -1183,7 +1166,6 @@ protected:
         ::InitializeCriticalSection(&m_csCanceled);
         m_fCanceled = FALSE;
         m_BalInfoCommand = { };
-        m_fApplying = FALSE;
         m_fRollingBack = FALSE;
 
         m_dwProgressPercentage = 0;
@@ -1212,7 +1194,6 @@ private:
     BOOTSTRAPPER_DISPLAY m_display;
     IBootstrapperEngine* m_pEngine;
 
-    BOOL m_fApplying;
     BOOL m_fRollingBack;
 
     DWORD m_dwProgressPercentage;
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
index e35678ad..92243540 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
@@ -63,15 +63,6 @@ static HRESULT BalBaseBAProcOnShutdown(
     return pBA->OnShutdown(&pResults->action);
 }
 
-static HRESULT BalBaseBAProcOnSystemShutdown(
-    __in IBootstrapperApplication* pBA,
-    __in BA_ONSYSTEMSHUTDOWN_ARGS* pArgs,
-    __inout BA_ONSYSTEMSHUTDOWN_RESULTS* pResults
-    )
-{
-    return pBA->OnSystemShutdown(pArgs->dwEndSession, &pResults->fCancel);
-}
-
 static HRESULT BalBaseBAProcOnDetectForwardCompatibleBundle(
     __in IBootstrapperApplication* pBA,
     __in BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_ARGS* pArgs,
@@ -803,9 +794,6 @@ static HRESULT WINAPI BalBaseBootstrapperApplicationProc(
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSHUTDOWN:
             hr = BalBaseBAProcOnShutdown(pBA, reinterpret_cast<BA_ONSHUTDOWN_ARGS*>(pvArgs), reinterpret_cast<BA_ONSHUTDOWN_RESULTS*>(pvResults));
             break;
-        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMSHUTDOWN:
-            hr = BalBaseBAProcOnSystemShutdown(pBA, reinterpret_cast<BA_ONSYSTEMSHUTDOWN_ARGS*>(pvArgs), reinterpret_cast<BA_ONSYSTEMSHUTDOWN_RESULTS*>(pvResults));
-            break;
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE:
             hr = BalBaseBAProcOnDetectForwardCompatibleBundle(pBA, reinterpret_cast<BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_ARGS*>(pvArgs), reinterpret_cast<BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_RESULTS*>(pvResults));
             break;
diff --git a/src/api/burn/balutil/inc/IBootstrapperApplication.h b/src/api/burn/balutil/inc/IBootstrapperApplication.h
index 6eca90fa..382d5aad 100644
--- a/src/api/burn/balutil/inc/IBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/IBootstrapperApplication.h
@@ -33,12 +33,6 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         __inout BOOTSTRAPPER_SHUTDOWN_ACTION* pAction
         ) = 0;
 
-    // OnSystemShutdown - called when the operating system is instructed to shutdown the machine.
-    STDMETHOD(OnSystemShutdown)(
-        __in DWORD dwEndSession,
-        __inout BOOL* pfCancel
-        ) = 0;
-
     // OnDetectBegin - called when the engine begins detection.
     STDMETHOD(OnDetectBegin)(
         __in BOOL fCached,
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
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index e93dbe90..c89d12c4 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -1265,9 +1265,6 @@ public: // IBootstrapperApplication
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSHUTDOWN:
             OnShutdownFallback(reinterpret_cast<BA_ONSHUTDOWN_ARGS*>(pvArgs), reinterpret_cast<BA_ONSHUTDOWN_RESULTS*>(pvResults));
             break;
-        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMSHUTDOWN:
-            OnSystemShutdownFallback(reinterpret_cast<BA_ONSYSTEMSHUTDOWN_ARGS*>(pvArgs), reinterpret_cast<BA_ONSYSTEMSHUTDOWN_RESULTS*>(pvResults));
-            break;
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE:
             OnDetectForwardCompatibleBundleFallback(reinterpret_cast<BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_ARGS*>(pvArgs), reinterpret_cast<BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_RESULTS*>(pvResults));
             break;
@@ -1544,14 +1541,6 @@ private: // privates
         m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONSHUTDOWN, pArgs, pResults, m_pvBAFunctionsProcContext);
     }
 
-    void OnSystemShutdownFallback(
-        __in BA_ONSYSTEMSHUTDOWN_ARGS* pArgs,
-        __inout BA_ONSYSTEMSHUTDOWN_RESULTS* pResults
-        )
-    {
-        m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONSYSTEMSHUTDOWN, pArgs, pResults, m_pvBAFunctionsProcContext);
-    }
-
     void OnDetectForwardCompatibleBundleFallback(
         __in BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_ARGS* pArgs,
         __inout BA_ONDETECTFORWARDCOMPATIBLEBUNDLE_RESULTS* pResults
@@ -2948,7 +2937,6 @@ private:
     {
 #pragma warning(suppress:4312)
         CWixStandardBootstrapperApplication* pBA = reinterpret_cast<CWixStandardBootstrapperApplication*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
-        BOOL fCancel = FALSE;
 
         switch (uMsg)
         {
@@ -2975,11 +2963,6 @@ private:
         case WM_THMUTIL_LOADED_CONTROL:
             return pBA->OnThemeLoadedControl(reinterpret_cast<THEME_LOADEDCONTROL_ARGS*>(wParam), reinterpret_cast<THEME_LOADEDCONTROL_RESULTS*>(lParam));
 
-        case WM_QUERYENDSESSION:
-            fCancel = true;
-            pBA->OnSystemShutdown(static_cast<DWORD>(lParam), &fCancel);
-            return !fCancel;
-
         case WM_CLOSE:
             // If the user chose not to close, do *not* let the default window proc handle the message.
             if (!pBA->OnClose())
diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index b4d74341..5f492efd 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -531,15 +531,6 @@ namespace WixToolset.Test.BA
             this.ShutdownUiThread();
         }
 
-        protected override void OnSystemShutdown(SystemShutdownEventArgs args)
-        {
-            // Always prevent shutdown.
-            this.Log("Disallowed system request to shut down the bootstrapper application.");
-            args.Cancel = true;
-
-            this.ShutdownUiThread();
-        }
-
         protected override void OnUnregisterBegin(UnregisterBeginEventArgs args)
         {
             if (this.forceKeepRegistration && args.RegistrationType == RegistrationType.None)
diff --git a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
index 7c533fea..3cf994fc 100644
--- a/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
+++ b/src/test/burn/TestData/Manual/BafThmutilTesting/BafThmUtilTesting.cpp
@@ -88,6 +88,56 @@ public: // IBAFunctions
         return hr;
     }
 
+    virtual STDMETHODIMP WndProc(
+        __in HWND hWnd,
+        __in UINT uMsg,
+        __in WPARAM /*wParam*/,
+        __in LPARAM lParam,
+        __inout BOOL* pfProcessed,
+        __inout LRESULT* plResult
+        )
+    {
+        switch (uMsg)
+        {
+        case WM_QUERYENDSESSION:
+            if (BOOTSTRAPPER_DISPLAY_FULL <= m_command.display)
+            {
+                DWORD dwEndSession = static_cast<DWORD>(lParam);
+                if (ENDSESSION_CRITICAL & dwEndSession)
+                {
+                    // Return false to get the WM_ENDSESSION message so that critical shutdowns can be delayed.
+                    *plResult = FALSE;
+                    *pfProcessed = TRUE;
+                }
+            }
+            break;
+        case WM_ENDSESSION:
+            if (BOOTSTRAPPER_DISPLAY_FULL <= m_command.display)
+            {
+                ::MessageBoxW(hWnd, L"WM_ENDSESSION", L"BAFunctions WndProc", MB_OK);
+            }
+            break;
+        }
+        return S_OK;
+    }
+
+public: //IBootstrapperApplication
+    virtual STDMETHODIMP OnExecuteBegin(
+        __in DWORD /*cExecutingPackages*/,
+        __inout BOOL* pfCancel
+        )
+    {
+        if (BOOTSTRAPPER_DISPLAY_FULL <= m_command.display)
+        {
+            if (IDCANCEL == ::MessageBoxW(m_hwndParent, L"Shutdown requests should be denied right now.", L"OnExecuteBegin", MB_OKCANCEL))
+            {
+                *pfCancel = TRUE;
+            }
+        }
+
+        return S_OK;
+    }
+
 private:
     HRESULT OnShowTheme()
     {
@@ -138,6 +188,8 @@ private:
 
         ThemeInitializeWindowClass(m_pBafTheme, &wc, CBafThmUtilTesting::TestingWndProc, m_hModule, BAFTHMUTILTESTING_WINDOW_CLASS);
 
+        Assert(wc.lpszClassName);
+
         // If the theme did not provide an icon, try using the icon from the bundle engine.
         if (!wc.hIcon)
         {
diff --git a/src/test/burn/TestData/Manual/BundleA/BundleA.wxs b/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
index 20706b6a..44abc645 100644
--- a/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
+++ b/src/test/burn/TestData/Manual/BundleA/BundleA.wxs
@@ -1,4 +1,4 @@
-<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
 
 
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
@@ -8,7 +8,10 @@
       <Payload SourceFile="$(var.BafThmUtilTesting.TargetPath)" bal:BAFunctions="yes" />
     </BootstrapperApplication>
     <PackageGroup Id="BundlePackages">
-      <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)" />
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)">
+        <MsiProperty Name="FORCERESTARTCA" Value="[FORCERESTARTCA]" />
+      </MsiPackage>
     </PackageGroup>
+    <Variable Name="FORCERESTARTCA" bal:Overridable="yes" />
   </Fragment>
 </Wix>
diff --git a/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj b/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj
index c9c81c3a..61dc27ab 100644
--- a/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj
+++ b/src/test/burn/TestData/Manual/PackageA/PackageA.wixproj
@@ -2,8 +2,12 @@
 <Project Sdk="WixToolset.Sdk">
   <PropertyGroup>
     <UpgradeCode>{0D803A6E-8090-4174-8DAC-810ECC2B1BBF}</UpgradeCode>
+    <ProductComponentsRef>true</ProductComponentsRef>
   </PropertyGroup>
   <ItemGroup>
     <Compile Include="..\..\Templates\Package.wxs" Link="Package.wxs" />
   </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.Util.wixext" />
+  </ItemGroup>
 </Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/Manual/PackageA/ProductA.wxs b/src/test/burn/TestData/Manual/PackageA/ProductA.wxs
new file mode 100644
index 00000000..bf1ef152
--- /dev/null
+++ b/src/test/burn/TestData/Manual/PackageA/ProductA.wxs
@@ -0,0 +1,15 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:util="http://wixtoolset.org/schemas/v4/wxs/util">
+    <Fragment>
+        <ComponentGroup Id="ProductComponents" />
+
+        <Property Id="FORCERESTARTCA" Secure="yes" />
+        <Property Id="ForceRestartCA" Value="&quot;shutdown.exe&quot; -r -f -t 0" />
+        <CustomAction Id="ForceRestartCA" DllEntry="WixQuietExec" BinaryRef="Wix4UtilCA_X86" Execute="deferred" Return="ignore" />
+
+        <InstallExecuteSequence>
+            <Custom Action="ForceRestartCA" After="InstallFiles" Condition="FORCERESTARTCA = 1" />
+        </InstallExecuteSequence>
+    </Fragment>
+</Wix>