From 14cdda3c489d6b9801f05939044e67b13939b42d Mon Sep 17 00:00:00 2001
From: Sean Hall <r.sean.hall@gmail.com>
Date: Sun, 25 Apr 2021 22:44:23 -0500
Subject: Set source of attached containers to WixBundleOriginalSource if set.
 Use file size when probing local files.

#5586
---
 src/engine/apply.cpp  | 50 +++++++++++++++++++++++++++++++++++++++++++++++---
 src/engine/engine.cpp | 27 +++++++++++++++++++++++++--
 2 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index 74f57f6a..58d41b28 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -133,6 +133,11 @@ static HRESULT AcquireContainerOrPayload(
     __in BURN_CACHE_PROGRESS_CONTEXT* pProgress,
     __out BOOL* pfRetry
     );
+static BOOL IsValidLocalFile(
+    __in_z LPCWSTR wzFilePath,
+    __in DWORD64 qwFileSize,
+    __in BOOL fMinimumFileSize
+    );
 static HRESULT LayoutOrCacheContainerOrPayload(
     __in BURN_CACHE_CONTEXT* pContext,
     __in_opt BURN_CONTAINER* pContainer,
@@ -1399,6 +1404,25 @@ static HRESULT AcquireContainerOrPayload(
     LPWSTR* pwzSourcePath = pContainer ? &pContainer->sczSourcePath : &pPayload->sczSourcePath;
     BOOL fFoundLocal = FALSE;
     BOOL fPreferExtract = FALSE;
+    DWORD64 qwFileSize = 0;
+    BOOL fMinimumFileSize = FALSE;
+
+    if (pContainer)
+    {
+        if (pContainer->fAttached)
+        {
+            fMinimumFileSize = TRUE;
+            qwFileSize = pContainer->qwAttachedOffset + pContainer->qwFileSize;
+        }
+        else if (pContainer->pbHash && pContext->wzLayoutDirectory)
+        {
+            qwFileSize = pContainer->qwFileSize;
+        }
+    }
+    else if (pPayload->pbHash)
+    {
+        qwFileSize = pPayload->qwFileSize;
+    }
 
     pContext->cSearchPaths = 0;
     *pfRetry = FALSE;
@@ -1427,7 +1451,7 @@ static HRESULT AcquireContainerOrPayload(
                 // When a payload comes from a container, the container has the highest chance of being correct.
                 // But we want to avoid extracting the container multiple times.
                 // So only consider the destination path, which means the container was already extracted.
-                if (FileExistsEx(pContext->rgSearchPaths[dwDestinationSearchPath], NULL))
+                if (IsValidLocalFile(pContext->rgSearchPaths[dwDestinationSearchPath], qwFileSize, fMinimumFileSize))
                 {
                     fFoundLocal = TRUE;
                     dwChosenSearchPath = dwDestinationSearchPath;
@@ -1442,8 +1466,8 @@ static HRESULT AcquireContainerOrPayload(
             {
                 for (DWORD i = 0; i < pContext->cSearchPaths; ++i)
                 {
-                    // If the file exists locally, choose it.
-                    if (FileExistsEx(pContext->rgSearchPaths[i], NULL))
+                    // If the file exists locally with the correct size, choose it.
+                    if (IsValidLocalFile(pContext->rgSearchPaths[i], qwFileSize, fMinimumFileSize))
                     {
                         dwChosenSearchPath = i;
 
@@ -1557,6 +1581,26 @@ LExit:
     return hr;
 }
 
+static BOOL IsValidLocalFile(
+    __in_z LPCWSTR wzFilePath,
+    __in DWORD64 qwFileSize,
+    __in BOOL fMinimumFileSize
+    )
+{
+    LONGLONG llFileSize = 0;
+
+    if (!qwFileSize)
+    {
+        return FileExistsEx(wzFilePath, NULL);
+    }
+    else
+    {
+        return SUCCEEDED(FileSize(wzFilePath, &llFileSize)) &&
+               (static_cast<DWORD64>(llFileSize) == qwFileSize ||
+                fMinimumFileSize && static_cast<DWORD64>(llFileSize) > qwFileSize);
+    }
+}
+
 static HRESULT LayoutOrCacheContainerOrPayload(
     __in BURN_CACHE_CONTEXT* pContext,
     __in_opt BURN_CONTAINER* pContainer,
diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp
index bb4061a7..e2728d7f 100644
--- a/src/engine/engine.cpp
+++ b/src/engine/engine.cpp
@@ -523,7 +523,8 @@ static HRESULT RunNormal(
     )
 {
     HRESULT hr = S_OK;
-    HANDLE hPipesCreatedEvent = NULL;
+    LPWSTR sczOriginalSource = NULL;
+    LPWSTR sczCopiedOriginalSource = NULL;
     BOOL fContinueExecution = TRUE;
     BOOL fReloadApp = FALSE;
     BOOL fSkipCleanup = FALSE;
@@ -558,6 +559,27 @@ static HRESULT RunNormal(
     hr = CoreQueryRegistration(pEngineState);
     ExitOnFailure(hr, "Failed to query registration.");
 
+    // Best effort to set the source of attached containers to BURN_BUNDLE_ORIGINAL_SOURCE.
+    hr = VariableGetString(&pEngineState->variables, BURN_BUNDLE_ORIGINAL_SOURCE, &sczOriginalSource);
+    if (SUCCEEDED(hr))
+    {
+        for (DWORD i = 0; i < pEngineState->containers.cContainers; ++i)
+        {
+            BURN_CONTAINER* pContainer = pEngineState->containers.rgContainers + i;
+            if (pContainer->fAttached)
+            {
+                hr = StrAllocString(&sczCopiedOriginalSource, sczOriginalSource, 0);
+                if (SUCCEEDED(hr))
+                {
+                    ReleaseNullStr(pContainer->sczSourcePath);
+                    pContainer->sczSourcePath = sczCopiedOriginalSource;
+                    sczCopiedOriginalSource = NULL;
+                }
+            }
+        }
+    }
+    hr = S_OK;
+
     // Set some built-in variables before loading the BA.
     hr = PlanSetVariables(pEngineState->command.action, &pEngineState->variables);
     ExitOnFailure(hr, "Failed to set action variables.");
@@ -613,7 +635,8 @@ LExit:
         ::PostMessageW(pEngineState->command.hwndSplashScreen, WM_CLOSE, 0, 0);
     }
 
-    ReleaseHandle(hPipesCreatedEvent);
+    ReleaseStr(sczOriginalSource);
+    ReleaseStr(sczCopiedOriginalSource);
 
     return hr;
 }
-- 
cgit v1.2.3-55-g6feb