Add BA events for setting the update bundle.

Fixes #6410

7 files changed, 147 insertions, 0 deletions

diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
index 0520463f..a78bf43f 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
@@ -253,6 +253,12 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public event EventHandler<CachePayloadExtractCompleteEventArgs> CachePayloadExtractComplete;
 
+        /// <inheritdoc/>
+        public event EventHandler<SetUpdateBeginEventArgs> SetUpdateBegin;
+
+        /// <inheritdoc/>
+        public event EventHandler<SetUpdateCompleteEventArgs> SetUpdateComplete;
+
         /// <summary>
         /// Entry point that is called when the bootstrapper application is ready to run.
         /// </summary>
@@ -1225,6 +1231,32 @@ namespace WixToolset.Mba.Core
             }
         }
 
+        /// <summary>
+        /// Called by the engine, raises the <see cref="SetUpdateBegin"/> event.
+        /// </summary>
+        /// <param name="args"></param>
+        protected virtual void OnSetUpdateBegin(SetUpdateBeginEventArgs args)
+        {
+            EventHandler<SetUpdateBeginEventArgs> handler = this.SetUpdateBegin;
+            if (null != handler)
+            {
+                handler(this, args);
+            }
+        }
+
+        /// <summary>
+        /// Called by the engine, raises the <see cref="SetUpdateComplete"/> event.
+        /// </summary>
+        /// <param name="args"></param>
+        protected virtual void OnSetUpdateComplete(SetUpdateCompleteEventArgs args)
+        {
+            EventHandler<SetUpdateCompleteEventArgs> handler = this.SetUpdateComplete;
+            if (null != handler)
+            {
+                handler(this, args);
+            }
+        }
+
         #region IBootstrapperApplication Members
 
         int IBootstrapperApplication.BAProc(int message, IntPtr pvArgs, IntPtr pvResults, IntPtr pvContext)
@@ -1895,6 +1927,22 @@ namespace WixToolset.Mba.Core
             return args.HResult;
         }
 
+        int IBootstrapperApplication.OnSetUpdateBegin()
+        {
+            SetUpdateBeginEventArgs args = new SetUpdateBeginEventArgs();
+            this.OnSetUpdateBegin(args);
+
+            return args.HResult;
+        }
+
+        int IBootstrapperApplication.OnSetUpdateComplete(int hrStatus, string wzPreviousPackageId, string wzNewPackageId)
+        {
+            SetUpdateCompleteEventArgs args = new SetUpdateCompleteEventArgs(hrStatus, wzPreviousPackageId, wzNewPackageId);
+            this.OnSetUpdateComplete(args);
+
+            return args.HResult;
+        }
+
         #endregion
     }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/BundleInfo.cs b/src/api/burn/WixToolset.Mba.Core/BundleInfo.cs
index 4a533bf9..ee751ebf 100644
--- a/src/api/burn/WixToolset.Mba.Core/BundleInfo.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BundleInfo.cs
@@ -41,6 +41,14 @@ namespace WixToolset.Mba.Core
             return package;
         }
 
+        /// <inheritdoc/>
+        public IPackageInfo AddUpdateBundleAsPackage(SetUpdateCompleteEventArgs e)
+        {
+            var package = PackageInfo.GetUpdateBundleAsPackage(e.NewPackageId);
+            this.Packages.Add(package.Id, package);
+            return package;
+        }
+
         /// <summary>
         /// Parses BA manifest from the given stream.
         /// </summary>
diff --git a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
index 556db821..55c9e74c 100644
--- a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
+++ b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
@@ -2230,4 +2230,41 @@ namespace WixToolset.Mba.Core
         /// </summary>
         public string PayloadId { get; private set; }
     }
