Add BA events for setting the update bundle.

Fixes #6410

24 files changed, 535 insertions, 37 deletions

diff --git a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index ad920577..2ffcf9d6 100644
--- a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -197,6 +197,8 @@ enum BOOTSTRAPPER_APPLICATION_MESSAGE
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPAYLOADEXTRACTCOMPLETE,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANROLLBACKBOUNDARY,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE,
 };
 
 enum BOOTSTRAPPER_APPLYCOMPLETE_ACTION
@@ -1228,6 +1230,29 @@ struct BA_ONROLLBACKMSITRANSACTIONCOMPLETE_RESULTS
     DWORD cbSize;
 };
 
+struct BA_ONSETUPDATEBEGIN_ARGS
+{
+    DWORD cbSize;
+};
+
+struct BA_ONSETUPDATEBEGIN_RESULTS
+{
+    DWORD cbSize;
+};
+
+struct BA_ONSETUPDATECOMPLETE_ARGS
+{
+    DWORD cbSize;
+    HRESULT hrStatus;
+    LPCWSTR wzPreviousPackageId;
+    LPCWSTR wzNewPackageId;
+};
+
+struct BA_ONSETUPDATECOMPLETE_RESULTS
+{
+    DWORD cbSize;
+};
+
 struct BA_ONSHUTDOWN_ARGS
 {
     DWORD cbSize;
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);
diff --git a/src/api/burn/balutil/balinfo.cpp b/src/api/burn/balutil/balinfo.cpp
index 2746f49e..d9cc9b76 100644
--- a/src/api/burn/balutil/balinfo.cpp
+++ b/src/api/burn/balutil/balinfo.cpp
@@ -166,7 +166,7 @@ LExit:
 
 DAPI_(HRESULT) BalInfoAddRelatedBundleAsPackage(
     __in BAL_INFO_PACKAGES* pPackages,
-    __in LPCWSTR wzId,
+    __in_z LPCWSTR wzId,
     __in BOOTSTRAPPER_RELATION_TYPE relationType,
     __in BOOL /*fPerMachine*/,
     __out_opt BAL_INFO_PACKAGE** ppPackage
@@ -228,6 +228,47 @@ LExit:
 }
 
 
