When rolling back a bundle failure, reinstall all upgrade related bundles.

Fixes #3421

11 files changed, 227 insertions, 26 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 5cc63d02..73f8fc72 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -210,6 +210,11 @@ static HRESULT ExecuteRelatedBundle(
     __out BOOL* pfSuspend,
     __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
+static HRESULT DoRestoreRelatedBundleActions(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
+    );
 static HRESULT ExecuteExePackage(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_EXECUTE_ACTION* pExecuteAction,
@@ -788,6 +793,9 @@ extern "C" HRESULT ApplyExecute(
                 {
                     if (pCheckpoint->pActiveRollbackBoundary->fVital)
                     {
+                        hrRollback = DoRestoreRelatedBundleActions(pEngineState, &context, pRestart);
+                        IgnoreRollbackError(hrRollback, "Failed rollback actions");
+
                         // If the rollback boundary is vital, end execution here.
                         break;
                     }
@@ -2590,6 +2598,48 @@ LExit:
     return hr;
 }
 
+static HRESULT DoRestoreRelatedBundleActions(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fRetryIgnored = FALSE;
+    BOOL fSuspendIgnored = FALSE;
+
+    // execute restore related bundle actions
+    for (DWORD i = 0; i < pEngineState->plan.cRestoreRelatedBundleActions; ++i)
+    {
+        BURN_EXECUTE_ACTION* pRestoreRelatedBundleAction = &pEngineState->plan.rgRestoreRelatedBundleActions[i];
+        if (pRestoreRelatedBundleAction->fDeleted)
+        {
+            continue;
+        }
+
+        BOOTSTRAPPER_APPLY_RESTART restart = BOOTSTRAPPER_APPLY_RESTART_NONE;
+        switch (pRestoreRelatedBundleAction->type)
+        {
+        case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
+            hr = ExecuteRelatedBundle(pEngineState, pRestoreRelatedBundleAction, pContext, TRUE, &fRetryIgnored, &fSuspendIgnored, &restart);
+            IgnoreRollbackError(hr, "Failed to restore related bundle package.");
+            break;
+
+        default:
+            hr = E_UNEXPECTED;
+            ExitOnFailure(hr, "Invalid restore related bundle action: %d.", pRestoreRelatedBundleAction->type);
+        }
+
+        if (*pRestart < restart)
+        {
+            *pRestart = restart;
+        }
+    }
+
+LExit:
+    return hr;
+}
+
 static HRESULT ExecuteExePackage(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_EXECUTE_ACTION* pExecuteAction,
diff --git a/src/burn/engine/bundlepackageengine.cpp b/src/burn/engine/bundlepackageengine.cpp
index cd84601f..b5fc51e5 100644
--- a/src/burn/engine/bundlepackageengine.cpp
+++ b/src/burn/engine/bundlepackageengine.cpp
@@ -51,7 +51,7 @@ extern "C" HRESULT BundlePackageEnginePlanCalculatePackage(
             execute = BOOTSTRAPPER_ACTION_STATE_NONE;
             break;
         case BOOTSTRAPPER_REQUEST_STATE_REPAIR:
-            execute = pPackage->Bundle.fRepairable ? BOOTSTRAPPER_ACTION_STATE_REPAIR : BOOTSTRAPPER_ACTION_STATE_NONE;
+            execute = BOOTSTRAPPER_ACTION_STATE_REPAIR;
             break;
         case BOOTSTRAPPER_REQUEST_STATE_ABSENT: __fallthrough;
         case BOOTSTRAPPER_REQUEST_STATE_CACHE:
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index 7f7a915e..8fac7bd0 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -558,7 +558,7 @@ extern "C" HRESULT CorePlan(
                 ExitOnFailure(hr, "Failed to plan packages.");
 
                 // Schedule the update of related bundles last.
-                hr = PlanRelatedBundlesComplete(&pEngineState->registration, &pEngineState->plan, &pEngineState->log, &pEngineState->variables, dwExecuteActionEarlyIndex);
+                hr = PlanRelatedBundlesComplete(&pEngineState->userExperience, &pEngineState->registration, &pEngineState->plan, &pEngineState->log, &pEngineState->variables, dwExecuteActionEarlyIndex);
                 ExitOnFailure(hr, "Failed to schedule related bundles.");
             }
         }
@@ -2309,7 +2309,7 @@ static void LogRelatedBundles(
 
             if (pRelatedBundle->fPlannable)
             {
-                LogId(REPORT_STANDARD, MSG_PLANNED_RELATED_BUNDLE, pPackage->sczId, LoggingRelationTypeToString(pRelatedBundle->relationType), LoggingRequestStateToString(pPackage->defaultRequested), LoggingRequestStateToString(pPackage->requested), LoggingActionStateToString(pPackage->execute), LoggingActionStateToString(pPackage->rollback), LoggingDependencyActionToString(pPackage->dependencyExecute));
+                LogId(REPORT_STANDARD, MSG_PLANNED_RELATED_BUNDLE, pPackage->sczId, LoggingRelationTypeToString(pRelatedBundle->relationType), LoggingRequestStateToString(pPackage->defaultRequested), LoggingRequestStateToString(pPackage->requested), LoggingActionStateToString(pPackage->execute), LoggingActionStateToString(pPackage->rollback), LoggingRequestStateToString(pRelatedBundle->defaultRequestedRestore), LoggingRequestStateToString(pRelatedBundle->requestedRestore), LoggingActionStateToString(pRelatedBundle->restore), LoggingDependencyActionToString(pPackage->dependencyExecute));
             }
         }
     }
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index b8c9a2d3..675a5644 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -352,13 +352,6 @@ Language=English
 Planned package: %1!ls!, state: %2!hs!, default requested: %3!hs!, ba requested: %4!hs!, execute: %5!hs!, rollback: %6!hs!, default cache strategy: %7!hs!, ba requested strategy: %8!hs!, cache: %9!hs!, uncache: %10!hs!, dependency: %11!hs!, expected install registration state: %12!hs!, expected cache registration state: %13!hs!
 .
 
-MessageId=202
-Severity=Success
-SymbolicName=MSG_PLANNED_BUNDLE_UX_CHANGED_REQUEST
-Language=English
-Planned bundle: %1!ls!, ba requested state: %2!hs! over default: %3!hs!
-.
-
 MessageId=203
 Severity=Success
 SymbolicName=MSG_PLANNED_MSI_FEATURE
@@ -391,7 +384,7 @@ MessageId=207
 Severity=Success
 SymbolicName=MSG_PLANNED_RELATED_BUNDLE
 Language=English
-Planned related bundle: %1!ls!, type: %2!hs!, default requested: %3!hs!, ba requested: %4!hs!, execute: %5!hs!, rollback: %6!hs!, dependency: %7!hs!
+Planned related bundle: %1!ls!, type: %2!hs!, default requested: %3!hs!, ba requested: %4!hs!, execute: %5!hs!, rollback: %6!hs!, default requested restore: %7!hs!, ba requested restore: %8!hs!, restore: %9!hs!, dependency: %10!hs!
 .
 
 MessageId=208
diff --git a/src/burn/engine/package.h b/src/burn/engine/package.h
index 6d1b5dd9..1e61d92b 100644
--- a/src/burn/engine/package.h
+++ b/src/burn/engine/package.h
@@ -309,7 +309,6 @@ typedef struct _BURN_PACKAGE
             LPCWSTR wzAncestors; // points directly into engine state.
             LPCWSTR wzEngineWorkingDirectory; // points directly into engine state.
 
-            BOOL fRepairable;
             BOOL fSupportsBurnProtocol;
 
             BURN_EXE_EXIT_CODE* rgExitCodes;
diff --git a/src/burn/engine/plan.cpp b/src/burn/engine/plan.cpp
index f850d49c..f1fb87b8 100644
--- a/src/burn/engine/plan.cpp
+++ b/src/burn/engine/plan.cpp
@@ -103,6 +103,10 @@ static HRESULT AppendCleanAction(
     __in BURN_PLAN* pPlan,
     __out BURN_CLEAN_ACTION** ppCleanAction
     );
+static HRESULT AppendRestoreRelatedBundleAction(
+    __in BURN_PLAN* pPlan,
+    __out BURN_EXECUTE_ACTION** ppExecuteAction
+    );
 static HRESULT ProcessPayloadGroup(
     __in BURN_PLAN* pPlan,
     __in BURN_PAYLOAD_GROUP* pPayloadGroup
@@ -196,6 +200,15 @@ extern "C" void PlanReset(
         MemFree(pPlan->rgRollbackActions);
     }
 
+    if (pPlan->rgRestoreRelatedBundleActions)
+    {
+        for (DWORD i = 0; i < pPlan->cRestoreRelatedBundleActions; ++i)
+        {
+            PlanUninitializeExecuteAction(&pPlan->rgRestoreRelatedBundleActions[i]);
+        }
+        MemFree(pPlan->rgRestoreRelatedBundleActions);
+    }
+
     if (pPlan->rgCleanActions)
     {
         // Nothing needs to be freed inside clean actions today.
@@ -1276,6 +1289,9 @@ extern "C" HRESULT PlanRelatedBundlesBegin(
             continue;
         }
 
+        pRelatedBundle->defaultRequestedRestore = BOOTSTRAPPER_REQUEST_STATE_NONE;
+        pRelatedBundle->requestedRestore = BOOTSTRAPPER_REQUEST_STATE_NONE;
+        pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_NONE;
         pRelatedBundle->package.defaultRequested = BOOTSTRAPPER_REQUEST_STATE_NONE;
         pRelatedBundle->package.requested = BOOTSTRAPPER_REQUEST_STATE_NONE;
 
@@ -1312,12 +1328,6 @@ extern "C" HRESULT PlanRelatedBundlesBegin(
         hr = UserExperienceOnPlanRelatedBundle(pUserExperience, pRelatedBundle->package.sczId, &pRelatedBundle->package.requested);
         ExitOnRootFailure(hr, "BA aborted plan related bundle.");
 
-        // Log when the BA changed the bundle state so the engine doesn't get blamed for planning the wrong thing.
-        if (pRelatedBundle->package.requested != pRelatedBundle->package.defaultRequested)
-        {
-            LogId(REPORT_STANDARD, MSG_PLANNED_BUNDLE_UX_CHANGED_REQUEST, pRelatedBundle->package.sczId, LoggingRequestStateToString(pRelatedBundle->package.requested), LoggingRequestStateToString(pRelatedBundle->package.defaultRequested));
-        }
-
         // If uninstalling and the dependent related bundle may be executed, ignore its provider key to allow for downgrades with ref-counting.
         if (BOOTSTRAPPER_ACTION_UNINSTALL == pPlan->action && BOOTSTRAPPER_RELATION_DEPENDENT == pRelatedBundle->relationType && BOOTSTRAPPER_REQUEST_STATE_NONE != pRelatedBundle->package.requested)
         {
@@ -1340,6 +1350,7 @@ LExit:
 }
 
 extern "C" HRESULT PlanRelatedBundlesComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_PLAN* pPlan,
     __in BURN_LOGGING* pLog,
@@ -1359,16 +1370,19 @@ extern "C" HRESULT PlanRelatedBundlesComplete(
     ExitOnFailure(hr, "Failed to create dictionary for planned packages.");
 
     BOOL fExecutingAnyPackage = FALSE;
+    BOOL fInstallingAnyPackage = FALSE;
 
     for (DWORD i = 0; i < pPlan->cExecuteActions; ++i)
     {
+        BOOTSTRAPPER_ACTION_STATE packageAction = BOOTSTRAPPER_ACTION_STATE_NONE;
+
         switch (pPlan->rgExecuteActions[i].type)
         {
         case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
-            if (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].relatedBundle.action)
-            {
-                fExecutingAnyPackage = TRUE;
+            packageAction = pPlan->rgExecuteActions[i].relatedBundle.action;
 
+            if (BOOTSTRAPPER_ACTION_STATE_NONE != packageAction)
+            {
                 BURN_PACKAGE* pPackage = &pPlan->rgExecuteActions[i].relatedBundle.pRelatedBundle->package;
                 if (pPackage->cDependencyProviders)
                 {
@@ -1380,21 +1394,24 @@ extern "C" HRESULT PlanRelatedBundlesComplete(
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_EXE_PACKAGE:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].exePackage.action);
+            packageAction = pPlan->rgExecuteActions[i].exePackage.action;
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_MSI_PACKAGE:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].msiPackage.action);
+            packageAction = pPlan->rgExecuteActions[i].msiPackage.action;
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_MSP_TARGET:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].mspTarget.action);
+            packageAction = pPlan->rgExecuteActions[i].mspTarget.action;
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_MSU_PACKAGE:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].msuPackage.action);
+            packageAction = pPlan->rgExecuteActions[i].msuPackage.action;
             break;
         }
+
+        fExecutingAnyPackage |= BOOTSTRAPPER_ACTION_STATE_NONE != packageAction;
+        fInstallingAnyPackage |= BOOTSTRAPPER_ACTION_STATE_INSTALL == packageAction || BOOTSTRAPPER_ACTION_STATE_MINOR_UPGRADE == packageAction;
     }
 
     for (DWORD i = 0; i < pRegistration->relatedBundles.cRelatedBundles; ++i)
@@ -1492,6 +1509,62 @@ extern "C" HRESULT PlanRelatedBundlesComplete(
             hr = DependencyPlanPackageComplete(&pRelatedBundle->package, pPlan);
             ExitOnFailure(hr, "Failed to complete plan dependency actions for related bundle package: %ls", pRelatedBundle->package.sczId);
         }
+
+        if (fInstallingAnyPackage && BOOTSTRAPPER_RELATION_UPGRADE == pRelatedBundle->relationType)
+        {
+            BURN_EXECUTE_ACTION* pAction = NULL;
+
+            pRelatedBundle->defaultRequestedRestore = pRelatedBundle->requestedRestore = BOOTSTRAPPER_REQUEST_STATE_FORCE_PRESENT;
+
+            hr = UserExperienceOnPlanRestoreRelatedBundle(pUserExperience, pRelatedBundle->package.sczId, &pRelatedBundle->requestedRestore);
+            ExitOnRootFailure(hr, "BA aborted plan restore related bundle.");
+
+            switch (pRelatedBundle->requestedRestore)
+            {
+            case BOOTSTRAPPER_REQUEST_STATE_REPAIR:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_REPAIR;
+                break;
+            case BOOTSTRAPPER_REQUEST_STATE_ABSENT: __fallthrough;
+            case BOOTSTRAPPER_REQUEST_STATE_CACHE: __fallthrough;
+            case BOOTSTRAPPER_REQUEST_STATE_FORCE_ABSENT:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_UNINSTALL;
+                break;
+            case BOOTSTRAPPER_REQUEST_STATE_FORCE_PRESENT:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_INSTALL;
+                break;
+            default:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_NONE;
+                break;
+            }
+
+            if (BOOTSTRAPPER_ACTION_STATE_NONE != pRelatedBundle->restore)
+            {
+                hr = AppendRestoreRelatedBundleAction(pPlan, &pAction);
+                ExitOnFailure(hr, "Failed to append restore related bundle action to plan.");
+
+                pAction->type = BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE;
+                pAction->relatedBundle.pRelatedBundle = pRelatedBundle;
+                pAction->relatedBundle.action = pRelatedBundle->restore;
+
+                if (pRelatedBundle->package.Bundle.sczIgnoreDependencies)
+                {
+                    hr = StrAllocString(&pAction->relatedBundle.sczIgnoreDependencies, pRelatedBundle->package.Bundle.sczIgnoreDependencies, 0);
+                    ExitOnFailure(hr, "Failed to allocate the list of dependencies to ignore.");
+                }
+
+                if (pRelatedBundle->package.Bundle.wzAncestors)
+                {
+                    hr = StrAllocString(&pAction->relatedBundle.sczAncestors, pRelatedBundle->package.Bundle.wzAncestors, 0);
+                    ExitOnFailure(hr, "Failed to allocate the list of ancestors.");
+                }
+
+                if (pRelatedBundle->package.Bundle.wzEngineWorkingDirectory)
+                {
+                    hr = StrAllocString(&pAction->relatedBundle.sczEngineWorkingDirectory, pRelatedBundle->package.Bundle.wzEngineWorkingDirectory, 0);
+                    ExitOnFailure(hr, "Failed to allocate the custom working directory.");
+                }
+            }
+        }
     }
 
 LExit:
@@ -2269,6 +2342,23 @@ LExit:
     return hr;
 }
 
