Share code to map stream name to payload when extracting containers.

4 files changed, 20 insertions, 52 deletions

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index b831be9c..4570ffb9 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -1113,7 +1113,8 @@ static HRESULT ExtractContainer(
     HRESULT hr = S_OK;
     BURN_CONTAINER_CONTEXT context = { };
     HANDLE hContainerHandle = INVALID_HANDLE_VALUE;
-    LPWSTR sczExtractPayloadId = NULL;
+    LPWSTR sczStreamName = NULL;
+    BURN_PAYLOAD* pExtract = NULL;
     BURN_CACHE_PROGRESS_CONTEXT progress = { };
 
     progress.pCacheContext = pContext;
@@ -1129,14 +1130,17 @@ static HRESULT ExtractContainer(
     hr = ContainerOpen(&context, pContainer, hContainerHandle, pContainer->sczUnverifiedPath);
     ExitOnFailure(hr, "Failed to open container: %ls.", pContainer->sczId);
 
-    while (S_OK == (hr = ContainerNextStream(&context, &sczExtractPayloadId)))
+    while (S_OK == (hr = ContainerNextStream(&context, &sczStreamName)))
     {
         BOOL fExtracted = FALSE;
 
-        for (DWORD iExtract = 0; iExtract < pContext->pPayloads->cPayloads; ++iExtract)
+        hr = PayloadFindEmbeddedBySourcePath(pContext->pPayloads, sczStreamName, &pExtract);
+        if (E_NOTFOUND != hr)
         {
-            BURN_PAYLOAD* pExtract = pContext->pPayloads->rgPayloads + iExtract;
-            if (pExtract->sczUnverifiedPath && pExtract->cRemainingInstances && CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, sczExtractPayloadId, -1, pExtract->sczSourcePath, -1))
+            ExitOnFailure(hr, "Failed to find embedded payload by source path: %ls container: %ls", sczStreamName, pContainer->sczId);
+
+            // Skip payloads that weren't planned or have already been cached.
+            if (pExtract->sczUnverifiedPath && pExtract->cRemainingInstances)
             {
                 progress.pPayload = pExtract;
 
@@ -1161,17 +1165,16 @@ static HRESULT ExtractContainer(
                 }
 
                 UserExperienceOnCachePayloadExtractComplete(pContext->pUX, pContainer->sczId, pExtract->sczKey, hr);
-                ExitOnFailure(hr, "Failed to extract payload: %ls from container: %ls", sczExtractPayloadId, pContainer->sczId);
+                ExitOnFailure(hr, "Failed to extract payload: %ls from container: %ls", sczStreamName, pContainer->sczId);
 
                 fExtracted = TRUE;
-                break;
             }
         }
 
         if (!fExtracted)
         {
             hr = ContainerSkipStream(&context);
-            ExitOnFailure(hr, "Failed to skip the extraction of payload: %ls from container: %ls", sczExtractPayloadId, pContainer->sczId);
+            ExitOnFailure(hr, "Failed to skip the extraction of payload: %ls from container: %ls", sczStreamName, pContainer->sczId);
         }
     }
 
@@ -1182,7 +1185,7 @@ static HRESULT ExtractContainer(
     ExitOnFailure(hr, "Failed to extract all payloads from container: %ls", pContainer->sczId);
 
 LExit:
-    ReleaseStr(sczExtractPayloadId);
+    ReleaseStr(sczStreamName);
     ContainerClose(&context);
 
     return hr;
diff --git a/src/engine/core.cpp b/src/engine/core.cpp
index aea614d2..969b94a0 100644
--- a/src/engine/core.cpp
+++ b/src/engine/core.cpp
@@ -155,7 +155,7 @@ extern "C" HRESULT CoreInitialize(
         hr = UserExperienceEnsureWorkingFolder(pEngineState->registration.sczId, &pEngineState->userExperience.sczTempDirectory);
         ExitOnFailure(hr, "Failed to get unique temporary folder for bootstrapper application.");
 
-        hr = PayloadExtractFromContainer(&pEngineState->userExperience.payloads, NULL, &containerContext, pEngineState->userExperience.sczTempDirectory);
+        hr = PayloadExtractUXContainer(&pEngineState->userExperience.payloads, &containerContext, pEngineState->userExperience.sczTempDirectory);
         ExitOnFailure(hr, "Failed to extract bootstrapper application payloads.");
 
         hr = PathConcat(pEngineState->userExperience.sczTempDirectory, L"BootstrapperApplicationData.xml", &pEngineState->command.wzBootstrapperApplicationDataPath);
diff --git a/src/engine/payload.cpp b/src/engine/payload.cpp
index f29fa2bd..72eb3476 100644
--- a/src/engine/payload.cpp
+++ b/src/engine/payload.cpp
@@ -190,9 +190,8 @@ extern "C" void PayloadsUninitialize(
     memset(pPayloads, 0, sizeof(BURN_PAYLOADS));
 }
 
-extern "C" HRESULT PayloadExtractFromContainer(
+extern "C" HRESULT PayloadExtractUXContainer(
     __in BURN_PAYLOADS* pPayloads,
-    __in_opt BURN_CONTAINER* pContainer,
     __in BURN_CONTAINER_CONTEXT* pContainerContext,
     __in_z LPCWSTR wzTargetDir
     )
@@ -215,7 +214,7 @@ extern "C" HRESULT PayloadExtractFromContainer(
         ExitOnFailure(hr, "Failed to get next stream.");
 
         // find payload by stream name
-        hr = FindEmbeddedBySourcePath(pPayloads, pContainer, sczStreamName, &pPayload);
+        hr = PayloadFindEmbeddedBySourcePath(pPayloads, sczStreamName, &pPayload);
         ExitOnFailure(hr, "Failed to find embedded payload: %ls", sczStreamName);
 
         // make file path
@@ -241,15 +240,11 @@ extern "C" HRESULT PayloadExtractFromContainer(
     {
         pPayload = &pPayloads->rgPayloads[i];
 
-        // if the payload is part of the container
-        if (!pContainer || pPayload->pContainer == pContainer)
+        // if the payload has not been acquired
+        if (BURN_PAYLOAD_STATE_ACQUIRED > pPayload->state)
         {
-            // if the payload has not been acquired
-            if (BURN_PAYLOAD_STATE_ACQUIRED > pPayload->state)
-            {
-                hr = E_INVALIDDATA;
-                ExitOnRootFailure(hr, "Payload was not found in container: %ls", pPayload->sczKey);
-            }
+            hr = E_INVALIDDATA;
+            ExitOnRootFailure(hr, "Payload was not found in container: %ls", pPayload->sczKey);
         }
     }
 
@@ -317,32 +312,3 @@ LExit:
 
 
 // internal function definitions
-
-static HRESULT FindEmbeddedBySourcePath(
-    __in BURN_PAYLOADS* pPayloads,
-    __in_opt BURN_CONTAINER* pContainer,
-    __in_z LPCWSTR wzStreamName,
-    __out BURN_PAYLOAD** ppPayload
-    )
-{
-    HRESULT hr = S_OK;
-
-    for (DWORD i = 0; i < pPayloads->cPayloads; ++i)
-    {
-        BURN_PAYLOAD* pPayload = &pPayloads->rgPayloads[i];
-
-        if (BURN_PAYLOAD_PACKAGING_EMBEDDED == pPayload->packaging && (!pContainer || pPayload->pContainer == pContainer))
-        {
-            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, pPayload->sczSourcePath, -1, wzStreamName, -1))
-            {
-                *ppPayload = pPayload;
-                ExitFunction1(hr = S_OK);
-            }
-        }
-    }
-
-    hr = E_NOTFOUND;
-
-LExit:
-    return hr;
-}
diff --git a/src/engine/payload.h b/src/engine/payload.h
index ad10509d..f28b437f 100644
--- a/src/engine/payload.h
+++ b/src/engine/payload.h
@@ -85,9 +85,8 @@ void PayloadUninitialize(
 void PayloadsUninitialize(
     __in BURN_PAYLOADS* pPayloads
     );
-HRESULT PayloadExtractFromContainer(
+HRESULT PayloadExtractUXContainer(
     __in BURN_PAYLOADS* pPayloads,
-    __in_opt BURN_CONTAINER* pContainer,
     __in BURN_CONTAINER_CONTEXT* pContainerContext,
     __in_z LPCWSTR wzTargetDir
     );