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