]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/CapsulePei/Common/CapsuleCoalesce.c
MdeModulePkg: Clean up source files
[mirror_edk2.git] / MdeModulePkg / Universal / CapsulePei / Common / CapsuleCoalesce.c
index a32b6d664582d4b70972df2c59a9822bdb4270f2..3575a94d0fb2f68c094165e29a65ab05a77a30d9 100644 (file)
@@ -1,7 +1,16 @@
 /** @file\r
   The logic to process capsule.\r
 \r
-Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>\r
+  Caution: This module requires additional review when modified.\r
+  This driver will have external input - capsule image.\r
+  This external input must be validated carefully to avoid security issue like\r
+  buffer overflow, integer overflow.\r
+\r
+  CapsuleDataCoalesce() will do basic validation before coalesce capsule data\r
+  into memory.\r
+\r
+(C) Copyright 2014 Hewlett-Packard Development Company, L.P.<BR>\r
+Copyright (c) 2011 - 2018, Intel Corporation. All rights reserved.<BR>\r
 This program and the accompanying materials\r
 are licensed and made available under the terms and conditions of the BSD License\r
 which accompanies this distribution.  The full text of the license may be found at\r
@@ -22,15 +31,9 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/PrintLib.h>\r
 #include <Library/BaseLib.h>\r
 \r
-#define MIN_COALESCE_ADDR                     (1024 * 1024)\r
-#define MAX_SUPPORT_CAPSULE_NUM               50\r
-\r
-#define EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE SIGNATURE_32 ('C', 'a', 'p', 'D')\r
+#include "CommonHeader.h"\r
 \r
-typedef struct {\r
-  UINT32  Signature;\r
-  UINT32  CapsuleSize;\r
-} EFI_CAPSULE_PEIM_PRIVATE_DATA;\r
+#define MIN_COALESCE_ADDR                     (1024 * 1024)\r
 \r
 /**\r
   Given a pointer to the capsule block list, info on the available system\r
@@ -42,7 +45,7 @@ typedef struct {
   @param MemSize     The size of the block of memory pointed to by MemBase\r
   @param DataSize    How big a free block we want to find\r
 \r
-  @return A pointer to a memory block of at least DataSize that lies somewhere \r
+  @return A pointer to a memory block of at least DataSize that lies somewhere\r
           between MemBase and (MemBase + MemSize). The memory pointed to does not\r
           contain any of the capsule block descriptors or capsule blocks pointed to\r
           by the BlockList.\r
@@ -56,20 +59,6 @@ FindFreeMem (
   UINTN                             DataSize\r
   );\r
 \r
-/**\r
-  Check the integrity of the capsule descriptors.\r
-\r
-  @param BlockList    Pointer to the capsule descriptors\r
-\r
-  @retval NULL           BlockList is not valid.\r
-  @retval LastBlockDesc  Last one Block in BlockList\r
-\r
-**/\r
-EFI_CAPSULE_BLOCK_DESCRIPTOR *\r
-ValidateCapsuleIntegrity (\r
-  IN EFI_CAPSULE_BLOCK_DESCRIPTOR    *BlockList\r
-  );\r
-\r
 /**\r
   The capsule block descriptors may be fragmented and spread all over memory.\r
   To simplify the coalescing of capsule blocks, first coalesce all the\r
@@ -79,19 +68,21 @@ ValidateCapsuleIntegrity (
   they are relocated into memory to turn them into a contiguous (null\r
   terminated) array.\r
 \r
-  @param PeiServices pointer to PEI services table\r
-  @param BlockList   pointer to the capsule block descriptors\r
-  @param MemBase     base of system memory in which we can work\r
-  @param MemSize     size of the system memory pointed to by MemBase\r
+  @param PeiServices    pointer to PEI services table\r
+  @param BlockList      pointer to the capsule block descriptors\r
+  @param NumDescriptors number of capsule data block descriptors, whose Length is non-zero.\r
+  @param MemBase        base of system memory in which we can work\r
+  @param MemSize        size of the system memory pointed to by MemBase\r
 \r
   @retval NULL    could not relocate the descriptors\r
-  @retval Pointer to the base of the successfully-relocated block descriptors. \r
+  @retval Pointer to the base of the successfully-relocated block descriptors.\r
 \r
 **/\r
 EFI_CAPSULE_BLOCK_DESCRIPTOR *\r
 RelocateBlockDescriptors (\r
   IN EFI_PEI_SERVICES                  **PeiServices,\r
   IN EFI_CAPSULE_BLOCK_DESCRIPTOR      *BlockList,\r
+  IN UINTN                              NumDescriptors,\r
   IN UINT8                             *MemBase,\r
   IN UINTN                             MemSize\r
   );\r
