From 5059f83d7dccd7221458fca414a0756bb83466a5 Mon Sep 17 00:00:00 2001
From: Bob Arnson <bob@firegiant.com>
Date: Fri, 9 Sep 2022 21:23:26 -0400
Subject: Support bundle updates in WixStdBA.

Fixes https://github.com/wixtoolset/issues/issues/6894.
---
 src/burn/engine/detect.cpp                         |   2 +-
 .../Bal/wixstdba/Resources/HyperlinkLargeTheme.xml |   5 +
 .../wixstdba/Resources/HyperlinkSidebarTheme.xml   |   5 +
 src/ext/Bal/wixstdba/Resources/HyperlinkTheme.wxl  |   2 +
 src/ext/Bal/wixstdba/Resources/HyperlinkTheme.xml  |   5 +
 src/ext/Bal/wixstdba/Resources/RtfLargeTheme.xml   |   5 +
 src/ext/Bal/wixstdba/Resources/RtfTheme.wxl        |   2 +
 src/ext/Bal/wixstdba/Resources/RtfTheme.xml        |   5 +
 .../WixStandardBootstrapperApplication.cpp         | 121 ++++++++++++++++++++-
 src/test/burn/TestBA/TestBA.cs                     |   2 +-
 .../burn/TestData/WixStdBaTests/BundleA/Bundle.wxs |   2 +
 11 files changed, 152 insertions(+), 4 deletions(-)

diff --git a/src/burn/engine/detect.cpp b/src/burn/engine/detect.cpp
index 25d184c1..fcd1c966 100644
--- a/src/burn/engine/detect.cpp
+++ b/src/burn/engine/detect.cpp
@@ -400,7 +400,7 @@ static HRESULT DetectAtomFeedUpdate(
             APPLICATION_UPDATE_ENTRY* pAppUpdateEntry = &pApupChain->rgEntries[i];
             APPLICATION_UPDATE_ENCLOSURE* pEnclosure = pAppUpdateEntry->rgEnclosures;
 
-            if (pEnclosure)
+            if (pEnclosure && pEnclosure->rgbDigest && *pEnclosure->rgbDigest)
             {
                 hr = StrAllocHexEncode(pEnclosure->rgbDigest, pEnclosure->cbDigest, &sczHash);
                 ExitOnFailure(hr, "Failed to encode hash as string.");
diff --git a/src/ext/Bal/wixstdba/Resources/HyperlinkLargeTheme.xml b/src/ext/Bal/wixstdba/Resources/HyperlinkLargeTheme.xml
index c018249f..6e7b69ee 100644
--- a/src/ext/Bal/wixstdba/Resources/HyperlinkLargeTheme.xml
+++ b/src/ext/Bal/wixstdba/Resources/HyperlinkLargeTheme.xml
@@ -12,6 +12,9 @@
         <ImageControl X="11" Y="11" Width="64" Height="64" ImageFile="logo.png" Visible="yes"/>
         <Label X="80" Y="11" Width="-11" Height="64" FontId="1" Visible="yes" DisablePrefix="yes">#(loc.Title)</Label>
 
+        <Page Name="Loading">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes" Visible="no" Name="CheckingForUpdatesLabel" />
+        </Page>
         <Page Name="Help">
             <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.HelpHeader)</Label>
             <Label X="11" Y="112" Width="-11" Height="-35" FontId="3" DisablePrefix="yes">#(loc.HelpText)</Label>
@@ -29,6 +32,7 @@
             <Hypertext Name="EulaHyperlink" X="11" Y="-107" Width="-11" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallLicenseLinkText)</Hypertext>
             <Label Name="InstallVersion" X="11" Y="-73" Width="246" Height="17" FontId="3" DisablePrefix="yes" VisibleCondition="WixStdBAShowVersion">#(loc.InstallVersion)</Label>
             <Checkbox Name="EulaAcceptCheckbox" X="-11" Y="-41" Width="260" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallAcceptCheckbox)</Checkbox>
+            <Button Name="InstallUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="OptionsButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" VisibleCondition="NOT WixStdBASuppressOptionsUI">
                 <Text>#(loc.InstallOptionsButton)</Text>
                 <ChangePageAction Page="Options" />
@@ -65,6 +69,7 @@
         </Page>
         <Page Name="Modify">
             <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ModifyHeader)</Label>