+DAPI_(HRESULT) BalInfoAddUpdateBundleAsPackage(
+    __in BAL_INFO_PACKAGES* pPackages,
+    __in_z LPCWSTR wzId,
+    __in_z LPCWSTR /*wzPreviousId*/,
+    __out_opt BAL_INFO_PACKAGE** ppPackage
+    )
+{
+    HRESULT hr = S_OK;
+    BAL_INFO_PACKAGE* pPackage = NULL;
+
+    // Check to see if the bundle is already in the list of packages.
+    for (DWORD i = 0; i < pPackages->cPackages; ++i)
+    {
+        if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzId, -1, pPackages->rgPackages[i].sczId, -1))
+        {
+            ExitFunction1(hr = HRESULT_FROM_WIN32(ERROR_ALREADY_EXISTS));
+        }
+    }
+
+    // Add the update bundle as a package.
+    hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&pPackages->rgPackages), pPackages->cPackages + 1, sizeof(BAL_INFO_PACKAGE), 2);
+    ExitOnFailure(hr, "Failed to allocate memory for update bundle package information.");
+
+    pPackage = pPackages->rgPackages + pPackages->cPackages;
+    ++pPackages->cPackages;
+
+    hr = StrAllocString(&pPackage->sczId, wzId, 0);
+    ExitOnFailure(hr, "Failed to copy update bundle package id.");
+
+    pPackage->type = BAL_INFO_PACKAGE_TYPE_BUNDLE_UPDATE;
+
+    if (ppPackage)
+    {
+        *ppPackage = pPackage;
+    }
+
+LExit:
+    return hr;
+}
+
+
 DAPI_(HRESULT) BalInfoFindPackageById(
     __in BAL_INFO_PACKAGES* pPackages,
     __in LPCWSTR wzId,
diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 21cace1f..2698a6e3 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -82,6 +82,8 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTCOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPAYLOADEXTRACTCOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS,
     BA_FUNCTIONS_MESSAGE_ONPLANROLLBACKBOUNDARY = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANROLLBACKBOUNDARY,
+    BA_FUNCTIONS_MESSAGE_ONSETUPDATEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN,
+    BA_FUNCTIONS_MESSAGE_ONSETUPDATECOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE,
 
     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 e45db2dc..8d9bddca 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -794,6 +794,20 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnSetUpdateBegin()
+    {
+        return S_OK;
+    }
+
+    virtual STDMETHODIMP OnSetUpdateComplete(
+        __in HRESULT /*hrStatus*/,
+        __in_z_opt LPCWSTR /*wzPreviousPackageId*/,
+        __in_z_opt LPCWSTR /*wzNewPackageId*/
+        )
+    {
+        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 e841c660..1ab0df59 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -153,6 +153,8 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTCOMPLETE:
         case BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS:
         case BA_FUNCTIONS_MESSAGE_ONPLANROLLBACKBOUNDARY:
+        case BA_FUNCTIONS_MESSAGE_ONSETUPDATEBEGIN:
+        case BA_FUNCTIONS_MESSAGE_ONSETUPDATECOMPLETE:
             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 c10b662c..4d043dfe 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -990,6 +990,20 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnSetUpdateBegin()
+    {
+        return S_OK;
+    }
+
+    virtual STDMETHODIMP OnSetUpdateComplete(
+        __in HRESULT /*hrStatus*/,
+        __in_z_opt LPCWSTR /*wzPreviousPackageId*/,
+        __in_z_opt LPCWSTR /*wzNewPackageId*/
+        )
+    {
+        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 f17e1fcb..59bfc1f8 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
@@ -666,6 +666,24 @@ static HRESULT BalBaseBAProcOnCachePayloadExtractComplete(
     return pBA->OnCachePayloadExtractComplete(pArgs->wzContainerId, pArgs->wzPayloadId, pArgs->hrStatus);
 }
 
+static HRESULT BalBaseBAProcOnSetUpdateBegin(
+    __in IBootstrapperApplication* pBA,
+    __in BA_ONSETUPDATEBEGIN_ARGS* /*pArgs*/,
+    __inout BA_ONSETUPDATEBEGIN_RESULTS* /*pResults*/
+    )
+{
+    return pBA->OnSetUpdateBegin();
+}
+
+static HRESULT BalBaseBAProcOnSetUpdateComplete(
+    __in IBootstrapperApplication* pBA,
+    __in BA_ONSETUPDATECOMPLETE_ARGS* pArgs,
+    __inout BA_ONSETUPDATECOMPLETE_RESULTS* /*pResults*/
+    )
+{
+    return pBA->OnSetUpdateComplete(pArgs->hrStatus, pArgs->wzPreviousPackageId, pArgs->wzNewPackageId);
+}
+
 /*******************************************************************
 BalBaseBootstrapperApplicationProc - requires pvContext to be of type IBootstrapperApplication.
                                      Provides a default mapping between the new message based BA interface and
@@ -904,6 +922,12 @@ static HRESULT WINAPI BalBaseBootstrapperApplicationProc(
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANROLLBACKBOUNDARY:
             hr = BalBaseBAProcOnPlanRollbackBoundary(pBA, reinterpret_cast<BA_ONPLANROLLBACKBOUNDARY_ARGS*>(pvArgs), reinterpret_cast<BA_ONPLANROLLBACKBOUNDARY_RESULTS*>(pvResults));
             break;
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN:
+            hr = BalBaseBAProcOnSetUpdateBegin(pBA, reinterpret_cast<BA_ONSETUPDATEBEGIN_ARGS*>(pvArgs), reinterpret_cast<BA_ONSETUPDATEBEGIN_RESULTS*>(pvResults));
+            break;
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE:
+            hr = BalBaseBAProcOnSetUpdateComplete(pBA, reinterpret_cast<BA_ONSETUPDATECOMPLETE_ARGS*>(pvArgs), reinterpret_cast<BA_ONSETUPDATECOMPLETE_RESULTS*>(pvResults));
+            break;
         }
     }
 
diff --git a/src/api/burn/balutil/inc/IBootstrapperApplication.h b/src/api/burn/balutil/inc/IBootstrapperApplication.h
index 1fc99988..be9b7b6e 100644
--- a/src/api/burn/balutil/inc/IBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/IBootstrapperApplication.h
@@ -658,4 +658,12 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         __in_z_opt LPCWSTR wzPayloadId,
         __in HRESULT hrStatus
         ) = 0;
+
+    STDMETHOD(OnSetUpdateBegin)() = 0;
+
+    STDMETHOD(OnSetUpdateComplete)(
+        __in HRESULT hrStatus,
+        __in_z_opt LPCWSTR wzPreviousPackageId,
+        __in_z_opt LPCWSTR wzNewPackageId
+        ) = 0;
 };
diff --git a/src/api/burn/balutil/inc/balinfo.h b/src/api/burn/balutil/inc/balinfo.h
index 769becb2..8f61685f 100644
--- a/src/api/burn/balutil/inc/balinfo.h
+++ b/src/api/burn/balutil/inc/balinfo.h
@@ -16,6 +16,7 @@ typedef enum BAL_INFO_PACKAGE_TYPE
     BAL_INFO_PACKAGE_TYPE_BUNDLE_UPGRADE,
     BAL_INFO_PACKAGE_TYPE_BUNDLE_ADDON,
     BAL_INFO_PACKAGE_TYPE_BUNDLE_PATCH,
+    BAL_INFO_PACKAGE_TYPE_BUNDLE_UPDATE,
 } BAL_INFO_PACKAGE_TYPE;
 
 typedef enum _BAL_INFO_RESTART
@@ -125,7 +126,7 @@ DAPI_(HRESULT) BalInfoParseFromXml(
  ********************************************************************/
 DAPI_(HRESULT) BalInfoAddRelatedBundleAsPackage(
     __in BAL_INFO_PACKAGES* pPackages,
-    __in LPCWSTR wzId,
+    __in_z LPCWSTR wzId,
     __in BOOTSTRAPPER_RELATION_TYPE relationType,
     __in BOOL fPerMachine,
     __out_opt BAL_INFO_PACKAGE** ppPackage
@@ -133,6 +134,18 @@ DAPI_(HRESULT) BalInfoAddRelatedBundleAsPackage(
 
 
 /*******************************************************************
+ BalInfoAddUpdateBundleAsPackage - adds an update bundle as a package.
+
+ ********************************************************************/
+DAPI_(HRESULT) BalInfoAddUpdateBundleAsPackage(
+    __in BAL_INFO_PACKAGES* pPackages,
+    __in_z LPCWSTR wzId,
+    __in_z LPCWSTR wzPreviousId,
+    __out_opt BAL_INFO_PACKAGE** ppPackage
+    );
+
+
+/*******************************************************************
  BalInfoFindPackageById - finds a package by its id.
 
  ********************************************************************/
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index bbd0ff96..812c7261 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -1193,7 +1193,7 @@ extern "C" HRESULT CoreCreateUpdateBundleCommandLine(
 {
     HRESULT hr = S_OK;
 
-    hr = CoreRecreateCommandLine(psczCommandLine, BOOTSTRAPPER_ACTION_INSTALL, pInternalCommand, pCommand, BOOTSTRAPPER_RELATION_NONE, FALSE);
+    hr = CoreRecreateCommandLine(psczCommandLine, BOOTSTRAPPER_ACTION_INSTALL, pInternalCommand, pCommand, BOOTSTRAPPER_RELATION_UPDATE, FALSE);
     ExitOnFailure(hr, "Failed to recreate update bundle command-line.");
 
 LExit:
diff --git a/src/burn/engine/externalengine.cpp b/src/burn/engine/externalengine.cpp
index da84c83d..abe9b8bc 100644
--- a/src/burn/engine/externalengine.cpp
+++ b/src/burn/engine/externalengine.cpp
@@ -269,64 +269,82 @@ HRESULT ExternalEngineSetUpdate(
     )
 {
     HRESULT hr = S_OK;
+    BOOL fLeaveCriticalSection = FALSE;
     LPWSTR sczFilePath = NULL;
     LPWSTR sczCommandline = NULL;
+    LPWSTR sczPreviousId = NULL;
+    LPCWSTR wzNewId = NULL;
     UUID guid = { };
     WCHAR wzGuid[39];
     RPC_STATUS rs = RPC_S_OK;
+    BOOL fRemove = (!wzLocalSource || !*wzLocalSource) && (!wzDownloadSource || !*wzDownloadSource);
+
+    UserExperienceOnSetUpdateBegin(&pEngineState->userExperience);
 
     ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
+    fLeaveCriticalSection = TRUE;
     hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
     ExitOnFailure(hr, "Engine is active, cannot change engine state.");
 
-    if ((!wzLocalSource || !*wzLocalSource) && (!wzDownloadSource || !*wzDownloadSource))
-    {
-        UpdateUninitialize(&pEngineState->update);
-    }
-    else if (BOOTSTRAPPER_UPDATE_HASH_TYPE_NONE == hashType && (0 != cbHash || rgbHash))
+    if (!fRemove)
     {
-        hr = E_INVALIDARG;
+        if (BOOTSTRAPPER_UPDATE_HASH_TYPE_NONE == hashType && (0 != cbHash || rgbHash))
+        {
+            ExitFunction1(hr = E_INVALIDARG);
+        }
+        else if (BOOTSTRAPPER_UPDATE_HASH_TYPE_SHA512 == hashType && (SHA512_HASH_LEN != cbHash || !rgbHash))
+        {
+            ExitFunction1(hr = E_INVALIDARG);
+        }
     }
-    else if (BOOTSTRAPPER_UPDATE_HASH_TYPE_SHA512 == hashType && (SHA512_HASH_LEN != cbHash || !rgbHash))
+
+    sczPreviousId = pEngineState->update.package.sczId;
+    pEngineState->update.package.sczId = NULL;
+    UpdateUninitialize(&pEngineState->update);
+
+    if (fRemove)
     {
-        hr = E_INVALIDARG;
+        ExitFunction();
     }
-    else
-    {
-        UpdateUninitialize(&pEngineState->update);
 
-        hr = CoreCreateUpdateBundleCommandLine(&sczCommandline, &pEngineState->internalCommand, &pEngineState->command);
-        ExitOnFailure(hr, "Failed to create command-line for update bundle.");
+    hr = CoreCreateUpdateBundleCommandLine(&sczCommandline, &pEngineState->internalCommand, &pEngineState->command);
+    ExitOnFailure(hr, "Failed to create command-line for update bundle.");
 
-        // Bundles would fail to use the downloaded update bundle, as the running bundle would be one of the search paths.
-        // Here I am generating a random guid, but in the future it would be nice if the feed would provide the ID of the update.
-        rs = ::UuidCreate(&guid);
-        hr = HRESULT_FROM_RPC(rs);
-        ExitOnFailure(hr, "Failed to create bundle update guid.");
+    // Bundles would fail to use the downloaded update bundle, as the running bundle would be one of the search paths.
+    // Here I am generating a random guid, but in the future it would be nice if the feed would provide the ID of the update.
+    rs = ::UuidCreate(&guid);
+    hr = HRESULT_FROM_RPC(rs);
+    ExitOnFailure(hr, "Failed to create bundle update guid.");
 
-        if (!::StringFromGUID2(guid, wzGuid, countof(wzGuid)))
-        {
-            hr = E_OUTOFMEMORY;
-            ExitOnRootFailure(hr, "Failed to convert bundle update guid into string.");
-        }
+    if (!::StringFromGUID2(guid, wzGuid, countof(wzGuid)))
+    {
+        hr = E_OUTOFMEMORY;
+        ExitOnRootFailure(hr, "Failed to convert bundle update guid into string.");
+    }
 
-        hr = StrAllocFormatted(&sczFilePath, L"%ls\\%ls", wzGuid, pEngineState->registration.sczExecutableName);
-        ExitOnFailure(hr, "Failed to build bundle update file path.");
+    hr = StrAllocFormatted(&sczFilePath, L"%ls\\%ls", wzGuid, pEngineState->registration.sczExecutableName);
+    ExitOnFailure(hr, "Failed to build bundle update file path.");
 
-        if (!wzLocalSource || !*wzLocalSource)
-        {
-            wzLocalSource = sczFilePath;
-        }
+    if (!wzLocalSource || !*wzLocalSource)
+    {
+        wzLocalSource = sczFilePath;
+    }
 
-        hr = PseudoBundleInitialize(FILEMAKEVERSION(rmj, rmm, rup, rpr), &pEngineState->update.package, FALSE, pEngineState->registration.sczId, BOOTSTRAPPER_RELATION_UPDATE, BOOTSTRAPPER_PACKAGE_STATE_ABSENT, FALSE, sczFilePath, wzLocalSource, wzDownloadSource, qwSize, TRUE, sczCommandline, NULL, NULL, NULL, rgbHash, cbHash);
-        ExitOnFailure(hr, "Failed to set update bundle.");
+    hr = PseudoBundleInitializeUpdateBundle(&pEngineState->update.package, wzGuid, pEngineState->registration.sczId, sczFilePath, wzLocalSource, wzDownloadSource, qwSize, sczCommandline, rgbHash, cbHash);
+    ExitOnFailure(hr, "Failed to set update bundle.");
 
-        pEngineState->update.fUpdateAvailable = TRUE;
-    }
+    pEngineState->update.fUpdateAvailable = TRUE;
+    wzNewId = wzGuid;
 
 LExit:
-    ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
+    if (fLeaveCriticalSection)
+    {
+        ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
+    }
+
+    UserExperienceOnSetUpdateComplete(&pEngineState->userExperience, hr, sczPreviousId, wzNewId);
 
+    ReleaseStr(sczPreviousId);
     ReleaseStr(sczCommandline);
     ReleaseStr(sczFilePath);
 
diff --git a/src/burn/engine/pseudobundle.cpp b/src/burn/engine/pseudobundle.cpp
index 52b7bd8a..91c6c14f 100644
--- a/src/burn/engine/pseudobundle.cpp
+++ b/src/burn/engine/pseudobundle.cpp
@@ -240,3 +240,80 @@ LExit:
     ReleaseStr(sczArguments);
     return hr;
 }
+
+extern "C" HRESULT PseudoBundleInitializeUpdateBundle(
+    __in BURN_PACKAGE* pPackage,
+    __in_z LPCWSTR wzId,
+    __in_z LPCWSTR wzCacheId,
+    __in_z LPCWSTR wzFilePath,
+    __in_z LPCWSTR wzLocalSource,
+    __in_z_opt LPCWSTR wzDownloadSource,
+    __in DWORD64 qwSize,
+    __in_z LPCWSTR wzInstallArguments,
+    __in_opt const BYTE* pbHash,
+    __in const DWORD cbHash
+    )
+{
+    HRESULT hr = S_OK;
+    BURN_PAYLOAD* pPayload = NULL;
+
+    // Initialize the single payload, and fill out all the necessary fields
+    pPackage->payloads.rgItems = (BURN_PAYLOAD_GROUP_ITEM*)MemAlloc(sizeof(BURN_PAYLOAD_GROUP_ITEM), TRUE);
+    ExitOnNull(pPackage->payloads.rgItems, hr, E_OUTOFMEMORY, "Failed to allocate space for burn payload group inside of update bundle struct");
+    pPackage->payloads.cItems = 1;
+
+    pPayload = (BURN_PAYLOAD*)MemAlloc(sizeof(BURN_PAYLOAD), TRUE);
+    ExitOnNull(pPayload, hr, E_OUTOFMEMORY, "Failed to allocate space for burn payload inside of update bundle struct");
+    pPackage->payloads.rgItems[0].pPayload = pPayload;
+    pPayload->packaging = BURN_PAYLOAD_PACKAGING_EXTERNAL;
+    pPayload->qwFileSize = qwSize;
+    pPayload->verification = BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE;
+
+    hr = StrAllocString(&pPayload->sczKey, wzId, 0);
+    ExitOnFailure(hr, "Failed to copy key for pseudo bundle payload.");
+
+    hr = StrAllocString(&pPayload->sczFilePath, wzFilePath, 0);
+    ExitOnFailure(hr, "Failed to copy filename for pseudo bundle.");
+
+    hr = StrAllocString(&pPayload->sczSourcePath, wzLocalSource, 0);
+    ExitOnFailure(hr, "Failed to copy local source path for pseudo bundle.");
+
+    if (wzDownloadSource && *wzDownloadSource)
+    {
+        hr = StrAllocString(&pPayload->downloadSource.sczUrl, wzDownloadSource, 0);
+        ExitOnFailure(hr, "Failed to copy download source for pseudo bundle.");
+    }
+
+    if (pbHash)
+    {
+        pPayload->pbHash = static_cast<BYTE*>(MemAlloc(cbHash, FALSE));
+        ExitOnNull(pPayload->pbHash, hr, E_OUTOFMEMORY, "Failed to allocate memory for update bundle payload hash.");
+
+        pPayload->cbHash = cbHash;
+        memcpy_s(pPayload->pbHash, pPayload->cbHash, pbHash, cbHash);
+    }
+
+    pPackage->Exe.fPseudoBundle = TRUE;
+
+    pPackage->type = BURN_PACKAGE_TYPE_EXE;
+    pPackage->currentState = BOOTSTRAPPER_PACKAGE_STATE_ABSENT;
+    pPackage->qwInstallSize = qwSize;
+    pPackage->qwSize = qwSize;
+    pPackage->fVital = TRUE;
+
+    hr = StrAllocString(&pPackage->sczId, wzId, 0);
+    ExitOnFailure(hr, "Failed to copy id for update bundle.");
+
+    hr = StrAllocString(&pPackage->sczCacheId, wzCacheId, 0);
+    ExitOnFailure(hr, "Failed to copy cache id for update bundle.");
+
+    hr = StrAllocString(&pPackage->Exe.sczInstallArguments, wzInstallArguments, 0);
+    ExitOnFailure(hr, "Failed to copy install arguments for update bundle package");
+
+    // Assume the update bundle has the same engine version as this one.
+    pPackage->Exe.protocol = BURN_EXE_PROTOCOL_TYPE_BURN;
+    pPackage->Exe.fSupportsAncestors = TRUE;
+
+LExit:
+    return hr;
+}
\ No newline at end of file
diff --git a/src/burn/engine/pseudobundle.h b/src/burn/engine/pseudobundle.h
index aa26d29d..0fd4cbdb 100644
--- a/src/burn/engine/pseudobundle.h
+++ b/src/burn/engine/pseudobundle.h
@@ -32,6 +32,18 @@ HRESULT PseudoBundleInitializePassthrough(
     __in BOOTSTRAPPER_COMMAND* pCommand,
     __in BURN_PACKAGE* pPackage
     );
+HRESULT PseudoBundleInitializeUpdateBundle(
+    __in BURN_PACKAGE* pPackage,
+    __in_z LPCWSTR wzId,
+    __in_z LPCWSTR wzCacheId,
+    __in_z LPCWSTR wzFilePath,
+    __in_z LPCWSTR wzLocalSource,
+    __in_z_opt LPCWSTR wzDownloadSource,
+    __in DWORD64 qwSize,
+    __in_z LPCWSTR wzInstallArguments,
+    __in_opt const BYTE* pbHash,
+    __in const DWORD cbHash
+    );
 
 #if defined(__cplusplus)
 }
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 2bd6ecaf..a6d670ea 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -2241,6 +2241,50 @@ LExit:
     return hr;
 }
 
+EXTERN_C BAAPI UserExperienceOnSetUpdateBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONSETUPDATEBEGIN_ARGS args = { };
+    BA_ONSETUPDATEBEGIN_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+
+    results.cbSize = sizeof(results);
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN, &args, &results);
+    ExitOnFailure(hr, "BA OnSetUpdateBegin failed.");
+
+LExit:
+    return hr;
+}
+
+EXTERN_C BAAPI UserExperienceOnSetUpdateComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus,
+    __in_z_opt LPCWSTR wzPreviousPackageId,
+    __in_z_opt LPCWSTR wzNewPackageId
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONSETUPDATECOMPLETE_ARGS args = { };
+    BA_ONSETUPDATECOMPLETE_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.hrStatus = hrStatus;
+    args.wzPreviousPackageId = wzPreviousPackageId;
+    args.wzNewPackageId = wzNewPackageId;
+
+    results.cbSize = sizeof(results);
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE, &args, &results);
+    ExitOnFailure(hr, "BA OnSetUpdateComplete failed.");
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnShutdown(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __inout BOOTSTRAPPER_SHUTDOWN_ACTION* pAction
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index 2493569b..f7ac962c 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -505,6 +505,15 @@ BAAPI UserExperienceOnRollbackMsiTransactionComplete(
     __in LPCWSTR wzTransactionId,
     __in HRESULT hrStatus
     );
+BAAPI UserExperienceOnSetUpdateBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    );
+BAAPI UserExperienceOnSetUpdateComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus,
+    __in_z_opt LPCWSTR wzPreviousPackageId,
+    __in_z_opt LPCWSTR wzNewPackageId
+    );
 BAAPI UserExperienceOnShutdown(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __inout BOOTSTRAPPER_SHUTDOWN_ACTION* pAction
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index abf5af83..ec3c268a 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -1169,6 +1169,23 @@ public: // IBootstrapperApplication
         return hr;
     }
 
+    virtual STDMETHODIMP OnSetUpdateComplete(
+        __in HRESULT hrStatus,
+        __in_z_opt LPCWSTR wzPreviousPackageId,
+        __in_z_opt LPCWSTR wzNewPackageId
+        )
+    {
+        BAL_INFO_PACKAGE* pPackage = NULL;
+
+        if (SUCCEEDED(hrStatus) && wzNewPackageId &&
+            SUCCEEDED(BalInfoAddUpdateBundleAsPackage(&m_Bundle.packages, wzNewPackageId, wzPreviousPackageId, &pPackage)))
+        {
+            InitializePackageInfoForPackage(pPackage);
+        }
+
+        return S_OK;
+    }
+
     virtual STDMETHODIMP_(void) BAProcFallback(
         __in BOOTSTRAPPER_APPLICATION_MESSAGE message,
         __in const LPVOID pvArgs,
@@ -1400,6 +1417,12 @@ public: // IBootstrapperApplication
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANROLLBACKBOUNDARY:
             OnPlanRollbackBoundaryFallback(reinterpret_cast<BA_ONPLANROLLBACKBOUNDARY_ARGS*>(pvArgs), reinterpret_cast<BA_ONPLANROLLBACKBOUNDARY_RESULTS*>(pvResults));
             break;
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN:
+            OnSetUpdateBeginFallback(reinterpret_cast<BA_ONSETUPDATEBEGIN_ARGS*>(pvArgs), reinterpret_cast<BA_ONSETUPDATEBEGIN_RESULTS*>(pvResults));
+            break;
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE:
+            OnSetUpdateCompleteFallback(reinterpret_cast<BA_ONSETUPDATECOMPLETE_ARGS*>(pvArgs), reinterpret_cast<BA_ONSETUPDATECOMPLETE_RESULTS*>(pvResults));
+            break;
         default:
 #ifdef DEBUG
             BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "WIXSTDBA: Forwarding unknown BA message: %d", message);
@@ -2003,6 +2026,22 @@ private: // privates
         BalLogId(BOOTSTRAPPER_LOG_LEVEL_STANDARD, MSG_WIXSTDBA_PLANNED_ROLLBACK_BOUNDARY, m_hModule, pArgs->wzRollbackBoundaryId, LoggingBoolToString(fTransaction), LoggingBoolToString(pResults->fTransaction));
     }
 
