11 files changed, 132 insertions, 0 deletions

diff --git a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index 659901be..5c6258d0 100644
--- a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -209,6 +209,7 @@ enum BOOTSTRAPPER_APPLICATION_MESSAGE
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE,
 };
 
 enum BOOTSTRAPPER_APPLYCOMPLETE_ACTION
@@ -1208,6 +1209,20 @@ struct BA_ONPLANRELATEDBUNDLE_RESULTS
     BOOTSTRAPPER_REQUEST_STATE requestedState;
 };
 
+struct BA_ONPLANRESTORERELATEDBUNDLE_ARGS
+{
+    DWORD cbSize;
+    LPCWSTR wzBundleId;
+    BOOTSTRAPPER_REQUEST_STATE recommendedState;
+};
+
+struct BA_ONPLANRESTORERELATEDBUNDLE_RESULTS
+{
+    DWORD cbSize;
+    BOOL fCancel;
+    BOOTSTRAPPER_REQUEST_STATE requestedState;
+};
+
 struct BA_ONPLANROLLBACKBOUNDARY_ARGS
 {
     DWORD cbSize;
diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
index f277425e..b08e66c0 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
@@ -271,6 +271,9 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public event EventHandler<SetUpdateCompleteEventArgs> SetUpdateComplete;
 
+        /// <inheritdoc/>
+        public event EventHandler<PlanRestoreRelatedBundleEventArgs> PlanRestoreRelatedBundle;
+
         /// <summary>
         /// Entry point that is called when the bootstrapper application is ready to run.
         /// </summary>
@@ -1321,6 +1324,19 @@ namespace WixToolset.Mba.Core
             }
         }
 
+        /// <summary>
+        /// Called by the engine, raises the <see cref="PlanRestoreRelatedBundle"/> event.
+        /// </summary>
+        /// <param name="args">Additional arguments for this event.</param>
+        protected virtual void OnPlanRestoreRelatedBundle(PlanRestoreRelatedBundleEventArgs args)
+        {
+            EventHandler<PlanRestoreRelatedBundleEventArgs> handler = this.PlanRestoreRelatedBundle;
+            if (null != handler)
+            {
+                handler(this, args);
+            }
+        }
+
         #region IBootstrapperApplication Members
 
         int IBootstrapperApplication.BAProc(int message, IntPtr pvArgs, IntPtr pvResults, IntPtr pvContext)
@@ -2042,6 +2058,16 @@ namespace WixToolset.Mba.Core
             return args.HResult;
         }
 
+        int IBootstrapperApplication.OnPlanRestoreRelatedBundle(string wzBundleId, RequestState recommendedState, ref RequestState pRequestedState, ref bool fCancel)
+        {
+            PlanRestoreRelatedBundleEventArgs args = new PlanRestoreRelatedBundleEventArgs(wzBundleId, recommendedState, pRequestedState, fCancel);
+            this.OnPlanRestoreRelatedBundle(args);
+
+            pRequestedState = args.State;
+            fCancel = args.Cancel;
+            return args.HResult;
+        }
+
         #endregion
     }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
index d4d70651..2e1e1be3 100644
--- a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
+++ b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
@@ -2402,4 +2402,35 @@ namespace WixToolset.Mba.Core
         /// </summary>
         public string NewPackageId { get; private set; }
     }
+
+    /// <summary>
+    /// Event arguments for <see cref="IDefaultBootstrapperApplication.PlanRestoreRelatedBundle"/>
+    /// </summary>
+    [Serializable]
+    public class PlanRestoreRelatedBundleEventArgs : CancellableHResultEventArgs
+    {
+        /// <summary />
+        public PlanRestoreRelatedBundleEventArgs(string bundleId, RequestState recommendedState, RequestState state, bool cancelRecommendation)
+            : base(cancelRecommendation)
+        {
+            this.BundleId = bundleId;
+            this.RecommendedState = recommendedState;
+            this.State = state;
+        }
+
+        /// <summary>
+        /// Gets the identity of the bundle to plan for.
+        /// </summary>
+        public string BundleId { get; private set; }
+
+        /// <summary>
+        /// Gets the recommended requested state for the bundle.
+        /// </summary>
+        public RequestState RecommendedState { get; private set; }
+
+        /// <summary>
+        /// Gets or sets the requested state for the bundle.
+        /// </summary>
+        public RequestState State { get; set; }
+    }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
index 05f96106..4fbe5e18 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
@@ -1136,6 +1136,18 @@ namespace WixToolset.Mba.Core
             [MarshalAs(UnmanagedType.LPWStr)] string wzPreviousPackageId,
             [MarshalAs(UnmanagedType.LPWStr)] string wzNewPackageId
             );
+
+        /// <summary>
+        /// See <see cref="IDefaultBootstrapperApplication.PlanRestoreRelatedBundle"/>.
+        /// </summary>
+        [PreserveSig]
+        [return: MarshalAs(UnmanagedType.I4)]
+        int OnPlanRestoreRelatedBundle(
+            [MarshalAs(UnmanagedType.LPWStr)] string wzBundleId,
+            [MarshalAs(UnmanagedType.U4)] RequestState recommendedState,
+            [MarshalAs(UnmanagedType.U4)] ref RequestState pRequestedState,
+            [MarshalAs(UnmanagedType.Bool)] ref bool fCancel
+            );
     }
 
     /// <summary>