+            <Button Name="ModifyUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="RepairButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.ModifyRepairButton)</Button>
             <Button Name="UninstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ModifyUninstallButton)</Button>
             <Button Name="ModifyCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
diff --git a/src/ext/Bal/wixstdba/Resources/HyperlinkSidebarTheme.xml b/src/ext/Bal/wixstdba/Resources/HyperlinkSidebarTheme.xml
index 3630b5b9..59a03ad0 100644
--- a/src/ext/Bal/wixstdba/Resources/HyperlinkSidebarTheme.xml
+++ b/src/ext/Bal/wixstdba/Resources/HyperlinkSidebarTheme.xml
@@ -19,6 +19,9 @@
                 <CloseWindowAction />
             </Button>
         </Page>
+        <Page Name="Loading">
+            <Label X="185" Y="50" Width="-11" Height="30" FontId="2" DisablePrefix="yes" Visible="no" Name="CheckingForUpdatesLabel" />
+        </Page>
         <Page Name="Install">
             <Label X="185" Y="11" Width="-11" Height="32" FontId="1" DisablePrefix="yes">#(loc.Title)</Label>
             <ImageControl X="11" Y="11" Width="165" Height="400" ImageFile="logoside.png"/>
@@ -30,6 +33,7 @@
             <Hypertext Name="EulaHyperlink" X="185" Y="-111" Width="-11" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallLicenseLinkText)</Hypertext>
             <Label Name="InstallVersion" X="185" Y="-81" Width="-11" Height="17" FontId="3" DisablePrefix="yes" VisibleCondition="WixStdBAShowVersion">#(loc.InstallVersion)</Label>
             <Checkbox Name="EulaAcceptCheckbox" X="185" Y="-51" Width="-11" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallAcceptCheckbox)</Checkbox>
+            <Button Name="InstallUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="OptionsButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" VisibleCondition="NOT WixStdBASuppressOptionsUI">
                 <Text>#(loc.InstallOptionsButton)</Text>
                 <ChangePageAction Page="Options" />
@@ -72,6 +76,7 @@
             <Label X="185" Y="11" Width="-11" Height="32" FontId="1" DisablePrefix="yes">#(loc.Title)</Label>
             <ImageControl X="11" Y="11" Width="165" Height="400" ImageFile="logoside.png"/>
             <Label X="185" Y="50" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ModifyHeader)</Label>
+            <Button Name="ModifyUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="RepairButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.ModifyRepairButton)</Button>
             <Button Name="UninstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ModifyUninstallButton)</Button>
             <Button Name="ModifyCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
diff --git a/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.wxl b/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.wxl
index 8d222a3b..1ce15453 100644
--- a/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.wxl
+++ b/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.wxl
@@ -5,6 +5,8 @@
 <WixLocalization Culture="en-us" Language="1033" xmlns="http://wixtoolset.org/schemas/v4/wxl">
   <String Id="Caption">[WixBundleName] Setup</String>
   <String Id="Title">[WixBundleName]</String>
+  <String Id="CheckingForUpdatesLabel">Checking for updates</String>
+  <String Id="UpdateButton">&amp;Update to version [WixStdBAUpdateAvailable]</String>
   <String Id="InstallHeader">Welcome</String>
   <String Id="InstallMessage">Setup will install [WixBundleName] on your computer. Click Install to continue or Cancel to exit.</String>
   <String Id="InstallMessageOptions">Setup will install [WixBundleName] on your computer. Click Install to continue, Options to set installation options, or Cancel to exit.</String>
diff --git a/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.xml b/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.xml
index e654fe23..c82e1b9f 100644
--- a/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.xml
+++ b/src/ext/Bal/wixstdba/Resources/HyperlinkTheme.xml
@@ -20,9 +20,13 @@
                 <CloseWindowAction />
             </Button>
         </Page>
+        <Page Name="Loading">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes" Visible="no" Name="CheckingForUpdatesLabel" />
+        </Page>
         <Page Name="Install">
             <Hypertext Name="EulaHyperlink" X="11" Y="121" Width="-11" Height="51" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallLicenseLinkText)</Hypertext>
             <Checkbox Name="EulaAcceptCheckbox" X="-11" Y="-41" Width="260" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallAcceptCheckbox)</Checkbox>
