5 files changed, 116 insertions, 9 deletions

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index 7e03ebf9..dffbc6d6 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -768,6 +768,12 @@ static HRESULT ApplyCachePackage(
     HRESULT hr = S_OK;
     BOOTSTRAPPER_CACHEPACKAGECOMPLETE_ACTION cachePackageCompleteAction = BOOTSTRAPPER_CACHEPACKAGECOMPLETE_ACTION_NONE;
 
+    if (!pPackage->sczCacheFolder && !pContext->wzLayoutDirectory)
+    {
+        hr = CacheGetCompletedPath(pPackage->fPerMachine, pPackage->sczCacheId, &pPackage->sczCacheFolder);
+        ExitOnFailure(hr, "Failed to get cached path for package with cache id: %ls", pPackage->sczCacheId);
+    }
+
     for (;;)
     {
         hr = UserExperienceOnCachePackageBegin(pContext->pUX, pPackage->sczId, pPackage->payloads.cPayloads, pPackage->payloads.qwTotalSize);
@@ -872,6 +878,14 @@ static HRESULT ApplyLayoutContainer(
 
     Assert(!pContainer->fAttached);
 
+    hr = CacheVerifyContainer(pContainer, pContext->wzLayoutDirectory);
+    if (SUCCEEDED(hr))
+    {
+        // TODO: send Acquire and Verify messages to BA?
+        pContext->qwSuccessfulCacheProgress += pContainer->qwFileSize * 2;
+        ExitFunction();
+    }
+
     for (;;)
     {
         fRetry = FALSE;
@@ -916,7 +930,20 @@ static HRESULT ApplyProcessPayload(
     DWORD cTryAgainAttempts = 0;
     BOOL fRetry = FALSE;
 
-    Assert(pContext->pPayloads || pContext->wzLayoutDirectory);
+    Assert(pContext->pPayloads && pPackage || pContext->wzLayoutDirectory);
+
+    if (pPayload->pContainer && pContext->wzLayoutDirectory)
+    {
+        ExitFunction();
+    }
+
+    hr = CacheVerifyPayload(pPayload, pContext->wzLayoutDirectory ? pContext->wzLayoutDirectory : pPackage->sczCacheFolder);
+    if (SUCCEEDED(hr))
+    {
+        // TODO: send Acquire and Verify messages to BA?
+        pContext->qwSuccessfulCacheProgress += pPayload->qwFileSize * 2;
+        ExitFunction();
+    }
 
     for (;;)
     {
@@ -924,12 +951,7 @@ static HRESULT ApplyProcessPayload(
 
         if (pPayload->pContainer)
         {
-            if (pContext->wzLayoutDirectory)
-            {
-                ExitFunction();
-            }
-
-            // TODO: only extract container if payload isn't already cached and isn't already extracted
+            // TODO: only extract container if payload isn't already extracted
             hr = ApplyExtractContainer(pContext, pPayload->pContainer);
             ExitOnFailure(hr, "Failed to extract container for payload: %ls", pPayload->sczKey);
         }
diff --git a/src/engine/cache.cpp b/src/engine/cache.cpp
index 6ddbfb50..6ac313e0 100644
--- a/src/engine/cache.cpp
+++ b/src/engine/cache.cpp
@@ -57,6 +57,10 @@ static HRESULT TransferWorkingPathToUnverifiedPath(
     __in_z LPCWSTR wzUnverifiedPayloadPath,
     __in BOOL fMove
     );
+static HRESULT VerifyFileAgainstContainer(
+    __in BURN_CONTAINER* pContainer,
+    __in_z LPCWSTR wzVerifyPath
+    );
 static HRESULT VerifyFileAgainstPayload(
     __in BURN_PAYLOAD* pPayload,
     __in_z LPCWSTR wzVerifyPath
@@ -837,7 +841,7 @@ LExit:
 extern "C" HRESULT CacheCompletePayload(
     __in BOOL fPerMachine,
     __in BURN_PAYLOAD* pPayload,
-    __in_z_opt LPCWSTR wzCacheId,
+    __in_z LPCWSTR wzCacheId,
     __in_z LPCWSTR wzWorkingPayloadPath,
     __in BOOL fMove
     )
@@ -910,6 +914,44 @@ LExit:
     return hr;
 }
 
+extern "C" HRESULT CacheVerifyContainer(
+    __in BURN_CONTAINER* pContainer,
+    __in_z LPCWSTR wzCachedDirectory
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczCachedPath = NULL;
+
+    hr = PathConcat(wzCachedDirectory, pContainer->sczFilePath, &sczCachedPath);
+    ExitOnFailure(hr, "Failed to concat complete cached path.");
+
+    hr = VerifyFileAgainstContainer(pContainer, sczCachedPath);
+
+LExit:
+    ReleaseStr(sczCachedPath);
+
+    return hr;
+}
+
+extern "C" HRESULT CacheVerifyPayload(
+    __in BURN_PAYLOAD* pPayload,
+    __in_z LPCWSTR wzCachedDirectory
+    )
+{
+    HRESULT hr = S_OK;
+    LPWSTR sczCachedPath = NULL;
+
+    hr = PathConcat(wzCachedDirectory, pPayload->sczFilePath, &sczCachedPath);
+    ExitOnFailure(hr, "Failed to concat complete cached path.");
+
+    hr = VerifyFileAgainstPayload(pPayload, sczCachedPath);
+
+LExit:
+    ReleaseStr(sczCachedPath);
+
+    return hr;
+}
+
 extern "C" HRESULT CacheRemoveWorkingFolder(
     __in_z_opt LPCWSTR wzBundleId
     )
@@ -1383,6 +1425,38 @@ LExit:
     return hr;
 }
 
+static HRESULT VerifyFileAgainstContainer(
+    __in BURN_CONTAINER* pContainer,
+    __in_z LPCWSTR wzVerifyPath
+    )
+{
+    HRESULT hr = S_OK;
+    HANDLE hFile = INVALID_HANDLE_VALUE;
+
+    // Get the container on disk actual hash.
+    hFile = ::CreateFileW(wzVerifyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_DELETE, NULL, OPEN_EXISTING, FILE_FLAG_SEQUENTIAL_SCAN, NULL);
+    if (INVALID_HANDLE_VALUE == hFile)
+    {
+        hr = HRESULT_FROM_WIN32(::GetLastError());
+        if (E_PATHNOTFOUND == hr || E_FILENOTFOUND == hr)
+        {
+            ExitFunction(); // do not log error when the file was not found.
+        }
+        ExitOnRootFailure(hr, "Failed to open container at path: %ls", wzVerifyPath);
+    }
+
+    if (pContainer->pbHash) // the container should have a hash we can use to verify it.
+    {
+        hr = VerifyHash(pContainer->pbHash, pContainer->cbHash, pContainer->qwFileSize, wzVerifyPath, hFile);
+        ExitOnFailure(hr, "Failed to verify hash of container: %ls", pContainer->sczId);
+    }
+
+LExit:
+    ReleaseFileHandle(hFile);
+
+    return hr;
+}
+
 static HRESULT VerifyFileAgainstPayload(
     __in BURN_PAYLOAD* pPayload,
     __in_z LPCWSTR wzVerifyPath
diff --git a/src/engine/cache.h b/src/engine/cache.h
index 52b111e9..acc7acb7 100644
--- a/src/engine/cache.h
+++ b/src/engine/cache.h
@@ -119,10 +119,18 @@ HRESULT CacheLayoutPayload(
 HRESULT CacheCompletePayload(
     __in BOOL fPerMachine,
     __in BURN_PAYLOAD* pPayload,
-    __in_z_opt LPCWSTR wzCacheId,
+    __in_z LPCWSTR wzCacheId,
     __in_z LPCWSTR wzUnverifiedPayloadPath,
     __in BOOL fMove
     );
+HRESULT CacheVerifyContainer(
+    __in BURN_CONTAINER* pContainer,
+    __in_z LPCWSTR wzCachedDirectory
+    );
+HRESULT CacheVerifyPayload(
+    __in BURN_PAYLOAD* pPayload,
+    __in_z LPCWSTR wzCachedDirectory
+    );
 HRESULT CacheRemoveWorkingFolder(
     __in_z_opt LPCWSTR wzBundleId
     );
diff --git a/src/engine/package.h b/src/engine/package.h
index 34a3af26..2091af94 100644
--- a/src/engine/package.h
+++ b/src/engine/package.h
@@ -246,6 +246,7 @@ typedef struct _BURN_PACKAGE
     BURN_DEPENDENCY_ACTION dependencyExecute;   // only valid during Plan.
     BURN_DEPENDENCY_ACTION dependencyRollback;  // only valid during Plan.
     BOOL fDependencyManagerWasHere;             // only valid during Plan.
+    LPWSTR sczCacheFolder;                      // only valid during Apply.
     HRESULT hrCacheResult;                      // only valid during Apply.
 
     BURN_PACKAGE_REGISTRATION_STATE cacheRegistrationState;          // initialized during Detect, updated during Apply.
diff --git a/src/engine/plan.cpp b/src/engine/plan.cpp
index ce577da5..cfe4893b 100644
--- a/src/engine/plan.cpp
+++ b/src/engine/plan.cpp
@@ -1863,6 +1863,8 @@ static void ResetPlannedPackageState(
     pPackage->expectedCacheRegistrationState = BURN_PACKAGE_REGISTRATION_STATE_UNKNOWN;
     pPackage->expectedInstallRegistrationState = BURN_PACKAGE_REGISTRATION_STATE_UNKNOWN;
 
+    ReleaseNullStr(pPackage->sczCacheFolder);
+
     if (BURN_PACKAGE_TYPE_MSI == pPackage->type)
     {
         for (DWORD i = 0; i < pPackage->Msi.cFeatures; ++i)