7 files changed, 107 insertions, 59 deletions

diff --git a/src/balutil/inc/BAFunctions.h b/src/balutil/inc/BAFunctions.h
index 66852efa..07f7a750 100644
--- a/src/balutil/inc/BAFunctions.h
+++ b/src/balutil/inc/BAFunctions.h
@@ -42,7 +42,7 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONCACHEPACKAGEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPACKAGEBEGIN,
     BA_FUNCTIONS_MESSAGE_ONCACHEACQUIREBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIREBEGIN,
     BA_FUNCTIONS_MESSAGE_ONCACHEACQUIREPROGRESS = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIREPROGRESS,
-    BA_FUNCTIONS_MESSAGE_ONRESOLVESOURCE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONRESOLVESOURCE,
+    BA_FUNCTIONS_MESSAGE_ONCACHEACQUIRERESOLVING = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIRERESOLVING,
     BA_FUNCTIONS_MESSAGE_ONCACHEACQUIRECOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIRECOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONCACHEVERIFYBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEVERIFYBEGIN,
     BA_FUNCTIONS_MESSAGE_ONCACHEVERIFYCOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEVERIFYCOMPLETE,
diff --git a/src/balutil/inc/BalBaseBAFunctions.h b/src/balutil/inc/BalBaseBAFunctions.h
index 0164269c..ca727f49 100644
--- a/src/balutil/inc/BalBaseBAFunctions.h
+++ b/src/balutil/inc/BalBaseBAFunctions.h
@@ -90,6 +90,7 @@ public: // IBootstrapperApplication
     }
 
     virtual STDMETHODIMP OnDetectBegin(
+        __in BOOL /*fCached*/,
         __in BOOL /*fInstalled*/,
         __in DWORD /*cPackages*/,
         __inout BOOL* /*pfCancel*/
@@ -385,8 +386,11 @@ public: // IBootstrapperApplication
     virtual STDMETHODIMP OnCacheAcquireBegin(
         __in_z LPCWSTR /*wzPackageOrContainerId*/,
         __in_z_opt LPCWSTR /*wzPayloadId*/,
-        __in BOOTSTRAPPER_CACHE_OPERATION /*operation*/,
         __in_z LPCWSTR /*wzSource*/,
+        __in_z_opt LPCWSTR /*wzDownloadUrl*/,
+        __in_z_opt LPCWSTR /*wzPayloadContainerId*/,
+        __in BOOTSTRAPPER_CACHE_OPERATION /*recommendation*/,
+        __inout BOOTSTRAPPER_CACHE_OPERATION* /*pAction*/,
         __inout BOOL* /*pfCancel*/
         )
     {
@@ -405,13 +409,18 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
-    virtual STDMETHODIMP OnResolveSource(
-        __in_z LPCWSTR /*wzPackageOrContainerId*/,
+    virtual STDMETHODIMP OnCacheAcquireResolving(
+        __in_z_opt LPCWSTR /*wzPackageOrContainerId*/,
         __in_z_opt LPCWSTR /*wzPayloadId*/,
-        __in_z LPCWSTR /*wzLocalSource*/,
-        __in_z_opt LPCWSTR /*wzDownloadSource*/,
-        __in BOOTSTRAPPER_RESOLVESOURCE_ACTION /*recommendation*/,
-        __inout BOOTSTRAPPER_RESOLVESOURCE_ACTION* /*pAction*/,
+        __in_z LPCWSTR* /*rgSearchPaths*/,
+        __in DWORD /*cSearchPaths*/,
+        __in BOOL /*fFoundLocal*/,
+        __in DWORD /*dwRecommendedSearchPath*/,
+        __in_z_opt LPCWSTR /*wzDownloadUrl*/,
+        __in_z_opt LPCWSTR /*wzPayloadContainerId*/,
+        __in BOOTSTRAPPER_CACHE_RESOLVE_OPERATION /*recommendation*/,
+        __inout DWORD* /*pdwChosenSearchPath*/,
+        __inout BOOTSTRAPPER_CACHE_RESOLVE_OPERATION* /*pAction*/,
         __inout BOOL* /*pfCancel*/
         )
     {
diff --git a/src/balutil/inc/BalBaseBAFunctionsProc.h b/src/balutil/inc/BalBaseBAFunctionsProc.h
index e1de800a..7e89fe83 100644
--- a/src/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/balutil/inc/BalBaseBAFunctionsProc.h
@@ -77,7 +77,7 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_ONCACHEPACKAGEBEGIN:
         case BA_FUNCTIONS_MESSAGE_ONCACHEACQUIREBEGIN:
         case BA_FUNCTIONS_MESSAGE_ONCACHEACQUIREPROGRESS:
-        case BA_FUNCTIONS_MESSAGE_ONRESOLVESOURCE:
+        case BA_FUNCTIONS_MESSAGE_ONCACHEACQUIRERESOLVING:
         case BA_FUNCTIONS_MESSAGE_ONCACHEACQUIRECOMPLETE:
         case BA_FUNCTIONS_MESSAGE_ONCACHEVERIFYBEGIN:
         case BA_FUNCTIONS_MESSAGE_ONCACHEVERIFYCOMPLETE:
diff --git a/src/balutil/inc/BalBaseBootstrapperApplication.h b/src/balutil/inc/BalBaseBootstrapperApplication.h
index c9211e0f..3d3e4ffa 100644
--- a/src/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/balutil/inc/BalBaseBootstrapperApplication.h
@@ -87,6 +87,7 @@ public: // IBootstrapperApplication
     }
 
     virtual STDMETHODIMP OnDetectBegin(
+        __in BOOL /*fCached*/,
         __in BOOL /*fInstalled*/,
         __in DWORD /*cPackages*/,
         __inout BOOL* pfCancel
@@ -442,14 +443,17 @@ public: // IBootstrapperApplication
     }
 
     virtual STDMETHODIMP OnCacheAcquireBegin(
-        __in_z LPCWSTR wzPackageOrContainerId,
+        __in_z_opt LPCWSTR wzPackageOrContainerId,
         __in_z_opt LPCWSTR wzPayloadId,
-        __in BOOTSTRAPPER_CACHE_OPERATION /*operation*/,
         __in_z LPCWSTR /*wzSource*/,
+        __in_z_opt LPCWSTR /*wzDownloadUrl*/,
+        __in_z_opt LPCWSTR /*wzPayloadContainerId*/,
+        __in BOOTSTRAPPER_CACHE_OPERATION /*recommendation*/,
+        __inout BOOTSTRAPPER_CACHE_OPERATION* /*pAction*/,
         __inout BOOL* pfCancel
         )
     {
-        BalRetryStartPackage(BALRETRY_TYPE_CACHE, wzPackageOrContainerId, wzPayloadId);
+        BalRetryStartContainerOrPayload(wzPackageOrContainerId, wzPayloadId);
         *pfCancel |= CheckCanceled();
         return S_OK;
     }
@@ -488,13 +492,18 @@ public: // IBootstrapperApplication
         return hr;
     }
 
-    virtual STDMETHODIMP OnResolveSource(
-        __in_z LPCWSTR /*wzPackageOrContainerId*/,
+    virtual STDMETHODIMP OnCacheAcquireResolving(
+        __in_z_opt LPCWSTR /*wzPackageOrContainerId*/,
         __in_z_opt LPCWSTR /*wzPayloadId*/,
-        __in_z LPCWSTR /*wzLocalSource*/,
-        __in_z_opt LPCWSTR /*wzDownloadSource*/,
-        __in BOOTSTRAPPER_RESOLVESOURCE_ACTION /*recommendation*/,
-        __inout BOOTSTRAPPER_RESOLVESOURCE_ACTION* /*pAction*/,
+        __in_z LPCWSTR* /*rgSearchPaths*/,
+        __in DWORD /*cSearchPaths*/,
+        __in BOOL /*fFoundLocal*/,
+        __in DWORD /*dwRecommendedSearchPath*/,
+        __in_z_opt LPCWSTR /*wzDownloadUrl*/,
+        __in_z_opt LPCWSTR /*wzPayloadContainerId*/,
+        __in BOOTSTRAPPER_CACHE_RESOLVE_OPERATION /*recommendation*/,
+        __inout DWORD* /*pdwChosenSearchPath*/,
+        __inout BOOTSTRAPPER_CACHE_RESOLVE_OPERATION* /*pAction*/,
         __inout BOOL* pfCancel
         )
     {
@@ -518,7 +527,7 @@ public: // IBootstrapperApplication
             ExitFunction1(hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT));
         }
 
-        hr = BalRetryEndPackage(BALRETRY_TYPE_CACHE, wzPackageOrContainerId, wzPayloadId, hrStatus, &fRetry);
+        hr = BalRetryEndContainerOrPayload(wzPackageOrContainerId, wzPayloadId, hrStatus, &fRetry);
         ExitOnFailure(hr, "BalRetryEndPackage for cache failed");
 
         if (fRetry)
@@ -599,7 +608,7 @@ public: // IBootstrapperApplication
         // Only track retry on execution (not rollback).
         if (fExecute)
         {
-            BalRetryStartPackage(BALRETRY_TYPE_EXECUTE, wzPackageId, NULL);
+            BalRetryStartPackage(wzPackageId);
         }
 
         m_fRollingBack = !fExecute;
@@ -700,7 +709,7 @@ public: // IBootstrapperApplication
             ExitFunction1(hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT));
         }
 
