Enforce payload and container verification.

9 files changed, 209 insertions, 92 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index b4fa9dac..4d527409 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -1410,12 +1410,12 @@ static HRESULT AcquireContainerOrPayload(
             fMinimumFileSize = TRUE;
             qwFileSize = pContainer->qwAttachedOffset + pContainer->qwFileSize;
         }
-        else if (pContainer->pbHash && pContext->wzLayoutDirectory)
+        else
         {
             qwFileSize = pContainer->qwFileSize;
         }
     }
-    else if (pPayload->pbHash)
+    else if (BURN_PAYLOAD_VERIFICATION_HASH == pPayload->verification || BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE == pPayload->verification)
     {
         qwFileSize = pPayload->qwFileSize;
     }
diff --git a/src/burn/engine/cache.cpp b/src/burn/engine/cache.cpp
index 02cad4a5..44267cbc 100644
--- a/src/burn/engine/cache.cpp
+++ b/src/burn/engine/cache.cpp
@@ -114,10 +114,16 @@ static HRESULT RemoveBundleOrPackage(
     __in_z LPCWSTR wzBundleOrPackageId,
     __in_z LPCWSTR wzCacheId
     );
+static HRESULT VerifyFileSize(
+    __in HANDLE hFile,
+    __in DWORD64 qwFileSize,
+    __in_z LPCWSTR wzUnverifiedPayloadPath
+    );
 static HRESULT VerifyHash(
     __in BYTE* pbHash,
     __in DWORD cbHash,
     __in DWORD64 qwFileSize,
+    __in BOOL fVerifyFileSize,
     __in_z LPCWSTR wzUnverifiedPayloadPath,
     __in HANDLE hFile,
     __in BURN_CACHE_STEP cacheStep,
@@ -1513,11 +1519,16 @@ static HRESULT VerifyThenTransferContainer(
         ExitWithLastError(hr, "Failed to open container in working path: %ls", wzUnverifiedContainerPath);
     }
 
-    // Container should have a hash we can use to verify with.
-    if (pContainer->pbHash)
+
+    switch (pContainer->verification)
     {
-        hr = VerifyHash(pContainer->pbHash, pContainer->cbHash, pContainer->qwFileSize, wzUnverifiedContainerPath, hFile, BURN_CACHE_STEP_HASH, pfnCacheMessageHandler, pfnProgress, pContext);
+    case BURN_CONTAINER_VERIFICATION_HASH:
+        hr = VerifyHash(pContainer->pbHash, pContainer->cbHash, pContainer->qwFileSize, TRUE, wzUnverifiedContainerPath, hFile, BURN_CACHE_STEP_HASH, pfnCacheMessageHandler, pfnProgress, pContext);
         ExitOnFailure(hr, "Failed to verify container hash: %ls", wzCachedPath);
+        break;
+    default:
+        ExitOnRootFailure(hr = E_INVALIDARG, "Container has no verification information: %ls", pContainer->sczId);
+        break;
     }
 
     LogStringLine(REPORT_STANDARD, "%ls container from working path '%ls' to path '%ls'", fMove ? L"Moving" : L"Copying", wzUnverifiedContainerPath, wzCachedPath);
@@ -1550,10 +1561,16 @@ static HRESULT VerifyThenTransferPayload(
         ExitWithLastError(hr, "Failed to open payload in working path: %ls", wzUnverifiedPayloadPath);
     }
 
-    if (pPayload->pbHash) // the payload should have a hash we can use to verify it.
+    switch (pPayload->verification)
     {
-        hr = VerifyHash(pPayload->pbHash, pPayload->cbHash, pPayload->qwFileSize, wzUnverifiedPayloadPath, hFile, BURN_CACHE_STEP_HASH, pfnCacheMessageHandler, pfnProgress, pContext);
+    case BURN_PAYLOAD_VERIFICATION_HASH:
+        hr = VerifyHash(pPayload->pbHash, pPayload->cbHash, pPayload->qwFileSize, TRUE, wzUnverifiedPayloadPath, hFile, BURN_CACHE_STEP_HASH, pfnCacheMessageHandler, pfnProgress, pContext);
         ExitOnFailure(hr, "Failed to verify payload hash: %ls", wzCachedPath);
+        break;
+    case BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE: __fallthrough;
+    default:
+        ExitOnRootFailure(hr = E_INVALIDARG, "Payload has no verification information: %ls", pPayload->sczKey);
+        break;
     }
 
     LogStringLine(REPORT_STANDARD, "%ls payload from working path '%ls' to path '%ls'", fMove ? L"Moving" : L"Copying", wzUnverifiedPayloadPath, wzCachedPath);
@@ -1627,10 +1644,15 @@ static HRESULT VerifyFileAgainstContainer(
         ExitOnRootFailure(hr, "Failed to open container at path: %ls", wzVerifyPath);
     }
 
-    if (pContainer->pbHash) // the container should have a hash we can use to verify it.
+    switch (pContainer->verification)
     {
-        hr = VerifyHash(pContainer->pbHash, pContainer->cbHash, pContainer->qwFileSize, wzVerifyPath, hFile, cacheStep, pfnCacheMessageHandler, pfnProgress, pContext);
+    case BURN_CONTAINER_VERIFICATION_HASH:
+        hr = VerifyHash(pContainer->pbHash, pContainer->cbHash, pContainer->qwFileSize, TRUE, wzVerifyPath, hFile, cacheStep, pfnCacheMessageHandler, pfnProgress, pContext);
         ExitOnFailure(hr, "Failed to verify hash of container: %ls", pContainer->sczId);
+        break;
+    default:
+        ExitOnRootFailure(hr = E_INVALIDARG, "Container has no verification information: %ls", pContainer->sczId);
+        break;
     }
 
     if (fAlreadyCached)
@@ -1667,6 +1689,7 @@ static HRESULT VerifyFileAgainstPayload(
 {
     HRESULT hr = S_OK;
     HANDLE hFile = INVALID_HANDLE_VALUE;
+    BOOL fVerifyFileSize = FALSE;
 
     // Get the payload on disk actual hash.
     hFile = ::CreateFileW(wzVerifyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_DELETE, NULL, OPEN_EXISTING, FILE_FLAG_SEQUENTIAL_SCAN, NULL);
@@ -1680,10 +1703,33 @@ static HRESULT VerifyFileAgainstPayload(
         ExitOnRootFailure(hr, "Failed to open payload at path: %ls", wzVerifyPath);
     }
 
-    if (pPayload->pbHash) // the payload should have a hash we can use to verify it.
+    switch (pPayload->verification)
     {
-        hr = VerifyHash(pPayload->pbHash, pPayload->cbHash, pPayload->qwFileSize, wzVerifyPath, hFile, cacheStep, pfnCacheMessageHandler, pfnProgress, pContext);
+    case BURN_PAYLOAD_VERIFICATION_HASH:
+        fVerifyFileSize = TRUE;
+
+        hr = VerifyHash(pPayload->pbHash, pPayload->cbHash, pPayload->qwFileSize, fVerifyFileSize, wzVerifyPath, hFile, cacheStep, pfnCacheMessageHandler, pfnProgress, pContext);
         ExitOnFailure(hr, "Failed to verify hash of payload: %ls", pPayload->sczKey);
+
+        break;
+    case BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE:
+        fVerifyFileSize = 0 != pPayload->qwFileSize;
+
+        if (pPayload->pbHash)
+        {
+            hr = VerifyHash(pPayload->pbHash, pPayload->cbHash, pPayload->qwFileSize, fVerifyFileSize, wzVerifyPath, hFile, cacheStep, pfnCacheMessageHandler, pfnProgress, pContext);
+            ExitOnFailure(hr, "Failed to verify hash of payload: %ls", pPayload->sczKey);
+        }
+        else if (fVerifyFileSize)
+        {
+            hr = VerifyFileSize(hFile, pPayload->qwFileSize, wzVerifyPath);
+            ExitOnFailure(hr, "Failed to verify file size for path: %ls", wzVerifyPath);
+        }
+
+        break;
+    default:
+        ExitOnRootFailure(hr = E_INVALIDARG, "Payload has no verification information: %ls", pPayload->sczKey);
+        break;
     }
 
     if (fAlreadyCached)
@@ -2015,10 +2061,32 @@ LExit:
     return hr;
 }
 
+static HRESULT VerifyFileSize(
+    __in HANDLE hFile,
+    __in DWORD64 qwFileSize,
+    __in_z LPCWSTR wzUnverifiedPayloadPath
+    )
+{
+    HRESULT hr = S_OK;
+    LONGLONG llSize = 0;
+
+    hr = FileSizeByHandle(hFile, &llSize);
+    ExitOnFailure(hr, "Failed to get file size for path: %ls", wzUnverifiedPayloadPath);
+
+    if (static_cast<DWORD64>(llSize) != qwFileSize)
+    {
+        ExitOnRootFailure(hr = ERROR_FILE_CORRUPT, "File size mismatch for path: %ls, expected: %llu, actual: %lld", wzUnverifiedPayloadPath, qwFileSize, llSize);
+    }
+
+LExit:
+    return hr;
+}
+
 static HRESULT VerifyHash(
     __in BYTE* pbHash,
     __in DWORD cbHash,
     __in DWORD64 qwFileSize,
+    __in BOOL fVerifyFileSize,
     __in_z LPCWSTR wzUnverifiedPayloadPath,
     __in HANDLE hFile,
     __in BURN_CACHE_STEP cacheStep,
@@ -2032,19 +2100,16 @@ static HRESULT VerifyHash(
     HRESULT hr = S_OK;
     BYTE rgbActualHash[SHA512_HASH_LEN] = { };
     DWORD64 qwHashedBytes = 0;
-    LONGLONG llSize = 0;
     LPWSTR pszExpected = NULL;
     LPWSTR pszActual = NULL;
 
     hr = SendCacheBeginMessage(pfnCacheMessageHandler, pContext, cacheStep);
     ExitOnFailure(hr, "Aborted cache verify hash begin.");
 
-    hr = FileSizeByHandle(hFile, &llSize);
-    ExitOnFailure(hr, "Failed to get file size for path: %ls", wzUnverifiedPayloadPath);
-
-    if (static_cast<DWORD64>(llSize) != qwFileSize)
+    if (fVerifyFileSize)
     {
-        ExitOnFailure(hr = ERROR_FILE_CORRUPT, "File size mismatch for path: %ls, expected: %llu, actual: %lld", wzUnverifiedPayloadPath, qwFileSize, llSize);
+        hr = VerifyFileSize(hFile, qwFileSize, wzUnverifiedPayloadPath);
+        ExitOnFailure(hr, "Failed to verify file size for path: %ls", wzUnverifiedPayloadPath);
     }
 
     // TODO: create a cryp hash file that sends progress.
diff --git a/src/burn/engine/container.cpp b/src/burn/engine/container.cpp
index 0cce3131..c6f2ada8 100644
--- a/src/burn/engine/container.cpp
+++ b/src/burn/engine/container.cpp
@@ -50,16 +50,9 @@ extern "C" HRESULT ContainersParseFromXml(
         hr = XmlGetAttributeEx(pixnNode, L"Id", &pContainer->sczId);
         ExitOnFailure(hr, "Failed to get @Id.");
 
-        // @Primary
-        hr = XmlGetYesNoAttribute(pixnNode, L"Primary", &pContainer->fPrimary);
-        if (E_NOTFOUND != hr)
-        {
-            ExitOnFailure(hr, "Failed to get @Primary.");
-        }
-
         // @Attached
         hr = XmlGetYesNoAttribute(pixnNode, L"Attached", &pContainer->fAttached);
-        if (E_NOTFOUND != hr || pContainer->fPrimary) // if it is a primary container, it has to be attached
+        if (E_NOTFOUND != hr)
         {
             ExitOnFailure(hr, "Failed to get @Attached.");
         }
@@ -87,10 +80,7 @@ extern "C" HRESULT ContainersParseFromXml(
         {
             // @FilePath
             hr = XmlGetAttributeEx(pixnNode, L"FilePath", &pContainer->sczFilePath);
-            if (E_NOTFOUND != hr)
-            {
-                ExitOnFailure(hr, "Failed to get @FilePath.");
-            }
+            ExitOnFailure(hr, "Failed to get @FilePath.");
         }
 
         // The source path starts as the file path.
@@ -99,23 +89,32 @@ extern "C" HRESULT ContainersParseFromXml(
 
         // @DownloadUrl
         hr = XmlGetAttributeEx(pixnNode, L"DownloadUrl", &pContainer->downloadSource.sczUrl);
-        if (E_NOTFOUND != hr || (!pContainer->fPrimary && !pContainer->sczSourcePath)) // if the package is not a primary package, it must have a source path or a download url
+        if (E_NOTFOUND != hr)
         {
-            ExitOnFailure(hr, "Failed to get @DownloadUrl. Either @SourcePath or @DownloadUrl needs to be provided.");
+            ExitOnFailure(hr, "Failed to get @DownloadUrl.");
         }
 
         // @Hash
         hr = XmlGetAttributeEx(pixnNode, L"Hash", &pContainer->sczHash);
-        if (SUCCEEDED(hr))
-        {
-            hr = StrAllocHexDecode(pContainer->sczHash, &pContainer->pbHash, &pContainer->cbHash);
-            ExitOnFailure(hr, "Failed to hex decode the Container/@Hash.");
-        }
-        else if (E_NOTFOUND != hr)
+        ExitOnFailure(hr, "Failed to get @Hash.");
+
+        hr = StrAllocHexDecode(pContainer->sczHash, &pContainer->pbHash, &pContainer->cbHash);
+        ExitOnFailure(hr, "Failed to hex decode the Container/@Hash.");
+
+        // @FileSize
+        hr = XmlGetAttributeEx(pixnNode, L"FileSize", &scz);
+        ExitOnFailure(hr, "Failed to get @FileSize.");
+
+        hr = StrStringToUInt64(scz, 0, &pContainer->qwFileSize);
+        ExitOnFailure(hr, "Failed to parse @FileSize.");
+
+        if (!pContainer->qwFileSize)
         {
-            ExitOnFailure(hr, "Failed to get @Hash.");
+            ExitOnRootFailure(hr = E_INVALIDDATA, "File size is required when verifying by hash for container: %ls", pContainer->sczId);
         }
 
+        pContainer->verification = BURN_CONTAINER_VERIFICATION_HASH;
+
         // prepare next iteration
         ReleaseNullObject(pixnNode);
     }
@@ -136,6 +135,7 @@ extern "C" HRESULT ContainersInitialize(
     )
 {
     HRESULT hr = S_OK;
+    DWORD64 qwSize = 0;
 
     if (pContainers->rgContainers)
     {
@@ -147,8 +147,13 @@ extern "C" HRESULT ContainersInitialize(
             // manifest contained and get the offset to the container.
             if (pContainer->fAttached)
             {
-                hr = SectionGetAttachedContainerInfo(pSection, pContainer->dwAttachedIndex, pContainer->type, &pContainer->qwAttachedOffset, &pContainer->qwFileSize, &pContainer->fActuallyAttached);
+                hr = SectionGetAttachedContainerInfo(pSection, pContainer->dwAttachedIndex, pContainer->type, &pContainer->qwAttachedOffset, &qwSize, &pContainer->fActuallyAttached);
                 ExitOnFailure(hr, "Failed to get attached container information.");
+
+                if (qwSize != pContainer->qwFileSize)
+                {
+                    ExitOnFailure(hr, "Attached container '%ls' size '%llu' didn't match size from manifest: '%llu'", pContainer->sczId, qwSize, pContainer->qwFileSize);
+                }
             }
         }
     }
@@ -195,7 +200,6 @@ extern "C" HRESULT ContainerOpenUX(
 
     // open attached container
     container.type = BURN_CONTAINER_TYPE_CABINET;
-    container.fPrimary = TRUE;
     container.fAttached = TRUE;
     container.dwAttachedIndex = 0;
 
diff --git a/src/burn/engine/container.h b/src/burn/engine/container.h
index c2c1c9a8..d454a248 100644
--- a/src/burn/engine/container.h
+++ b/src/burn/engine/container.h
@@ -52,6 +52,12 @@ enum BURN_CAB_OPERATION
     BURN_CAB_OPERATION_CLOSE,
 };
 
+enum BURN_CONTAINER_VERIFICATION
+{
+    BURN_CONTAINER_VERIFICATION_NONE,
+    BURN_CONTAINER_VERIFICATION_HASH,
+};
+
 
 // structs
 
@@ -59,7 +65,6 @@ typedef struct _BURN_CONTAINER
 {
     LPWSTR sczId;
     BURN_CONTAINER_TYPE type;
-    BOOL fPrimary;
     BOOL fAttached;
     DWORD dwAttachedIndex;
     DWORD64 qwFileSize;
@@ -69,6 +74,7 @@ typedef struct _BURN_CONTAINER
 
     BYTE* pbHash;
     DWORD cbHash;
+    BURN_CONTAINER_VERIFICATION verification;
     DWORD64 qwAttachedOffset;
     BOOL fActuallyAttached;     // indicates whether an attached container is attached or missing.
 
diff --git a/src/burn/engine/payload.cpp b/src/burn/engine/payload.cpp
index 72eb3476..392a3dd4 100644
--- a/src/burn/engine/payload.cpp
+++ b/src/burn/engine/payload.cpp
@@ -27,6 +27,8 @@ extern "C" HRESULT PayloadsParseFromXml(
     IXMLDOMNode* pixnNode = NULL;
     DWORD cNodes = 0;
     LPWSTR scz = NULL;
+    BOOL fChainPayload = pContainers && pLayoutPayloads; // These are required when parsing chain payloads.
+    BOOL fValidFileSize = FALSE;
 
     // select payload nodes
     hr = XmlSelectNodes(pixnBundle, L"Payload", &pixnNodes);
@@ -51,6 +53,7 @@ extern "C" HRESULT PayloadsParseFromXml(
     for (DWORD i = 0; i < cNodes; ++i)
     {
         BURN_PAYLOAD* pPayload = &pPayloads->rgPayloads[i];
+        fValidFileSize = FALSE;
 
         hr = XmlNextElement(pixnNodes, &pixnNode, NULL);
         ExitOnFailure(hr, "Failed to get next node.");
@@ -63,27 +66,36 @@ extern "C" HRESULT PayloadsParseFromXml(
         hr = XmlGetAttributeEx(pixnNode, L"FilePath", &pPayload->sczFilePath);
         ExitOnFailure(hr, "Failed to get @FilePath.");
 
-        // @Packaging
-        hr = XmlGetAttributeEx(pixnNode, L"Packaging", &scz);
-        ExitOnFailure(hr, "Failed to get @Packaging.");
+        // @SourcePath
+        hr = XmlGetAttributeEx(pixnNode, L"SourcePath", &pPayload->sczSourcePath);
+        ExitOnFailure(hr, "Failed to get @SourcePath.");
 
-        if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"embedded", -1))
+        if (!fChainPayload)
         {
+            // All non-chain payloads are embedded in the UX container.
             pPayload->packaging = BURN_PAYLOAD_PACKAGING_EMBEDDED;
         }
-        else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"external", -1))
-        {
-            pPayload->packaging = BURN_PAYLOAD_PACKAGING_EXTERNAL;
-        }
         else
         {
-            hr = E_INVALIDARG;
-            ExitOnFailure(hr, "Invalid value for @Packaging: %ls", scz);
-        }
+            // @Packaging
+            hr = XmlGetAttributeEx(pixnNode, L"Packaging", &scz);
+            ExitOnFailure(hr, "Failed to get @Packaging.");
 
-        // @Container
-        if (pContainers)
-        {
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"embedded", -1))
+            {
+                pPayload->packaging = BURN_PAYLOAD_PACKAGING_EMBEDDED;
+            }
+            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, scz, -1, L"external", -1))
+            {
+                pPayload->packaging = BURN_PAYLOAD_PACKAGING_EXTERNAL;
+            }
+            else
+            {
+                hr = E_INVALIDARG;
+                ExitOnFailure(hr, "Invalid value for @Packaging: %ls", scz);
+            }
+
+            // @Container
             hr = XmlGetAttributeEx(pixnNode, L"Container", &scz);
             if (E_NOTFOUND != hr || BURN_PAYLOAD_PACKAGING_EMBEDDED == pPayload->packaging)
             {
@@ -93,52 +105,67 @@ extern "C" HRESULT PayloadsParseFromXml(
                 hr = ContainerFindById(pContainers, scz, &pPayload->pContainer);
                 ExitOnFailure(hr, "Failed to to find container: %ls", scz);
             }
-        }
 
-        // @LayoutOnly
-        hr = XmlGetYesNoAttribute(pixnNode, L"LayoutOnly", &pPayload->fLayoutOnly);
-        if (E_NOTFOUND != hr)
-        {
-            ExitOnFailure(hr, "Failed to get @LayoutOnly.");
-        }
+            // @LayoutOnly
+            hr = XmlGetYesNoAttribute(pixnNode, L"LayoutOnly", &pPayload->fLayoutOnly);
+            if (E_NOTFOUND != hr)
+            {
+                ExitOnFailure(hr, "Failed to get @LayoutOnly.");
+            }
 
-        // @SourcePath
-        hr = XmlGetAttributeEx(pixnNode, L"SourcePath", &pPayload->sczSourcePath);
-        ExitOnFailure(hr, "Failed to get @SourcePath.");
+            // @DownloadUrl
+            hr = XmlGetAttributeEx(pixnNode, L"DownloadUrl", &pPayload->downloadSource.sczUrl);
+            if (E_NOTFOUND != hr)
+            {
+                ExitOnFailure(hr, "Failed to get @DownloadUrl.");
+            }
 
-        // @DownloadUrl
-        hr = XmlGetAttributeEx(pixnNode, L"DownloadUrl", &pPayload->downloadSource.sczUrl);
-        if (E_NOTFOUND != hr)
-        {
-            ExitOnFailure(hr, "Failed to get @DownloadUrl.");
-        }
+            // @FileSize
+            hr = XmlGetAttributeEx(pixnNode, L"FileSize", &scz);
+            if (E_NOTFOUND != hr)
+            {
+                ExitOnFailure(hr, "Failed to get @FileSize.");
 
-        // @FileSize
-        hr = XmlGetAttributeEx(pixnNode, L"FileSize", &scz);
-        if (E_NOTFOUND != hr)
-        {
-            ExitOnFailure(hr, "Failed to get @FileSize.");
+                hr = StrStringToUInt64(scz, 0, &pPayload->qwFileSize);
+                ExitOnFailure(hr, "Failed to parse @FileSize.");
 
-            hr = StrStringToUInt64(scz, 0, &pPayload->qwFileSize);
-            ExitOnFailure(hr, "Failed to parse @FileSize.");
-        }
+                fValidFileSize = TRUE;
+            }
+
+            // @Hash
+            hr = XmlGetAttributeEx(pixnNode, L"Hash", &scz);
+            if (E_NOTFOUND != hr)
+            {
+                ExitOnFailure(hr, "Failed to get @Hash.");
 
-        // @Hash
-        hr = XmlGetAttributeEx(pixnNode, L"Hash", &scz);
-        ExitOnFailure(hr, "Failed to get @Hash.");
+                hr = StrAllocHexDecode(scz, &pPayload->pbHash, &pPayload->cbHash);
+                ExitOnFailure(hr, "Failed to hex decode the Payload/@Hash.");
 
-        hr = StrAllocHexDecode(scz, &pPayload->pbHash, &pPayload->cbHash);
-        ExitOnFailure(hr, "Failed to hex decode the Payload/@Hash.");
+                if (BURN_PAYLOAD_VERIFICATION_NONE == pPayload->verification)
+                {
+                    pPayload->verification = BURN_PAYLOAD_VERIFICATION_HASH;
+                }
+            }
 
-        if (pPayload->fLayoutOnly && pLayoutPayloads)
-        {
-            hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&pLayoutPayloads->rgItems), pLayoutPayloads->cItems + 1, sizeof(BURN_PAYLOAD_GROUP_ITEM), 5);
-            ExitOnFailure(hr, "Failed to allocate memory for layout payloads.");
+            if (BURN_PAYLOAD_VERIFICATION_NONE == pPayload->verification)
+            {
+                ExitOnRootFailure(hr = E_INVALIDDATA, "There was no verification information for payload: %ls", pPayload->sczKey);
+            }
+            else if (BURN_PAYLOAD_VERIFICATION_HASH == pPayload->verification && !fValidFileSize)
+            {
+                ExitOnRootFailure(hr = E_INVALIDDATA, "File size is required when verifying by hash for payload: %ls", pPayload->sczKey);
+            }
 
-            pLayoutPayloads->rgItems[pLayoutPayloads->cItems].pPayload = pPayload;
-            ++pLayoutPayloads->cItems;
+            if (pPayload->fLayoutOnly)
+            {
+                hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&pLayoutPayloads->rgItems), pLayoutPayloads->cItems + 1, sizeof(BURN_PAYLOAD_GROUP_ITEM), 5);
+                ExitOnFailure(hr, "Failed to allocate memory for layout payloads.");
 
-            pLayoutPayloads->qwTotalSize += pPayload->qwFileSize;
+                pLayoutPayloads->rgItems[pLayoutPayloads->cItems].pPayload = pPayload;
+                ++pLayoutPayloads->cItems;
+
+                pLayoutPayloads->qwTotalSize += pPayload->qwFileSize;
+            }
         }
 
         // prepare next iteration
diff --git a/src/burn/engine/payload.h b/src/burn/engine/payload.h
index f28b437f..6fc6de5e 100644
--- a/src/burn/engine/payload.h
+++ b/src/burn/engine/payload.h
@@ -23,6 +23,13 @@ enum BURN_PAYLOAD_STATE
     BURN_PAYLOAD_STATE_CACHED,
 };
 
