6 files changed, 195 insertions, 35 deletions

diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index bbd0ff96..812c7261 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -1193,7 +1193,7 @@ extern "C" HRESULT CoreCreateUpdateBundleCommandLine(
 {
     HRESULT hr = S_OK;
 
-    hr = CoreRecreateCommandLine(psczCommandLine, BOOTSTRAPPER_ACTION_INSTALL, pInternalCommand, pCommand, BOOTSTRAPPER_RELATION_NONE, FALSE);
+    hr = CoreRecreateCommandLine(psczCommandLine, BOOTSTRAPPER_ACTION_INSTALL, pInternalCommand, pCommand, BOOTSTRAPPER_RELATION_UPDATE, FALSE);
     ExitOnFailure(hr, "Failed to recreate update bundle command-line.");
 
 LExit:
diff --git a/src/burn/engine/externalengine.cpp b/src/burn/engine/externalengine.cpp
index da84c83d..abe9b8bc 100644
--- a/src/burn/engine/externalengine.cpp
+++ b/src/burn/engine/externalengine.cpp
@@ -269,64 +269,82 @@ HRESULT ExternalEngineSetUpdate(
     )
 {
     HRESULT hr = S_OK;
+    BOOL fLeaveCriticalSection = FALSE;
     LPWSTR sczFilePath = NULL;
     LPWSTR sczCommandline = NULL;
+    LPWSTR sczPreviousId = NULL;
+    LPCWSTR wzNewId = NULL;
     UUID guid = { };
     WCHAR wzGuid[39];
     RPC_STATUS rs = RPC_S_OK;
+    BOOL fRemove = (!wzLocalSource || !*wzLocalSource) && (!wzDownloadSource || !*wzDownloadSource);
+
+    UserExperienceOnSetUpdateBegin(&pEngineState->userExperience);
 
     ::EnterCriticalSection(&pEngineState->userExperience.csEngineActive);
+    fLeaveCriticalSection = TRUE;
     hr = UserExperienceEnsureEngineInactive(&pEngineState->userExperience);
     ExitOnFailure(hr, "Engine is active, cannot change engine state.");
 
-    if ((!wzLocalSource || !*wzLocalSource) && (!wzDownloadSource || !*wzDownloadSource))
-    {
-        UpdateUninitialize(&pEngineState->update);
-    }
-    else if (BOOTSTRAPPER_UPDATE_HASH_TYPE_NONE == hashType && (0 != cbHash || rgbHash))
+    if (!fRemove)
     {
-        hr = E_INVALIDARG;
+        if (BOOTSTRAPPER_UPDATE_HASH_TYPE_NONE == hashType && (0 != cbHash || rgbHash))
+        {
+            ExitFunction1(hr = E_INVALIDARG);
+        }
+        else if (BOOTSTRAPPER_UPDATE_HASH_TYPE_SHA512 == hashType && (SHA512_HASH_LEN != cbHash || !rgbHash))
+        {
+            ExitFunction1(hr = E_INVALIDARG);
+        }
     }
-    else if (BOOTSTRAPPER_UPDATE_HASH_TYPE_SHA512 == hashType && (SHA512_HASH_LEN != cbHash || !rgbHash))
+
+    sczPreviousId = pEngineState->update.package.sczId;
+    pEngineState->update.package.sczId = NULL;
+    UpdateUninitialize(&pEngineState->update);
+
+    if (fRemove)
     {
-        hr = E_INVALIDARG;
+        ExitFunction();
     }
-    else
-    {
-        UpdateUninitialize(&pEngineState->update);
 
-        hr = CoreCreateUpdateBundleCommandLine(&sczCommandline, &pEngineState->internalCommand, &pEngineState->command);
-        ExitOnFailure(hr, "Failed to create command-line for update bundle.");
+    hr = CoreCreateUpdateBundleCommandLine(&sczCommandline, &pEngineState->internalCommand, &pEngineState->command);
+    ExitOnFailure(hr, "Failed to create command-line for update bundle.");
 
-        // Bundles would fail to use the downloaded update bundle, as the running bundle would be one of the search paths.
-        // Here I am generating a random guid, but in the future it would be nice if the feed would provide the ID of the update.
-        rs = ::UuidCreate(&guid);
-        hr = HRESULT_FROM_RPC(rs);
-        ExitOnFailure(hr, "Failed to create bundle update guid.");
+    // Bundles would fail to use the downloaded update bundle, as the running bundle would be one of the search paths.
+    // Here I am generating a random guid, but in the future it would be nice if the feed would provide the ID of the update.
+    rs = ::UuidCreate(&guid);
+    hr = HRESULT_FROM_RPC(rs);
+    ExitOnFailure(hr, "Failed to create bundle update guid.");
 
-        if (!::StringFromGUID2(guid, wzGuid, countof(wzGuid)))
-        {
-            hr = E_OUTOFMEMORY;
-            ExitOnRootFailure(hr, "Failed to convert bundle update guid into string.");
-        }
+    if (!::StringFromGUID2(guid, wzGuid, countof(wzGuid)))
+    {
+        hr = E_OUTOFMEMORY;
+        ExitOnRootFailure(hr, "Failed to convert bundle update guid into string.");
+    }
 
-        hr = StrAllocFormatted(&sczFilePath, L"%ls\\%ls", wzGuid, pEngineState->registration.sczExecutableName);
-        ExitOnFailure(hr, "Failed to build bundle update file path.");
+    hr = StrAllocFormatted(&sczFilePath, L"%ls\\%ls", wzGuid, pEngineState->registration.sczExecutableName);
+    ExitOnFailure(hr, "Failed to build bundle update file path.");
 
-        if (!wzLocalSource || !*wzLocalSource)
-        {
-            wzLocalSource = sczFilePath;
-        }
+    if (!wzLocalSource || !*wzLocalSource)
+    {
+        wzLocalSource = sczFilePath;
+    }
 
-        hr = PseudoBundleInitialize(FILEMAKEVERSION(rmj, rmm, rup, rpr), &pEngineState->update.package, FALSE, pEngineState->registration.sczId, BOOTSTRAPPER_RELATION_UPDATE, BOOTSTRAPPER_PACKAGE_STATE_ABSENT, FALSE, sczFilePath, wzLocalSource, wzDownloadSource, qwSize, TRUE, sczCommandline, NULL, NULL, NULL, rgbHash, cbHash);
-        ExitOnFailure(hr, "Failed to set update bundle.");
+    hr = PseudoBundleInitializeUpdateBundle(&pEngineState->update.package, wzGuid, pEngineState->registration.sczId, sczFilePath, wzLocalSource, wzDownloadSource, qwSize, sczCommandline, rgbHash, cbHash);
+    ExitOnFailure(hr, "Failed to set update bundle.");
 
-        pEngineState->update.fUpdateAvailable = TRUE;
-    }
+    pEngineState->update.fUpdateAvailable = TRUE;
+    wzNewId = wzGuid;
 
 LExit:
-    ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
+    if (fLeaveCriticalSection)
+    {
+        ::LeaveCriticalSection(&pEngineState->userExperience.csEngineActive);
+    }
+
+    UserExperienceOnSetUpdateComplete(&pEngineState->userExperience, hr, sczPreviousId, wzNewId);
 
+    ReleaseStr(sczPreviousId);
     ReleaseStr(sczCommandline);
     ReleaseStr(sczFilePath);
 
diff --git a/src/burn/engine/pseudobundle.cpp b/src/burn/engine/pseudobundle.cpp
index 52b7bd8a..91c6c14f 100644
--- a/src/burn/engine/pseudobundle.cpp
+++ b/src/burn/engine/pseudobundle.cpp
@@ -240,3 +240,80 @@ LExit:
     ReleaseStr(sczArguments);
     return hr;
 }
+
+extern "C" HRESULT PseudoBundleInitializeUpdateBundle(
+    __in BURN_PACKAGE* pPackage,
+    __in_z LPCWSTR wzId,
+    __in_z LPCWSTR wzCacheId,
+    __in_z LPCWSTR wzFilePath,
+    __in_z LPCWSTR wzLocalSource,
+    __in_z_opt LPCWSTR wzDownloadSource,
+    __in DWORD64 qwSize,
+    __in_z LPCWSTR wzInstallArguments,
+    __in_opt const BYTE* pbHash,
+    __in const DWORD cbHash
+    )
+{
+    HRESULT hr = S_OK;
+    BURN_PAYLOAD* pPayload = NULL;
+
+    // Initialize the single payload, and fill out all the necessary fields
+    pPackage->payloads.rgItems = (BURN_PAYLOAD_GROUP_ITEM*)MemAlloc(sizeof(BURN_PAYLOAD_GROUP_ITEM), TRUE);
+    ExitOnNull(pPackage->payloads.rgItems, hr, E_OUTOFMEMORY, "Failed to allocate space for burn payload group inside of update bundle struct");
+    pPackage->payloads.cItems = 1;
+
+    pPayload = (BURN_PAYLOAD*)MemAlloc(sizeof(BURN_PAYLOAD), TRUE);
+    ExitOnNull(pPayload, hr, E_OUTOFMEMORY, "Failed to allocate space for burn payload inside of update bundle struct");
+    pPackage->payloads.rgItems[0].pPayload = pPayload;
+    pPayload->packaging = BURN_PAYLOAD_PACKAGING_EXTERNAL;
+    pPayload->qwFileSize = qwSize;
+    pPayload->verification = BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE;
+
+    hr = StrAllocString(&pPayload->sczKey, wzId, 0);
+    ExitOnFailure(hr, "Failed to copy key for pseudo bundle payload.");
+
+    hr = StrAllocString(&pPayload->sczFilePath, wzFilePath, 0);
+    ExitOnFailure(hr, "Failed to copy filename for pseudo bundle.");
+
+    hr = StrAllocString(&pPayload->sczSourcePath, wzLocalSource, 0);
+    ExitOnFailure(hr, "Failed to copy local source path for pseudo bundle.");
+
+    if (wzDownloadSource && *wzDownloadSource)
+    {
+        hr = StrAllocString(&pPayload->downloadSource.sczUrl, wzDownloadSource, 0);
+        ExitOnFailure(hr, "Failed to copy download source for pseudo bundle.");
+    }
+
+    if (pbHash)
+    {
+        pPayload->pbHash = static_cast<BYTE*>(MemAlloc(cbHash, FALSE));
+        ExitOnNull(pPayload->pbHash, hr, E_OUTOFMEMORY, "Failed to allocate memory for update bundle payload hash.");
+
+        pPayload->cbHash = cbHash;
+        memcpy_s(pPayload->pbHash, pPayload->cbHash, pbHash, cbHash);
+    }
+
+    pPackage->Exe.fPseudoBundle = TRUE;
+
+    pPackage->type = BURN_PACKAGE_TYPE_EXE;
+    pPackage->currentState = BOOTSTRAPPER_PACKAGE_STATE_ABSENT;
+    pPackage->qwInstallSize = qwSize;
+    pPackage->qwSize = qwSize;
+    pPackage->fVital = TRUE;
+
+    hr = StrAllocString(&pPackage->sczId, wzId, 0);
+    ExitOnFailure(hr, "Failed to copy id for update bundle.");
+
+    hr = StrAllocString(&pPackage->sczCacheId, wzCacheId, 0);
+    ExitOnFailure(hr, "Failed to copy cache id for update bundle.");
+
+    hr = StrAllocString(&pPackage->Exe.sczInstallArguments, wzInstallArguments, 0);
+    ExitOnFailure(hr, "Failed to copy install arguments for update bundle package");
+
+    // Assume the update bundle has the same engine version as this one.
+    pPackage->Exe.protocol = BURN_EXE_PROTOCOL_TYPE_BURN;
+    pPackage->Exe.fSupportsAncestors = TRUE;
+
+LExit:
+    return hr;
+}
\ No newline at end of file
diff --git a/src/burn/engine/pseudobundle.h b/src/burn/engine/pseudobundle.h
index aa26d29d..0fd4cbdb 100644
--- a/src/burn/engine/pseudobundle.h
+++ b/src/burn/engine/pseudobundle.h
@@ -32,6 +32,18 @@ HRESULT PseudoBundleInitializePassthrough(
     __in BOOTSTRAPPER_COMMAND* pCommand,
     __in BURN_PACKAGE* pPackage
     );
+HRESULT PseudoBundleInitializeUpdateBundle(
+    __in BURN_PACKAGE* pPackage,
+    __in_z LPCWSTR wzId,
+    __in_z LPCWSTR wzCacheId,
+    __in_z LPCWSTR wzFilePath,
+    __in_z LPCWSTR wzLocalSource,
+    __in_z_opt LPCWSTR wzDownloadSource,
+    __in DWORD64 qwSize,
+    __in_z LPCWSTR wzInstallArguments,
+    __in_opt const BYTE* pbHash,
+    __in const DWORD cbHash
+    );
 
 #if defined(__cplusplus)
 }
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 2bd6ecaf..a6d670ea 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -2241,6 +2241,50 @@ LExit:
     return hr;
 }
 
