Reset the package's cache folder only once.

This behavior makes it work similar to the root package cache folder and avoids performance problems when there are many payloads in the package.

#4774

6 files changed, 145 insertions, 24 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 58d41b28..f0b78e81 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -549,12 +549,19 @@ extern "C" HRESULT ApplyCache(
         case BURN_CACHE_ACTION_TYPE_PACKAGE:
             pPackage = pCacheAction->package.pPackage;
 
-            if (!pPackage->fPerMachine && !cacheContext.wzLayoutDirectory)
+            if (!cacheContext.wzLayoutDirectory)
             {
-                hr = CacheGetCompletedPath(FALSE, pPackage->sczCacheId, &pPackage->sczCacheFolder);
-                ExitOnFailure(hr, "Failed to get cached path for package with cache id: %ls", pPackage->sczCacheId);
+                if (!pPackage->fPerMachine || INVALID_HANDLE_VALUE == cacheContext.hPipe)
+                {
+                    hr = CachePreparePackage(pPackage);
 
-                cacheContext.hPipe = INVALID_HANDLE_VALUE;
+                    cacheContext.hPipe = INVALID_HANDLE_VALUE;
+                }
+                else
+                {
+                    hr = ElevationCachePreparePackage(hPipe, pPackage);
+                }
+                LogExitOnFailure(hr, MSG_CACHE_PREPARE_PACKAGE_FAILED, "Cache prepare package failed: %ls", pPackage->sczId, NULL, NULL);
             }
 
             hr = ApplyCachePackage(&cacheContext, pPackage);
diff --git a/src/burn/engine/cache.cpp b/src/burn/engine/cache.cpp
index d1999a0d..02cad4a5 100644
--- a/src/burn/engine/cache.cpp
+++ b/src/burn/engine/cache.cpp
@@ -29,7 +29,8 @@ static HRESULT SecurePerMachineCacheRoot();
 static HRESULT CreateCompletedPath(
     __in BOOL fPerMachine,
     __in LPCWSTR wzCacheId,
-    __out_z LPWSTR* psczCacheDirectory
+    __in LPCWSTR wzFilePath,
+    __out_z LPWSTR* psczCachePath
     );
 static HRESULT CreateUnverifiedPath(
     __in BOOL fPerMachine,
@@ -750,6 +751,20 @@ extern "C" BOOL CacheBundleRunningFromCache()
     return vfRunningFromCache;
 }
 
+HRESULT CachePreparePackage(
+    __in BURN_PACKAGE* pPackage
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (!pPackage->sczCacheFolder)
+    {
+        hr = CreateCompletedPath(pPackage->fPerMachine, pPackage->sczCacheId, NULL, &pPackage->sczCacheFolder);
+    }
+
+    return hr;
+}
+
 extern "C" HRESULT CacheBundleToCleanRoom(
     __in BURN_SECTION* pSection,
     __deref_out_z_opt LPWSTR* psczCleanRoomBundlePath
@@ -852,7 +867,7 @@ extern "C" HRESULT CacheCompleteBundle(
     LPWSTR sczSourceDirectory = NULL;
     LPWSTR sczPayloadSourcePath = NULL;
 
-    hr = CreateCompletedPath(fPerMachine, wzBundleId, &sczTargetDirectory);
+    hr = CreateCompletedPath(fPerMachine, wzBundleId, NULL, &sczTargetDirectory);
     ExitOnFailure(hr, "Failed to create completed cache path for bundle.");
 
     hr = PathConcat(sczTargetDirectory, wzExecutableName, &sczTargetPath);
@@ -957,16 +972,12 @@ extern "C" HRESULT CacheCompletePayload(
     )
 {
     HRESULT hr = S_OK;
-    LPWSTR sczCachedDirectory = NULL;
     LPWSTR sczCachedPath = NULL;
     LPWSTR sczUnverifiedPayloadPath = NULL;
 
-    hr = CreateCompletedPath(fPerMachine, wzCacheId, &sczCachedDirectory);
+    hr = CreateCompletedPath(fPerMachine, wzCacheId, pPayload->sczFilePath, &sczCachedPath);
     ExitOnFailure(hr, "Failed to get cached path for package with cache id: %ls", wzCacheId);
 
-    hr = PathConcat(sczCachedDirectory, pPayload->sczFilePath, &sczCachedPath);
-    ExitOnFailure(hr, "Failed to concat complete cached path.");
-
     // If the cached file matches what we expected, we're good.
     hr = VerifyFileAgainstPayload(pPayload, sczCachedPath, TRUE, BURN_CACHE_STEP_HASH_TO_SKIP_VERIFY, pfnCacheMessageHandler, pfnProgress, pContext);
     if (SUCCEEDED(hr))
@@ -1016,7 +1027,6 @@ extern "C" HRESULT CacheCompletePayload(
 LExit:
     ReleaseStr(sczUnverifiedPayloadPath);
     ReleaseStr(sczCachedPath);
-    ReleaseStr(sczCachedDirectory);
 
     return hr;
 }
@@ -1403,11 +1413,13 @@ LExit:
 static HRESULT CreateCompletedPath(
     __in BOOL fPerMachine,
     __in LPCWSTR wzId,
-    __out_z LPWSTR* psczCacheDirectory
+    __in LPCWSTR wzFilePath,
+    __out_z LPWSTR* psczCachePath
     )
 {
     HRESULT hr = S_OK;
     LPWSTR sczCacheDirectory = NULL;
+    LPWSTR sczCacheFile = NULL;
 
     if (fPerMachine)
     {
@@ -1415,22 +1427,38 @@ static HRESULT CreateCompletedPath(
         ExitOnFailure(hr, "Failed to secure per-machine cache root.");
     }
 
-    // Get the cache completed path, ensure it exists, and reset any permissions people
-    // might have tried to set on the directory so we inherit the (correct!) security
-    // permissions from the parent directory.
+    // Get the cache completed path.
     hr = CacheGetCompletedPath(fPerMachine, wzId, &sczCacheDirectory);
     ExitOnFailure(hr, "Failed to get cache directory.");
 
+    // Ensure it exists.
     hr = DirEnsureExists(sczCacheDirectory, NULL);
     ExitOnFailure(hr, "Failed to create cache directory: %ls", sczCacheDirectory);
 
-    ResetPathPermissions(fPerMachine, sczCacheDirectory);
+    if (!wzFilePath)
+    {
+        // Reset any permissions people might have tried to set on the directory
+        // so we inherit the (correct!) security permissions from the parent directory.
+        ResetPathPermissions(fPerMachine, sczCacheDirectory);
 
-    *psczCacheDirectory = sczCacheDirectory;
-    sczCacheDirectory = NULL;
+        *psczCachePath = sczCacheDirectory;
+        sczCacheDirectory = NULL;
+    }
+    else
+    {
+        // Get the cache completed file path.
+        hr = PathConcat(sczCacheDirectory, wzFilePath, &sczCacheFile);
+        ExitOnFailure(hr, "Failed to construct cache file.");
+
+        // Don't reset permissions here. The payload's package must reset its cache folder when it starts caching.
+
+        *psczCachePath = sczCacheFile;
+        sczCacheFile = NULL;
+    }
 
 LExit:
     ReleaseStr(sczCacheDirectory);
+    ReleaseStr(sczCacheFile);
     return hr;
 }
 
diff --git a/src/burn/engine/cache.h b/src/burn/engine/cache.h
index a300e99d..dfb4f69f 100644
--- a/src/burn/engine/cache.h
+++ b/src/burn/engine/cache.h
@@ -122,6 +122,9 @@ void CacheSendErrorCallback(
     __out_opt BOOL* pfRetry
     );
 BOOL CacheBundleRunningFromCache();
+HRESULT CachePreparePackage(
+    __in BURN_PACKAGE* pPackage
+    );
 HRESULT CacheBundleToCleanRoom(
     __in BURN_SECTION* pSection,
     __deref_out_z_opt LPWSTR* psczCleanRoomBundlePath
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 9d1b8fc7..85e4b692 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -14,6 +14,7 @@ typedef enum _BURN_ELEVATION_MESSAGE_TYPE
     BURN_ELEVATION_MESSAGE_TYPE_SESSION_RESUME,
     BURN_ELEVATION_MESSAGE_TYPE_SESSION_END,
     BURN_ELEVATION_MESSAGE_TYPE_SAVE_STATE,
+    BURN_ELEVATION_MESSAGE_TYPE_CACHE_PREPARE_PACKAGE,
     BURN_ELEVATION_MESSAGE_TYPE_CACHE_COMPLETE_PAYLOAD,
     BURN_ELEVATION_MESSAGE_TYPE_CACHE_VERIFY_PAYLOAD,
     BURN_ELEVATION_MESSAGE_TYPE_CACHE_CLEANUP,
@@ -186,6 +187,11 @@ static HRESULT OnSaveState(
     __in BYTE* pbData,
     __in SIZE_T cbData
     );
+static HRESULT OnCachePreparePackage(
+    __in BURN_PACKAGES* pPackages,
+    __in BYTE* pbData,
+    __in SIZE_T cbData
+    );
 static HRESULT OnCacheCompletePayload(
     __in HANDLE hPipe,
     __in BURN_PACKAGES* pPackages,
@@ -612,6 +618,32 @@ LExit:
     return hr;
 }
 
+extern "C" HRESULT ElevationCachePreparePackage(
+    __in HANDLE hPipe,
+    __in BURN_PACKAGE* pPackage
+    )
+{
+    HRESULT hr = S_OK;
+    BYTE* pbData = NULL;
+    SIZE_T cbData = 0;
+    DWORD dwResult = 0;
+
+    // Serialize message data.
+    hr = BuffWriteString(&pbData, &cbData, pPackage ? pPackage->sczId : NULL);
+    ExitOnFailure(hr, "Failed to write package id to message buffer.");
+
+    // Send message.
+    hr = PipeSendMessage(hPipe, BURN_ELEVATION_MESSAGE_TYPE_CACHE_PREPARE_PACKAGE, pbData, cbData, NULL, NULL, &dwResult);
+    ExitOnFailure(hr, "Failed to send BURN_ELEVATION_MESSAGE_TYPE_CACHE_PREPARE_PACKAGE message to per-machine process.");
+
+    hr = (HRESULT)dwResult;
+
+LExit:
+    ReleaseBuffer(pbData);
+
+    return hr;
+}
+
 /*******************************************************************
  ElevationCacheCompletePayload - 
 
@@ -1856,6 +1888,10 @@ static HRESULT ProcessElevatedChildCacheMessage(
 
     switch (pMsg->dwMessage)
     {
+    case BURN_ELEVATION_MESSAGE_TYPE_CACHE_PREPARE_PACKAGE:
+        hrResult = OnCachePreparePackage(pContext->pPackages, (BYTE*)pMsg->pvData, pMsg->cbData);
+        break;
+
     case BURN_ELEVATION_MESSAGE_TYPE_CACHE_COMPLETE_PAYLOAD:
         hrResult = OnCacheCompletePayload(pContext->hPipe, pContext->pPackages, pContext->pPayloads, (BYTE*)pMsg->pvData, pMsg->cbData);
         break;
@@ -2153,6 +2189,41 @@ LExit:
     return hr;
 }
 
+static HRESULT OnCachePreparePackage(
+    __in BURN_PACKAGES* pPackages,
+    __in BYTE* pbData,
+    __in SIZE_T cbData
+    )
+{
+    HRESULT hr = S_OK;
+    SIZE_T iData = 0;
+    LPWSTR scz = NULL;
+    BURN_PACKAGE* pPackage = NULL;
+
+    // Deserialize message data.
+    hr = BuffReadString(pbData, cbData, &iData, &scz);
+    ExitOnFailure(hr, "Failed to read package id.");
+
+    if (scz && *scz)
+    {
+        hr = PackageFindById(pPackages, scz, &pPackage);
+        ExitOnFailure(hr, "Failed to find package: %ls", scz);
+    }
+    else
+    {
+        hr = E_INVALIDARG;
+        ExitOnRootFailure(hr, "Invalid data passed to cache prepare package.");
+    }
+
+    hr = CachePreparePackage(pPackage);
+    ExitOnFailure(hr, "Failed to prepare cache package.");
+
+LExit:
+    ReleaseStr(scz);
+
+    return hr;
+}
+
 static HRESULT OnCacheCompletePayload(
     __in HANDLE hPipe,
     __in BURN_PACKAGES* pPackages,
@@ -2225,7 +2296,6 @@ static HRESULT OnCacheVerifyPayload(
     LPWSTR scz = NULL;
     BURN_PACKAGE* pPackage = NULL;
     BURN_PAYLOAD* pPayload = NULL;
-    LPWSTR sczCacheDirectory = NULL;
 
     // Deserialize message data.
     hr = BuffReadString(pbData, cbData, &iData, &scz);
@@ -2248,10 +2318,13 @@ static HRESULT OnCacheVerifyPayload(
 
     if (pPackage && pPayload)
     {
-        hr = CacheGetCompletedPath(TRUE, pPackage->sczCacheId, &sczCacheDirectory);
-        ExitOnFailure(hr, "Failed to get cached path for package with cache id: %ls", pPackage->sczCacheId);
+        if (!pPackage->sczCacheFolder)
+        {
+            hr = E_INVALIDSTATE;
+            ExitOnRootFailure(hr, "Cache verify payload called without starting its package.");
+        }
 
-        hr = CacheVerifyPayload(pPayload, sczCacheDirectory, BurnCacheMessageHandler, ElevatedProgressRoutine, hPipe);
+        hr = CacheVerifyPayload(pPayload, pPackage->sczCacheFolder, BurnCacheMessageHandler, ElevatedProgressRoutine, hPipe);
     }
     else
     {
@@ -2261,7 +2334,6 @@ static HRESULT OnCacheVerifyPayload(
     // Nothing should be logged on failure.
 
 LExit:
-    ReleaseStr(sczCacheDirectory);
     ReleaseStr(scz);
 
     return hr;
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index 9244f36c..af2dec8b 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -50,6 +50,10 @@ HRESULT ElevationSaveState(
     __in_bcount(cbBuffer) BYTE* pbBuffer,
     __in SIZE_T cbBuffer
     );
+HRESULT ElevationCachePreparePackage(
+    __in HANDLE hPipe,
+    __in BURN_PACKAGE* pPackage
+    );
 HRESULT ElevationCacheCompletePayload(
     __in HANDLE hPipe,
     __in BURN_PACKAGE* pPackage,
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index 25d5b4e4..84046314 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -717,6 +717,13 @@ Language=English
 Acquiring container: %1!ls!, %3!hs! from: %4!ls!
 .
 
+MessageId=337
+Severity=Error
+SymbolicName=MSG_CACHE_PREPARE_PACKAGE_FAILED
+Language=English
+Failed to prepare package: %2!ls!, error: 0x%1!ls!
+.
+
 MessageId=338
 Severity=Success
 SymbolicName=MSG_ACQUIRE_PACKAGE_PAYLOAD