4 files changed, 74 insertions, 0 deletions

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;