+            <Button Name="InstallUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="OptionsButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" VisibleCondition="NOT WixStdBASuppressOptionsUI">
                 <Text>#(loc.InstallOptionsButton)</Text>
                 <ChangePageAction Page="Options" />
@@ -59,6 +63,7 @@
         </Page>
         <Page Name="Modify">
             <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ModifyHeader)</Label>
+            <Button Name="ModifyUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="RepairButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.ModifyRepairButton)</Button>
             <Button Name="UninstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ModifyUninstallButton)</Button>
             <Button Name="ModifyCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
diff --git a/src/ext/Bal/wixstdba/Resources/RtfLargeTheme.xml b/src/ext/Bal/wixstdba/Resources/RtfLargeTheme.xml
index e5d18ec7..b30fe3cd 100644
--- a/src/ext/Bal/wixstdba/Resources/RtfLargeTheme.xml
+++ b/src/ext/Bal/wixstdba/Resources/RtfLargeTheme.xml
@@ -20,11 +20,15 @@
                 <CloseWindowAction />
             </Button>
         </Page>
+        <Page Name="Loading">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes" Visible="no" Name="CheckingForUpdatesLabel" />
+        </Page>
         <Page Name="Install">
             <Label X="11" Y="80" Width="-11" Height="-70" TabStop="no" FontId="2" HexStyle="800000" DisablePrefix="yes" />
             <Richedit Name="EulaRichedit" X="12" Y="81" Width="-12" Height="-71" TabStop="yes" FontId="0" />
             <Label Name="InstallVersion" X="11" Y="-41" Width="210" Height="17" FontId="3" DisablePrefix="yes" VisibleCondition="WixStdBAShowVersion">#(loc.InstallVersion)</Label>
             <Checkbox Name="EulaAcceptCheckbox" X="-11" Y="-41" Width="260" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallAcceptCheckbox)</Checkbox>
+            <Button Name="InstallUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="OptionsButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" VisibleCondition="NOT WixStdBASuppressOptionsUI">
                 <Text>#(loc.InstallOptionsButton)</Text>
                 <ChangePageAction Page="Options" />
@@ -61,6 +65,7 @@
         </Page>
         <Page Name="Modify">
             <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ModifyHeader)</Label>
+            <Button Name="ModifyUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="RepairButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.ModifyRepairButton)</Button>
             <Button Name="UninstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ModifyUninstallButton)</Button>
             <Button Name="ModifyCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
diff --git a/src/ext/Bal/wixstdba/Resources/RtfTheme.wxl b/src/ext/Bal/wixstdba/Resources/RtfTheme.wxl
index efe66f41..8125f17e 100644
--- a/src/ext/Bal/wixstdba/Resources/RtfTheme.wxl
+++ b/src/ext/Bal/wixstdba/Resources/RtfTheme.wxl
@@ -5,6 +5,8 @@
 <WixLocalization Culture="en-us" Language="1033" xmlns="http://wixtoolset.org/schemas/v4/wxl">
   <String Id="Caption">[WixBundleName] Setup</String>
   <String Id="Title">[WixBundleName]</String>
+  <String Id="CheckingForUpdatesLabel">Checking for updates</String>
+  <String Id="UpdateButton">&amp;Update to version [WixStdBAUpdateAvailable]</String>
   <String Id="InstallVersion">Version [WixBundleVersion]</String>
   <String Id="ConfirmCancelMessage">Are you sure you want to cancel?</String>
   <String Id="ExecuteUpgradeRelatedBundleMessage">Previous version</String>
diff --git a/src/ext/Bal/wixstdba/Resources/RtfTheme.xml b/src/ext/Bal/wixstdba/Resources/RtfTheme.xml
index 72345343..e1a486c8 100644
--- a/src/ext/Bal/wixstdba/Resources/RtfTheme.xml
+++ b/src/ext/Bal/wixstdba/Resources/RtfTheme.xml
@@ -20,9 +20,13 @@
                 <CloseWindowAction />
             </Button>
         </Page>
+        <Page Name="Loading">
+            <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes" Visible="no" Name="CheckingForUpdatesLabel" />
+        </Page>
         <Page Name="Install">
             <Richedit Name="EulaRichedit" X="11" Y="80" Width="-11" Height="-70" TabStop="yes" FontId="0" HexStyle="800000" />
             <Checkbox Name="EulaAcceptCheckbox" X="-11" Y="-41" Width="260" Height="17" TabStop="yes" FontId="3" HideWhenDisabled="yes">#(loc.InstallAcceptCheckbox)</Checkbox>
