Add BA events for setting the update bundle.

Fixes #6410

1 files changed, 39 insertions, 0 deletions

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(