+enum BURN_PAYLOAD_VERIFICATION
+{
+    BURN_PAYLOAD_VERIFICATION_NONE,
+    BURN_PAYLOAD_VERIFICATION_HASH,
+    BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE,
+};
+
 
 // structs
 
@@ -36,6 +43,7 @@ typedef struct _BURN_PAYLOAD
 
     BYTE* pbHash;
     DWORD cbHash;
+    BURN_PAYLOAD_VERIFICATION verification;
 
     LPWSTR sczSourcePath;
     BURN_CONTAINER* pContainer;
diff --git a/src/burn/engine/pseudobundle.cpp b/src/burn/engine/pseudobundle.cpp
index 180cc621..1b2eae75 100644
--- a/src/burn/engine/pseudobundle.cpp
+++ b/src/burn/engine/pseudobundle.cpp
@@ -69,6 +69,11 @@ extern "C" HRESULT PseudoBundleInitialize(
         memcpy_s(pPayload->pbHash, pPayload->cbHash, pbHash, cbHash);
     }
 
+    if (BOOTSTRAPPER_RELATION_UPDATE == relationType)
+    {
+        pPayload->verification = BURN_PAYLOAD_VERIFICATION_UPDATE_BUNDLE;
+    }
+
     pPackage->Exe.fPseudoBundle = TRUE;
 
     pPackage->type = BURN_PACKAGE_TYPE_EXE;
