Allow BA to update feed source

Fixes wixtoolset/issues#5568

6 files changed, 77 insertions, 2 deletions

diff --git a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperEngine.h b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperEngine.h
index 99e9b9f0..0a974563 100644
--- a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperEngine.h
+++ b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperEngine.h
@@ -112,6 +112,7 @@ enum BOOTSTRAPPER_ENGINE_MESSAGE
     BOOTSTRAPPER_ENGINE_MESSAGE_APPLY,
     BOOTSTRAPPER_ENGINE_MESSAGE_QUIT,
     BOOTSTRAPPER_ENGINE_MESSAGE_LAUNCHAPPROVEDEXE,
+    BOOTSTRAPPER_ENGINE_MESSAGE_SETUPDATESOURCE,
     BOOTSTRAPPER_ENGINE_MESSAGE_COMPAREVERSIONS,
 };
 
@@ -276,6 +277,17 @@ typedef struct _BAENGINE_LAUNCHAPPROVEDEXE_RESULTS
     DWORD cbSize;
 } BAENGINE_LAUNCHAPPROVEDEXE_RESULTS;
 
+typedef struct _BAENGINE_SETUPDATESOURCE_ARGS
+{
+    DWORD cbSize;
+    LPCWSTR wzUrl;
+} BAENGINE_SETUPDATESOURCE_ARGS;
+
+typedef struct _BAENGINE_SETUPDATESOURCE_RESULTS
+{
+    DWORD cbSize;
+} BAENGINE_SETUPDATESOURCE_RESULTS;
+
 typedef struct _BAENGINE_LOG_ARGS
 {
     DWORD cbSize;
diff --git a/src/engine/EngineForApplication.cpp b/src/engine/EngineForApplication.cpp
index 361e0f4e..83d88ba1 100644
--- a/src/engine/EngineForApplication.cpp
+++ b/src/engine/EngineForApplication.cpp
@@ -411,6 +411,22 @@ LExit:
     return hr;
 }
 