@@ -102,7 +93,7 @@ RelocateBlockDescriptors (
   @param CapsuleHeader   The pointer to EFI_CAPSULE_HEADER\r
 \r
   @retval FALSE  Capsule is OK\r
-  @retval TRUE   Capsule is corrupted \r
+  @retval TRUE   Capsule is corrupted\r
 \r
 **/\r
 BOOLEAN\r
@@ -136,10 +127,9 @@ IsOverlapped (
   refers to is.\r
 \r
   @param Desc            Pointer to the capsule block descriptors\r
-                         NumDescriptors  - optional pointer to where to return the number of descriptors\r
-                         CapsuleSize     - optional pointer to where to return the capsule size\r
-  @param NumDescriptors  Optional pointer to where to return the number of descriptors\r
-  @param CapsuleSize     Optional pointer to where to return the capsule size\r
+  @param NumDescriptors  Optional pointer to where to return the number of capsule data descriptors, whose Length is non-zero.\r
+  @param CapsuleSize     Optional pointer to where to return the capsule image size\r
+  @param CapsuleNumber   Optional pointer to where to return the number of capsule\r
 \r
   @retval EFI_NOT_FOUND   No descriptors containing data in the list\r
   @retval EFI_SUCCESS     Return data is valid\r
@@ -149,7 +139,8 @@ EFI_STATUS
 GetCapsuleInfo (\r
   IN EFI_CAPSULE_BLOCK_DESCRIPTOR   *Desc,\r
   IN OUT UINTN                      *NumDescriptors OPTIONAL,\r
-  IN OUT UINTN                      *CapsuleSize OPTIONAL\r
+  IN OUT UINTN                      *CapsuleSize OPTIONAL,\r
+  IN OUT UINTN                      *CapsuleNumber OPTIONAL\r
   );\r
 \r
 /**\r
@@ -162,7 +153,7 @@ GetCapsuleInfo (
   @param MemSize     The size of the block of memory pointed to by MemBase\r
   @param DataSize    How big a free block we want to find\r
 \r
-  @return A pointer to a memory block of at least DataSize that lies somewhere \r
+  @return A pointer to a memory block of at least DataSize that lies somewhere\r
           between MemBase and (MemBase + MemSize). The memory pointed to does not\r
           contain any of the capsule block descriptors or capsule blocks pointed to\r
           by the BlockList.\r
@@ -242,10 +233,69 @@ FindFreeMem (
   return MemBase;\r
 }\r
 \r
+/**\r
+  Validate capsule by MemoryResource.\r
+\r
+  @param MemoryResource  Pointer to the buffer of memory resource descriptor.\r
+  @param Address         Address to be validated.\r
+  @param Size            Size to be validated.\r
+\r
+  @retval TRUE  No memory resource descriptor reported in HOB list before capsule Coalesce,\r
+                or it is valid in one MemoryResource.\r
+          FALSE It is not in any MemoryResource.\r
+\r
+**/\r
+BOOLEAN\r
+ValidateCapsuleByMemoryResource (\r
+  IN MEMORY_RESOURCE_DESCRIPTOR     *MemoryResource,\r
+  IN EFI_PHYSICAL_ADDRESS           Address,\r
+  IN UINT64                         Size\r
+  )\r
+{\r
+  UINTN             Index;\r
+\r
+  //\r
+  // Sanity Check\r
+  //\r
+  if (Size > MAX_ADDRESS) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: Size(0x%lx) > MAX_ADDRESS\n", Size));\r
+    return FALSE;\r
+  }\r
+\r
+  //\r
+  // Sanity Check\r
+  //\r
+  if (Address > (MAX_ADDRESS - Size)) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: Address(0x%lx) > (MAX_ADDRESS - Size(0x%lx))\n", Address, Size));\r
+    return FALSE;\r
+  }\r
+\r
+  if (MemoryResource == NULL) {\r
+    //\r
+    // No memory resource descriptor reported in HOB list before capsule Coalesce.\r
+    //\r
+    return TRUE;\r
+  }\r
+\r
+  for (Index = 0; MemoryResource[Index].ResourceLength != 0; Index++) {\r
+    if ((Address >= MemoryResource[Index].PhysicalStart) &&\r
+        ((Address + Size) <= (MemoryResource[Index].PhysicalStart + MemoryResource[Index].ResourceLength))) {\r
+      DEBUG ((EFI_D_INFO, "Address(0x%lx) Size(0x%lx) in MemoryResource[0x%x] - Start(0x%lx) Length(0x%lx)\n",\r
+                          Address, Size,\r
+                          Index, MemoryResource[Index].PhysicalStart, MemoryResource[Index].ResourceLength));\r
+      return TRUE;\r
+    }\r
+  }\r
+\r
+  DEBUG ((EFI_D_ERROR, "ERROR: Address(0x%lx) Size(0x%lx) not in any MemoryResource\n", Address, Size));\r
+  return FALSE;\r
+}\r
+\r
 /**\r
   Check the integrity of the capsule descriptors.\r
 \r
-  @param BlockList    Pointer to the capsule descriptors\r
+  @param BlockList       Pointer to the capsule descriptors\r
+  @param MemoryResource  Pointer to the buffer of memory resource descriptor.\r
 \r
   @retval NULL           BlockList is not valid.\r
   @retval LastBlockDesc  Last one Block in BlockList\r
@@ -253,29 +303,45 @@ FindFreeMem (
 **/\r
 EFI_CAPSULE_BLOCK_DESCRIPTOR *\r
 ValidateCapsuleIntegrity (\r
-  IN EFI_CAPSULE_BLOCK_DESCRIPTOR    *BlockList\r
+  IN EFI_CAPSULE_BLOCK_DESCRIPTOR    *BlockList,\r
+  IN MEMORY_RESOURCE_DESCRIPTOR      *MemoryResource\r
   )\r
 {\r
   EFI_CAPSULE_HEADER             *CapsuleHeader;\r
   UINT64                         CapsuleSize;\r
-  UINT32                         CapsuleCount;\r
+  UINT                         CapsuleCount;\r
   EFI_CAPSULE_BLOCK_DESCRIPTOR   *Ptr;\r
 \r
+  DEBUG ((EFI_D_INFO, "ValidateCapsuleIntegrity\n"));\r
+\r
   //\r
   // Go through the list to look for inconsistencies. Check for:\r
   //   * misaligned block descriptors.\r
   //   * The first capsule header guid\r
   //   * The first capsule header flag\r
-  //   * Data + Length < Data (wrap)\r
+  //   * The first capsule header HeaderSize\r
+  //   * Below check will be done in ValidateCapsuleByMemoryResource()\r
+  //     Length > MAX_ADDRESS\r
+  //     Ptr + sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR) > MAX_ADDRESS\r
+  //     DataBlock + Length > MAX_ADDRESS\r
+  //\r
   CapsuleSize  = 0;\r
   CapsuleCount = 0;\r
   Ptr = BlockList;\r
+\r
+  if (!ValidateCapsuleByMemoryResource (MemoryResource, (EFI_PHYSICAL_ADDRESS) (UINTN) Ptr, sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {\r
+    return NULL;\r
+  }\r
+\r
+  DEBUG ((EFI_D_INFO, "Ptr - 0x%x\n", Ptr));\r
+  DEBUG ((EFI_D_INFO, "Ptr->Length - 0x%x\n", Ptr->Length));\r
+  DEBUG ((EFI_D_INFO, "Ptr->Union - 0x%x\n", Ptr->Union.ContinuationPointer));\r
   while ((Ptr->Length != 0) || (Ptr->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
     //\r
     // Make sure the descriptor is aligned at UINT64 in memory\r
     //\r
-    if ((UINTN) Ptr & 0x07) {\r
-      DEBUG ((EFI_D_ERROR, "BlockList address failed alignment check\n"));\r
+    if ((UINTN) Ptr & (sizeof(UINT64) - 1)) {\r
+      DEBUG ((EFI_D_ERROR, "ERROR: BlockList address failed alignment check\n"));\r
       return NULL;\r
     }\r
 \r
@@ -285,7 +351,17 @@ ValidateCapsuleIntegrity (
       // else.\r
       //\r
       Ptr = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) Ptr->Union.ContinuationPointer;\r
+      if (!ValidateCapsuleByMemoryResource (MemoryResource, (EFI_PHYSICAL_ADDRESS) (UINTN) Ptr, sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {\r
+        return NULL;\r
+      }\r
+      DEBUG ((EFI_D_INFO, "Ptr(C) - 0x%x\n", Ptr));\r
+      DEBUG ((EFI_D_INFO, "Ptr->Length - 0x%x\n", Ptr->Length));\r
+      DEBUG ((EFI_D_INFO, "Ptr->Union - 0x%x\n", Ptr->Union.ContinuationPointer));\r
     } else {\r
+      if (!ValidateCapsuleByMemoryResource (MemoryResource, Ptr->Union.DataBlock, Ptr->Length)) {\r
+        return NULL;\r
+      }\r
+\r
       //\r
       //To enhance the reliability of check-up, the first capsule's header is checked here.\r
       //More reliabilities check-up will do later.\r
@@ -295,29 +371,63 @@ ValidateCapsuleIntegrity (
         //Move to the first capsule to check its header.\r
         //\r
         CapsuleHeader = (EFI_CAPSULE_HEADER*)((UINTN)Ptr->Union.DataBlock);\r
+        //\r
+        // Sanity check\r
+        //\r
+        if (Ptr->Length < sizeof(EFI_CAPSULE_HEADER)) {\r
+          DEBUG ((EFI_D_ERROR, "ERROR: Ptr->Length(0x%lx) < sizeof(EFI_CAPSULE_HEADER)\n", Ptr->Length));\r
+          return NULL;\r
+        }\r
+        //\r
+        // Make sure HeaderSize field is valid\r
+        //\r
+        if (CapsuleHeader->HeaderSize > CapsuleHeader->CapsuleImageSize) {\r
+          DEBUG ((EFI_D_ERROR, "ERROR: CapsuleHeader->HeaderSize(0x%x) > CapsuleHeader->CapsuleImageSize(0x%x)\n", CapsuleHeader->HeaderSize, CapsuleHeader->CapsuleImageSize));\r
+          return NULL;\r
+        }\r
         if (IsCapsuleCorrupted (CapsuleHeader)) {\r
           return NULL;\r
         }\r
         CapsuleCount ++;\r
         CapsuleSize = CapsuleHeader->CapsuleImageSize;\r
+      }\r
+\r
+      if (CapsuleSize >= Ptr->Length) {\r
+        CapsuleSize = CapsuleSize - Ptr->Length;\r
       } else {\r
-        if (CapsuleSize >= Ptr->Length) {\r
-          CapsuleSize = CapsuleSize - Ptr->Length;\r
-        } else {\r
-          CapsuleSize = 0;\r
-        }\r
+        DEBUG ((EFI_D_ERROR, "ERROR: CapsuleSize(0x%lx) < Ptr->Length(0x%lx)\n", CapsuleSize, Ptr->Length));\r
+        //\r
+        // Sanity check\r
+        //\r
+        return NULL;\r
       }\r
+\r
       //\r
       // Move to next BLOCK descriptor\r
       //\r
       Ptr++;\r
+      if (!ValidateCapsuleByMemoryResource (MemoryResource, (EFI_PHYSICAL_ADDRESS) (UINTN) Ptr, sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {\r
+        return NULL;\r
+      }\r
+      DEBUG ((EFI_D_INFO, "Ptr(B) - 0x%x\n", Ptr));\r
+      DEBUG ((EFI_D_INFO, "Ptr->Length - 0x%x\n", Ptr->Length));\r
+      DEBUG ((EFI_D_INFO, "Ptr->Union - 0x%x\n", Ptr->Union.ContinuationPointer));\r
     }\r
   }\r
 \r
   if (CapsuleCount == 0) {\r
     //\r
-    // No any capsule is found in BlockList.\r
+    // No any capsule is found in BlockList\r
+    //\r
+    DEBUG ((EFI_D_ERROR, "ERROR: CapsuleCount(0x%x) == 0\n", CapsuleCount));\r
+    return NULL;\r
+  }\r
+\r
+  if (CapsuleSize != 0) {\r
+    //\r
+    // Capsule data is incomplete.\r
     //\r
+    DEBUG ((EFI_D_ERROR, "ERROR: CapsuleSize(0x%lx) != 0\n", CapsuleSize));\r
     return NULL;\r
   }\r
 \r
@@ -333,19 +443,21 @@ ValidateCapsuleIntegrity (
   they are relocated into memory to turn them into a contiguous (null\r
   terminated) array.\r
 \r
-  @param PeiServices pointer to PEI services table\r
-  @param BlockList   pointer to the capsule block descriptors\r
-  @param MemBase     base of system memory in which we can work\r
-  @param MemSize     size of the system memory pointed to by MemBase\r
+  @param PeiServices    pointer to PEI services table\r
+  @param BlockList      pointer to the capsule block descriptors\r
+  @param NumDescriptors number of capsule data block descriptors, whose Length is non-zero.\r
+  @param MemBase        base of system memory in which we can work\r
+  @param MemSize        size of the system memory pointed to by MemBase\r
 \r
   @retval NULL    could not relocate the descriptors\r
-  @retval Pointer to the base of the successfully-relocated block descriptors. \r
+  @retval Pointer to the base of the successfully-relocated block descriptors.\r
 \r
 **/\r
 EFI_CAPSULE_BLOCK_DESCRIPTOR  *\r
 RelocateBlockDescriptors (\r
   IN EFI_PEI_SERVICES                   **PeiServices,\r
   IN EFI_CAPSULE_BLOCK_DESCRIPTOR       *BlockList,\r
+  IN UINTN                              NumDescriptors,\r
   IN UINT8                              *MemBase,\r
   IN UINTN                              MemSize\r
   )\r
@@ -354,21 +466,17 @@ RelocateBlockDescriptors (
   EFI_CAPSULE_BLOCK_DESCRIPTOR   *CurrBlockDescHead;\r
   EFI_CAPSULE_BLOCK_DESCRIPTOR   *TempBlockDesc;\r
   EFI_CAPSULE_BLOCK_DESCRIPTOR   *PrevBlockDescTail;\r
-  UINTN                          NumDescriptors;\r
   UINTN                          BufferSize;\r
   UINT8                          *RelocBuffer;\r
   UINTN                          BlockListSize;\r
+\r
   //\r
   // Get the info on the blocks and descriptors. Since we're going to move\r
   // the descriptors low in memory, adjust the base/size values accordingly here.\r
-  // GetCapsuleInfo() returns the number of legit descriptors, so add one for\r
-  // a terminator.\r
+  // NumDescriptors is the number of legit data descriptors, so add one for\r
+  // a terminator. (Already done by caller, no check is needed.)\r
   //\r
-  if (GetCapsuleInfo (BlockList, &NumDescriptors, NULL) != EFI_SUCCESS) {\r
-    return NULL;\r
-  }\r
 \r
-  NumDescriptors++;\r
   BufferSize    = NumDescriptors * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
   NewBlockList  = (EFI_CAPSULE_BLOCK_DESCRIPTOR *) MemBase;\r
   if (MemSize < BufferSize) {\r
@@ -407,12 +515,11 @@ RelocateBlockDescriptors (
         }\r
 \r
         CopyMem ((VOID *) RelocBuffer, (VOID *) (UINTN) TempBlockDesc->Union.DataBlock, (UINTN) TempBlockDesc->Length);\r
+        DEBUG ((EFI_D_INFO, "Capsule relocate descriptors from/to/size  0x%lX 0x%lX 0x%lX\n", TempBlockDesc->Union.DataBlock, (UINT64)(UINTN)RelocBuffer, TempBlockDesc->Length));\r
         TempBlockDesc->Union.DataBlock = (EFI_PHYSICAL_ADDRESS) (UINTN) RelocBuffer;\r
-\r
-        DEBUG ((EFI_D_INFO, "Capsule relocate descriptors from/to/size  0x%X 0x%X 0x%X\n", (UINT32)(UINTN)TempBlockDesc->Union.DataBlock, (UINT32)(UINTN)RelocBuffer, (UINT32)(UINTN)TempBlockDesc->Length));\r
       }\r
+      TempBlockDesc++;\r
     }\r
-    TempBlockDesc++;\r
   }\r
   //\r
   // Now go through all the block descriptors to make sure that they're not\r
@@ -522,10 +629,9 @@ IsOverlapped (
   refers to is.\r
 \r
   @param Desc            Pointer to the capsule block descriptors\r
-                         NumDescriptors  - optional pointer to where to return the number of descriptors\r
-                         CapsuleSize     - optional pointer to where to return the capsule size\r
-  @param NumDescriptors  Optional pointer to where to return the number of descriptors\r
-  @param CapsuleSize     Optional pointer to where to return the capsule size\r
+  @param NumDescriptors  Optional pointer to where to return the number of capsule data descriptors, whose Length is non-zero.\r
+  @param CapsuleSize     Optional pointer to where to return the capsule image size\r
+  @param CapsuleNumber   Optional pointer to where to return the number of capsule\r
 \r
   @retval EFI_NOT_FOUND   No descriptors containing data in the list\r
   @retval EFI_SUCCESS     Return data is valid\r
@@ -535,16 +641,24 @@ EFI_STATUS
 GetCapsuleInfo (\r
   IN EFI_CAPSULE_BLOCK_DESCRIPTOR   *Desc,\r
   IN OUT UINTN                      *NumDescriptors OPTIONAL,\r
-  IN OUT UINTN                      *CapsuleSize OPTIONAL\r
+  IN OUT UINTN                      *CapsuleSize OPTIONAL,\r
+  IN OUT UINTN                      *CapsuleNumber OPTIONAL\r
   )\r
 {\r
-  UINTN Count;\r
-  UINTN Size;\r
+  UINTN                          Count;\r
+  UINTN                          Size;\r
+  UINTN                          Number;\r
+  UINTN                          ThisCapsuleImageSize;\r
+  EFI_CAPSULE_HEADER             *CapsuleHeader;\r
+\r
+  DEBUG ((EFI_D_INFO, "GetCapsuleInfo enter\n"));\r
 \r
   ASSERT (Desc != NULL);\r
 \r
   Count = 0;\r
   Size  = 0;\r
+  Number = 0;\r
+  ThisCapsuleImageSize = 0;\r
 \r
   while (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL) {\r
     if (Desc->Length == 0) {\r
@@ -553,8 +667,39 @@ GetCapsuleInfo (
       //\r
       Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) Desc->Union.ContinuationPointer;\r
     } else {\r
+      //\r
+      // Sanity Check\r
+      // It is needed, because ValidateCapsuleIntegrity() only validate one individual capsule Size.\r
+      // While here we need check all capsules size.\r
+      //\r
+      if (Desc->Length >= (MAX_ADDRESS - Size)) {\r
+        DEBUG ((EFI_D_ERROR, "ERROR: Desc->Length(0x%lx) >= (MAX_ADDRESS - Size(0x%x))\n", Desc->Length, Size));\r
+        return EFI_OUT_OF_RESOURCES;\r
+      }\r
       Size += (UINTN) Desc->Length;\r
       Count++;\r
+\r
+      //\r
+      // See if this is first capsule's header\r
+      //\r
+      if (ThisCapsuleImageSize == 0) {\r
+        CapsuleHeader = (EFI_CAPSULE_HEADER*)((UINTN)Desc->Union.DataBlock);\r
+        //\r
+        // This has been checked in ValidateCapsuleIntegrity()\r
+        //\r
+        Number ++;\r
+        ThisCapsuleImageSize = CapsuleHeader->CapsuleImageSize;\r
+      }\r
+\r
+      //\r
+      // This has been checked in ValidateCapsuleIntegrity()\r
+      //\r
+      ASSERT (ThisCapsuleImageSize >= Desc->Length);\r
+      ThisCapsuleImageSize = (UINTN)(ThisCapsuleImageSize - Desc->Length);\r
+\r
+      //\r
+      // Move to next\r
+      //\r
       Desc++;\r
     }\r
   }\r
@@ -562,9 +707,15 @@ GetCapsuleInfo (
   // If no descriptors, then fail\r
   //\r
   if (Count == 0) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: Count == 0\n"));\r
     return EFI_NOT_FOUND;\r
   }\r
 \r
+  //\r
+  // checked in ValidateCapsuleIntegrity()\r
+  //\r
+  ASSERT (ThisCapsuleImageSize == 0);\r
+\r
   if (NumDescriptors != NULL) {\r
     *NumDescriptors = Count;\r
   }\r
@@ -573,6 +724,10 @@ GetCapsuleInfo (
     *CapsuleSize = Size;\r
   }\r
 \r
+  if (CapsuleNumber != NULL) {\r
+    *CapsuleNumber = Number;\r
+  }\r
+\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -582,7 +737,7 @@ GetCapsuleInfo (
   @param CapsuleHeader   The pointer to EFI_CAPSULE_HEADER\r
 \r
   @retval FALSE  Capsule is OK\r
-  @retval TRUE   Capsule is corrupted \r
+  @retval TRUE   Capsule is corrupted\r
 \r
 **/\r
 BOOLEAN\r
@@ -614,7 +769,7 @@ IsCapsuleCorrupted (
   capsule gets coalesced. This can be useful in narrowing down\r
   where capsule data corruption occurs.\r
 \r
-  The test pattern mode fills in memory with a counting UINT32 value. \r
+  The test pattern mode fills in memory with a counting UINT32 value.\r
   If the capsule is not divided up in a multiple of 4-byte blocks, then\r
   things get messy doing the check. Therefore there are some cases\r
   here where we just give up and skip the pre-coalesce check.\r
@@ -631,6 +786,9 @@ CapsuleTestPatternPreCoalesce (
   UINT32  *TestPtr;\r
   UINT32  TestCounter;\r
   UINT32  TestSize;\r
+\r
+  DEBUG ((EFI_D_INFO, "CapsuleTestPatternPreCoalesce\n"));\r
+\r
   //\r
   // Find first data descriptor\r
   //\r
@@ -697,6 +855,7 @@ CapsuleTestPatternPreCoalesce (
   Get capsule descriptors from variable CapsuleUpdateData, CapsuleUpdateData1, CapsuleUpdateData2...\r
 \r
   @param BlockListBuffer            Pointer to the buffer of capsule descriptors variables\r
+  @param MemoryResource             Pointer to the buffer of memory resource descriptor.\r
   @param BlockDescriptorList        Pointer to the capsule descriptors list\r
 \r
   @retval EFI_SUCCESS               a valid capsule is present\r
@@ -705,7 +864,8 @@ CapsuleTestPatternPreCoalesce (
 EFI_STATUS\r
 BuildCapsuleDescriptors (\r
   IN  EFI_PHYSICAL_ADDRESS            *BlockListBuffer,\r
-  OUT EFI_CAPSULE_BLOCK_DESCRIPTOR    **BlockDescriptorList \r
+  IN  MEMORY_RESOURCE_DESCRIPTOR      *MemoryResource,\r
+  OUT EFI_CAPSULE_BLOCK_DESCRIPTOR    **BlockDescriptorList\r
   )\r
 {\r
   UINTN                            Index;\r
@@ -713,42 +873,42 @@ BuildCapsuleDescriptors (
   EFI_CAPSULE_BLOCK_DESCRIPTOR     *TempBlock;\r
   EFI_CAPSULE_BLOCK_DESCRIPTOR     *HeadBlock;\r
 \r
+  DEBUG ((EFI_D_INFO, "BuildCapsuleDescriptors enter\n"));\r
+\r
   LastBlock         = NULL;\r
   HeadBlock         = NULL;\r
   TempBlock         = NULL;\r
   Index             = 0;\r
 \r
   while (BlockListBuffer[Index] != 0) {\r
-    if (Index == 0) {\r
-      //\r
-      // For the first Capsule Image, test integrity of descriptors.\r
-      //\r
-      LastBlock = ValidateCapsuleIntegrity ((EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index]);\r
-      if (LastBlock == NULL) {\r
-        return EFI_NOT_FOUND;\r
-      }\r
-      //\r
-      // Return the base of the block descriptors\r
-      //\r
-      HeadBlock = (EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index];\r
-    } else {        \r
-      //\r
-      // Test integrity of descriptors.\r
-      //\r
-      TempBlock = ValidateCapsuleIntegrity ((EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index]);\r
-      if (TempBlock == NULL) {\r
-        return EFI_NOT_FOUND;\r
+    //\r
+    // Test integrity of descriptors.\r
+    //\r
+    if (BlockListBuffer[Index] < MAX_ADDRESS) {\r
+      TempBlock = ValidateCapsuleIntegrity ((EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index], MemoryResource);\r
+      if (TempBlock != NULL) {\r
+        if (LastBlock == NULL) {\r
+          LastBlock = TempBlock;\r
+\r
+          //\r
+          // Return the base of the block descriptors\r
+          //\r
+          HeadBlock = (EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index];\r
+        } else {\r
+          //\r
+          // Combine the different BlockList into single BlockList.\r
+          //\r
+          LastBlock->Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)BlockListBuffer[Index];\r
+          LastBlock->Length          = 0;\r
+          LastBlock                  = TempBlock;\r
+        }\r
       }\r
-      //\r
-      // Combine the different BlockList into single BlockList.\r
-      //\r
-      LastBlock->Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)BlockListBuffer[Index];\r
-      LastBlock->Length          = 0;\r
-      LastBlock                  = TempBlock;\r
+    } else {\r
+      DEBUG ((EFI_D_ERROR, "ERROR: BlockListBuffer[Index](0x%lx) < MAX_ADDRESS\n", BlockListBuffer[Index]));\r
     }\r
     Index ++;\r
   }\r
-  \r
+\r
   if (HeadBlock != NULL) {\r
     *BlockDescriptorList = HeadBlock;\r
     return EFI_SUCCESS;\r
@@ -761,26 +921,32 @@ BuildCapsuleDescriptors (
 \r
   Memory Map for coalesced capsule:\r
   MemBase +   ---->+---------------------------+<-----------+\r
-  MemSize          |    CapsuleOffset[49]      |            |\r
-                   +---------------------------+            |\r
-                   |    ................       |            |\r
-                   +---------------------------+            |\r
-                   |    CapsuleOffset[2]       |            |\r
-                   +---------------------------+            |\r
-                   |    CapsuleOffset[1]       |            |\r
-                   +---------------------------+            |\r
-                   |    CapsuleOffset[0]       |       CapsuleSize     \r
-                   +---------------------------+            |\r
-                   |    CapsuleNumber          |            |\r
-                   +---------------------------+            |\r
-                   |                           |            |       \r
-                   |                           |            |       \r
-                   |    Capsule Image          |            |   \r
-                   |                           |            |       \r
-                   |                           |            |       \r
-                   +---------------------------+            |\r
+  MemSize          | ------------------------- |            |\r
+                   | |  Capsule [Num-1]      | |            |\r
+                   | ------------------------- |            |\r
+                   | |  ................     | |            |\r
+                   | ------------------------- |            |\r
+                   | |  Capsule [1]          | |            |\r
+                   | ------------------------- |            |\r
+                   | |  Capsule [0]          | |            |\r
+                   | ------------------------- |            |\r
+                   |    Capsule Image          |            |\r
+CapsuleImageBase-->+---------------------------+\r
+                   | ------------------------- |            |\r
+                   | |  CapsuleOffset[Num-1] | |            |\r
+                   | ------------------------- |            |\r
+                   | |  ................     | |        CapsuleSize\r
+                   | ------------------------- |            |\r
+                   | |  CapsuleOffset[1]     | |            |\r
+                   | ------------------------- |            |\r
+                   | |  CapsuleOffset[0]     | |            |\r
+                   |---------------------------|            |\r
+                   | |  CapsuleNumber        | |            |\r
+                   | ------------------------- |            |\r
+                   | |  CapsuleAllImageSize  | |            |\r
+                   | ------------------------- |            |\r
                    |    PrivateData            |            |\r
-   DestPtr  ---->  +---------------------------+<-----------+\r
+     DestPtr  ---->+---------------------------+<-----------+\r
                    |                           |            |\r
                    |     FreeMem               |        FreeMemSize\r
                    |                           |            |\r
@@ -798,8 +964,13 @@ BuildCapsuleDescriptors (
                    |    PrivateDataDesc 0      |\r
       MemBase ---->+---------------------------+<----- BlockList\r
 \r
+  Caution: This function may receive untrusted input.\r
+  The capsule data is external input, so this routine will do basic validation before\r
+  coalesce capsule data into memory.\r
+\r
   @param PeiServices        General purpose services available to every PEIM.\r
-  @param BlockListBuffer    Point to the buffer of Capsule Descriptor Variables.\r
+  @param BlockListBuffer    Pointer to the buffer of Capsule Descriptor Variables.\r
+  @param MemoryResource     Pointer to the buffer of memory resource descriptor.\r
   @param MemoryBase         Pointer to the base of a block of memory that we can walk\r
                             all over while trying to coalesce our buffers.\r
                             On output, this variable will hold the base address of\r
@@ -821,26 +992,26 @@ EFIAPI
 CapsuleDataCoalesce (\r
   IN EFI_PEI_SERVICES                **PeiServices,\r
   IN EFI_PHYSICAL_ADDRESS            *BlockListBuffer,\r
+  IN MEMORY_RESOURCE_DESCRIPTOR      *MemoryResource,\r
   IN OUT VOID                        **MemoryBase,\r
   IN OUT UINTN                       *MemorySize\r
   )\r
 {\r
   VOID                           *NewCapsuleBase;\r
-  VOID                           *DataPtr;\r
-  UINT8                          CapsuleIndex;\r
+  VOID                           *CapsuleImageBase;\r
+  UINTN                          CapsuleIndex;\r
   UINT8                          *FreeMemBase;\r
   UINT8                          *DestPtr;\r
+  UINTN                          DestLength;\r
   UINT8                          *RelocPtr;\r
-  UINT32                         CapsuleOffset[MAX_SUPPORT_CAPSULE_NUM]; \r
-  UINT32                         *AddDataPtr;\r
-  UINT32                         CapsuleTimes; \r
-  UINT64                         SizeLeft; \r
-  UINT64                         CapsuleImageSize; \r
+  UINTN                          CapsuleTimes;\r
+  UINT64                         SizeLeft;\r
+  UINT64                         CapsuleImageSize;\r
   UINTN                          CapsuleSize;\r
+  UINTN                          CapsuleNumber;\r
   UINTN                          DescriptorsSize;\r
   UINTN                          FreeMemSize;\r
   UINTN                          NumDescriptors;\r
-  BOOLEAN                        IsCorrupted;\r
   BOOLEAN                        CapsuleBeginFlag;\r
   EFI_STATUS                     Status;\r
   EFI_CAPSULE_HEADER             *CapsuleHeader;\r
@@ -851,22 +1022,22 @@ CapsuleDataCoalesce (
   EFI_CAPSULE_BLOCK_DESCRIPTOR   *TempBlockDesc;\r
   EFI_CAPSULE_BLOCK_DESCRIPTOR   PrivateDataDesc[2];\r
 \r
+  DEBUG ((EFI_D_INFO, "CapsuleDataCoalesce enter\n"));\r
+\r
   CapsuleIndex     = 0;\r
   SizeLeft         = 0;\r
   CapsuleTimes     = 0;\r
   CapsuleImageSize = 0;\r
   PrivateDataPtr   = NULL;\r
-  AddDataPtr       = NULL;\r
   CapsuleHeader    = NULL;\r
   CapsuleBeginFlag = TRUE;\r
-  IsCorrupted      = TRUE;\r
   CapsuleSize      = 0;\r
   NumDescriptors   = 0;\r
-  \r
+\r
   //\r
   // Build capsule descriptors list\r
   //\r
-  Status = BuildCapsuleDescriptors (BlockListBuffer, &BlockList);\r
+  Status = BuildCapsuleDescriptors (BlockListBuffer, MemoryResource, &BlockList);\r
   if (EFI_ERROR (Status)) {\r
     return Status;\r
   }\r
@@ -880,30 +1051,71 @@ CapsuleDataCoalesce (
   // returns the number of descriptors that actually point to data, so add\r
   // one for a terminator. Do that below.\r
   //\r
-  GetCapsuleInfo (BlockList, &NumDescriptors, &CapsuleSize);\r
-  if ((CapsuleSize == 0) || (NumDescriptors == 0)) {\r
+  Status = GetCapsuleInfo (BlockList, &NumDescriptors, &CapsuleSize, &CapsuleNumber);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+  DEBUG ((EFI_D_INFO, "CapsuleSize - 0x%x\n", CapsuleSize));\r
+  DEBUG ((EFI_D_INFO, "CapsuleNumber - 0x%x\n", CapsuleNumber));\r
+  DEBUG ((EFI_D_INFO, "NumDescriptors - 0x%x\n", NumDescriptors));\r
+  if ((CapsuleSize == 0) || (NumDescriptors == 0) || (CapsuleNumber == 0)) {\r
     return EFI_NOT_FOUND;\r
   }\r
 \r
+  if (CapsuleNumber - 1 >= (MAX_ADDRESS - (sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA)  + sizeof(UINT64))) / sizeof(UINT64)) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: CapsuleNumber - 0x%x\n", CapsuleNumber));\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
+\r
   //\r
   // Initialize our local copy of private data. When we're done, we'll create a\r
   // descriptor for it as well so that it can be put into free memory without\r
   // trashing anything.\r
   //\r
-  PrivateData.Signature     = EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE;\r
-  PrivateData.CapsuleSize   = (UINT32) CapsuleSize;\r
+  PrivateData.Signature           = EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE;\r
+  PrivateData.CapsuleAllImageSize = (UINT64) CapsuleSize;\r
+  PrivateData.CapsuleNumber       = (UINT64) CapsuleNumber;\r
+  PrivateData.CapsuleOffset[0]    = 0;\r
+  //\r
+  // NOTE: Only data in sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) is valid, CapsuleOffset field is uninitialized at this moment.\r
+  // The code sets partial length here for Descriptor.Length check, but later it will use full length to reserve those PrivateData region.\r
+  //\r
   PrivateDataDesc[0].Union.DataBlock  = (EFI_PHYSICAL_ADDRESS) (UINTN) &PrivateData;\r
   PrivateDataDesc[0].Length           = sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA);\r
   PrivateDataDesc[1].Union.DataBlock  = (EFI_PHYSICAL_ADDRESS) (UINTN) BlockList;\r
   PrivateDataDesc[1].Length           = 0;\r
   //\r
-  // In addition to PrivateDataDesc[1:0], one terminator is added\r
-  // See below RelocateBlockDescriptors()\r
+  // Add PrivateDataDesc[0] in beginning, as it is new descriptor. PrivateDataDesc[1] is NOT needed.\r
+  // In addition, one NULL terminator is added in the end. See RelocateBlockDescriptors().\r
+  //\r
+  NumDescriptors  += 2;\r
+  //\r
+  // Sanity check\r
+  //\r
+  if (CapsuleSize >= (MAX_ADDRESS - (sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64) + sizeof(UINT64)))) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: CapsuleSize - 0x%x\n", CapsuleSize));\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
+  //\r
+  // Need add sizeof(UINT64) for PrivateData alignment\r
   //\r
-  NumDescriptors  += 3;\r
-  CapsuleSize     += sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + sizeof(CapsuleOffset) + sizeof(UINT32);\r
+  CapsuleSize     += sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64) + sizeof(UINT64);\r
   BlockList        = PrivateDataDesc;\r
+  //\r
+  // Sanity check\r
+  //\r
+  if (NumDescriptors >= (MAX_ADDRESS / sizeof(EFI_CAPSULE_BLOCK_DESCRIPTOR))) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: NumDescriptors - 0x%x\n", NumDescriptors));\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
   DescriptorsSize  = NumDescriptors * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
+  //\r
+  // Sanity check\r
+  //\r
+  if (DescriptorsSize >= (MAX_ADDRESS - CapsuleSize)) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: DescriptorsSize - 0x%lx, CapsuleSize - 0x%lx\n", (UINT64)DescriptorsSize, (UINT64)CapsuleSize));\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
 \r
   //\r
   // Don't go below some min address. If the base is below it,\r
@@ -912,6 +1124,7 @@ CapsuleDataCoalesce (
   DEBUG ((EFI_D_INFO, "Capsule Memory range from 0x%8X to 0x%8X\n", (UINTN) *MemoryBase, (UINTN)*MemoryBase + *MemorySize));\r
   if ((UINTN)*MemoryBase < (UINTN) MIN_COALESCE_ADDR) {\r
     if (((UINTN)*MemoryBase + *MemorySize) < (UINTN) MIN_COALESCE_ADDR) {\r
+      DEBUG ((EFI_D_ERROR, "ERROR: *MemoryBase + *MemorySize - 0x%x\n", (UINTN)*MemoryBase + *MemorySize));\r
       return EFI_BUFFER_TOO_SMALL;\r
     } else {\r
       *MemorySize = *MemorySize - ((UINTN) MIN_COALESCE_ADDR - (UINTN) *MemoryBase);\r
@@ -920,6 +1133,7 @@ CapsuleDataCoalesce (
   }\r
 \r
   if (*MemorySize <= (CapsuleSize + DescriptorsSize)) {\r
+    DEBUG ((EFI_D_ERROR, "ERROR: CapsuleSize + DescriptorsSize - 0x%x\n", CapsuleSize + DescriptorsSize));\r
     return EFI_BUFFER_TOO_SMALL;\r
   }\r
 \r
@@ -931,7 +1145,7 @@ CapsuleDataCoalesce (
   // Relocate all the block descriptors to low memory to make further\r
   // processing easier.\r
   //\r
-  BlockList = RelocateBlockDescriptors (PeiServices, BlockList, FreeMemBase, FreeMemSize);\r
+  BlockList = RelocateBlockDescriptors (PeiServices, BlockList, NumDescriptors, FreeMemBase, FreeMemSize);\r
   if (BlockList == NULL) {\r
     //\r
     // Not enough room to relocate the descriptors\r
@@ -940,13 +1154,16 @@ CapsuleDataCoalesce (
   }\r
 \r
   //\r
-  // Take the top of memory for the capsule. Naturally align.\r
+  // Take the top of memory for the capsule. UINT64 align up.\r
   //\r
   DestPtr         = FreeMemBase + FreeMemSize - CapsuleSize;\r
-  DestPtr         = (UINT8 *) ((UINTN) DestPtr &~ (UINTN) (sizeof (UINTN) - 1));\r
+  DestPtr         = (UINT8 *) (((UINTN)DestPtr + sizeof (UINT64) - 1) & ~(sizeof (UINT64) - 1));\r
   FreeMemBase     = (UINT8 *) BlockList + DescriptorsSize;\r
-  FreeMemSize     = FreeMemSize - DescriptorsSize - CapsuleSize;\r
+  FreeMemSize     = (UINTN) DestPtr - (UINTN) FreeMemBase;\r
   NewCapsuleBase  = (VOID *) DestPtr;\r
+  CapsuleImageBase = (UINT8 *)NewCapsuleBase + sizeof(EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64);\r
+\r
+  PrivateDataPtr = (EFI_CAPSULE_PEIM_PRIVATE_DATA *) NewCapsuleBase;\r
 \r
   //\r
   // Move all the blocks to the top (high) of memory.\r
@@ -955,6 +1172,16 @@ CapsuleDataCoalesce (
   //\r
   CurrentBlockDesc = BlockList;\r
   while ((CurrentBlockDesc->Length != 0) || (CurrentBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
+    if (CapsuleTimes == 0) {\r
+      //\r
+      // The first entry is the block descriptor for EFI_CAPSULE_PEIM_PRIVATE_DATA.\r
+      // CapsuleOffset field is uninitialized at this time. No need copy it, but need to reserve for future use.\r
+      //\r
+      ASSERT (CurrentBlockDesc->Union.DataBlock == (UINT64)(UINTN)&PrivateData);\r
+      DestLength = sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64);\r
+    } else {\r
+      DestLength = (UINTN)CurrentBlockDesc->Length;\r
+    }\r
     //\r
     // See if any of the remaining capsule blocks are in the way\r
     //\r
@@ -965,7 +1192,7 @@ CapsuleDataCoalesce (
       //\r
       if (IsOverlapped (\r
             (UINT8 *) DestPtr,\r
-            (UINTN) CurrentBlockDesc->Length,\r
+            (UINTN) DestLength,\r
             (UINT8 *) (UINTN) TempBlockDesc->Union.DataBlock,\r
             (UINTN) TempBlockDesc->Length\r
             )) {\r
@@ -1005,77 +1232,50 @@ CapsuleDataCoalesce (
         CapsuleBeginFlag  = FALSE;\r
         CapsuleHeader     = (EFI_CAPSULE_HEADER*)(UINTN)CurrentBlockDesc->Union.DataBlock;\r
         SizeLeft          = CapsuleHeader->CapsuleImageSize;\r
-        if (!IsCapsuleCorrupted (CapsuleHeader)) {\r
 \r
-          if (CapsuleIndex > (MAX_SUPPORT_CAPSULE_NUM - 1)) {\r
-            DEBUG ((EFI_D_ERROR, "Capsule number exceeds the max number of %d!\n", MAX_SUPPORT_CAPSULE_NUM));\r
-            return  EFI_BUFFER_TOO_SMALL;\r
-          }\r
+        //\r
+        // No more check here is needed, because IsCapsuleCorrupted() already in ValidateCapsuleIntegrity()\r
+        //\r
+        ASSERT (CapsuleIndex < CapsuleNumber);\r
 \r
-          //\r
-          // Relocate this valid capsule\r
-          //\r
-          IsCorrupted  = FALSE;\r
-          CapsuleImageSize += SizeLeft;\r
-          CopyMem ((VOID *) DestPtr, (VOID *) (UINTN) CurrentBlockDesc->Union.DataBlock, (UINTN) CurrentBlockDesc->Length);\r
-          DEBUG ((EFI_D_INFO, "Capsule coalesce block no.0x%8X from 0x%8lX to 0x%8lX with size 0x%8X\n",CapsuleTimes,\r
-                 (UINTN)CurrentBlockDesc->Union.DataBlock, (UINTN)DestPtr, (UINTN)CurrentBlockDesc->Length));\r
-          //\r
-          // Cache the begin offset of this capsule\r
-          //\r
-          CapsuleOffset[CapsuleIndex++] = (UINT32) (UINTN) DestPtr - (UINT32)(UINTN)NewCapsuleBase - (UINT32)sizeof(EFI_CAPSULE_PEIM_PRIVATE_DATA);\r
-          DestPtr += CurrentBlockDesc->Length;\r
-        }\r
         //\r
-        // If the current block length is greater than or equal to SizeLeft, this is the \r
-        // start of the next capsule\r
+        // Relocate this capsule\r
         //\r
-        if (CurrentBlockDesc->Length < SizeLeft) {\r
-          SizeLeft -= CurrentBlockDesc->Length;\r
-        } else {\r
-          //\r
-          // Start the next cycle\r
-          //\r
-          SizeLeft         = 0;\r
-          IsCorrupted      = TRUE;\r
-          CapsuleBeginFlag = TRUE;          \r
-        }\r
-      } else {\r
+        CapsuleImageSize += SizeLeft;\r
         //\r
-        //Go on relocating the current capule image.\r
+        // Cache the begin offset of this capsule\r
         //\r
-        if (CurrentBlockDesc->Length < SizeLeft) {\r
-          if (!IsCorrupted) {\r
-            CopyMem ((VOID *) DestPtr, (VOID *) (UINTN) (CurrentBlockDesc->Union.DataBlock), (UINTN)CurrentBlockDesc->Length);\r
-            DEBUG ((EFI_D_INFO, "Capsule coalesce block no.0x%8X from 0x%8lX to 0x%8lX with size 0x%8X\n",CapsuleTimes,\r
-                   (UINTN)CurrentBlockDesc->Union.DataBlock, (UINTN)DestPtr, (UINTN)CurrentBlockDesc->Length));\r
-            DestPtr += CurrentBlockDesc->Length;\r
-          }\r
-          SizeLeft -= CurrentBlockDesc->Length;\r
-        } else {\r
-          //\r
-          //Here is the end of the current capsule image.\r
-          //\r
-          if (!IsCorrupted) {\r
-            CopyMem ((VOID *) DestPtr, (VOID *)(UINTN)(CurrentBlockDesc->Union.DataBlock), (UINTN)CurrentBlockDesc->Length);\r
-            DEBUG ((EFI_D_INFO, "Capsule coalesce block no.0x%8X from 0x%8lX to 0x%8lX with size 0x%8X\n",CapsuleTimes,\r
-                   (UINTN)CurrentBlockDesc->Union.DataBlock, (UINTN)DestPtr, (UINTN)CurrentBlockDesc->Length));\r
-            DestPtr += CurrentBlockDesc->Length;\r
-          }\r
-          //\r
-          // Start the next cycle\r
-          //\r
-          SizeLeft = 0;\r
-          IsCorrupted = TRUE;\r
-          CapsuleBeginFlag = TRUE; \r
-        }\r
+        ASSERT (PrivateDataPtr->Signature == EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE);\r
+        ASSERT ((UINTN)DestPtr >= (UINTN)CapsuleImageBase);\r
+        PrivateDataPtr->CapsuleOffset[CapsuleIndex++] = (UINTN)DestPtr - (UINTN)CapsuleImageBase;\r
+      }\r
+\r
+      //\r
+      // Below ASSERT is checked in ValidateCapsuleIntegrity()\r
+      //\r
+      ASSERT (CurrentBlockDesc->Length <= SizeLeft);\r
+\r
+      CopyMem ((VOID *) DestPtr, (VOID *) (UINTN) (CurrentBlockDesc->Union.DataBlock), (UINTN)CurrentBlockDesc->Length);\r
+      DEBUG ((EFI_D_INFO, "Capsule coalesce block no.0x%lX from 0x%lX to 0x%lX with size 0x%lX\n",(UINT64)CapsuleTimes,\r
+             CurrentBlockDesc->Union.DataBlock, (UINT64)(UINTN)DestPtr, CurrentBlockDesc->Length));\r
+      DestPtr += CurrentBlockDesc->Length;\r
+      SizeLeft -= CurrentBlockDesc->Length;\r
+\r
+      if (SizeLeft == 0) {\r
+        //\r
+        //Here is the end of the current capsule image.\r
+        //\r
+        CapsuleBeginFlag = TRUE;\r
       }\r
     } else {\r
       //\r
-      //The first entry is the block descriptor for EFI_CAPSULE_PEIM_PRIVATE_DATA.\r
+      // The first entry is the block descriptor for EFI_CAPSULE_PEIM_PRIVATE_DATA.\r
+      // CapsuleOffset field is uninitialized at this time. No need copy it, but need to reserve for future use.\r
       //\r
+      ASSERT (CurrentBlockDesc->Length == sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA));\r
+      ASSERT ((UINTN)DestPtr == (UINTN)NewCapsuleBase);\r
       CopyMem ((VOID *) DestPtr, (VOID *) (UINTN) CurrentBlockDesc->Union.DataBlock, (UINTN) CurrentBlockDesc->Length);\r
-      DestPtr += CurrentBlockDesc->Length;\r
+      DestPtr += sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64);\r
     }\r
     //\r
     //Walk through the block descriptor list.\r
@@ -1086,21 +1286,12 @@ CapsuleDataCoalesce (
   // We return the base of memory we want reserved, and the size.\r
   // The memory peim should handle it appropriately from there.\r
   //\r
-  *MemorySize = (UINTN) CapsuleImageSize;\r
+  *MemorySize = (UINTN) CapsuleSize;\r
   *MemoryBase = (VOID *) NewCapsuleBase;\r
 \r
-  //\r
-  //Append the offsets of mutiply capsules to the continous buffer\r
-  //\r
-  DataPtr    = (VOID*)((UINTN)NewCapsuleBase + sizeof(EFI_CAPSULE_PEIM_PRIVATE_DATA) + (UINTN)CapsuleImageSize);\r
-  AddDataPtr = (UINT32*)(((UINTN) DataPtr + sizeof(UINT32) - 1) &~ (UINT32) (sizeof (UINT32) - 1));\r
-\r
-  *AddDataPtr++ = CapsuleIndex;\r
-\r
-  CopyMem (AddDataPtr, &CapsuleOffset[0], sizeof (UINT32) * CapsuleIndex);\r
-\r
-  PrivateDataPtr = (EFI_CAPSULE_PEIM_PRIVATE_DATA *) NewCapsuleBase;\r
-  PrivateDataPtr->CapsuleSize = (UINT32) CapsuleImageSize;\r
+  ASSERT (PrivateDataPtr->Signature == EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE);\r
+  ASSERT (PrivateDataPtr->CapsuleAllImageSize == CapsuleImageSize);\r
+  ASSERT (PrivateDataPtr->CapsuleNumber == CapsuleIndex);\r
 \r
   return EFI_SUCCESS;\r
 }\r