Add BA events for setting the update bundle.

Fixes #6410

16 files changed, 292 insertions, 2 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.
 
  ********************************************************************/