Add BA events for setting the update bundle.

Fixes #6410

8 files changed, 120 insertions, 2 deletions

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.
 
  ********************************************************************/