+    void OnSetUpdateBeginFallback(
+        __in BA_ONSETUPDATEBEGIN_ARGS* pArgs,
+        __inout BA_ONSETUPDATEBEGIN_RESULTS* pResults
+        )
+    {
+        m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONSETUPDATEBEGIN, pArgs, pResults, m_pvBAFunctionsProcContext);
+    }
+
+    void OnSetUpdateCompleteFallback(
+        __in BA_ONSETUPDATECOMPLETE_ARGS* pArgs,
+        __inout BA_ONSETUPDATECOMPLETE_RESULTS* pResults
+        )
+    {
+        m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONSETUPDATECOMPLETE, pArgs, pResults, m_pvBAFunctionsProcContext);
+    }
+
 
 public: //CBalBaseBootstrapperApplication
     virtual STDMETHODIMP Initialize(
diff --git a/src/test/burn/WixToolset.WixBA/InstallationViewModel.cs b/src/test/burn/WixToolset.WixBA/InstallationViewModel.cs
index 14ea561f..1846d51b 100644
--- a/src/test/burn/WixToolset.WixBA/InstallationViewModel.cs
+++ b/src/test/burn/WixToolset.WixBA/InstallationViewModel.cs
@@ -85,6 +85,7 @@ namespace WixToolset.WixBA
             WixBA.Model.Bootstrapper.ExecutePackageComplete += this.ExecutePackageComplete;
             WixBA.Model.Bootstrapper.Error += this.ExecuteError;
             WixBA.Model.Bootstrapper.ApplyComplete += this.ApplyComplete;
+            WixBA.Model.Bootstrapper.SetUpdateComplete += this.SetUpdateComplete;
         }
 
         void RootPropertyChanged(object sender, PropertyChangedEventArgs e)
@@ -409,6 +410,14 @@ namespace WixToolset.WixBA
             }
         }
 
+        private void SetUpdateComplete(object sender, SetUpdateCompleteEventArgs e)
+        {
+            if (!String.IsNullOrEmpty(e.NewPackageId) && !WixBA.Model.BAManifest.Bundle.Packages.ContainsKey(e.NewPackageId))
+            {
+                WixBA.Model.BAManifest.Bundle.AddUpdateBundleAsPackage(e);
+            }
+        }
+
         private void DetectComplete(object sender, DetectCompleteEventArgs e)
         {
             // Parse the command line string before any planning.