Add support for downloading embedded payloads.

#5253

4 files changed, 21 insertions, 22 deletions

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index 4570ffb9..74f57f6a 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -1422,19 +1422,19 @@ static HRESULT AcquireContainerOrPayload(
             hr = CacheGetLocalSourcePaths(wzRelativePath, *pwzSourcePath, wzDestinationPath, pContext->wzLayoutDirectory, pContext->pVariables, &pContext->rgSearchPaths, &pContext->cSearchPaths, &dwChosenSearchPath, &dwDestinationSearchPath);
             ExitOnFailure(hr, "Failed to search local source.");
 
-            // When a payload comes from a container, the container has the highest chance of being correct.
-            // But we want to avoid extracting the container multiple times.
-            // So only consider the destination path, which means the container was already extracted.
             if (wzPayloadContainerId)
             {
+                // When a payload comes from a container, the container has the highest chance of being correct.
+                // But we want to avoid extracting the container multiple times.
+                // So only consider the destination path, which means the container was already extracted.
                 if (FileExistsEx(pContext->rgSearchPaths[dwDestinationSearchPath], NULL))
                 {
                     fFoundLocal = TRUE;
                     dwChosenSearchPath = dwDestinationSearchPath;
                 }
-                else
+                else // don't prefer the container if extracting it already failed.
                 {
-                    fPreferExtract = TRUE;
+                    fPreferExtract = SUCCEEDED(pPayload->pContainer->hrExtract);
                 }
             }
 
@@ -1470,14 +1470,14 @@ static HRESULT AcquireContainerOrPayload(
                 {
                     resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_LOCAL;
                 }
-                else if (wzPayloadContainerId)
-                {
-                    resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_CONTAINER;
-                }
                 else if (*pwzDownloadUrl && **pwzDownloadUrl)
                 {
                     resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_DOWNLOAD;
                 }
+                else if (wzPayloadContainerId)
+                {
+                    resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_CONTAINER;
+                }
             }
 
             // Let the BA have a chance to override the source.
@@ -1526,7 +1526,7 @@ static HRESULT AcquireContainerOrPayload(
 
         break;
     case BOOTSTRAPPER_CACHE_OPERATION_EXTRACT:
-        Assert(pPayload->pContainer);
+        Assert(pPayload && pPayload->pContainer);
 
         hr = ApplyExtractContainer(pContext, pPayload->pContainer);
         ExitOnFailure(hr, "Failed to extract container for payload: %ls", wzPayloadId);
@@ -1541,6 +1541,15 @@ static HRESULT AcquireContainerOrPayload(
     hr = CompleteCacheProgress(pProgress, pContainer ? pContainer->qwFileSize : pPayload->qwFileSize);
 
 LExit:
+    if (BOOTSTRAPPER_CACHE_OPERATION_EXTRACT == cacheOperation)
+    {
+        if (FAILED(hr) && SUCCEEDED(pPayload->pContainer->hrExtract) &&
+            (fFoundLocal || pPayload->downloadSource.sczUrl && *pPayload->downloadSource.sczUrl))
+        {
+            *pfRetry = TRUE;
+        }
+        pPayload->pContainer->hrExtract = hr;
+    }
     UserExperienceOnCacheAcquireComplete(pContext->pUX, wzPackageOrContainerId, wzPayloadId, hr, pfRetry);
 
     pContext->cSearchPathsMax = max(pContext->cSearchPaths, pContext->cSearchPathsMax);
diff --git a/src/engine/container.h b/src/engine/container.h
index 3174eb38..88fe48f1 100644
--- a/src/engine/container.h
+++ b/src/engine/container.h
@@ -79,6 +79,7 @@ typedef struct _BURN_CONTAINER
     DWORD64 qwExtractSizeTotal;
     DWORD64 qwCommittedCacheProgress;
     DWORD64 qwCommittedExtractProgress;
+    HRESULT hrExtract;
 } BURN_CONTAINER;
 
 typedef struct _BURN_CONTAINERS
diff --git a/src/engine/externalengine.cpp b/src/engine/externalengine.cpp
index d6c44736..63177722 100644
--- a/src/engine/externalengine.cpp
+++ b/src/engine/externalengine.cpp
@@ -357,12 +357,6 @@ HRESULT ExternalEngineSetLocalSource(
         hr = PayloadFindById(&pEngineState->payloads, wzPayloadId, &pPayload);
         ExitOnFailure(hr, "BA requested unknown payload with id: %ls", wzPayloadId);
 
-        if (BURN_PAYLOAD_PACKAGING_EMBEDDED == pPayload->packaging)
-        {
-            hr = HRESULT_FROM_WIN32(ERROR_INVALID_OPERATION);
-            ExitOnFailure(hr, "BA denied while trying to set source on embedded payload: %ls", wzPayloadId);
-        }
-
         hr = StrAllocString(&pPayload->sczSourcePath, wzPath, 0);
         ExitOnFailure(hr, "Failed to set source path for payload.");
     }
@@ -408,12 +402,6 @@ HRESULT ExternalEngineSetDownloadSource(
         hr = PayloadFindById(&pEngineState->payloads, wzPayloadId, &pPayload);
         ExitOnFailure(hr, "BA requested unknown payload with id: %ls", wzPayloadId);
 
-        if (BURN_PAYLOAD_PACKAGING_EMBEDDED == pPayload->packaging)
-        {
-            hr = HRESULT_FROM_WIN32(ERROR_INVALID_OPERATION);
-            ExitOnFailure(hr, "BA denied while trying to set download URL on embedded payload: %ls", wzPayloadId);
-        }
-
         pDownloadSource = &pPayload->downloadSource;
     }
     else if (wzPackageOrContainerId && *wzPackageOrContainerId)
diff --git a/src/engine/plan.cpp b/src/engine/plan.cpp
index f662c445..e5c1ee36 100644
--- a/src/engine/plan.cpp
+++ b/src/engine/plan.cpp
@@ -1832,6 +1832,7 @@ static void ResetPlannedContainerState(
     pContainer->qwExtractSizeTotal = 0;
     pContainer->qwCommittedCacheProgress = 0;
     pContainer->qwCommittedExtractProgress = 0;
+    pContainer->hrExtract = S_OK;
 }
 
 static void ResetPlannedPayloadsState(