+static HRESULT BAEngineSetUpdateSource(
+    __in BOOTSTRAPPER_ENGINE_CONTEXT* pContext,
+    __in const LPVOID pvArgs,
+    __inout LPVOID pvResults
+    )
+{
+    HRESULT hr = S_OK;
+    ValidateMessageArgs(hr, pvArgs, BAENGINE_SETUPDATESOURCE_ARGS, pArgs);
+    ValidateMessageResults(hr, pvResults, BAENGINE_SETUPDATESOURCE_RESULTS, pResults);
+
+    hr = ExternalEngineSetUpdateSource(pContext->pEngineState, pArgs->wzUrl);
+
+LExit:
+    return hr;
+}
+
 HRESULT WINAPI EngineForApplicationProc(
     __in BOOTSTRAPPER_ENGINE_MESSAGE message,
     __in const LPVOID pvArgs,
@@ -497,6 +513,9 @@ HRESULT WINAPI EngineForApplicationProc(
     case BOOTSTRAPPER_ENGINE_MESSAGE_LAUNCHAPPROVEDEXE:
         hr = BAEngineLaunchApprovedExe(pContext, pvArgs, pvResults);
         break;
+    case BOOTSTRAPPER_ENGINE_MESSAGE_SETUPDATESOURCE:
+        hr = BAEngineSetUpdateSource(pContext, pvArgs, pvResults);
+        break;
     case BOOTSTRAPPER_ENGINE_MESSAGE_COMPAREVERSIONS:
         hr = BAEngineCompareVersions(pContext, pvArgs, pvResults);
         break;
diff --git a/src/engine/detect.cpp b/src/engine/detect.cpp
index 8ca74986..3a58f3dc 100644
--- a/src/engine/detect.cpp
+++ b/src/engine/detect.cpp
@@ -258,6 +258,7 @@ extern "C" HRESULT DetectUpdate(
     BOOL fBeginCalled = FALSE;
     BOOL fSkip = TRUE;
     BOOL fIgnoreError = FALSE;
+    LPWSTR sczOriginalSource = NULL;
 
     // If no update source was specified, skip update detection.
     if (!pUpdate->sczUpdateSource || !*pUpdate->sczUpdateSource)
@@ -266,7 +267,11 @@ extern "C" HRESULT DetectUpdate(
     }
 
     fBeginCalled = TRUE;
-    hr = UserExperienceOnDetectUpdateBegin(pUX, pUpdate->sczUpdateSource, &fSkip);
+
+    hr = StrAllocString(&sczOriginalSource, pUpdate->sczUpdateSource, 0);
+    ExitOnFailure(hr, "Failed to duplicate update feed source.");
+
+    hr = UserExperienceOnDetectUpdateBegin(pUX, sczOriginalSource, &fSkip);
     ExitOnRootFailure(hr, "BA aborted detect update begin.");
 
     if (!fSkip)
@@ -276,6 +281,8 @@ extern "C" HRESULT DetectUpdate(
     }
 
 LExit:
+    ReleaseStr(sczOriginalSource);
+
     if (fBeginCalled)
     {
         UserExperienceOnDetectUpdateComplete(pUX, hr, &fIgnoreError);
diff --git a/src/engine/externalengine.cpp b/src/engine/externalengine.cpp
index d881544c..fffd96bf 100644
--- a/src/engine/externalengine.cpp
+++ b/src/engine/externalengine.cpp
@@ -733,6 +733,38 @@ LExit:
     return hr;
 }
 
+HRESULT ExternalEngineSetUpdateSource(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in_z LPCWSTR wzUrl
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fLeaveCriticalSection = FALSE;
+
+    ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
+    fLeaveCriticalSection = TRUE;
+    hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
+    ExitOnFailure(hr, "Engine is active, cannot change engine state.");
+
+    if (wzUrl && *wzUrl)
+    {
+        hr = StrAllocString(&pEngineState->update.sczUpdateSource, wzUrl, 0);
+        ExitOnFailure(hr, "Failed to set feed download URL.");
+    }
+    else // no URL provided means clear out the whole download source.
+    {
+        ReleaseNullStr(pEngineState->update.sczUpdateSource);
+    }
+
+LExit:
+    if (fLeaveCriticalSection)
+    {
+        ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
+    }
+
+    return hr;
+}
+
 // TODO: callers need to provide the original size (at the time of first public release) of the struct instead of the current size.
 HRESULT WINAPI ExternalEngineValidateMessageParameter(
     __in_opt const LPVOID pv,
diff --git a/src/engine/externalengine.h b/src/engine/externalengine.h
index 3f7bc8c8..a007e5b2 100644
--- a/src/engine/externalengine.h
+++ b/src/engine/externalengine.h
@@ -165,6 +165,11 @@ HRESULT ExternalEngineLaunchApprovedExe(
     __in const DWORD dwWaitForInputIdleTimeout
     );
 
+HRESULT ExternalEngineSetUpdateSource(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in_z LPCWSTR wzUrl
+    );
+
 HRESULT WINAPI ExternalEngineValidateMessageParameter(
     __in_opt const LPVOID pv,
     __in SIZE_T cbSizeOffset,
diff --git a/src/engine/userexperience.cpp b/src/engine/userexperience.cpp
index 5a225651..b81dbfb2 100644
--- a/src/engine/userexperience.cpp
+++ b/src/engine/userexperience.cpp
@@ -111,7 +111,7 @@ extern "C" HRESULT UserExperienceLoad(
     args.pCommand = pCommand;
     args.pfnBootstrapperEngineProc = EngineForApplicationProc;
     args.pvBootstrapperEngineProcContext = pEngineContext;
-    args.qwEngineAPIVersion = MAKEQWORDVERSION(2021, 3, 2, 0);
+    args.qwEngineAPIVersion = MAKEQWORDVERSION(2021, 4, 14, 0);
 
     results.cbSize = sizeof(BOOTSTRAPPER_CREATE_RESULTS);