+            <Button Name="InstallUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="OptionsButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" VisibleCondition="NOT WixStdBASuppressOptionsUI">
                 <Text>#(loc.InstallOptionsButton)</Text>
                 <ChangePageAction Page="Options" />
@@ -59,6 +63,7 @@
         </Page>
         <Page Name="Modify">
             <Label X="11" Y="80" Width="-11" Height="30" FontId="2" DisablePrefix="yes">#(loc.ModifyHeader)</Label>
+            <Button Name="ModifyUpdateButton" X="11" Y="-11" Width="200" Height="23" TabStop="yes" FontId="0" EnableCondition="WixStdBAUpdateAvailable" HideWhenDisabled="yes">#(loc.UpdateButton)</Button>
             <Button Name="RepairButton" X="-171" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0" HideWhenDisabled="yes">#(loc.ModifyRepairButton)</Button>
             <Button Name="UninstallButton" X="-91" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">#(loc.ModifyUninstallButton)</Button>
             <Button Name="ModifyCancelButton" X="-11" Y="-11" Width="75" Height="23" TabStop="yes" FontId="0">
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 50de2c7f..3539450e 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -23,6 +23,7 @@ static const LPCWSTR WIXSTDBA_VARIABLE_LANGUAGE_ID = L"WixStdBALanguageId";
 static const LPCWSTR WIXSTDBA_VARIABLE_RESTART_REQUIRED = L"WixStdBARestartRequired";
 static const LPCWSTR WIXSTDBA_VARIABLE_SHOW_VERSION = L"WixStdBAShowVersion";
 static const LPCWSTR WIXSTDBA_VARIABLE_SUPPRESS_OPTIONS_UI = L"WixStdBASuppressOptionsUI";
+static const LPCWSTR WIXSTDBA_VARIABLE_UPDATE_AVAILABLE = L"WixStdBAUpdateAvailable";
 
 enum WIXSTDBA_STATE
 {
@@ -98,6 +99,11 @@ enum WIXSTDBA_CONTROL
     WIXSTDBA_CONTROL_REPAIR_BUTTON,
     WIXSTDBA_CONTROL_UNINSTALL_BUTTON,
 
+    // Updates
+    WIXSTDBA_CONTROL_CHECKING_FOR_UPDATES_LABEL,
+    WIXSTDBA_CONTROL_INSTALL_UPDATE_BUTTON,
+    WIXSTDBA_CONTROL_MODIFY_UPDATE_BUTTON,
+
     // Progress page
     WIXSTDBA_CONTROL_CACHE_PROGRESS_PACKAGE_TEXT,
     WIXSTDBA_CONTROL_CACHE_PROGRESS_BAR,
@@ -342,6 +348,70 @@ public: // IBootstrapperApplication
     }
 
 