+EXTERN_C BAAPI UserExperienceOnSetUpdateBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONSETUPDATEBEGIN_ARGS args = { };
+    BA_ONSETUPDATEBEGIN_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+
+    results.cbSize = sizeof(results);
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN, &args, &results);
+    ExitOnFailure(hr, "BA OnSetUpdateBegin failed.");
+
+LExit:
+    return hr;
+}
+
+EXTERN_C BAAPI UserExperienceOnSetUpdateComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus,
+    __in_z_opt LPCWSTR wzPreviousPackageId,
+    __in_z_opt LPCWSTR wzNewPackageId
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONSETUPDATECOMPLETE_ARGS args = { };
+    BA_ONSETUPDATECOMPLETE_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.hrStatus = hrStatus;
+    args.wzPreviousPackageId = wzPreviousPackageId;
+    args.wzNewPackageId = wzNewPackageId;
+
+    results.cbSize = sizeof(results);
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE, &args, &results);
+    ExitOnFailure(hr, "BA OnSetUpdateComplete failed.");
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnShutdown(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __inout BOOTSTRAPPER_SHUTDOWN_ACTION* pAction
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index 2493569b..f7ac962c 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -505,6 +505,15 @@ BAAPI UserExperienceOnRollbackMsiTransactionComplete(
     __in LPCWSTR wzTransactionId,
     __in HRESULT hrStatus
     );
+BAAPI UserExperienceOnSetUpdateBegin(
+    __in BURN_USER_EXPERIENCE* pUserExperience
+    );
+BAAPI UserExperienceOnSetUpdateComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in HRESULT hrStatus,
+    __in_z_opt LPCWSTR wzPreviousPackageId,
+    __in_z_opt LPCWSTR wzNewPackageId
+    );
 BAAPI UserExperienceOnShutdown(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __inout BOOTSTRAPPER_SHUTDOWN_ACTION* pAction