Allow setting source from OnCacheAcquireResolving.

6 files changed, 92 insertions, 43 deletions

diff --git a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index d8994c26..603df890 100644
--- a/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -64,6 +64,20 @@ enum BOOTSTRAPPER_CACHE_OPERATION
     BOOTSTRAPPER_CACHE_OPERATION_EXTRACT,
 };
 
+enum BOOTSTRAPPER_CACHE_RESOLVE_OPERATION
+{
+    // There is no source available.
+    BOOTSTRAPPER_CACHE_RESOLVE_NONE,
+    // Copy the payload or container from the chosen local source.
+    BOOTSTRAPPER_CACHE_RESOLVE_LOCAL,
+    // Download the payload or container from the download URL.
+    BOOTSTRAPPER_CACHE_RESOLVE_DOWNLOAD,
+    // Extract the payload from the container.
+    BOOTSTRAPPER_CACHE_RESOLVE_CONTAINER,
+    // Look again for the payload or container locally.
+    BOOTSTRAPPER_CACHE_RESOLVE_RETRY,
+};
+
 enum BOOTSTRAPPER_CACHE_VERIFY_STEP
 {
     BOOTSTRAPPER_CACHE_VERIFY_STEP_STAGE,
@@ -379,14 +393,14 @@ struct BA_ONCACHEACQUIRERESOLVING_ARGS
     DWORD dwRecommendedSearchPath;
     LPCWSTR wzDownloadUrl;
     LPCWSTR wzPayloadContainerId;
-    BOOTSTRAPPER_CACHE_OPERATION recommendation;
+    BOOTSTRAPPER_CACHE_RESOLVE_OPERATION recommendation;
 };
 
 struct BA_ONCACHEACQUIRERESOLVING_RESULTS
 {
     DWORD cbSize;
     DWORD dwChosenSearchPath;
-    BOOTSTRAPPER_CACHE_OPERATION action;
+    BOOTSTRAPPER_CACHE_RESOLVE_OPERATION action;
     BOOL fCancel;
 };
 
diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index 1bc01d44..9cba0483 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -1390,6 +1390,7 @@ static HRESULT AcquireContainerOrPayload(
     LPCWSTR wzRelativePath = pContainer ? pContainer->sczFilePath : pPayload->sczFilePath;
     DWORD dwChosenSearchPath = 0;
     BOOTSTRAPPER_CACHE_OPERATION cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_NONE;
+    BOOTSTRAPPER_CACHE_RESOLVE_OPERATION resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_NONE;
     LPWSTR* pwzDownloadUrl = pContainer ? &pContainer->downloadSource.sczUrl : &pPayload->downloadSource.sczUrl;
     LPWSTR* pwzSourcePath = pContainer ? &pContainer->sczSourcePath : &pPayload->sczSourcePath;
     BOOL fFoundLocal = FALSE;
@@ -1405,47 +1406,70 @@ static HRESULT AcquireContainerOrPayload(
     if (BOOTSTRAPPER_CACHE_OPERATION_DOWNLOAD != cacheOperation &&
         BOOTSTRAPPER_CACHE_OPERATION_EXTRACT != cacheOperation)
     {
-        hr = CacheGetLocalSourcePaths(wzRelativePath, *pwzSourcePath, wzDestinationPath, pContext->wzLayoutDirectory, pContext->pVariables, &pContext->rgSearchPaths, &pContext->cSearchPaths);
-        ExitOnFailure(hr, "Failed to search local source.");
-
-        for (DWORD i = 0; i < pContext->cSearchPaths; ++i)
+        do
         {
-            // If the file exists locally, choose it.
-            if (FileExistsEx(pContext->rgSearchPaths[i], NULL))
-            {
-                dwChosenSearchPath = i;
+            fFoundLocal = FALSE;
+            resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_NONE;
+            dwChosenSearchPath = 0;
 
-                fFoundLocal = TRUE;
-                break;
-            }
-        }
+            hr = CacheGetLocalSourcePaths(wzRelativePath, *pwzSourcePath, wzDestinationPath, pContext->wzLayoutDirectory, pContext->pVariables, &pContext->rgSearchPaths, &pContext->cSearchPaths, &dwChosenSearchPath);
+            ExitOnFailure(hr, "Failed to search local source.");
 
-        if (BOOTSTRAPPER_CACHE_OPERATION_COPY == cacheOperation)
-        {
-            if (!fFoundLocal)
+            for (DWORD i = 0; i < pContext->cSearchPaths; ++i)
             {
-                cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_NONE;
+                // If the file exists locally, choose it.
+                if (FileExistsEx(pContext->rgSearchPaths[i], NULL))
+                {
+                    dwChosenSearchPath = i;
+
+                    fFoundLocal = TRUE;
+                    break;
+                }
             }
-        }
-        else
-        {
-            if (fFoundLocal) // the file exists locally, so copy it.
+
+            if (BOOTSTRAPPER_CACHE_OPERATION_COPY == cacheOperation)
             {
-                cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_COPY;
+                if (fFoundLocal)
+                {
+                    resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_LOCAL;
+                }
             }
-            else if (wzPayloadContainerId)
+            else
             {
-                cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_EXTRACT;
+                if (fFoundLocal) // the file exists locally, so copy it.
+                {
+                    resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_LOCAL;
+                }
+                else if (wzPayloadContainerId)
+                {
+                    resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_CONTAINER;
+                }
+                else if (*pwzDownloadUrl && **pwzDownloadUrl)
+                {
+                    resolveOperation = BOOTSTRAPPER_CACHE_RESOLVE_DOWNLOAD;
+                }
             }
-            else if (*pwzDownloadUrl && **pwzDownloadUrl)
+
+            // Let the BA have a chance to override the source.
+            hr = UserExperienceOnCacheAcquireResolving(pContext->pUX, wzPackageOrContainerId, wzPayloadId, pContext->rgSearchPaths, pContext->cSearchPaths, fFoundLocal, &dwChosenSearchPath, pwzDownloadUrl, wzPayloadContainerId, &resolveOperation);
+            ExitOnRootFailure(hr, "BA aborted cache acquire resolving.");
+
+            switch (resolveOperation)
             {
+            case BOOTSTRAPPER_CACHE_RESOLVE_LOCAL:
+                cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_COPY;
+                break;
+            case BOOTSTRAPPER_CACHE_RESOLVE_DOWNLOAD:
                 cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_DOWNLOAD;
+                break;
+            case BOOTSTRAPPER_CACHE_RESOLVE_CONTAINER:
+                cacheOperation = BOOTSTRAPPER_CACHE_OPERATION_EXTRACT;
+                break;
+            case BOOTSTRAPPER_CACHE_RESOLVE_RETRY:
+                pContext->cSearchPathsMax = max(pContext->cSearchPaths, pContext->cSearchPathsMax);
+                break;
             }
-        }
-
-        // Let the BA have a chance to override the action, but their chance to change the source is during begin or complete.
-        hr = UserExperienceOnCacheAcquireResolving(pContext->pUX, wzPackageOrContainerId, wzPayloadId, pContext->rgSearchPaths, pContext->cSearchPaths, fFoundLocal, &dwChosenSearchPath, *pwzDownloadUrl, wzPayloadContainerId, &cacheOperation);
-        ExitOnRootFailure(hr, "BA aborted cache acquire resolving.");
+        } while (BOOTSTRAPPER_CACHE_RESOLVE_RETRY == resolveOperation);
     }
 
     switch (cacheOperation)
diff --git a/src/engine/cache.cpp b/src/engine/cache.cpp
index 16f2c1e5..764de065 100644
--- a/src/engine/cache.cpp
+++ b/src/engine/cache.cpp
@@ -440,7 +440,8 @@ extern "C" HRESULT CacheGetLocalSourcePaths(
     __in_z_opt LPCWSTR wzLayoutDirectory,
     __in BURN_VARIABLES* pVariables,
     __inout LPWSTR** prgSearchPaths,
-    __out DWORD* pcSearchPaths
+    __out DWORD* pcSearchPaths,
+    __out DWORD* pdwLikelySearchPath
     )
 {
     HRESULT hr = S_OK;
@@ -452,6 +453,7 @@ extern "C" HRESULT CacheGetLocalSourcePaths(
     BOOL fTryRelativePath = FALSE;
     BOOL fSourceIsAbsolute = FALSE;
     DWORD cSearchPaths = 0;
+    DWORD dwLikelySearchPath = 0;
 
     AssertSz(vfInitializedCache, "Cache wasn't initialized");
 
@@ -473,6 +475,12 @@ extern "C" HRESULT CacheGetLocalSourcePaths(
         hr = StrAllocString(psczPath, wzSourcePath, 0);
         ExitOnFailure(hr, "Failed to copy absolute source path.");
     }
+    else
+    {
+        // If none of the paths exist, then most BAs will want to prompt the user with a possible path.
+        // The destination path is a temporary location and so not really a possible path.
+        dwLikelySearchPath = 1;
+    }
 
     // Try the destination path next.
     hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(prgSearchPaths), cSearchPaths + 1, sizeof(LPWSTR), BURN_CACHE_MAX_SEARCH_PATHS);
@@ -593,6 +601,7 @@ LExit:
 
     AssertSz(cSearchPaths <= BURN_CACHE_MAX_SEARCH_PATHS, "Got more than BURN_CACHE_MAX_SEARCH_PATHS search paths");
     *pcSearchPaths = cSearchPaths;
+    *pdwLikelySearchPath = dwLikelySearchPath;
 
     return hr;
 }
diff --git a/src/engine/cache.h b/src/engine/cache.h
index cd964649..afa18e47 100644
--- a/src/engine/cache.h
+++ b/src/engine/cache.h
@@ -101,7 +101,8 @@ HRESULT CacheGetLocalSourcePaths(
     __in_z_opt LPCWSTR wzLayoutDirectory,
     __in BURN_VARIABLES* pVariables,
     __inout LPWSTR** prgSearchPaths,
-    __out DWORD* pcSearchPaths
+    __out DWORD* pcSearchPaths,
+    __out DWORD* pdwLikelySearchPath
     );
 HRESULT CacheSetLastUsedSource(
     __in BURN_VARIABLES* pVariables,
diff --git a/src/engine/userexperience.cpp b/src/engine/userexperience.cpp
index 8c6f29f3..7e68d664 100644
--- a/src/engine/userexperience.cpp
+++ b/src/engine/userexperience.cpp
@@ -515,9 +515,9 @@ EXTERN_C BAAPI UserExperienceOnCacheAcquireResolving(
     __in DWORD cSearchPaths,
     __in BOOL fFoundLocal,
     __in DWORD* pdwChosenSearchPath,
-    __in_z_opt LPCWSTR wzDownloadUrl,
+    __in_z_opt LPWSTR* pwzDownloadUrl,
     __in_z_opt LPCWSTR wzPayloadContainerId,
-    __inout BOOTSTRAPPER_CACHE_OPERATION* pCacheOperation
+    __inout BOOTSTRAPPER_CACHE_RESOLVE_OPERATION* pCacheOperation
     )
 {
     HRESULT hr = S_OK;
@@ -531,14 +531,14 @@ EXTERN_C BAAPI UserExperienceOnCacheAcquireResolving(
     args.cSearchPaths = cSearchPaths;
     args.fFoundLocal = fFoundLocal;
     args.dwRecommendedSearchPath = *pdwChosenSearchPath;
-    args.wzDownloadUrl = wzDownloadUrl;
+    args.wzDownloadUrl = *pwzDownloadUrl;
     args.recommendation = *pCacheOperation;
 
     results.cbSize = sizeof(results);
     results.dwChosenSearchPath = *pdwChosenSearchPath;
     results.action = *pCacheOperation;
 
-    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIRERESOLVING, &args, &results);
+    hr = SendBAMessageFromInactiveEngine(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIRERESOLVING, &args, &results);
     ExitOnFailure(hr, "BA OnCacheAcquireResolving failed.");
 
     if (results.fCancel)
@@ -548,13 +548,14 @@ EXTERN_C BAAPI UserExperienceOnCacheAcquireResolving(
     else
     {
         // Verify the BA requested an action that is possible.
-        if (BOOTSTRAPPER_CACHE_OPERATION_DOWNLOAD == results.action && wzDownloadUrl && *wzDownloadUrl ||
-            BOOTSTRAPPER_CACHE_OPERATION_EXTRACT == results.action && wzPayloadContainerId ||
-            BOOTSTRAPPER_CACHE_OPERATION_NONE == results.action)
+        if (BOOTSTRAPPER_CACHE_RESOLVE_DOWNLOAD == results.action && *pwzDownloadUrl && **pwzDownloadUrl ||
+            BOOTSTRAPPER_CACHE_RESOLVE_CONTAINER == results.action && wzPayloadContainerId ||
+            BOOTSTRAPPER_CACHE_RESOLVE_RETRY == results.action ||
+            BOOTSTRAPPER_CACHE_RESOLVE_NONE == results.action)
         {
             *pCacheOperation = results.action;
         }
-        else if (BOOTSTRAPPER_CACHE_OPERATION_COPY == results.action && results.dwChosenSearchPath < cSearchPaths)
+        else if (BOOTSTRAPPER_CACHE_RESOLVE_LOCAL == results.action && results.dwChosenSearchPath < cSearchPaths)
         {
             *pdwChosenSearchPath = results.dwChosenSearchPath;
             *pCacheOperation = results.action;
diff --git a/src/engine/userexperience.h b/src/engine/userexperience.h
index 58e0431e..e1041624 100644
--- a/src/engine/userexperience.h
+++ b/src/engine/userexperience.h
@@ -154,9 +154,9 @@ BAAPI UserExperienceOnCacheAcquireResolving(
     __in DWORD cSearchPaths,
     __in BOOL fFoundLocal,
     __in DWORD* pdwChosenSearchPath,
-    __in_z_opt LPCWSTR wzDownloadUrl,
+    __in_z_opt LPWSTR* pwzDownloadUrl,
     __in_z_opt LPCWSTR wzPayloadContainerId,
-    __inout BOOTSTRAPPER_CACHE_OPERATION* pCacheOperation
+    __inout BOOTSTRAPPER_CACHE_RESOLVE_OPERATION* pCacheOperation
     );
 BAAPI UserExperienceOnCacheBegin(
     __in BURN_USER_EXPERIENCE* pUserExperience