diff --git a/src/burn/test/BurnUnitTest/CacheTest.cpp b/src/burn/test/BurnUnitTest/CacheTest.cpp
index d0cc237f..e9ad555b 100644
--- a/src/burn/test/BurnUnitTest/CacheTest.cpp
+++ b/src/burn/test/BurnUnitTest/CacheTest.cpp
@@ -71,6 +71,8 @@ namespace Bootstrapper
                 payload.sczFilePath = L"CacheSignatureTest.File";
                 payload.pbHash = pb;
                 payload.cbHash = cb;
+                payload.qwFileSize = 27;
+                payload.verification = BURN_PAYLOAD_VERIFICATION_HASH;
 
                 hr = CacheCompletePayload(package.fPerMachine, &payload, package.sczCacheId, sczPayloadPath, FALSE, CacheTestEventRoutine, CacheTestProgressRoutine, &context);
                 Assert::Equal(S_OK, hr);
diff --git a/src/burn/test/BurnUnitTest/PlanTest.cpp b/src/burn/test/BurnUnitTest/PlanTest.cpp
index 001acaee..045c510e 100644
--- a/src/burn/test/BurnUnitTest/PlanTest.cpp
+++ b/src/burn/test/BurnUnitTest/PlanTest.cpp
@@ -71,7 +71,7 @@ namespace Bootstrapper
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
             Assert::Equal(107082ull, pPlan->qwEstimatedSize);
-            Assert::Equal(506145ull, pPlan->qwCacheSizeTotal);
+            Assert::Equal(522548ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
             dwIndex = 0;