From 6c740c1070475a8a894393a186cdcdd8a8a81112 Mon Sep 17 00:00:00 2001
From: Rob Mensching <rob@firegiant.com>
Date: Wed, 14 Apr 2021 12:18:39 -0700
Subject: Allow BA to update feed source

Fixes wixtoolset/issues#5568
---
 src/WixToolset.Mba.Core/Engine.cs              |  6 ++++++
 src/WixToolset.Mba.Core/IBootstrapperEngine.cs |  8 ++++++++
 src/WixToolset.Mba.Core/IEngine.cs             |  6 ++++++
 src/balutil/BalBootstrapperEngine.cpp          | 15 +++++++++++++++
 src/balutil/inc/IBootstrapperEngine.h          |  4 ++++
 5 files changed, 39 insertions(+)

diff --git a/src/WixToolset.Mba.Core/Engine.cs b/src/WixToolset.Mba.Core/Engine.cs
index d5c43a53..e07ecd8b 100644
--- a/src/WixToolset.Mba.Core/Engine.cs
+++ b/src/WixToolset.Mba.Core/Engine.cs
@@ -246,6 +246,12 @@ namespace WixToolset.Mba.Core
             this.engine.SetUpdate(localSource, downloadSource, size, hashType, hash, null == hash ? 0 : hash.Length);
         }
 
+        /// <inheritdoc/>
+        public void SetUpdateSource(string url)
+        {
+            this.engine.SetUpdateSource(url);
+        }
+
         /// <inheritdoc/>
         public void SetLocalSource(string packageOrContainerId, string payloadId, string path)
         {
diff --git a/src/WixToolset.Mba.Core/IBootstrapperEngine.cs b/src/WixToolset.Mba.Core/IBootstrapperEngine.cs
index af3101dc..78753a42 100644
--- a/src/WixToolset.Mba.Core/IBootstrapperEngine.cs
+++ b/src/WixToolset.Mba.Core/IBootstrapperEngine.cs
@@ -277,6 +277,14 @@ namespace WixToolset.Mba.Core
             [MarshalAs(UnmanagedType.U4)] int dwWaitForInputIdleTimeout
             );
 
+        /// <summary>
+        /// Sets the URL to the update feed.
+        /// </summary>
+        /// <param name="url">URL of the update feed.</param>
+        void SetUpdateSource(
+            [MarshalAs(UnmanagedType.LPWStr)] string url
+            );
+
         /// <summary>
         /// See <see cref="IEngine.CompareVersions(string, string)"/>.
         /// </summary>
diff --git a/src/WixToolset.Mba.Core/IEngine.cs b/src/WixToolset.Mba.Core/IEngine.cs
index 0899ec43..3e636961 100644
--- a/src/WixToolset.Mba.Core/IEngine.cs
+++ b/src/WixToolset.Mba.Core/IEngine.cs
@@ -144,6 +144,12 @@ namespace WixToolset.Mba.Core
         /// <param name="hash">Optional hash expected for the update.</param>
         void SetUpdate(string localSource, string downloadSource, long size, UpdateHashType hashType, byte[] hash);
 
+        /// <summary>
+        /// Sets the URL to the update feed.
+        /// </summary>
+        /// <param name="url">URL of the update feed.</param>
+        void SetUpdateSource(string url);
+
         /// <summary>
         /// Set the local source for a package or container.
         /// </summary>
diff --git a/src/balutil/BalBootstrapperEngine.cpp b/src/balutil/BalBootstrapperEngine.cpp
index 8e40f7ae..dda98cb9 100644
--- a/src/balutil/BalBootstrapperEngine.cpp
+++ b/src/balutil/BalBootstrapperEngine.cpp
@@ -535,6 +535,21 @@ public: // IBootstrapperEngine
         return m_pfnBAEngineProc(BOOTSTRAPPER_ENGINE_MESSAGE_LAUNCHAPPROVEDEXE, &args, &results, m_pvBAEngineProcContext);
     }
 
+    virtual STDMETHODIMP SetUpdateSource(
+        __in_z LPCWSTR wzUrl
+        )
+    {
+        BAENGINE_SETUPDATESOURCE_ARGS args = { };
+        BAENGINE_SETUPDATESOURCE_RESULTS results = { };
+
+        args.cbSize = sizeof(args);
+        args.wzUrl = wzUrl;
+
+        results.cbSize = sizeof(results);
+
+        return m_pfnBAEngineProc(BOOTSTRAPPER_ENGINE_MESSAGE_SETUPDATESOURCE, &args, &results, m_pvBAEngineProcContext);
+    }
+
     virtual STDMETHODIMP CompareVersions(
         __in_z LPCWSTR wzVersion1,
         __in_z LPCWSTR wzVersion2,
diff --git a/src/balutil/inc/IBootstrapperEngine.h b/src/balutil/inc/IBootstrapperEngine.h
index cd89a9f0..af6379f4 100644
--- a/src/balutil/inc/IBootstrapperEngine.h
+++ b/src/balutil/inc/IBootstrapperEngine.h
@@ -128,6 +128,10 @@ DECLARE_INTERFACE_IID_(IBootstrapperEngine, IUnknown, "6480D616-27A0-44D7-905B-8
         __in DWORD dwWaitForInputIdleTimeout
         ) = 0;
 
+    STDMETHOD(SetUpdateSource)(
+        __in_z LPCWSTR wzUrl
+        ) = 0;
+
     STDMETHOD(CompareVersions)(
         __in_z LPCWSTR wzVersion1,
         __in_z LPCWSTR wzVersion2,
-- 
cgit v1.2.3-55-g6feb