+static HRESULT AppendRestoreRelatedBundleAction(
+    __in BURN_PLAN* pPlan,
+    __out BURN_EXECUTE_ACTION** ppExecuteAction
+    )
+{
+    HRESULT hr = S_OK;
+
+    hr = MemEnsureArraySizeForNewItems(reinterpret_cast<LPVOID*>(&pPlan->rgRestoreRelatedBundleActions), pPlan->cRestoreRelatedBundleActions, 1, sizeof(BURN_EXECUTE_ACTION), 5);
+    ExitOnFailure(hr, "Failed to grow plan's array of restore related bundle actions.");
+
+    *ppExecuteAction = pPlan->rgRestoreRelatedBundleActions + pPlan->cRestoreRelatedBundleActions;
+    ++pPlan->cRestoreRelatedBundleActions;
+
+LExit:
+    return hr;
+}
+
 static HRESULT ProcessPayloadGroup(
     __in BURN_PLAN* pPlan,
     __in BURN_PAYLOAD_GROUP* pPayloadGroup
@@ -2725,6 +2815,28 @@ static void ExecuteActionLog(
     }
 }
 
+static void RestoreRelatedBundleActionLog(
+    __in DWORD iAction,
+    __in BURN_EXECUTE_ACTION* pAction
+    )
+{
+    switch (pAction->type)
+    {
+    case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
+        LogStringLine(PlanDumpLevel, "Restore action[%u]: RELATED_BUNDLE package id: %ls, action: %hs, ignore dependencies: %ls", iAction, pAction->relatedBundle.pRelatedBundle->package.sczId, LoggingActionStateToString(pAction->relatedBundle.action), pAction->relatedBundle.sczIgnoreDependencies);
+        break;
+
+    default:
+        AssertSz(FALSE, "Unknown execute action type.");
+        break;
+    }
+
+    if (pAction->fDeleted)
+    {
+        LogStringLine(PlanDumpLevel, "      (deleted action)");
+    }
+}
+
 static void CleanActionLog(
     __in DWORD iAction,
     __in BURN_CLEAN_ACTION* pAction
@@ -2784,6 +2896,11 @@ extern "C" void PlanDump(
         ExecuteActionLog(i, pPlan->rgRollbackActions + i, TRUE);
     }
 
+    for (DWORD i = 0; i < pPlan->cRestoreRelatedBundleActions; ++i)
+    {
+        RestoreRelatedBundleActionLog(i, pPlan->rgRestoreRelatedBundleActions + i);
+    }
+
     for (DWORD i = 0; i < pPlan->cCleanActions; ++i)
     {
         CleanActionLog(i, pPlan->rgCleanActions + i);
diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h
index 0734e39f..3dce8e5d 100644
--- a/src/burn/engine/plan.h
+++ b/src/burn/engine/plan.h
@@ -280,6 +280,9 @@ typedef struct _BURN_PLAN
     BURN_EXECUTE_ACTION* rgRollbackActions;
     DWORD cRollbackActions;
 
+    BURN_EXECUTE_ACTION* rgRestoreRelatedBundleActions;
+    DWORD cRestoreRelatedBundleActions;
+
     BURN_CLEAN_ACTION* rgCleanActions;
     DWORD cCleanActions;
 
@@ -394,6 +397,7 @@ HRESULT PlanRelatedBundlesBegin(
     __in BURN_PLAN* pPlan
     );
 HRESULT PlanRelatedBundlesComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_PLAN* pPlan,
     __in BURN_LOGGING* pLog,
diff --git a/src/burn/engine/pseudobundle.cpp b/src/burn/engine/pseudobundle.cpp
index 94b095c5..082c4487 100644
--- a/src/burn/engine/pseudobundle.cpp
+++ b/src/burn/engine/pseudobundle.cpp
@@ -51,7 +51,6 @@ extern "C" HRESULT PseudoBundleInitializeRelated(
     pPackage->fVital = FALSE;
 
     pPackage->fPermanent = FALSE;
-    pPackage->Bundle.fRepairable = TRUE;
     pPackage->Bundle.fSupportsBurnProtocol = fSupportsBurnProtocol;
 
     hr = StrAllocString(&pPackage->sczId, wzId, 0);
diff --git a/src/burn/engine/registration.h b/src/burn/engine/registration.h
index 0ae61974..64191828 100644
--- a/src/burn/engine/registration.h
+++ b/src/burn/engine/registration.h
@@ -61,6 +61,10 @@ typedef struct _BURN_RELATED_BUNDLE
     BOOL fPlannable;
 
     BURN_PACKAGE package;
+
+    BOOTSTRAPPER_REQUEST_STATE defaultRequestedRestore;
+    BOOTSTRAPPER_REQUEST_STATE requestedRestore;
+    BOOTSTRAPPER_ACTION_STATE restore;
 } BURN_RELATED_BUNDLE;
 
 typedef struct _BURN_RELATED_BUNDLES
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 1439f5f2..59988bef 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -2176,6 +2176,36 @@ LExit:
     return hr;
 }
 
+EXTERN_C BAAPI UserExperienceOnPlanRestoreRelatedBundle(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in_z LPCWSTR wzBundleId,
+    __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONPLANRESTORERELATEDBUNDLE_ARGS args = { };
+    BA_ONPLANRESTORERELATEDBUNDLE_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.wzBundleId = wzBundleId;
+    args.recommendedState = *pRequestedState;
+
+    results.cbSize = sizeof(results);
+    results.requestedState = *pRequestedState;
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE, &args, &results);
+    ExitOnFailure(hr, "BA OnPlanRestoreRelatedBundle failed.");
+
+    if (results.fCancel)
+    {
+        hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT);
+    }
+    *pRequestedState = results.requestedState;
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnPlanRollbackBoundary(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in_z LPCWSTR wzRollbackBoundaryId,
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index e8341120..8106d7f7 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -497,6 +497,11 @@ BAAPI UserExperienceOnPlanRelatedBundle(
     __in_z LPCWSTR wzBundleId,
     __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState
     );
+BAAPI UserExperienceOnPlanRestoreRelatedBundle(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in_z LPCWSTR wzBundleId,
+    __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState
+    );
 BAAPI UserExperienceOnPlanRollbackBoundary(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in_z LPCWSTR wzRollbackBoundaryId,