+    virtual STDMETHODIMP OnDetectUpdateBegin(
+        __in_z LPCWSTR wzUpdateLocation,
+        __inout BOOL* pfCancel,
+        __inout BOOL* pfSkip
+    )
+    {
+#ifdef DEBUG
+        BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "WIXSTDBA: OnDetectUpdateBegin() - update location: %ls", wzUpdateLocation);
+#endif
+
+        // Try update detection only if we have a potential update source and are in full UI mode.
+        *pfSkip = !wzUpdateLocation
+            || !*wzUpdateLocation
+            || BOOTSTRAPPER_DISPLAY_FULL != m_command.display;
+
+        ThemeShowControl(m_pControlCheckingForUpdatesLabel, *pfSkip ? SW_HIDE : SW_SHOW);
+
+        return __super::OnDetectUpdateBegin(wzUpdateLocation, pfCancel, pfSkip);
+    }
+
+
+    virtual STDMETHODIMP OnDetectUpdate(
+        __in_z LPCWSTR wzUpdateLocation,
+        __in DWORD64 dw64Size,
+        __in_z_opt LPCWSTR wzHash,
+        __in BOOTSTRAPPER_UPDATE_HASH_TYPE hashAlgorithm,
+        __in LPCWSTR wzUpdateVersion,
+        __in_z LPCWSTR wzTitle,
+        __in_z LPCWSTR wzSummary,
+        __in_z LPCWSTR wzContentType,
+        __in_z LPCWSTR wzContent,
+        __inout BOOL* pfCancel,
+        __inout BOOL* pfStopProcessingUpdates
+    )
+    {
+        HRESULT hr = S_OK;
+        int nResult = 0;
+
+        hr = VerCompareStringVersions(m_sczBundleVersion, wzUpdateVersion, TRUE/*fStrict*/, &nResult);
+        BalExitOnFailure(hr, "Failed to compare bundle version: %ls to update version: %ls.", m_sczBundleVersion, wzUpdateVersion);
+
+        // Burn sends the feed in descending version order so we need only the first one.
+        *pfStopProcessingUpdates = TRUE;
+
+        if (0 <= nResult)
+        {
+            BalLog(BOOTSTRAPPER_LOG_LEVEL_VERBOSE, "WIXSTDBA: Update version: %ls is a match or downgrade for bundle version: %ls.", wzUpdateVersion, m_sczBundleVersion);
+        }
+        else
+        {
+            BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "WIXSTDBA: Update v%ls for bundle v%ls available from: %ls.", wzUpdateVersion, m_sczBundleVersion, wzUpdateLocation);
+
+            hr = BalSetVersionVariable(WIXSTDBA_VARIABLE_UPDATE_AVAILABLE, wzUpdateVersion);
+            BalExitOnFailure(hr, "Failed to set WixStdBAUpdateAvailable value: %ls.", wzUpdateVersion);
+
+            hr = m_pEngine->SetUpdate(NULL, wzUpdateLocation, dw64Size, hashAlgorithm, wzHash);
+            BalExitOnFailure(hr, "Failed to set update location: %ls.", wzUpdateLocation);
+        }
+
+    LExit:
+        return __super::OnDetectUpdate(wzUpdateLocation, dw64Size, wzHash, hashAlgorithm, wzUpdateVersion, wzTitle, wzSummary, wzContentType, wzContent, pfCancel, pfStopProcessingUpdates);
+    }
+
+
     virtual STDMETHODIMP OnDetectComplete(
         __in HRESULT hrStatus,
         __in BOOL /*fEligibleForCleanup*/
@@ -3741,10 +3811,10 @@ private:
 
         m_state = state;
 
-        // If our install is at the end (success or failure) and we're not showing full UI or
+        // If our install is at the end (success or failure) and we're not showing full UI or not updating or
         // we successfully installed the prerequisite(s) and either no restart is required or can automatically restart
         // then exit.
-        if ((WIXSTDBA_STATE_APPLIED <= m_state && BOOTSTRAPPER_DISPLAY_FULL > m_command.display) ||
+        if ((WIXSTDBA_STATE_APPLIED <= m_state && (BOOTSTRAPPER_DISPLAY_FULL > m_command.display || BOOTSTRAPPER_ACTION_UPDATE_REPLACE == m_plannedAction)) ||
             (WIXSTDBA_STATE_APPLIED == m_state && m_fPrereq && (!m_fRestartRequired || m_fShouldRestart && m_fAllowRestart)))
         {
             // Quietly exit.
@@ -4037,6 +4107,15 @@ private:
     }
 
 
+    //
+    // OnClickUpdateButton - start the update process.
+    //
+    void OnClickUpdateButton()
+    {
+        this->OnPlan(BOOTSTRAPPER_ACTION_UPDATE_REPLACE);
+    }
+
+
     //
     // OnClickCloseButton - close the application.
     //
@@ -4263,6 +4342,13 @@ private:
                 pResults->lResult = 0;
                 ExitFunction();
 
+            case WIXSTDBA_CONTROL_INSTALL_UPDATE_BUTTON:
+            case WIXSTDBA_CONTROL_MODIFY_UPDATE_BUTTON:
+                OnClickUpdateButton();
+                fProcessed = TRUE;
+                pResults->lResult = 0;
+                ExitFunction();
+
             case WIXSTDBA_CONTROL_LAUNCH_BUTTON:
                 OnClickLaunchButton();
                 fProcessed = TRUE;
@@ -4381,6 +4467,10 @@ LExit:
         __in HRESULT hrStatus
         )
     {
+#ifdef DEBUG
+        BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "WIXSTDBA: SetState() - setting state to %u with status 0x%08x.", state, hrStatus);
+#endif
+
         if (FAILED(hrStatus))
         {
             m_hrFinal = hrStatus;
@@ -4786,6 +4876,27 @@ public:
         m_pControlUninstallButton = NULL;
         ++pAssignControl;
 
+        pAssignControl->wId = WIXSTDBA_CONTROL_CHECKING_FOR_UPDATES_LABEL;
+        pAssignControl->wzName = L"CheckingForUpdatesLabel";
+        pAssignControl->ppControl = &m_pControlCheckingForUpdatesLabel;
+        pAssignControl->dwAutomaticBehaviorType = dwAutomaticBehaviorType;
+        m_pControlCheckingForUpdatesLabel = NULL;
+        ++pAssignControl;
+
+        pAssignControl->wId = WIXSTDBA_CONTROL_INSTALL_UPDATE_BUTTON;
+        pAssignControl->wzName = L"InstallUpdateButton";
+        pAssignControl->ppControl = &m_pControlInstallUpdateButton;
+        pAssignControl->dwAutomaticBehaviorType = THEME_CONTROL_AUTOMATIC_BEHAVIOR_ALL;
+        m_pControlInstallUpdateButton = NULL;
+        ++pAssignControl;
+
+        pAssignControl->wId = WIXSTDBA_CONTROL_MODIFY_UPDATE_BUTTON;
+        pAssignControl->wzName = L"ModifyUpdateButton";
+        pAssignControl->ppControl = &m_pControlModifyUpdateButton;
+        pAssignControl->dwAutomaticBehaviorType = THEME_CONTROL_AUTOMATIC_BEHAVIOR_ALL;
+        m_pControlModifyUpdateButton = NULL;
+        ++pAssignControl;
+
         pAssignControl->wId = WIXSTDBA_CONTROL_CACHE_PROGRESS_PACKAGE_TEXT;
         pAssignControl->wzName = L"CacheProgressPackageText";
         pAssignControl->ppControl = &m_pControlCacheProgressPackageText;
@@ -4929,6 +5040,7 @@ public:
         ReleaseStr(m_sczLicenseUrl);
         ReleaseStr(m_sczBundleVersion);
         ReleaseStr(m_sczAfterForcedRestartPackage);
+
         ReleaseNullObject(m_pEngine);
     }
 
@@ -4966,6 +5078,11 @@ private:
     const THEME_CONTROL* m_pControlRepairButton;
     const THEME_CONTROL* m_pControlUninstallButton;
 
+    // Update/loading pages
+    const THEME_CONTROL* m_pControlCheckingForUpdatesLabel;
+    const THEME_CONTROL* m_pControlInstallUpdateButton;
+    const THEME_CONTROL* m_pControlModifyUpdateButton;
+
     // Progress page
     const THEME_CONTROL* m_pControlCacheProgressPackageText;
     const THEME_CONTROL* m_pControlCacheProgressbar;
diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index 087219af..01109911 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -220,7 +220,7 @@ namespace WixToolset.Test.BA
             if (!this.UpdateAvailable && this.Engine.CompareVersions(e.Version, this.Version) > 0)
             {
                 this.Log(String.Format("Selected update v{0}", e.Version));
-                this.Engine.SetUpdate(null, e.UpdateLocation, e.Size, UpdateHashType.None, null);
+                this.Engine.SetUpdate(null, e.UpdateLocation, e.Size, e.HashAlgorithm, e.Hash);
                 this.UpdateAvailable = true;
             }
         }
diff --git a/src/test/burn/TestData/WixStdBaTests/BundleA/Bundle.wxs b/src/test/burn/TestData/WixStdBaTests/BundleA/Bundle.wxs
index 99320b35..98a00344 100644
--- a/src/test/burn/TestData/WixStdBaTests/BundleA/Bundle.wxs
+++ b/src/test/burn/TestData/WixStdBaTests/BundleA/Bundle.wxs
@@ -4,6 +4,8 @@
   <Bundle Name="~$(var.TestGroupName) - $(var.BundleName)" Version="$(var.Version)" UpgradeCode="$(var.UpgradeCode)" Compressed="yes" bal:CommandLineVariables="caseSensitive">
     <Log Prefix="~$(var.TestGroupName)_$(var.BundleName)" />
 
+    <Update Location='http://wixtoolset.org/releases/feed/v3.14' />
+    
     <Variable Name="TestGroupName" Value="$(var.TestGroupName)" />
 
     <BootstrapperApplication>
-- 
cgit v1.2.3-55-g6feb