-        hr = BalRetryEndPackage(BALRETRY_TYPE_EXECUTE, wzPackageId, NULL, hrStatus, &fRetry);
+        hr = BalRetryEndPackage(wzPackageId, hrStatus, &fRetry);
         ExitOnFailure(hr, "BalRetryEndPackage for execute failed");
 
         if (fRetry)
@@ -951,6 +960,9 @@ protected:
         m_fApplying = FALSE;
         m_fRollingBack = FALSE;
 
+        m_dwProgressPercentage = 0;
+        m_dwOverallProgressPercentage = 0;
+
         BalRetryInitialize(dwRetryCount, dwRetryTimeout);
     }
 
diff --git a/src/balutil/inc/BalBaseBootstrapperApplicationProc.h b/src/balutil/inc/BalBaseBootstrapperApplicationProc.h
index d8a6590b..42ffeb79 100644
--- a/src/balutil/inc/BalBaseBootstrapperApplicationProc.h
+++ b/src/balutil/inc/BalBaseBootstrapperApplicationProc.h
@@ -15,7 +15,7 @@ static HRESULT BalBaseBAProcOnDetectBegin(
     __inout BA_ONDETECTBEGIN_RESULTS* pResults
     )
 {
-    return pBA->OnDetectBegin(pArgs->fInstalled, pArgs->cPackages, &pResults->fCancel);
+    return pBA->OnDetectBegin(pArgs->fCached, pArgs->fInstalled, pArgs->cPackages, &pResults->fCancel);
 }
 
 static HRESULT BalBaseBAProcOnDetectComplete(
@@ -303,7 +303,7 @@ static HRESULT BalBaseBAProcOnCacheAcquireBegin(
     __inout BA_ONCACHEACQUIREBEGIN_RESULTS* pResults
     )
 {
-    return pBA->OnCacheAcquireBegin(pArgs->wzPackageOrContainerId, pArgs->wzPayloadId, pArgs->operation, pArgs->wzSource, &pResults->fCancel);
+    return pBA->OnCacheAcquireBegin(pArgs->wzPackageOrContainerId, pArgs->wzPayloadId, pArgs->wzSource, pArgs->wzDownloadUrl, pArgs->wzPayloadContainerId, pArgs->recommendation, &pResults->action, &pResults->fCancel);
 }
 
 static HRESULT BalBaseBAProcOnCacheAcquireProgress(
@@ -315,13 +315,13 @@ static HRESULT BalBaseBAProcOnCacheAcquireProgress(
     return pBA->OnCacheAcquireProgress(pArgs->wzPackageOrContainerId, pArgs->wzPayloadId, pArgs->dw64Progress, pArgs->dw64Total, pArgs->dwOverallPercentage, &pResults->fCancel);
 }
 
-static HRESULT BalBaseBAProcOnResolveSource(
+static HRESULT BalBaseBAProcOnCacheAcquireResolving(
     __in IBootstrapperApplication* pBA,
-    __in BA_ONRESOLVESOURCE_ARGS* pArgs,
-    __inout BA_ONRESOLVESOURCE_RESULTS* pResults
+    __in BA_ONCACHEACQUIRERESOLVING_ARGS* pArgs,
+    __inout BA_ONCACHEACQUIRERESOLVING_RESULTS* pResults
     )
 {
-    return pBA->OnResolveSource(pArgs->wzPackageOrContainerId, pArgs->wzPayloadId, pArgs->wzLocalSource, pArgs->wzDownloadSource, pArgs->recommendation, &pResults->action, &pResults->fCancel);
+    return pBA->OnCacheAcquireResolving(pArgs->wzPackageOrContainerId, pArgs->wzPayloadId, pArgs->rgSearchPaths, pArgs->cSearchPaths, pArgs->fFoundLocal, pArgs->dwRecommendedSearchPath, pArgs->wzDownloadUrl, pArgs->wzPayloadContainerId, pArgs->recommendation, &pResults->dwChosenSearchPath, &pResults->action, &pResults->fCancel);
 }
 
 static HRESULT BalBaseBAProcOnCacheAcquireComplete(
@@ -713,8 +713,8 @@ static HRESULT WINAPI BalBaseBootstrapperApplicationProc(
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIREPROGRESS:
             hr = BalBaseBAProcOnCacheAcquireProgress(pBA, reinterpret_cast<BA_ONCACHEACQUIREPROGRESS_ARGS*>(pvArgs), reinterpret_cast<BA_ONCACHEACQUIREPROGRESS_RESULTS*>(pvResults));
             break;
-        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONRESOLVESOURCE:
-            hr = BalBaseBAProcOnResolveSource(pBA, reinterpret_cast<BA_ONRESOLVESOURCE_ARGS*>(pvArgs), reinterpret_cast<BA_ONRESOLVESOURCE_RESULTS*>(pvResults));
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIRERESOLVING:
+            hr = BalBaseBAProcOnCacheAcquireResolving(pBA, reinterpret_cast<BA_ONCACHEACQUIRERESOLVING_ARGS*>(pvArgs), reinterpret_cast<BA_ONCACHEACQUIRERESOLVING_RESULTS*>(pvResults));
             break;
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEACQUIRECOMPLETE:
             hr = BalBaseBAProcOnCacheAcquireComplete(pBA, reinterpret_cast<BA_ONCACHEACQUIRECOMPLETE_ARGS*>(pvArgs), reinterpret_cast<BA_ONCACHEACQUIRECOMPLETE_RESULTS*>(pvResults));
diff --git a/src/balutil/inc/IBootstrapperApplication.h b/src/balutil/inc/IBootstrapperApplication.h
index ed70d8fe..8fcdd318 100644
--- a/src/balutil/inc/IBootstrapperApplication.h
+++ b/src/balutil/inc/IBootstrapperApplication.h
@@ -21,6 +21,7 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
 
     // OnDetectBegin - called when the engine begins detection.
     STDMETHOD(OnDetectBegin)(
+        __in BOOL fCached,
         __in BOOL fInstalled,
         __in DWORD cPackages,
         __inout BOOL* pfCancel
@@ -279,14 +280,20 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         __inout BOOL* pfCancel
         )  = 0;
 
-    // OnCacheAcquireBegin - called when the engine begins copying or
-    //                       downloading a payload to the working folder.
+    // OnCacheAcquireBegin - called when the engine begins acquiring a payload or container.
+    //
+    // Notes:
+    //  It is expected the BA may call IBootstrapperEngine::SetLocalSource() or IBootstrapperEngine::SetDownloadSource()
+    //  to update the source location before returning.
     //
     STDMETHOD(OnCacheAcquireBegin)(
         __in_z_opt LPCWSTR wzPackageOrContainerId,
         __in_z_opt LPCWSTR wzPayloadId,
-        __in BOOTSTRAPPER_CACHE_OPERATION operation,
         __in_z LPCWSTR wzSource,
+        __in_z_opt LPCWSTR wzDownloadUrl,
+        __in_z_opt LPCWSTR wzPayloadContainerId,
+        __in BOOTSTRAPPER_CACHE_OPERATION recommendation,
+        __inout BOOTSTRAPPER_CACHE_OPERATION* pAction,
         __inout BOOL* pfCancel
         ) = 0;
 
@@ -302,27 +309,38 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         __inout BOOL* pfCancel
         ) = 0;
 
-    // OnResolveSource - called when a payload or container cannot be found locally.
+    // OnCacheAcquireResolving - called to allow the BA to override the acquisition action for the payload or container.
     //
     // Parameters:
+    //  wzPackageOrContainerId will be NULL when resolving a layout-only payload.
     //  wzPayloadId will be NULL when resolving a container.
-    //  wzDownloadSource will be NULL if the container or payload does not provide a DownloadURL.
+    //  wzDownloadUrl will be NULL if the container or payload does not provide a DownloadURL.
+    //  wzPayloadContainerId will not be NULL if acquiring a payload that is in a container.
     //
-    // Notes:
-    //  It is expected the BA may call IBootstrapperEngine::SetLocalSource() or IBootstrapperEngine::SetDownloadSource()
-    //  to update the source location before returning BOOTSTRAPPER_RESOLVESOURCE_ACTION_RETRY or BOOTSTRAPPER_RESOLVESOURCE_ACTION_DOWNLOAD.
-    STDMETHOD(OnResolveSource)(
-        __in_z LPCWSTR wzPackageOrContainerId,
+    //  rgSearchPaths are the search paths used for source resolution.
+    //  fFoundLocal is TRUE when dwRecommendedSearchPath indicates that the file was found.
+    //  dwRecommendedSearchPath is the index into rgSearchPaths for the recommended local file.
+    //
+    STDMETHOD(OnCacheAcquireResolving)(
+        __in_z_opt LPCWSTR wzPackageOrContainerId,
         __in_z_opt LPCWSTR wzPayloadId,
-        __in_z LPCWSTR wzLocalSource,
-        __in_z_opt LPCWSTR wzDownloadSource,
-        __in BOOTSTRAPPER_RESOLVESOURCE_ACTION recommendation,
-        __inout BOOTSTRAPPER_RESOLVESOURCE_ACTION* pAction,
+        __in_z LPCWSTR* rgSearchPaths,
+        __in DWORD cSearchPaths,
+        __in BOOL fFoundLocal,
+        __in DWORD dwRecommendedSearchPath,
+        __in_z_opt LPCWSTR wzDownloadUrl,
+        __in_z_opt LPCWSTR wzPayloadContainerId,
+        __in BOOTSTRAPPER_CACHE_RESOLVE_OPERATION recommendation,
+        __inout DWORD* pdwChosenSearchPath,
+        __inout BOOTSTRAPPER_CACHE_RESOLVE_OPERATION* pAction,
         __inout BOOL* pfCancel
         ) = 0;
 
-    // OnCacheAcquireComplete - called after the engine copied or downloaded
-    //                          a payload to the working folder.
+    // OnCacheAcquireComplete - called after the engine acquired the payload or container.
+    //
+    // Notes:
+    //  It is expected the BA may call IBootstrapperEngine::SetLocalSource() or IBootstrapperEngine::SetDownloadSource()
+    //  to update the source location before returning BOOTSTRAPPER_CACHEACQUIRECOMPLETE_ACTION_RETRY.
     //
     STDMETHOD(OnCacheAcquireComplete)(
         __in_z_opt LPCWSTR wzPackageOrContainerId,
diff --git a/src/balutil/inc/balretry.h b/src/balutil/inc/balretry.h
index 040ab4ae..35282a7e 100644
--- a/src/balutil/inc/balretry.h
+++ b/src/balutil/inc/balretry.h
@@ -6,12 +6,6 @@
 extern "C" {
 #endif
 
-typedef enum BALRETRY_TYPE
-{
-    BALRETRY_TYPE_CACHE,
-    BALRETRY_TYPE_EXECUTE,
-} BALRETRY_TYPE;
-
 /*******************************************************************
  BalRetryInitialize - initialize the retry count and timeout between
                       retries (in milliseconds).
@@ -33,9 +27,7 @@ DAPI_(void) BalRetryUninitialize();
                         wait the specified timeout.
 ********************************************************************/
 DAPI_(void) BalRetryStartPackage(
-    __in BALRETRY_TYPE type,
-    __in_z_opt LPCWSTR wzPackageId,
-    __in_z_opt LPCWSTR wzPayloadId
+    __in_z LPCWSTR wzPackageId
     );
 
 /*******************************************************************
@@ -43,17 +35,34 @@ DAPI_(void) BalRetryStartPackage(
                         to consider.
 ********************************************************************/
 DAPI_(void) BalRetryErrorOccurred(
-    __in_z_opt LPCWSTR wzPackageId,
+    __in_z LPCWSTR wzPackageId,
     __in DWORD dwError
     );
 
 /*******************************************************************
- BalRetryEndPackage - returns IDRETRY is a retry is recommended or 
-                      IDNOACTION if a retry is not recommended.
+ BalRetryEndPackage - returns TRUE if a retry is recommended.
 ********************************************************************/
 DAPI_(HRESULT) BalRetryEndPackage(
-    __in BALRETRY_TYPE type,
-    __in_z_opt LPCWSTR wzPackageId,
+    __in_z LPCWSTR wzPackageId,
+    __in HRESULT hrError,
+    __inout BOOL* pfRetry
+    );
+
+/*******************************************************************
+ BalRetryStartContainerOrPayload - call when a container or payload
+        begins to be acquired. If the target is being retried,
+        the function will wait the specified timeout.
+********************************************************************/
+DAPI_(void) BalRetryStartContainerOrPayload(
+    __in_z_opt LPCWSTR wzContainerOrPackageId,
+    __in_z_opt LPCWSTR wzPayloadId
+    );
+
+/*******************************************************************
+ BalRetryEndContainerOrPayload - returns TRUE if a retry is recommended.
+********************************************************************/
+DAPI_(HRESULT) BalRetryEndContainerOrPayload(
+    __in_z_opt LPCWSTR wzContainerOrPackageId,
     __in_z_opt LPCWSTR wzPayloadId,
     __in HRESULT hrError,
     __inout BOOL* pfRetry