+
+    /// <summary>
+    /// EventArgs for <see cref="IDefaultBootstrapperApplication.SetUpdateBegin"/>.
+    /// </summary>
+    [Serializable]
+    public class SetUpdateBeginEventArgs : HResultEventArgs
+    {
+        /// <summary />
+        public SetUpdateBeginEventArgs()
+        {
+        }
+    }
+
+    /// <summary>
+    /// Event arguments for <see cref="IDefaultBootstrapperApplication.SetUpdateComplete"/>
+    /// </summary>
+    [Serializable]
+    public class SetUpdateCompleteEventArgs : StatusEventArgs
+    {
+        /// <summary />
+        public SetUpdateCompleteEventArgs(int hrStatus, string previousPackageId, string newPackageId)
+            : base(hrStatus)
+        {
+            this.PreviousPackageId = previousPackageId;
+            this.NewPackageId = newPackageId;
+        }
+
+        /// <summary>
+        /// Gets the identifier of the update package that was removed.
+        /// </summary>
+        public string PreviousPackageId { get; private set; }
+
+        /// <summary>
+        /// Gets the identifier of the update package that was added.
+        /// </summary>
+        public string NewPackageId { get; private set; }
+    }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
index 259c407f..f0c0b7ec 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
@@ -1097,6 +1097,24 @@ namespace WixToolset.Mba.Core
             [MarshalAs(UnmanagedType.LPWStr)] string wzPayloadId,
             int hrStatus
             );
+
+        /// <summary>
+        /// See <see cref="IDefaultBootstrapperApplication.SetUpdateBegin"/>.
+        /// </summary>
+        [PreserveSig]
+        [return: MarshalAs(UnmanagedType.I4)]
+        int OnSetUpdateBegin();
+
+        /// <summary>
+        /// See <see cref="IDefaultBootstrapperApplication.SetUpdateComplete"/>.
+        /// </summary>
+        [PreserveSig]
+        [return: MarshalAs(UnmanagedType.I4)]
+        int OnSetUpdateComplete(
+            int hrStatus,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzPreviousPackageId,
+            [MarshalAs(UnmanagedType.LPWStr)] string wzNewPackageId
+            );
     }
 
     /// <summary>
diff --git a/src/api/burn/WixToolset.Mba.Core/IBundleInfo.cs b/src/api/burn/WixToolset.Mba.Core/IBundleInfo.cs
index 35decc88..3227b72d 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBundleInfo.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBundleInfo.cs
@@ -40,5 +40,12 @@ namespace WixToolset.Mba.Core
         /// <param name="e"></param>
         /// <returns>The created <see cref="IPackageInfo"/>.</returns>
         IPackageInfo AddRelatedBundleAsPackage(DetectRelatedBundleEventArgs e);
+
+        /// <summary>
+        /// Adds an update bundle as a package.
+        /// </summary>
+        /// <param name="e"></param>
+        /// <returns>The created <see cref="IPackageInfo"/>.</returns>
+        IPackageInfo AddUpdateBundleAsPackage(SetUpdateCompleteEventArgs e);
     }
 }
\ No newline at end of file
diff --git a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
index 20ce9f88..e809a965 100644
--- a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
@@ -344,6 +344,16 @@ namespace WixToolset.Mba.Core
         event EventHandler<RollbackMsiTransactionCompleteEventArgs> RollbackMsiTransactionComplete;
 
         /// <summary>
+        /// Fired when the engine has begun to setup the update package.
+        /// </summary>
+        event EventHandler<SetUpdateBeginEventArgs> SetUpdateBegin;
+
+        /// <summary>
+        /// Fired when the engine has completed setting up the update package.
+        /// </summary>
+        event EventHandler<SetUpdateCompleteEventArgs> SetUpdateComplete;
+
+        /// <summary>
         /// Fired when the engine is shutting down the bootstrapper application.
         /// </summary>
         event EventHandler<ShutdownEventArgs> Shutdown;
diff --git a/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs b/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs
index 567a7cdd..3681a497 100644
--- a/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs
+++ b/src/api/burn/WixToolset.Mba.Core/PackageInfo.cs
@@ -51,6 +51,11 @@ namespace WixToolset.Mba.Core
         /// 
         /// </summary>
         PatchBundle,
+
+        /// <summary>
+        /// 
+        /// </summary>
+        UpdateBundle,
     }
 
     /// <summary>
@@ -269,6 +274,20 @@ namespace WixToolset.Mba.Core
             return package;
         }
 
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="id"></param>
+        /// <returns></returns>
+        public static IPackageInfo GetUpdateBundleAsPackage(string id)
+        {
+            PackageInfo package = new PackageInfo();
+            package.Id = id;
+            package.Type = PackageType.UpdateBundle;
+
+            return package;
+        }
+
         internal static void ParseBalPackageInfoFromXml(XPathNavigator root, XmlNamespaceManager namespaceManager, Dictionary<string, IPackageInfo> packagesById)
         {
             XPathNodeIterator nodes = root.Select("/p:BootstrapperApplicationData/p:WixBalPackageInfo", namespaceManager);