Use MSI transaction end result to detect whether reboot is needed

7 files changed, 167 insertions, 69 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 8fbaa76e..9a1a7bf9 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -300,12 +300,14 @@ static HRESULT ExecuteMsiBeginTransaction(
 static HRESULT ExecuteMsiCommitTransaction(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
-    __in BURN_EXECUTE_CONTEXT* pContext
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     );
 static HRESULT ExecuteMsiRollbackTransaction(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
-    __in BURN_EXECUTE_CONTEXT* pContext
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     );
 static void ResetTransactionRegistrationState(
     __in BURN_ENGINE_STATE* pEngineState,
@@ -645,7 +647,7 @@ extern "C" HRESULT ApplyCache(
             Assert(pPlan->sczLayoutDirectory);
             hr = ApplyLayoutContainer(&cacheContext, pCacheAction->container.pContainer);
             ExitOnFailure(hr, "Failed cache action: %ls", L"layout container");
-            
+
             break;
 
         case BURN_CACHE_ACTION_TYPE_SIGNAL_SYNCPOINT:
@@ -793,7 +795,7 @@ extern "C" HRESULT ApplyExecute(
 
                 if (pCheckpoint && pCheckpoint->pActiveRollbackBoundary && pCheckpoint->pActiveRollbackBoundary->fActiveTransaction)
                 {
-                    hrRollback = ExecuteMsiCommitTransaction(pEngineState, pCheckpoint->pActiveRollbackBoundary, &context);
+                    hrRollback = ExecuteMsiCommitTransaction(pEngineState, pCheckpoint->pActiveRollbackBoundary, &context, pRestart);
                     IgnoreRollbackError(hrRollback, "Failed commit transaction from disable rollback");
                 }
 
@@ -806,7 +808,7 @@ extern "C" HRESULT ApplyExecute(
                 // If inside a MSI transaction, roll it back.
                 if (pCheckpoint->pActiveRollbackBoundary && pCheckpoint->pActiveRollbackBoundary->fActiveTransaction)
                 {
-                    hrRollback = ExecuteMsiRollbackTransaction(pEngineState, pCheckpoint->pActiveRollbackBoundary, &context);
+                    hrRollback = ExecuteMsiRollbackTransaction(pEngineState, pCheckpoint->pActiveRollbackBoundary, &context, pRestart);
                     IgnoreRollbackError(hrRollback, "Failed rolling back transaction");
                 }
 
@@ -2058,13 +2060,13 @@ static HRESULT DownloadPayload(
     cacheCallback.pfnProgress = CacheProgressRoutine;
     cacheCallback.pfnCancel = NULL; // TODO: set this
     cacheCallback.pv = pProgress;
-   
+
     authenticationData.pUX = pProgress->pCacheContext->pUX;
     authenticationData.wzPackageOrContainerId = wzPackageOrContainerId;
     authenticationData.wzPayloadId = wzPayloadId;
     authenticationCallback.pv =  static_cast<LPVOID>(&authenticationData);
     authenticationCallback.pfnAuthenticate = &AuthenticationRequired;
-        
+
     hr = DownloadUrl(pDownloadSource, qwDownloadSize, wzDestinationPath, &cacheCallback, &authenticationCallback);
     ExitOnFailure(hr, "Failed attempt to download URL: '%ls' to: '%ls'", pDownloadSource->sczUrl, wzDestinationPath);
 
@@ -2526,7 +2528,7 @@ static HRESULT DoExecuteAction(
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_COMMIT_MSI_TRANSACTION:
-            hr = ExecuteMsiCommitTransaction(pEngineState, pExecuteAction->msiTransaction.pRollbackBoundary, pContext);
+            hr = ExecuteMsiCommitTransaction(pEngineState, pExecuteAction->msiTransaction.pRollbackBoundary, pContext, &restart);
             ExitOnFailure(hr, "Failed to execute commit MSI transaction action.");
             break;
 
@@ -3381,11 +3383,13 @@ LExit:
 static HRESULT ExecuteMsiCommitTransaction(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
-    __in BURN_EXECUTE_CONTEXT* /*pContext*/
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     )
 {
     HRESULT hr = S_OK;
     BOOL fCommitBeginCalled = FALSE;
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION action = BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION_NONE;
 
     if (!pRollbackBoundary->fActiveTransaction)
     {
@@ -3398,12 +3402,12 @@ static HRESULT ExecuteMsiCommitTransaction(
 
     if (pEngineState->plan.fPerMachine)
     {
-        hr = ElevationMsiCommitTransaction(pEngineState->companionConnection.hPipe, pRollbackBoundary);
+        hr = ElevationMsiCommitTransaction(pEngineState->companionConnection.hPipe, pRollbackBoundary, MsiExecuteMessageHandler, pContext, pRestart);
         ExitOnFailure(hr, "Failed to commit an elevated MSI transaction.");
     }
     else
     {
-        hr = MsiEngineCommitTransaction(pRollbackBoundary);
+        hr = MsiEngineCommitTransaction(pRollbackBoundary, pRestart);
     }
 
     // Assume that MsiEndTransaction can only be called once for each MsiBeginTransaction.
@@ -3414,7 +3418,12 @@ static HRESULT ExecuteMsiCommitTransaction(
 LExit:
     if (fCommitBeginCalled)
     {
-        UserExperienceOnCommitMsiTransactionComplete(&pEngineState->userExperience, pRollbackBoundary->sczId, hr);
+        UserExperienceOnCommitMsiTransactionComplete(&pEngineState->userExperience, pRollbackBoundary->sczId, hr, *pRestart, &action);
+
+        if (action == BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION_RESTART)
+        {
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_INITIATED;
+        }
     }
 
     return hr;
@@ -3423,11 +3432,13 @@ LExit:
 static HRESULT ExecuteMsiRollbackTransaction(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
-    __in BURN_EXECUTE_CONTEXT* /*pContext*/
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     )
 {
     HRESULT hr = S_OK;
     BOOL fRollbackBeginCalled = FALSE;
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION action = BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION_NONE;
 
     if (!pRollbackBoundary->fActiveTransaction)
     {
@@ -3439,12 +3450,12 @@ static HRESULT ExecuteMsiRollbackTransaction(
 
     if (pEngineState->plan.fPerMachine)
     {
-        hr = ElevationMsiRollbackTransaction(pEngineState->companionConnection.hPipe, pRollbackBoundary);
+        hr = ElevationMsiRollbackTransaction(pEngineState->companionConnection.hPipe, pRollbackBoundary, MsiExecuteMessageHandler, pContext, pRestart);
         ExitOnFailure(hr, "Failed to rollback an elevated MSI transaction.");
     }
     else
     {
-        hr = MsiEngineRollbackTransaction(pRollbackBoundary);
+        hr = MsiEngineRollbackTransaction(pRollbackBoundary, pRestart);
     }
 
 LExit:
@@ -3454,7 +3465,12 @@ LExit:
 
     if (fRollbackBeginCalled)
     {
-        UserExperienceOnRollbackMsiTransactionComplete(&pEngineState->userExperience, pRollbackBoundary->sczId, hr);
+        UserExperienceOnRollbackMsiTransactionComplete(&pEngineState->userExperience, pRollbackBoundary->sczId, hr, *pRestart, &action);
+
+        if (action == BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION_RESTART)
+        {
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_INITIATED;
+        }
     }
 
     return hr;
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 154c407d..63a76c2c 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -371,12 +371,14 @@ static HRESULT OnMsiBeginTransaction(
 static HRESULT OnMsiCommitTransaction(
     __in BURN_PACKAGES* pPackages,
     __in BYTE* pbData,
-    __in SIZE_T cbData
+    __in SIZE_T cbData,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     );
 static HRESULT OnMsiRollbackTransaction(
     __in BURN_PACKAGES* pPackages,
     __in BYTE* pbData,
-    __in SIZE_T cbData
+    __in SIZE_T cbData,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     );
 static HRESULT ElevatedOnPauseAUBegin(
     __in HANDLE hPipe
@@ -496,7 +498,7 @@ extern "C" HRESULT ElevationApplyInitialize(
 
     hr = BuffWriteNumber(&pbData, &cbData, (DWORD)!pPlan->pInternalCommand->fDisableSystemRestore);
     ExitOnFailure(hr, "Failed to write system restore point action to message buffer.");
-    
+
     hr = VariableSerialize(pVariables, FALSE, &pbData, &cbData);
     ExitOnFailure(hr, "Failed to write variables.");
 
@@ -544,7 +546,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationSessionBegin - 
+ ElevationSessionBegin -
 
 *******************************************************************/
 extern "C" HRESULT ElevationSessionBegin(
@@ -602,7 +604,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationSessionEnd - 
+ ElevationSessionEnd -
 
 *******************************************************************/
 extern "C" HRESULT ElevationSessionEnd(
@@ -648,7 +650,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationSaveState - 
+ ElevationSaveState -
 
 *******************************************************************/
 HRESULT ElevationSaveState(
@@ -697,7 +699,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationCacheCompletePayload - 
+ ElevationCacheCompletePayload -
 
 *******************************************************************/
 extern "C" HRESULT ElevationCacheCompletePayload(
@@ -785,7 +787,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationCacheCleanup - 
+ ElevationCacheCleanup -
 
 *******************************************************************/
 extern "C" HRESULT ElevationCacheCleanup(
@@ -838,7 +840,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationExecuteRelatedBundle - 
+ ElevationExecuteRelatedBundle -
 
 *******************************************************************/
 extern "C" HRESULT ElevationExecuteRelatedBundle(
@@ -899,7 +901,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationExecuteBundlePackage - 
+ ElevationExecuteBundlePackage -
 
 *******************************************************************/
 extern "C" HRESULT ElevationExecuteBundlePackage(
@@ -963,7 +965,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationExecuteExePackage - 
+ ElevationExecuteExePackage -
 
 *******************************************************************/
 extern "C" HRESULT ElevationExecuteExePackage(
@@ -1047,12 +1049,16 @@ LExit:
 
 extern "C" HRESULT ElevationMsiCommitTransaction(
     __in HANDLE hPipe,
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __in PFN_MSIEXECUTEMESSAGEHANDLER pfnMessageHandler,
+    __in LPVOID pvContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     )
 {
     HRESULT hr = S_OK;
     BYTE* pbData = NULL;
     SIZE_T cbData = 0;
+    BURN_ELEVATION_MSI_MESSAGE_CONTEXT context = { };
     DWORD dwResult = ERROR_SUCCESS;
 
     // serialize message data
@@ -1062,10 +1068,15 @@ extern "C" HRESULT ElevationMsiCommitTransaction(
     hr = BuffWriteString(&pbData, &cbData, pRollbackBoundary->sczLogPath);
     ExitOnFailure(hr, "Failed to write transaction log path to message buffer.");
 
-    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_COMMIT_MSI_TRANSACTION, pbData, cbData, NULL, NULL, &dwResult);
+    // send message
+    context.pfnMessageHandler = pfnMessageHandler;
+    context.pvContext = pvContext;
+
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_COMMIT_MSI_TRANSACTION, pbData, cbData, ProcessMsiPackageMessages, &context, &dwResult);
     ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_COMMIT_MSI_TRANSACTION message to per-machine process.");
 
     hr = static_cast<HRESULT>(dwResult);
+    *pRestart = context.restart;
 
 LExit:
     ReleaseBuffer(pbData);
@@ -1075,12 +1086,16 @@ LExit:
 
 extern "C" HRESULT ElevationMsiRollbackTransaction(
     __in HANDLE hPipe,
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __in PFN_MSIEXECUTEMESSAGEHANDLER pfnMessageHandler,
+    __in LPVOID pvContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     )
 {
     HRESULT hr = S_OK;
     BYTE* pbData = NULL;
     SIZE_T cbData = 0;
+    BURN_ELEVATION_MSI_MESSAGE_CONTEXT context = { };
     DWORD dwResult = ERROR_SUCCESS;
 
     // serialize message data
@@ -1090,10 +1105,15 @@ extern "C" HRESULT ElevationMsiRollbackTransaction(
     hr = BuffWriteString(&pbData, &cbData, pRollbackBoundary->sczLogPath);
     ExitOnFailure(hr, "Failed to write transaction log path to message buffer.");
 
-    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_ROLLBACK_MSI_TRANSACTION, pbData, cbData, NULL, NULL, &dwResult);
+    // send message
+    context.pfnMessageHandler = pfnMessageHandler;
+    context.pvContext = pvContext;
+
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_ROLLBACK_MSI_TRANSACTION, pbData, cbData, ProcessMsiPackageMessages, &context, &dwResult);
     ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_ROLLBACK_MSI_TRANSACTION message to per-machine process.");
 
     hr = static_cast<HRESULT>(dwResult);
+    *pRestart = context.restart;
 
 LExit:
     ReleaseBuffer(pbData);
@@ -1104,7 +1124,7 @@ LExit:
 
 
 /*******************************************************************
- ElevationExecuteMsiPackage - 
+ ElevationExecuteMsiPackage -
 
 *******************************************************************/
 extern "C" HRESULT ElevationExecuteMsiPackage(
@@ -1189,7 +1209,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationExecuteMspPackage - 
+ ElevationExecuteMspPackage -
 
 *******************************************************************/
 extern "C" HRESULT ElevationExecuteMspPackage(
@@ -1269,7 +1289,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationExecuteMsuPackage - 
+ ElevationExecuteMsuPackage -
 
 *******************************************************************/
 extern "C" HRESULT ElevationExecuteMsuPackage(
@@ -1480,7 +1500,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationCleanPackage - 
+ ElevationCleanPackage -
 
 *******************************************************************/
 extern "C" HRESULT ElevationCleanPackage(
@@ -1545,7 +1565,7 @@ LExit:
 }
 
 /*******************************************************************
- ElevationChildPumpMessages - 
+ ElevationChildPumpMessages -
 
 *******************************************************************/
 extern "C" HRESULT ElevationChildPumpMessages(
@@ -1875,7 +1895,7 @@ static HRESULT ProcessGenericExecuteMessages(
 
     hr = BuffReadNumber((BYTE*)pMsg->pvData, pMsg->cbData, &iData, &message.dwUIHint);
     ExitOnFailure(hr, "Failed to allowed results.");
-    
+
     // Process the message.
     switch (pMsg->dwMessage)
     {
@@ -2052,7 +2072,7 @@ static HRESULT ProcessMsiPackageMessages(
         ExitOnRootFailure(hr, "Invalid package message.");
         break;
     }
-    
+
     // send message
     *pdwResult = (DWORD)pContext->pfnMessageHandler(&message, pContext->pvContext);
 
@@ -2154,11 +2174,13 @@ static HRESULT ProcessElevatedChildMessage(
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_COMMIT_MSI_TRANSACTION:
-        hrResult = OnMsiCommitTransaction(pContext->pPackages, (BYTE*)pMsg->pvData, pMsg->cbData);
+        hrResult = OnMsiCommitTransaction(pContext->pPackages, (BYTE*)pMsg->pvData, pMsg->cbData, &restart);
+        fSendRestart = TRUE;
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_ROLLBACK_MSI_TRANSACTION:
-        hrResult = OnMsiRollbackTransaction(pContext->pPackages, (BYTE*)pMsg->pvData, pMsg->cbData);
+        hrResult = OnMsiRollbackTransaction(pContext->pPackages, (BYTE*)pMsg->pvData, pMsg->cbData, &restart);
+        fSendRestart = TRUE;
         break;
 
     case BURN_ELEVATION_MESSAGE_TYPE_APPLY_INITIALIZE:
@@ -3922,7 +3944,8 @@ LExit:
 static HRESULT OnMsiCommitTransaction(
     __in BURN_PACKAGES* pPackages,
     __in BYTE* pbData,
-    __in SIZE_T cbData
+    __in SIZE_T cbData,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     )
 {
     HRESULT hr = S_OK;
@@ -3946,7 +3969,7 @@ static HRESULT OnMsiCommitTransaction(
         pRollbackBoundary->sczLogPath = sczLogPath;
     }
 
-    hr = MsiEngineCommitTransaction(pRollbackBoundary);
+    hr = MsiEngineCommitTransaction(pRollbackBoundary, pRestart);
 
 LExit:
     ReleaseStr(sczId);
@@ -3963,7 +3986,8 @@ LExit:
 static HRESULT OnMsiRollbackTransaction(
     __in BURN_PACKAGES* pPackages,
     __in BYTE* pbData,
-    __in SIZE_T cbData
+    __in SIZE_T cbData,
+    __out BOOTSTRAPPER_APPLY_RESTART *pRestart
     )
 {
     HRESULT hr = S_OK;
@@ -3987,7 +4011,7 @@ static HRESULT OnMsiRollbackTransaction(
         pRollbackBoundary->sczLogPath = sczLogPath;
     }
 
-    hr = MsiEngineRollbackTransaction(pRollbackBoundary);
+    hr = MsiEngineRollbackTransaction(pRollbackBoundary, pRestart);
 
 LExit:
     ReleaseStr(sczId);
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index 810287a3..02019977 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -191,11 +191,17 @@ HRESULT ElevationMsiBeginTransaction(
     );
 HRESULT ElevationMsiCommitTransaction(
     __in HANDLE hPipe,
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __in PFN_MSIEXECUTEMESSAGEHANDLER pfnMessageHandler,
+    __in LPVOID pvContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
 HRESULT ElevationMsiRollbackTransaction(
     __in HANDLE hPipe,
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __in PFN_MSIEXECUTEMESSAGEHANDLER pfnMessageHandler,
+    __in LPVOID pvContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
 
 #ifdef __cplusplus
diff --git a/src/burn/engine/msiengine.cpp b/src/burn/engine/msiengine.cpp
index 2c1b0d4f..fec995e0 100644
--- a/src/burn/engine/msiengine.cpp
+++ b/src/burn/engine/msiengine.cpp
@@ -1139,33 +1139,65 @@ LExit:
 }
 
 extern "C" HRESULT MsiEngineCommitTransaction(
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     )
 {
     HRESULT hr = S_OK;
+    WIU_RESTART restart = WIU_RESTART_NONE;
 
     LogId(REPORT_STANDARD, MSG_MSI_TRANSACTION_COMMIT, pRollbackBoundary->sczId);
 
-    hr = WiuEndTransaction(MSITRANSACTIONSTATE_COMMIT, WIU_LOG_DEFAULT | INSTALLLOGMODE_VERBOSE, pRollbackBoundary->sczLogPath);
+    hr = WiuEndTransaction(MSITRANSACTIONSTATE_COMMIT, WIU_LOG_DEFAULT | INSTALLLOGMODE_VERBOSE, pRollbackBoundary->sczLogPath, &restart);
     ExitOnFailure(hr, "Failed to commit the MSI transaction");
 
 LExit:
+    switch (restart)
+    {
+        case WIU_RESTART_NONE:
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_NONE;
+            break;
+
+        case WIU_RESTART_REQUIRED:
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_REQUIRED;
+            break;
+
+        case WIU_RESTART_INITIATED:
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_INITIATED;
+            break;
+    }
 
     return hr;
 }
 
 extern "C" HRESULT MsiEngineRollbackTransaction(
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     )
 {
     HRESULT hr = S_OK;
+    WIU_RESTART restart = WIU_RESTART_NONE;
 
     LogId(REPORT_WARNING, MSG_MSI_TRANSACTION_ROLLBACK, pRollbackBoundary->sczId);
 
-    hr = WiuEndTransaction(MSITRANSACTIONSTATE_ROLLBACK, WIU_LOG_DEFAULT | INSTALLLOGMODE_VERBOSE, pRollbackBoundary->sczLogPath);
+    hr = WiuEndTransaction(MSITRANSACTIONSTATE_ROLLBACK, WIU_LOG_DEFAULT | INSTALLLOGMODE_VERBOSE, pRollbackBoundary->sczLogPath, &restart);
     ExitOnFailure(hr, "Failed to rollback the MSI transaction");
 
 LExit:
+    switch (restart)
+    {
+        case WIU_RESTART_NONE:
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_NONE;
+            break;
+
+        case WIU_RESTART_REQUIRED:
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_REQUIRED;
+            break;
+
+        case WIU_RESTART_INITIATED:
+            *pRestart = BOOTSTRAPPER_APPLY_RESTART_INITIATED;
+            break;
+    }
 
     return hr;
 }
@@ -1246,7 +1278,7 @@ extern "C" HRESULT MsiEngineExecutePackage(
 
     // Best effort to set the execute package action variable.
     VariableSetNumeric(pVariables, BURN_BUNDLE_EXECUTE_PACKAGE_ACTION, pExecuteAction->msiPackage.action, TRUE);
-    
+
     // Wire up the external UI handler and logging.
     if (pExecuteAction->msiPackage.fDisableExternalUiHandler)
     {
diff --git a/src/burn/engine/msiengine.h b/src/burn/engine/msiengine.h
index bc356fab..862c4f6a 100644
--- a/src/burn/engine/msiengine.h
+++ b/src/burn/engine/msiengine.h
@@ -60,10 +60,12 @@ HRESULT MsiEngineBeginTransaction(
     __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
     );
 HRESULT MsiEngineCommitTransaction(
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
 HRESULT MsiEngineRollbackTransaction(
-    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary
+    __in BURN_ROLLBACK_BOUNDARY* pRollbackBoundary,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
 HRESULT MsiEngineExecutePackage(
     __in_opt HWND hwndParent,
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 28429394..372ca901 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -35,7 +35,7 @@ static HRESULT SendBAMessageFromInactiveEngine(
 // function definitions
 
 /*******************************************************************
- UserExperienceParseFromXml - 
+ UserExperienceParseFromXml -
 
 *******************************************************************/
 extern "C" HRESULT UserExperienceParseFromXml(
@@ -72,7 +72,7 @@ LExit:
 }
 
 /*******************************************************************
- UserExperienceUninitialize - 
+ UserExperienceUninitialize -
 
 *******************************************************************/
 extern "C" void UserExperienceUninitialize(
@@ -87,7 +87,7 @@ extern "C" void UserExperienceUninitialize(
 }
 
 /*******************************************************************
- UserExperienceLoad - 
+ UserExperienceLoad -
 
 *******************************************************************/
 extern "C" HRESULT UserExperienceLoad(
@@ -129,7 +129,7 @@ LExit:
 }
 
 /*******************************************************************
- UserExperienceUnload - 
+ UserExperienceUnload -
 
 *******************************************************************/
 extern "C" HRESULT UserExperienceUnload(
@@ -380,7 +380,7 @@ EXTERN_C BAAPI UserExperienceOnBeginMsiTransactionBegin(
     args.wzTransactionId = wzTransactionId;
 
     results.cbSize = sizeof(results);
-    
+
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONBEGINMSITRANSACTIONBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnBeginMsiTransactionBegin failed.");
 
@@ -1017,7 +1017,7 @@ EXTERN_C BAAPI UserExperienceOnCommitMsiTransactionBegin(
     args.wzTransactionId = wzTransactionId;
 
     results.cbSize = sizeof(results);
-    
+
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONCOMMITMSITRANSACTIONBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnCommitMsiTransactionBegin failed.");
 
@@ -1033,8 +1033,10 @@ LExit:
 EXTERN_C BAAPI UserExperienceOnCommitMsiTransactionComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in LPCWSTR wzTransactionId,
-    __in HRESULT hrStatus
-    )
+    __in HRESULT hrStatus,
+    __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* pAction
+)
 {
     HRESULT hr = S_OK;
     BA_ONCOMMITMSITRANSACTIONCOMPLETE_ARGS args = { };
@@ -1043,12 +1045,17 @@ EXTERN_C BAAPI UserExperienceOnCommitMsiTransactionComplete(
     args.cbSize = sizeof(args);
     args.wzTransactionId = wzTransactionId;
     args.hrStatus = hrStatus;
+    args.restart = restart;
+    args.recommendation = *pAction;
 
     results.cbSize = sizeof(results);
+    results.action = *pAction;
 
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONCOMMITMSITRANSACTIONCOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnCommitMsiTransactionComplete failed.");
 
+    *pAction = results.action;
+
 LExit:
     return hr;
 }
@@ -1908,7 +1915,7 @@ EXTERN_C BAAPI UserExperienceOnPauseAUBegin(
     args.cbSize = sizeof(args);
 
     results.cbSize = sizeof(results);
-    
+
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONPAUSEAUTOMATICUPDATESBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnPauseAUBegin failed.");
 
@@ -2523,7 +2530,7 @@ EXTERN_C BAAPI UserExperienceOnRollbackMsiTransactionBegin(
     args.wzTransactionId = wzTransactionId;
 
     results.cbSize = sizeof(results);
-    
+
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONROLLBACKMSITRANSACTIONBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnRollbackMsiTransactionBegin failed.");
 
@@ -2534,7 +2541,9 @@ LExit:
 EXTERN_C BAAPI UserExperienceOnRollbackMsiTransactionComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in LPCWSTR wzTransactionId,
-    __in HRESULT hrStatus
+    __in HRESULT hrStatus,
+    __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION *pAction
     )
 {
     HRESULT hr = S_OK;
@@ -2544,12 +2553,17 @@ EXTERN_C BAAPI UserExperienceOnRollbackMsiTransactionComplete(
     args.cbSize = sizeof(args);
     args.wzTransactionId = wzTransactionId;
     args.hrStatus = hrStatus;
+    args.restart = restart;
+    args.recommendation = *pAction;
 
     results.cbSize = sizeof(results);
+    results.action = *pAction;
 
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONROLLBACKMSITRANSACTIONCOMPLETE, &args, &results);
     ExitOnFailure(hr, "BA OnRollbackMsiTransactionComplete failed.");
 
+    *pAction = results.action;
+
 LExit:
     return hr;
 }
@@ -2651,7 +2665,7 @@ EXTERN_C BAAPI UserExperienceOnSystemRestorePointBegin(
     args.cbSize = sizeof(args);
 
     results.cbSize = sizeof(results);
-    
+
     hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSYSTEMRESTOREPOINTBEGIN, &args, &results);
     ExitOnFailure(hr, "BA OnSystemRestorePointBegin failed.");
 
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index e342da03..4f15c5d7 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -251,8 +251,10 @@ BAAPI UserExperienceOnCommitMsiTransactionBegin(
 BAAPI UserExperienceOnCommitMsiTransactionComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in LPCWSTR wzTransactionId,
-    __in HRESULT hrStatus
-    );
+    __in HRESULT hrStatus,
+    __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* pAction
+);
 BAAPI UserExperienceOnDetectBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in BOOL fCached,
@@ -567,8 +569,10 @@ BAAPI UserExperienceOnRollbackMsiTransactionBegin(
 BAAPI UserExperienceOnRollbackMsiTransactionComplete(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in LPCWSTR wzTransactionId,
-    __in HRESULT hrStatus
-    );
+    __in HRESULT hrStatus,
+    __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* pAction
+);
 BAAPI UserExperienceOnSetUpdateBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience
     );