diff --git a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
index ce06408e..c237cb9d 100644
--- a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
@@ -334,6 +334,11 @@ namespace WixToolset.Mba.Core
         event EventHandler<PlanRelatedBundleEventArgs> PlanRelatedBundle;
 
         /// <summary>
+        /// Fired when the engine has begun planning an upgrade related bundle for restoring in case of failure.
+        /// </summary>
+        event EventHandler<PlanRestoreRelatedBundleEventArgs> PlanRestoreRelatedBundle;
+
+        /// <summary>
         /// Fired when the engine is planning a rollback boundary.
         /// </summary>
         event EventHandler<PlanRollbackBoundaryEventArgs> PlanRollbackBoundary;
diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 84359d65..cbef88df 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -88,6 +88,7 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN,
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE,
+    BA_FUNCTIONS_MESSAGE_ONPLANRESTORERELATEDBUNDLE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE,
 
     BA_FUNCTIONS_MESSAGE_ONTHEMELOADED = 1024,
     BA_FUNCTIONS_MESSAGE_WNDPROC,
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctions.h b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
index c6d0924f..e98ebc9f 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -849,6 +849,16 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnPlanRestoreRelatedBundle(
+        __in_z LPCWSTR /*wzBundleId*/,
+        __in BOOTSTRAPPER_REQUEST_STATE /*recommendedState*/,
+        __inout BOOTSTRAPPER_REQUEST_STATE* /*pRequestedState*/,
+        __inout BOOL* /*pfCancel*/
+        )
+    {
+        return S_OK;
+    }
+
 public: // IBAFunctions
     virtual STDMETHODIMP OnPlan(
         )
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
index 5d5ff098..09cc189e 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -159,6 +159,7 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN:
         case BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE:
         case BA_FUNCTIONS_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE:
+        case BA_FUNCTIONS_MESSAGE_ONPLANRESTORERELATEDBUNDLE:
             hr = BalBaseBootstrapperApplicationProc((BOOTSTRAPPER_APPLICATION_MESSAGE)message, pvArgs, pvResults, pvContext);
             break;
         case BA_FUNCTIONS_MESSAGE_ONTHEMELOADED:
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
index e1a36fdf..6a24f24b 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -1047,6 +1047,17 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnPlanRestoreRelatedBundle(
+        __in_z LPCWSTR /*wzBundleId*/,
+        __in BOOTSTRAPPER_REQUEST_STATE /*recommendedState*/,
+        __inout BOOTSTRAPPER_REQUEST_STATE* /*pRequestedState*/,
+        __inout BOOL* pfCancel
+        )
+    {
+        *pfCancel |= CheckCanceled();
+        return S_OK;
+    }
+
 public: //CBalBaseBootstrapperApplication
     virtual STDMETHODIMP Initialize(
         __in const BOOTSTRAPPER_CREATE_ARGS* pCreateArgs
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
index 1ee5258e..d40390e5 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
@@ -720,6 +720,15 @@ static HRESULT BalBaseBAProcOnSetUpdateComplete(
     return pBA->OnSetUpdateComplete(pArgs->hrStatus, pArgs->wzPreviousPackageId, pArgs->wzNewPackageId);
 }
 
+static HRESULT BalBaseBAProcOnPlanRestoreRelatedBundle(
+    __in IBootstrapperApplication* pBA,
+    __in BA_ONPLANRESTORERELATEDBUNDLE_ARGS* pArgs,
+    __inout BA_ONPLANRESTORERELATEDBUNDLE_RESULTS* pResults
+    )
+{
+    return pBA->OnPlanRestoreRelatedBundle(pArgs->wzBundleId, pArgs->recommendedState, &pResults->requestedState, &pResults->fCancel);
+}
+
 /*******************************************************************
 BalBaseBootstrapperApplicationProc - requires pvContext to be of type IBootstrapperApplication.
                                      Provides a default mapping between the new message based BA interface and
@@ -976,6 +985,9 @@ static HRESULT WINAPI BalBaseBootstrapperApplicationProc(
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE:
             hr = BalBaseBAProcOnPlannedCompatiblePackage(pBA, reinterpret_cast<BA_ONPLANNEDCOMPATIBLEPACKAGE_ARGS*>(pvArgs), reinterpret_cast<BA_ONPLANNEDCOMPATIBLEPACKAGE_RESULTS*>(pvResults));
             break;
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE:
+            hr = BalBaseBAProcOnPlanRestoreRelatedBundle(pBA, reinterpret_cast<BA_ONPLANRESTORERELATEDBUNDLE_ARGS*>(pvArgs), reinterpret_cast<BA_ONPLANRESTORERELATEDBUNDLE_RESULTS*>(pvResults));
+            break;
         }
     }
 
diff --git a/src/api/burn/balutil/inc/IBootstrapperApplication.h b/src/api/burn/balutil/inc/IBootstrapperApplication.h
index 640f609d..5932c06e 100644
--- a/src/api/burn/balutil/inc/IBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/IBootstrapperApplication.h
@@ -691,4 +691,12 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         __in_z_opt LPCWSTR wzPreviousPackageId,
         __in_z_opt LPCWSTR wzNewPackageId
         ) = 0;
+
+    // OnPlanRestoreRelatedBundle - called when the engine begins planning an upgrade related bundle for restoring in case of failure.
+    STDMETHOD(OnPlanRestoreRelatedBundle)(
+        __in_z LPCWSTR wzBundleId,
+        __in BOOTSTRAPPER_REQUEST_STATE recommendedState,
+        __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState,
+        __inout BOOL* pfCancel
+        ) = 0;
 };