]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
MdeModulePkg: Replace BSD License with BSD+Patent License
[mirror_edk2.git] / MdeModulePkg / Universal / CapsulePei / UefiCapsule.c
index 7c10a3c56887a711c80bf149d1d1ba91a4a9bb83..36b04e34064b1caec15d86b4fffb23e9bb6532af 100644 (file)
 /** @file\r
   Capsule update PEIM for UEFI2.0\r
 \r
-Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2006 - 2018, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>\r
 \r
-This program and the accompanying materials\r
-are licensed and made available under the terms and conditions\r
-of the BSD License which accompanies this distribution.  The\r
-full text of the license may be found at\r
-http://opensource.org/licenses/bsd-license.php\r
-\r
-THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
 #include "Capsule.h"\r
 \r
-EFI_PHYSICAL_ADDRESS           *mBufferAddress;\r
+#ifdef MDE_CPU_IA32\r
+//\r
+// Global Descriptor Table (GDT)\r
+//\r
+GLOBAL_REMOVE_IF_UNREFERENCED IA32_SEGMENT_DESCRIPTOR mGdtEntries[] = {\r
+/* selector { Global Segment Descriptor                              } */\r
+/* 0x00 */  {{0,      0,  0,  0,    0,  0,  0,  0,    0,  0, 0,  0,  0}}, //null descriptor\r
+/* 0x08 */  {{0xffff, 0,  0,  0x3,  1,  0,  1,  0xf,  0,  0, 1,  1,  0}}, //linear data segment descriptor\r
+/* 0x10 */  {{0xffff, 0,  0,  0xf,  1,  0,  1,  0xf,  0,  0, 1,  1,  0}}, //linear code segment descriptor\r
+/* 0x18 */  {{0xffff, 0,  0,  0x3,  1,  0,  1,  0xf,  0,  0, 1,  1,  0}}, //system data segment descriptor\r
+/* 0x20 */  {{0xffff, 0,  0,  0xb,  1,  0,  1,  0xf,  0,  0, 1,  1,  0}}, //system code segment descriptor\r
+/* 0x28 */  {{0,      0,  0,  0,    0,  0,  0,  0,    0,  0, 0,  0,  0}}, //spare segment descriptor\r
+/* 0x30 */  {{0xffff, 0,  0,  0x3,  1,  0,  1,  0xf,  0,  0, 1,  1,  0}}, //system data segment descriptor\r
+/* 0x38 */  {{0xffff, 0,  0,  0xb,  1,  0,  1,  0xf,  0,  1, 0,  1,  0}}, //system code segment descriptor\r
+/* 0x40 */  {{0,      0,  0,  0,    0,  0,  0,  0,    0,  0, 0,  0,  0}}, //spare segment descriptor\r
+};\r
 \r
-/**\r
-  Check every capsule header.\r
+//\r
+// IA32 Gdt register\r
+//\r
+GLOBAL_REMOVE_IF_UNREFERENCED CONST IA32_DESCRIPTOR mGdt = {\r
+  sizeof (mGdtEntries) - 1,\r
+  (UINTN) mGdtEntries\r
+  };\r
 \r
-  @param CapsuleHeader   The pointer to EFI_CAPSULE_HEADER\r
 \r
-  @retval FALSE  Capsule is OK\r
-  @retval TRUE   Capsule is corrupted \r
+/**\r
+  The function will check if 1G page is supported.\r
+\r
+  @retval TRUE   1G page is supported.\r
+  @retval FALSE  1G page is not supported.\r
 \r
 **/\r
 BOOLEAN\r
-IsCapsuleCorrupted (\r
-  IN EFI_CAPSULE_HEADER       *CapsuleHeader\r
+IsPage1GSupport (\r
+  VOID\r
   )\r
 {\r
+  UINT32                                        RegEax;\r
+  UINT32                                        RegEdx;\r
+  BOOLEAN                                       Page1GSupport;\r
+\r
+  Page1GSupport = FALSE;\r
+  if (PcdGetBool(PcdUse1GPageTable)) {\r
+    AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);\r
+    if (RegEax >= 0x80000001) {\r
+      AsmCpuid (0x80000001, NULL, NULL, NULL, &RegEdx);\r
+      if ((RegEdx & BIT26) != 0) {\r
+        Page1GSupport = TRUE;\r
+      }\r
+    }\r
+  }\r
+\r
+  return Page1GSupport;\r
+}\r
+\r
+/**\r
+  Calculate the total size of page table.\r
+\r
+  @param[in] Page1GSupport      1G page support or not.\r
+\r
+  @return The size of page table.\r
+\r
+**/\r
+UINTN\r
+CalculatePageTableSize (\r
+  IN BOOLEAN                                    Page1GSupport\r
+  )\r
+{\r
+  UINTN                                         ExtraPageTablePages;\r
+  UINTN                                         TotalPagesNum;\r
+  UINT8                                         PhysicalAddressBits;\r
+  UINT32                                        NumberOfPml4EntriesNeeded;\r
+  UINT32                                        NumberOfPdpEntriesNeeded;\r
+\r
   //\r
-  //A capsule to be updated across a system reset should contain CAPSULE_FLAGS_PERSIST_ACROSS_RESET.\r
+  // Create 4G page table by default,\r
+  // and let PF handler to handle > 4G request.\r
   //\r
-  if ((CapsuleHeader->Flags & CAPSULE_FLAGS_PERSIST_ACROSS_RESET) == 0) {\r
-    return TRUE;\r
-  }\r
+  PhysicalAddressBits = 32;\r
+  ExtraPageTablePages = EXTRA_PAGE_TABLE_PAGES;\r
+\r
   //\r
-  //Make sure the flags combination is supported by the platform.\r
+  // Calculate the table entries needed.\r
   //\r
-  if ((CapsuleHeader->Flags & (CAPSULE_FLAGS_PERSIST_ACROSS_RESET | CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE)) == CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE) {\r
-    return TRUE;\r
+  if (PhysicalAddressBits <= 39 ) {\r
+    NumberOfPml4EntriesNeeded = 1;\r
+    NumberOfPdpEntriesNeeded = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 30));\r
+  } else {\r
+    NumberOfPml4EntriesNeeded = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 39));\r
+    NumberOfPdpEntriesNeeded = 512;\r
   }\r
-  if ((CapsuleHeader->Flags & (CAPSULE_FLAGS_PERSIST_ACROSS_RESET | CAPSULE_FLAGS_INITIATE_RESET)) == CAPSULE_FLAGS_INITIATE_RESET) {\r
-    return TRUE;\r
+\r
+  if (!Page1GSupport) {\r
+    TotalPagesNum = (NumberOfPdpEntriesNeeded + 1) * NumberOfPml4EntriesNeeded + 1;\r
+  } else {\r
+    TotalPagesNum = NumberOfPml4EntriesNeeded + 1;\r
   }\r
+  TotalPagesNum += ExtraPageTablePages;\r
 \r
-  return FALSE;\r
+  return EFI_PAGES_TO_SIZE (TotalPagesNum);\r
 }\r
 \r
 /**\r
-  Check the integrity of the capsule descriptors.\r
+  Allocates and fills in the Page Directory and Page Table Entries to\r
+  establish a 4G page table.\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
+  @param[in] PageTablesAddress  The base address of page table.\r
+  @param[in] Page1GSupport      1G page support or not.\r
 \r
 **/\r
-EFI_CAPSULE_BLOCK_DESCRIPTOR *\r
-ValidateCapsuleIntegrity (\r
-  IN EFI_CAPSULE_BLOCK_DESCRIPTOR    *BlockList\r
+VOID\r
+Create4GPageTables (\r
+  IN EFI_PHYSICAL_ADDRESS   PageTablesAddress,\r
+  IN BOOLEAN                Page1GSupport\r
   )\r
 {\r
-  EFI_CAPSULE_HEADER             *CapsuleHeader;\r
-  UINT64                         CapsuleSize;\r
-  UINT32                         CapsuleCount;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR   *Ptr;\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
-  CapsuleSize  = 0;\r
-  CapsuleCount = 0;\r
-  Ptr = BlockList;\r
-  while ((Ptr->Length != 0) || (Ptr->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
+  UINT8                                         PhysicalAddressBits;\r
+  EFI_PHYSICAL_ADDRESS                          PageAddress;\r
+  UINTN                                         IndexOfPml4Entries;\r
+  UINTN                                         IndexOfPdpEntries;\r
+  UINTN                                         IndexOfPageDirectoryEntries;\r
+  UINT32                                        NumberOfPml4EntriesNeeded;\r
+  UINT32                                        NumberOfPdpEntriesNeeded;\r
+  PAGE_MAP_AND_DIRECTORY_POINTER                *PageMapLevel4Entry;\r
+  PAGE_MAP_AND_DIRECTORY_POINTER                *PageMap;\r
+  PAGE_MAP_AND_DIRECTORY_POINTER                *PageDirectoryPointerEntry;\r
+  PAGE_TABLE_ENTRY                              *PageDirectoryEntry;\r
+  UINTN                                         BigPageAddress;\r
+  PAGE_TABLE_1G_ENTRY                           *PageDirectory1GEntry;\r
+  UINT64                                        AddressEncMask;\r
+\r
+  //\r
+  // Make sure AddressEncMask is contained to smallest supported address field.\r
+  //\r
+  AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;\r
+\r
+  //\r
+  // Create 4G page table by default,\r
+  // and let PF handler to handle > 4G request.\r
+  //\r
+  PhysicalAddressBits = 32;\r
+\r
+  //\r
+  // Calculate the table entries needed.\r
+  //\r
+  if (PhysicalAddressBits <= 39 ) {\r
+    NumberOfPml4EntriesNeeded = 1;\r
+    NumberOfPdpEntriesNeeded = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 30));\r
+  } else {\r
+    NumberOfPml4EntriesNeeded = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 39));\r
+    NumberOfPdpEntriesNeeded = 512;\r
+  }\r
+\r
+  //\r
+  // Pre-allocate big pages to avoid later allocations.\r
+  //\r
+  BigPageAddress = (UINTN) PageTablesAddress;\r
+\r
+  //\r
+  // By architecture only one PageMapLevel4 exists - so lets allocate storage for it.\r
+  //\r
+  PageMap         = (VOID *) BigPageAddress;\r
+  BigPageAddress += SIZE_4KB;\r
+\r
+  PageMapLevel4Entry = PageMap;\r
+  PageAddress        = 0;\r
+  for (IndexOfPml4Entries = 0; IndexOfPml4Entries < NumberOfPml4EntriesNeeded; IndexOfPml4Entries++, PageMapLevel4Entry++) {\r
     //\r
-    // Make sure the descriptor is aligned at UINT64 in memory\r
+    // Each PML4 entry points to a page of Page Directory Pointer entires.\r
+    // So lets allocate space for them and fill them in in the IndexOfPdpEntries loop.\r
     //\r
-    if ((UINTN) Ptr & 0x07) {\r
-      DEBUG ((EFI_D_ERROR, "BlockList address failed alignment check\n"));\r
-      return NULL;\r
-    }\r
+    PageDirectoryPointerEntry = (VOID *) BigPageAddress;\r
+    BigPageAddress += SIZE_4KB;\r
 \r
-    if (Ptr->Length == 0) {\r
-      //\r
-      // Descriptor points to another list of block descriptors somewhere\r
-      // else.\r
-      //\r
-      Ptr = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) Ptr->Union.ContinuationPointer;\r
+    //\r
+    // Make a PML4 Entry\r
+    //\r
+    PageMapLevel4Entry->Uint64 = (UINT64)(UINTN)PageDirectoryPointerEntry | AddressEncMask;\r
+    PageMapLevel4Entry->Bits.ReadWrite = 1;\r
+    PageMapLevel4Entry->Bits.Present = 1;\r
+\r
+    if (Page1GSupport) {\r
+      PageDirectory1GEntry = (VOID *) PageDirectoryPointerEntry;\r
+\r
+      for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectory1GEntry++, PageAddress += SIZE_1GB) {\r
+        //\r
+        // Fill in the Page Directory entries\r
+        //\r
+        PageDirectory1GEntry->Uint64 = (UINT64)PageAddress | AddressEncMask;\r
+        PageDirectory1GEntry->Bits.ReadWrite = 1;\r
+        PageDirectory1GEntry->Bits.Present = 1;\r
+        PageDirectory1GEntry->Bits.MustBe1 = 1;\r
+      }\r
     } else {\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
-      //\r
-      if (CapsuleSize == 0) {\r
+      for (IndexOfPdpEntries = 0; IndexOfPdpEntries < NumberOfPdpEntriesNeeded; IndexOfPdpEntries++, PageDirectoryPointerEntry++) {\r
         //\r
-        //Move to the first capsule to check its header.\r
+        // Each Directory Pointer entries points to a page of Page Directory entires.\r
+        // So allocate space for them and fill them in in the IndexOfPageDirectoryEntries loop.\r
         //\r
-        CapsuleHeader = (EFI_CAPSULE_HEADER*)((UINTN)Ptr->Union.DataBlock);\r
-        if (IsCapsuleCorrupted (CapsuleHeader)) {\r
-          return NULL;\r
-        }\r
-        CapsuleCount ++;\r
-        CapsuleSize = CapsuleHeader->CapsuleImageSize;\r
-      } else {\r
-        if (CapsuleSize >= Ptr->Length) {\r
-          CapsuleSize = CapsuleSize - Ptr->Length;\r
-        } else {\r
-          CapsuleSize = 0;\r
+        PageDirectoryEntry = (VOID *) BigPageAddress;\r
+        BigPageAddress += SIZE_4KB;\r
+\r
+        //\r
+        // Fill in a Page Directory Pointer Entries\r
+        //\r
+        PageDirectoryPointerEntry->Uint64 = (UINT64)(UINTN)PageDirectoryEntry | AddressEncMask;\r
+        PageDirectoryPointerEntry->Bits.ReadWrite = 1;\r
+        PageDirectoryPointerEntry->Bits.Present = 1;\r
+\r
+        for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectoryEntry++, PageAddress += SIZE_2MB) {\r
+          //\r
+          // Fill in the Page Directory entries\r
+          //\r
+          PageDirectoryEntry->Uint64 = (UINT64)PageAddress | AddressEncMask;\r
+          PageDirectoryEntry->Bits.ReadWrite = 1;\r
+          PageDirectoryEntry->Bits.Present = 1;\r
+          PageDirectoryEntry->Bits.MustBe1 = 1;\r
         }\r
       }\r
-      //\r
-      // Move to next BLOCK descriptor\r
-      //\r
-      Ptr++;\r
+\r
+      for (; IndexOfPdpEntries < 512; IndexOfPdpEntries++, PageDirectoryPointerEntry++) {\r
+        ZeroMem (\r
+          PageDirectoryPointerEntry,\r
+          sizeof(PAGE_MAP_AND_DIRECTORY_POINTER)\r
+          );\r
+      }\r
     }\r
   }\r
 \r
-  if (CapsuleCount == 0) {\r
-    //\r
-    // No any capsule is found in BlockList.\r
-    //\r
-    return NULL;\r
+  //\r
+  // For the PML4 entries we are not using fill in a null entry.\r
+  //\r
+  for (; IndexOfPml4Entries < 512; IndexOfPml4Entries++, PageMapLevel4Entry++) {\r
+    ZeroMem (\r
+      PageMapLevel4Entry,\r
+      sizeof (PAGE_MAP_AND_DIRECTORY_POINTER)\r
+      );\r
   }\r
-\r
-  return Ptr;\r
 }\r
 \r
+/**\r
+  Return function from long mode to 32-bit mode.\r
+\r
+  @param  EntrypointContext  Context for mode switching\r
+  @param  ReturnContext      Context for mode switching\r
+\r
+**/\r
+VOID\r
+ReturnFunction (\r
+  SWITCH_32_TO_64_CONTEXT  *EntrypointContext,\r
+  SWITCH_64_TO_32_CONTEXT  *ReturnContext\r
+  )\r
+{\r
+  //\r
+  // Restore original GDT\r
+  //\r
+  AsmWriteGdtr (&ReturnContext->Gdtr);\r
+\r
+  //\r
+  // return to original caller\r
+  //\r
+  LongJump ((BASE_LIBRARY_JUMP_BUFFER  *)(UINTN)EntrypointContext->JumpBuffer, 1);\r
+\r
+  //\r
+  // never be here\r
+  //\r
+  ASSERT (FALSE);\r
+}\r
 \r
 /**\r
-  Checks for the presence of capsule descriptors.\r
-  Get capsule descriptors from variable CapsuleUpdateData, CapsuleUpdateData1, CapsuleUpdateData2...\r
+  Thunk function from 32-bit protection mode to long mode.\r
 \r
-  @param BlockList        Pointer to the capsule descriptors\r
+  @param  PageTableAddress  Page table base address\r
+  @param  Context           Context for mode switching\r
+  @param  ReturnContext     Context for mode switching\r
+\r
+  @retval EFI_SUCCESS  Function successfully executed.\r
 \r
-  @retval EFI_SUCCESS     a valid capsule is present\r
-  @retval EFI_NOT_FOUND   if a valid capsule is not present\r
 **/\r
 EFI_STATUS\r
-GetCapsuleDescriptors (\r
-  IN OUT EFI_CAPSULE_BLOCK_DESCRIPTOR    **BlockList OPTIONAL\r
+Thunk32To64 (\r
+  EFI_PHYSICAL_ADDRESS          PageTableAddress,\r
+  SWITCH_32_TO_64_CONTEXT       *Context,\r
+  SWITCH_64_TO_32_CONTEXT       *ReturnContext\r
   )\r
 {\r
-  EFI_STATUS                       Status;\r
-  UINTN                            Size;\r
-  UINTN                            Index;\r
-  UINTN                            TempIndex;\r
-  UINTN                            ValidIndex;\r
-  BOOLEAN                          Flag;\r
-  CHAR16                           CapsuleVarName[30];\r
-  CHAR16                           *TempVarName;\r
-  EFI_PHYSICAL_ADDRESS             CapsuleDataPtr64;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR     *LastBlock;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR     *TempBlock;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR     *HeadBlock;\r
-  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;\r
+  UINTN                       SetJumpFlag;\r
+  EFI_STATUS                  Status;\r
 \r
-  LastBlock         = NULL;\r
-  HeadBlock         = NULL;\r
-  TempBlock         = NULL;\r
-  Index             = 0;\r
-  TempVarName       = NULL;\r
-  CapsuleVarName[0] = 0;\r
-  ValidIndex        = 0;\r
-  \r
-  Status = PeiServicesLocatePpi (\r
-              &gEfiPeiReadOnlyVariable2PpiGuid,\r
-              0,\r
-              NULL,\r
-              (VOID **) &PPIVariableServices\r
-              );\r
-  if (Status == EFI_SUCCESS) {\r
-    StrCpy (CapsuleVarName, EFI_CAPSULE_VARIABLE_NAME);\r
-    TempVarName = CapsuleVarName + StrLen (CapsuleVarName);\r
-    Size = sizeof (CapsuleDataPtr64);\r
-    while (1) {\r
-      if (Index == 0) {\r
-        //\r
-        // For the first Capsule Image\r
-        //\r
-        Status = PPIVariableServices->GetVariable (\r
-                                        PPIVariableServices,\r
-                                        CapsuleVarName,\r
-                                        &gEfiCapsuleVendorGuid,\r
-                                        NULL,\r
-                                        &Size,\r
-                                        (VOID *) &CapsuleDataPtr64\r
-                                        );\r
-        if (EFI_ERROR (Status)) {\r
-          DEBUG ((EFI_D_ERROR, "Capsule -- capsule variable not set\n"));\r
-          return EFI_NOT_FOUND;\r
-        }\r
-        //\r
-        // We have a chicken/egg situation where the memory init code needs to\r
-        // know the boot mode prior to initializing memory. For this case, our\r
-        // validate function will fail. We can detect if this is the case if blocklist\r
-        // pointer is null. In that case, return success since we know that the\r
-        // variable is set.\r
-        //\r
-        if (BlockList == NULL) {\r
-          return EFI_SUCCESS;\r
-        }\r
-        //\r
-        // Test integrity of descriptors.\r
-        //\r
-        LastBlock = ValidateCapsuleIntegrity ((EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)CapsuleDataPtr64);\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)CapsuleDataPtr64;\r
-      } else {\r
-        UnicodeValueToString (TempVarName, 0, Index, 0);\r
-        Status = PPIVariableServices->GetVariable (\r
-                                        PPIVariableServices,\r
-                                        CapsuleVarName,\r
-                                        &gEfiCapsuleVendorGuid,\r
-                                        NULL,\r
-                                        &Size,\r
-                                        (VOID *) &CapsuleDataPtr64\r
-                                        );\r
-        if (EFI_ERROR (Status)) {\r
-          break;\r
-        }\r
-        \r
-        //\r
-        // If this BlockList has been linked before, skip this variable\r
-        //\r
-        Flag = FALSE;\r
-        for (TempIndex = 0; TempIndex < ValidIndex; TempIndex++) {\r
-          if (mBufferAddress[TempIndex] == CapsuleDataPtr64) {\r
-            Flag = TRUE;\r
-            break;\r
-          }\r
-        }\r
-        if (Flag) {\r
-          Index ++;\r
-          continue;\r
-        }\r
+  //\r
+  // Save return address, LongJump will return here then\r
+  //\r
+  SetJumpFlag = SetJump ((BASE_LIBRARY_JUMP_BUFFER  *) (UINTN) Context->JumpBuffer);\r
 \r
-        //\r
-        // Test integrity of descriptors.\r
-        //\r
-        TempBlock = ValidateCapsuleIntegrity ((EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)CapsuleDataPtr64);\r
-        if (TempBlock == NULL) {\r
-          return EFI_NOT_FOUND;\r
-        }\r
-        //\r
-        // Combine the different BlockList into single BlockList.\r
-        //\r
-        LastBlock->Union.DataBlock = CapsuleDataPtr64;\r
-        LastBlock->Length          = 0;\r
-        LastBlock                  = TempBlock;\r
-      }\r
-      \r
-      //\r
-      // Cache BlockList which has been processed\r
-      //\r
-      mBufferAddress[ValidIndex++] = CapsuleDataPtr64;\r
-      Index ++;\r
-    }\r
+  if (SetJumpFlag == 0) {\r
+\r
+    //\r
+    // Build 4G Page Tables.\r
+    //\r
+    Create4GPageTables (PageTableAddress, Context->Page1GSupport);\r
+\r
+    //\r
+    // Create 64-bit GDT\r
+    //\r
+    AsmWriteGdtr (&mGdt);\r
+\r
+    //\r
+    // Write CR3\r
+    //\r
+    AsmWriteCr3 ((UINTN) PageTableAddress);\r
+\r
+    DEBUG ((\r
+      DEBUG_INFO,\r
+      "%a() Stack Base: 0x%lx, Stack Size: 0x%lx\n",\r
+      __FUNCTION__,\r
+      Context->StackBufferBase,\r
+      Context->StackBufferLength\r
+      ));\r
+\r
+    //\r
+    // Disable interrupt of Debug timer, since the IDT table cannot work in long mode\r
+    //\r
+    SaveAndSetDebugTimerInterrupt (FALSE);\r
+    //\r
+    // Transfer to long mode\r
+    //\r
+    AsmEnablePaging64 (\r
+       0x38,\r
+      (UINT64) Context->EntryPoint,\r
+      (UINT64)(UINTN) Context,\r
+      (UINT64)(UINTN) ReturnContext,\r
+      Context->StackBufferBase + Context->StackBufferLength\r
+      );\r
   }\r
-  \r
-  if (HeadBlock != NULL) {\r
-    *BlockList = HeadBlock;\r
-    return EFI_SUCCESS;\r
+\r
+  //\r
+  // Convert to 32-bit Status and return\r
+  //\r
+  Status = EFI_SUCCESS;\r
+  if ((UINTN) ReturnContext->ReturnStatus != 0) {\r
+    Status = ENCODE_ERROR ((UINTN) ReturnContext->ReturnStatus);\r
   }\r
-  return EFI_NOT_FOUND;\r
+\r
+  return Status;\r
 }\r
 \r
 /**\r
-  Given a pointer to a capsule block descriptor, traverse the list to figure\r
-  out how many legitimate descriptors there are, and how big the capsule it\r
-  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
-\r
-  @retval EFI_NOT_FOUND   No descriptors containing data in the list\r
-  @retval EFI_SUCCESS     Return data is valid\r
+  If in 32 bit protection mode, and coalesce image is of X64, switch to long mode.\r
+\r
+  @param  LongModeBuffer            The context of long mode.\r
+  @param  CoalesceEntry             Entry of coalesce image.\r
+  @param  BlockListAddr             Address of block list.\r
+  @param  MemoryResource            Pointer to the buffer of memory resource descriptor.\r
+  @param  MemoryBase                Base of memory range.\r
+  @param  MemorySize                Size of memory range.\r
+\r
+  @retval EFI_SUCCESS               Successfully switched to long mode and execute coalesce.\r
+  @retval Others                    Failed to execute coalesce in long mode.\r
+\r
 **/\r
 EFI_STATUS\r
-GetCapsuleInfo (\r
-  IN EFI_CAPSULE_BLOCK_DESCRIPTOR   *Desc,\r
-  IN OUT UINTN                      *NumDescriptors OPTIONAL,\r
-  IN OUT UINTN                      *CapsuleSize OPTIONAL\r
+ModeSwitch (\r
+  IN EFI_CAPSULE_LONG_MODE_BUFFER   *LongModeBuffer,\r
+  IN COALESCE_ENTRY                 CoalesceEntry,\r
+  IN EFI_PHYSICAL_ADDRESS           BlockListAddr,\r
+  IN MEMORY_RESOURCE_DESCRIPTOR     *MemoryResource,\r
+  IN OUT VOID                       **MemoryBase,\r
+  IN OUT UINTN                      *MemorySize\r
   )\r
 {\r
-  UINTN Count;\r
-  UINTN Size;\r
-\r
-  ASSERT (Desc != NULL);\r
-\r
-  Count = 0;\r
-  Size  = 0;\r
+  EFI_STATUS                           Status;\r
+  EFI_PHYSICAL_ADDRESS                 MemoryBase64;\r
+  UINT64                               MemorySize64;\r
+  EFI_PHYSICAL_ADDRESS                 MemoryEnd64;\r
+  SWITCH_32_TO_64_CONTEXT              Context;\r
+  SWITCH_64_TO_32_CONTEXT              ReturnContext;\r
+  BASE_LIBRARY_JUMP_BUFFER             JumpBuffer;\r
+  EFI_PHYSICAL_ADDRESS                 ReservedRangeBase;\r
+  EFI_PHYSICAL_ADDRESS                 ReservedRangeEnd;\r
+  BOOLEAN                              Page1GSupport;\r
+\r
+  ZeroMem (&Context, sizeof (SWITCH_32_TO_64_CONTEXT));\r
+  ZeroMem (&ReturnContext, sizeof (SWITCH_64_TO_32_CONTEXT));\r
+\r
+  MemoryBase64  = (UINT64) (UINTN) *MemoryBase;\r
+  MemorySize64  = (UINT64) (UINTN) *MemorySize;\r
+  MemoryEnd64   = MemoryBase64 + MemorySize64;\r
+\r
+  Page1GSupport = IsPage1GSupport ();\r
+\r
+  //\r
+  // Merge memory range reserved for stack and page table\r
+  //\r
+  if (LongModeBuffer->StackBaseAddress < LongModeBuffer->PageTableAddress) {\r
+    ReservedRangeBase = LongModeBuffer->StackBaseAddress;\r
+    ReservedRangeEnd  = LongModeBuffer->PageTableAddress + CalculatePageTableSize (Page1GSupport);\r
+  } else {\r
+    ReservedRangeBase = LongModeBuffer->PageTableAddress;\r
+    ReservedRangeEnd  = LongModeBuffer->StackBaseAddress + LongModeBuffer->StackSize;\r
+  }\r
 \r
-  while (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL) {\r
-    if (Desc->Length == 0) {\r
-      //\r
-      // Descriptor points to another list of block descriptors somewhere\r
-      //\r
-      Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) Desc->Union.ContinuationPointer;\r
+  //\r
+  // Check if memory range reserved is overlap with MemoryBase ~ MemoryBase + MemorySize.\r
+  // If they are overlapped, get a larger range to process capsule data.\r
+  //\r
+  if (ReservedRangeBase <= MemoryBase64) {\r
+    if (ReservedRangeEnd < MemoryEnd64) {\r
+      MemoryBase64 = ReservedRangeEnd;\r
     } else {\r
-      Size += (UINTN) Desc->Length;\r
-      Count++;\r
-      Desc++;\r
+      DEBUG ((EFI_D_ERROR, "Memory is not enough to process capsule!\n"));\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
+  } else if (ReservedRangeBase < MemoryEnd64) {\r
+    if (ReservedRangeEnd < MemoryEnd64   &&\r
+        ReservedRangeBase - MemoryBase64 < MemoryEnd64 - ReservedRangeEnd) {\r
+      MemoryBase64 = ReservedRangeEnd;\r
+    } else {\r
+      MemorySize64 = (UINT64)(UINTN)(ReservedRangeBase - MemoryBase64);\r
     }\r
   }\r
+\r
   //\r
-  // If no descriptors, then fail\r
+  // Initialize context jumping to 64-bit enviroment\r
   //\r
-  if (Count == 0) {\r
-    return EFI_NOT_FOUND;\r
-  }\r
+  Context.JumpBuffer            = (EFI_PHYSICAL_ADDRESS)(UINTN)&JumpBuffer;\r
+  Context.StackBufferBase       = LongModeBuffer->StackBaseAddress;\r
+  Context.StackBufferLength     = LongModeBuffer->StackSize;\r
+  Context.EntryPoint            = (EFI_PHYSICAL_ADDRESS)(UINTN)CoalesceEntry;\r
+  Context.BlockListAddr         = BlockListAddr;\r
+  Context.MemoryResource        = (EFI_PHYSICAL_ADDRESS)(UINTN)MemoryResource;\r
+  Context.MemoryBase64Ptr       = (EFI_PHYSICAL_ADDRESS)(UINTN)&MemoryBase64;\r
+  Context.MemorySize64Ptr       = (EFI_PHYSICAL_ADDRESS)(UINTN)&MemorySize64;\r
+  Context.Page1GSupport         = Page1GSupport;\r
+  Context.AddressEncMask        = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;\r
 \r
-  if (NumDescriptors != NULL) {\r
-    *NumDescriptors = Count;\r
+  //\r
+  // Prepare data for return back\r
+  //\r
+  ReturnContext.ReturnCs           = 0x10;\r
+  ReturnContext.ReturnEntryPoint   = (EFI_PHYSICAL_ADDRESS)(UINTN)ReturnFunction;\r
+  //\r
+  // Will save the return status of processing capsule\r
+  //\r
+  ReturnContext.ReturnStatus       = 0;\r
+\r
+  //\r
+  // Save original GDT\r
+  //\r
+  AsmReadGdtr ((IA32_DESCRIPTOR *)&ReturnContext.Gdtr);\r
+\r
+  Status = Thunk32To64 (LongModeBuffer->PageTableAddress, &Context, &ReturnContext);\r
+\r
+  if (!EFI_ERROR (Status)) {\r
+    *MemoryBase = (VOID *) (UINTN) MemoryBase64;\r
+    *MemorySize = (UINTN) MemorySize64;\r
   }\r
 \r
-  if (CapsuleSize != NULL) {\r
-    *CapsuleSize = Size;\r
+  return Status;\r
+\r
+}\r
+\r
+/**\r
+  Locates the coalesce image entry point, and detects its machine type.\r
+\r
+  @param CoalesceImageEntryPoint   Pointer to coalesce image entry point for output.\r
+  @param CoalesceImageMachineType  Pointer to machine type of coalesce image.\r
+\r
+  @retval EFI_SUCCESS     Coalesce image successfully located.\r
+  @retval Others          Failed to locate the coalesce image.\r
+\r
+**/\r
+EFI_STATUS\r
+FindCapsuleCoalesceImage (\r
+  OUT EFI_PHYSICAL_ADDRESS    *CoalesceImageEntryPoint,\r
+  OUT UINT16                  *CoalesceImageMachineType\r
+  )\r
+{\r
+  EFI_STATUS                           Status;\r
+  UINTN                                Instance;\r
+  EFI_PEI_LOAD_FILE_PPI                *LoadFile;\r
+  EFI_PEI_FV_HANDLE                    VolumeHandle;\r
+  EFI_PEI_FILE_HANDLE                  FileHandle;\r
+  EFI_PHYSICAL_ADDRESS                 CoalesceImageAddress;\r
+  UINT64                               CoalesceImageSize;\r
+  UINT32                               AuthenticationState;\r
+\r
+  Instance = 0;\r
+\r
+  while (TRUE) {\r
+    Status = PeiServicesFfsFindNextVolume (Instance++, &VolumeHandle);\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+    Status = PeiServicesFfsFindFileByName (PcdGetPtr(PcdCapsuleCoalesceFile), VolumeHandle, &FileHandle);\r
+    if (!EFI_ERROR (Status)) {\r
+      Status = PeiServicesLocatePpi (&gEfiPeiLoadFilePpiGuid, 0, NULL, (VOID **) &LoadFile);\r
+      ASSERT_EFI_ERROR (Status);\r
+\r
+      Status = LoadFile->LoadFile (\r
+                           LoadFile,\r
+                           FileHandle,\r
+                           &CoalesceImageAddress,\r
+                           &CoalesceImageSize,\r
+                           CoalesceImageEntryPoint,\r
+                           &AuthenticationState\r
+                           );\r
+      if (EFI_ERROR (Status)) {\r
+        DEBUG ((EFI_D_ERROR, "Unable to find PE32 section in CapsuleX64 image ffs %r!\n", Status));\r
+        return Status;\r
+      }\r
+      *CoalesceImageMachineType = PeCoffLoaderGetMachineType ((VOID *) (UINTN) CoalesceImageAddress);\r
+      break;\r
+    } else {\r
+      continue;\r
+    }\r
   }\r
 \r
-  return EFI_SUCCESS;\r
+  return Status;\r
 }\r
 \r
+/**\r
+  Gets the reserved long mode buffer.\r
 \r
+  @param  LongModeBuffer  Pointer to the long mode buffer for output.\r
+\r
+  @retval EFI_SUCCESS     Long mode buffer successfully retrieved.\r
+  @retval Others          Variable storing long mode buffer not found.\r
+\r
+**/\r
+EFI_STATUS\r
+GetLongModeContext (\r
+  OUT EFI_CAPSULE_LONG_MODE_BUFFER *LongModeBuffer\r
+  )\r
+{\r
+  EFI_STATUS   Status;\r
+  UINTN        Size;\r
+  EFI_PEI_READ_ONLY_VARIABLE2_PPI *PPIVariableServices;\r
+\r
+  Status = PeiServicesLocatePpi (\r
+             &gEfiPeiReadOnlyVariable2PpiGuid,\r
+             0,\r
+             NULL,\r
+             (VOID **) &PPIVariableServices\r
+             );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  Size = sizeof (EFI_CAPSULE_LONG_MODE_BUFFER);\r
+  Status = PPIVariableServices->GetVariable (\r
+                                  PPIVariableServices,\r
+                                  EFI_CAPSULE_LONG_MODE_BUFFER_NAME,\r
+                                  &gEfiCapsuleVendorGuid,\r
+                                  NULL,\r
+                                  &Size,\r
+                                  LongModeBuffer\r
+                                  );\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG (( EFI_D_ERROR, "Error Get LongModeBuffer variable %r!\n", Status));\r
+  }\r
+  return Status;\r
+}\r
+#endif\r
+\r
+#if defined (MDE_CPU_IA32) || defined (MDE_CPU_X64)\r
 /**\r
-  Try to verify the integrity of a capsule test pattern before the\r
-  capsule gets coalesced. This can be useful in narrowing down\r
-  where capsule data corruption occurs.\r
+  Get physical address bits.\r
 \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
+  @return Physical address bits.\r
 \r
-  @param PeiServices  PEI services table\r
-  @param Desc         Pointer to capsule descriptors\r
 **/\r
-VOID\r
-CapsuleTestPatternPreCoalesce (\r
-  IN EFI_PEI_SERVICES              **PeiServices,\r
-  IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *Desc\r
+UINT8\r
+GetPhysicalAddressBits (\r
+  VOID\r
   )\r
 {\r
-  UINT32  *TestPtr;\r
-  UINT32  TestCounter;\r
-  UINT32  TestSize;\r
+  UINT32                        RegEax;\r
+  UINT8                         PhysicalAddressBits;\r
+  VOID                          *Hob;\r
+\r
   //\r
-  // Find first data descriptor\r
+  // Get physical address bits supported.\r
   //\r
-  while ((Desc->Length == 0) && (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
-    Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) Desc->Union.ContinuationPointer;\r
+  Hob = GetFirstHob (EFI_HOB_TYPE_CPU);\r
+  if (Hob != NULL) {\r
+    PhysicalAddressBits = ((EFI_HOB_CPU *) Hob)->SizeOfMemorySpace;\r
+  } else {\r
+    AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);\r
+    if (RegEax >= 0x80000008) {\r
+      AsmCpuid (0x80000008, &RegEax, NULL, NULL, NULL);\r
+      PhysicalAddressBits = (UINT8) RegEax;\r
+    } else {\r
+      PhysicalAddressBits = 36;\r
+    }\r
   }\r
 \r
-  if (Desc->Union.ContinuationPointer == 0) {\r
-    return ;\r
-  }\r
   //\r
-  // First one better be long enough to at least hold the test signature\r
+  // IA-32e paging translates 48-bit linear addresses to 52-bit physical addresses.\r
   //\r
-  if (Desc->Length < sizeof (UINT32)) {\r
-    DEBUG ((EFI_D_INFO, "Capsule test pattern pre-coalesce punted #1\n"));\r
-    return ;\r
+  ASSERT (PhysicalAddressBits <= 52);\r
+  if (PhysicalAddressBits > 48) {\r
+    PhysicalAddressBits = 48;\r
   }\r
 \r
-  TestPtr = (UINT32 *) (UINTN) Desc->Union.DataBlock;\r
-  if (*TestPtr != CAPSULE_TEST_SIGNATURE) {\r
-    return ;\r
-  }\r
+  return PhysicalAddressBits;\r
+}\r
+#endif\r
 \r
-  TestCounter = 0;\r
-  TestSize    = (UINT32) Desc->Length - 2 * sizeof (UINT32);\r
-  //\r
-  // Skip over the signature and the size fields in the pattern data header\r
-  //\r
-  TestPtr += 2;\r
-  while (1) {\r
-    if ((TestSize & 0x03) != 0) {\r
-      DEBUG ((EFI_D_INFO, "Capsule test pattern pre-coalesce punted #2\n"));\r
-      return ;\r
-    }\r
+/**\r
+  Sort memory resource entries based upon PhysicalStart, from low to high.\r
 \r
-    while (TestSize > 0) {\r
-      if (*TestPtr != TestCounter) {\r
-        DEBUG ((EFI_D_INFO, "Capsule test pattern pre-coalesce failed data corruption check\n"));\r
-        return ;\r
-      }\r
+  @param[in, out] MemoryResource    A pointer to the memory resource entry buffer.\r
 \r
-      TestSize -= sizeof (UINT32);\r
-      TestCounter++;\r
-      TestPtr++;\r
-    }\r
-    Desc++;\r
-    while ((Desc->Length == 0) && (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
-      Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) Desc->Union.ContinuationPointer;\r
-    }\r
+**/\r
+VOID\r
+SortMemoryResourceDescriptor (\r
+  IN OUT MEMORY_RESOURCE_DESCRIPTOR *MemoryResource\r
+  )\r
+{\r
+  MEMORY_RESOURCE_DESCRIPTOR        *MemoryResourceEntry;\r
+  MEMORY_RESOURCE_DESCRIPTOR        *NextMemoryResourceEntry;\r
+  MEMORY_RESOURCE_DESCRIPTOR        TempMemoryResource;\r
+\r
+  MemoryResourceEntry = MemoryResource;\r
+  NextMemoryResourceEntry = MemoryResource + 1;\r
+  while (MemoryResourceEntry->ResourceLength != 0) {\r
+    while (NextMemoryResourceEntry->ResourceLength != 0) {\r
+      if (MemoryResourceEntry->PhysicalStart > NextMemoryResourceEntry->PhysicalStart) {\r
+        CopyMem (&TempMemoryResource, MemoryResourceEntry, sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
+        CopyMem (MemoryResourceEntry, NextMemoryResourceEntry, sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
+        CopyMem (NextMemoryResourceEntry, &TempMemoryResource, sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
+      }\r
 \r
-    if (Desc->Union.ContinuationPointer == (EFI_PHYSICAL_ADDRESS) (UINTN) NULL) {\r
-      return ;\r
+      NextMemoryResourceEntry = NextMemoryResourceEntry + 1;\r
     }\r
-    TestSize = (UINT32) Desc->Length;\r
-    TestPtr  = (UINT32 *) (UINTN) Desc->Union.DataBlock;\r
+\r
+    MemoryResourceEntry     = MemoryResourceEntry + 1;\r
+    NextMemoryResourceEntry = MemoryResourceEntry + 1;\r
   }\r
 }\r
 \r
-\r
 /**\r
-  Determine if two buffers overlap in memory.\r
+  Merge continous memory resource entries.\r
 \r
-  @param Buff1   pointer to first buffer\r
-  @param Size1   size of Buff1\r
-  @param Buff2   pointer to second buffer\r
-  @param Size2   size of Buff2\r
+  @param[in, out] MemoryResource    A pointer to the memory resource entry buffer.\r
 \r
-  @retval TRUE    Buffers overlap in memory.\r
-  @retval FALSE   Buffer doesn't overlap.\r
 **/\r
-BOOLEAN\r
-IsOverlapped (\r
-  UINT8     *Buff1,\r
-  UINTN     Size1,\r
-  UINT8     *Buff2,\r
-  UINTN     Size2\r
+VOID\r
+MergeMemoryResourceDescriptor (\r
+  IN OUT MEMORY_RESOURCE_DESCRIPTOR *MemoryResource\r
   )\r
 {\r
-  //\r
-  // If buff1's end is less than the start of buff2, then it's ok.\r
-  // Also, if buff1's start is beyond buff2's end, then it's ok.\r
-  //\r
-  if (((Buff1 + Size1) <= Buff2) || (Buff1 >= (Buff2 + Size2))) {\r
-    return FALSE;\r
+  MEMORY_RESOURCE_DESCRIPTOR        *MemoryResourceEntry;\r
+  MEMORY_RESOURCE_DESCRIPTOR        *NewMemoryResourceEntry;\r
+  MEMORY_RESOURCE_DESCRIPTOR        *NextMemoryResourceEntry;\r
+  MEMORY_RESOURCE_DESCRIPTOR        *MemoryResourceEnd;\r
+\r
+  MemoryResourceEntry = MemoryResource;\r
+  NewMemoryResourceEntry = MemoryResource;\r
+  while (MemoryResourceEntry->ResourceLength != 0) {\r
+    CopyMem (NewMemoryResourceEntry, MemoryResourceEntry, sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
+    NextMemoryResourceEntry = MemoryResourceEntry + 1;\r
+\r
+    while ((NextMemoryResourceEntry->ResourceLength != 0) &&\r
+           (NextMemoryResourceEntry->PhysicalStart == (MemoryResourceEntry->PhysicalStart + MemoryResourceEntry->ResourceLength))) {\r
+      MemoryResourceEntry->ResourceLength += NextMemoryResourceEntry->ResourceLength;\r
+      if (NewMemoryResourceEntry != MemoryResourceEntry) {\r
+        NewMemoryResourceEntry->ResourceLength += NextMemoryResourceEntry->ResourceLength;\r
+      }\r
+\r
+      NextMemoryResourceEntry = NextMemoryResourceEntry + 1;\r
+    }\r
+\r
+    MemoryResourceEntry = NextMemoryResourceEntry;\r
+    NewMemoryResourceEntry = NewMemoryResourceEntry + 1;\r
   }\r
 \r
-  return TRUE;\r
+  //\r
+  // Set NULL terminate memory resource descriptor after merging.\r
+  //\r
+  MemoryResourceEnd = NewMemoryResourceEntry;\r
+  ZeroMem (MemoryResourceEnd, sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
 }\r
 \r
 /**\r
-  Given a pointer to the capsule block list, info on the available system\r
-  memory, and the size of a buffer, find a free block of memory where a\r
-  buffer of the given size can be copied to safely.\r
-\r
-  @param BlockList   Pointer to head of capsule block descriptors\r
-  @param MemBase     Pointer to the base of memory in which we want to find free space\r
-  @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
-  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
+  Build memory resource descriptor from resource descriptor in HOB list.\r
+\r
+  @return Pointer to the buffer of memory resource descriptor.\r
+          NULL if no memory resource descriptor reported in HOB list\r
+          before capsule Coalesce.\r
+\r
 **/\r
-UINT8 *\r
-FindFreeMem (\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR      *BlockList,\r
-  UINT8                             *MemBase,\r
-  UINTN                             MemSize,\r
-  UINTN                             DataSize\r
+MEMORY_RESOURCE_DESCRIPTOR *\r
+BuildMemoryResourceDescriptor (\r
+  VOID\r
   )\r
 {\r
-  UINTN                           Size;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR    *CurrDesc;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR    *TempDesc;\r
-  UINT8                           *MemEnd;\r
-  BOOLEAN                         Failed;\r
+  EFI_PEI_HOB_POINTERS          Hob;\r
+  UINTN                         Index;\r
+  EFI_HOB_RESOURCE_DESCRIPTOR   *ResourceDescriptor;\r
+  MEMORY_RESOURCE_DESCRIPTOR    *MemoryResource;\r
+  EFI_STATUS                    Status;\r
 \r
   //\r
-  // Need at least enough to copy the data to at the end of the buffer, so\r
-  // say the end is less the data size for easy comparisons here.\r
+  // Get the count of memory resource descriptor.\r
   //\r
-  MemEnd    = MemBase + MemSize - DataSize;\r
-  CurrDesc  = BlockList;\r
-  //\r
-  // Go through all the descriptor blocks and see if any obstruct the range\r
-  //\r
-  while (CurrDesc != NULL) {\r
-    //\r
-    // Get the size of this block list and see if it's in the way\r
-    //\r
-    Failed    = FALSE;\r
-    TempDesc  = CurrDesc;\r
-    Size      = sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
-    while (TempDesc->Length != 0) {\r
-      Size += sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
-      TempDesc++;\r
+  Index = 0;\r
+  Hob.Raw = GetFirstHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR);\r
+  while (Hob.Raw != NULL) {\r
+    ResourceDescriptor = (EFI_HOB_RESOURCE_DESCRIPTOR *) Hob.Raw;\r
+    if (ResourceDescriptor->ResourceType == EFI_RESOURCE_SYSTEM_MEMORY) {\r
+      Index++;\r
     }\r
+    Hob.Raw = GET_NEXT_HOB (Hob);\r
+    Hob.Raw = GetNextHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR, Hob.Raw);\r
+  }\r
 \r
-    if (IsOverlapped (MemBase, DataSize, (UINT8 *) CurrDesc, Size)) {\r
-      //\r
-      // Set our new base to the end of this block list and start all over\r
-      //\r
-      MemBase   = (UINT8 *) CurrDesc + Size;\r
-      CurrDesc  = BlockList;\r
-      if (MemBase > MemEnd) {\r
-        return NULL;\r
-      }\r
-\r
-      Failed = TRUE;\r
-    }\r
-    //\r
-    // Now go through all the blocks and make sure none are in the way\r
-    //\r
-    while ((CurrDesc->Length != 0) && (!Failed)) {\r
-      if (IsOverlapped (MemBase, DataSize, (UINT8 *) (UINTN) CurrDesc->Union.DataBlock, (UINTN) CurrDesc->Length)) {\r
-        //\r
-        // Set our new base to the end of this block and start all over\r
-        //\r
-        Failed    = TRUE;\r
-        MemBase   = (UINT8 *) ((UINTN) CurrDesc->Union.DataBlock) + CurrDesc->Length;\r
-        CurrDesc  = BlockList;\r
-        if (MemBase > MemEnd) {\r
-          return NULL;\r
-        }\r
-      }\r
-      CurrDesc++;\r
-    }\r
+  if (Index == 0) {\r
+    DEBUG ((EFI_D_INFO | EFI_D_WARN, "No memory resource descriptor reported in HOB list before capsule Coalesce\n"));\r
+#if defined (MDE_CPU_IA32) || defined (MDE_CPU_X64)\r
     //\r
-    // Normal continuation -- jump to next block descriptor list\r
+    // Allocate memory to hold memory resource descriptor,\r
+    // include extra one NULL terminate memory resource descriptor.\r
     //\r
-    if (!Failed) {\r
-      CurrDesc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) CurrDesc->Union.ContinuationPointer;\r
+    Status = PeiServicesAllocatePool ((1 + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR), (VOID **) &MemoryResource);\r
+    ASSERT_EFI_ERROR (Status);\r
+    ZeroMem (MemoryResource, (1 + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
+\r
+    MemoryResource[0].PhysicalStart = 0;\r
+    MemoryResource[0].ResourceLength = LShiftU64 (1, GetPhysicalAddressBits ());\r
+    DEBUG ((EFI_D_INFO, "MemoryResource[0x0] - Start(0x%0lx) Length(0x%0lx)\n",\r
+                        MemoryResource[0x0].PhysicalStart, MemoryResource[0x0].ResourceLength));\r
+    return MemoryResource;\r
+#else\r
+    return NULL;\r
+#endif\r
+  }\r
+\r
+  //\r
+  // Allocate memory to hold memory resource descriptor,\r
+  // include extra one NULL terminate memory resource descriptor.\r
+  //\r
+  Status = PeiServicesAllocatePool ((Index + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR), (VOID **) &MemoryResource);\r
+  ASSERT_EFI_ERROR (Status);\r
+  ZeroMem (MemoryResource, (Index + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR));\r
+\r
+  //\r
+  // Get the content of memory resource descriptor.\r
+  //\r
+  Index = 0;\r
+  Hob.Raw = GetFirstHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR);\r
+  while (Hob.Raw != NULL) {\r
+    ResourceDescriptor = (EFI_HOB_RESOURCE_DESCRIPTOR *) Hob.Raw;\r
+    if (ResourceDescriptor->ResourceType == EFI_RESOURCE_SYSTEM_MEMORY) {\r
+      DEBUG ((EFI_D_INFO, "MemoryResource[0x%x] - Start(0x%0lx) Length(0x%0lx)\n",\r
+                          Index, ResourceDescriptor->PhysicalStart, ResourceDescriptor->ResourceLength));\r
+      MemoryResource[Index].PhysicalStart = ResourceDescriptor->PhysicalStart;\r
+      MemoryResource[Index].ResourceLength = ResourceDescriptor->ResourceLength;\r
+      Index++;\r
     }\r
+    Hob.Raw = GET_NEXT_HOB (Hob);\r
+    Hob.Raw = GetNextHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR, Hob.Raw);\r
+  }\r
+\r
+  SortMemoryResourceDescriptor (MemoryResource);\r
+  MergeMemoryResourceDescriptor (MemoryResource);\r
+\r
+  DEBUG ((DEBUG_INFO, "Dump MemoryResource[] after sorted and merged\n"));\r
+  for (Index = 0; MemoryResource[Index].ResourceLength != 0; Index++) {\r
+    DEBUG ((\r
+      DEBUG_INFO,\r
+      "  MemoryResource[0x%x] - Start(0x%0lx) Length(0x%0lx)\n",\r
+      Index,\r
+      MemoryResource[Index].PhysicalStart,\r
+      MemoryResource[Index].ResourceLength\r
+      ));\r
   }\r
-  return MemBase;\r
+\r
+  return MemoryResource;\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
-  capsule block descriptors low in memory.\r
-\r
-  The descriptors passed in can be fragmented throughout memory. Here\r
-  they are relocated into memory to turn them into a contiguous (null\r
-  terminated) array.\r
+  Checks for the presence of capsule descriptors.\r
+  Get capsule descriptors from variable CapsuleUpdateData, CapsuleUpdateData1, CapsuleUpdateData2...\r
+  and save to DescriptorBuffer.\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 DescriptorBuffer        Pointer to the capsule descriptors\r
 \r
-  @retval NULL    could not relocate the descriptors\r
-  @retval Pointer to the base of the successfully-relocated block descriptors. \r
+  @retval EFI_SUCCESS     a valid capsule is present\r
+  @retval EFI_NOT_FOUND   if a valid capsule is not present\r
 **/\r
-EFI_CAPSULE_BLOCK_DESCRIPTOR  *\r
-RelocateBlockDescriptors (\r
-  IN EFI_PEI_SERVICES                   **PeiServices,\r
-  IN EFI_CAPSULE_BLOCK_DESCRIPTOR       *BlockList,\r
-  IN UINT8                              *MemBase,\r
-  IN UINTN                              MemSize\r
+EFI_STATUS\r
+GetCapsuleDescriptors (\r
+  IN EFI_PHYSICAL_ADDRESS     *DescriptorBuffer\r
   )\r
 {\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR   *NewBlockList;\r
-  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
-  // 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
-  //\r
-  if (GetCapsuleInfo (BlockList, &NumDescriptors, NULL) != EFI_SUCCESS) {\r
-    return NULL;\r
-  }\r
+  EFI_STATUS                       Status;\r
+  UINTN                            Size;\r
+  UINTN                            Index;\r
+  UINTN                            TempIndex;\r
+  UINTN                            ValidIndex;\r
+  BOOLEAN                          Flag;\r
+  CHAR16                           CapsuleVarName[30];\r
+  CHAR16                           *TempVarName;\r
+  EFI_PHYSICAL_ADDRESS             CapsuleDataPtr64;\r
+  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;\r
 \r
-  NumDescriptors++;\r
-  BufferSize    = NumDescriptors * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
-  NewBlockList  = (EFI_CAPSULE_BLOCK_DESCRIPTOR *) MemBase;\r
-  if (MemSize < BufferSize) {\r
-    return NULL;\r
-  }\r
+  Index             = 0;\r
+  TempVarName       = NULL;\r
+  CapsuleVarName[0] = 0;\r
+  ValidIndex        = 0;\r
+  CapsuleDataPtr64  = 0;\r
 \r
-  MemSize -= BufferSize;\r
-  MemBase += BufferSize;\r
-  //\r
-  // Go through all the blocks and make sure none are in the way\r
-  //\r
-  TempBlockDesc = BlockList;\r
-  while (TempBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL) {\r
-    if (TempBlockDesc->Length == 0) {\r
-      //\r
-      // Next block of descriptors\r
-      //\r
-      TempBlockDesc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) TempBlockDesc->Union.ContinuationPointer;\r
-    } else {\r
-      //\r
-      // If the capsule data pointed to by this descriptor is in the way,\r
-      // move it.\r
-      //\r
-      if (IsOverlapped (\r
-            (UINT8 *) NewBlockList,\r
-            BufferSize,\r
-            (UINT8 *) (UINTN) TempBlockDesc->Union.DataBlock,\r
-            (UINTN) TempBlockDesc->Length\r
-            )) {\r
+  Status = PeiServicesLocatePpi (\r
+              &gEfiPeiReadOnlyVariable2PpiGuid,\r
+              0,\r
+              NULL,\r
+              (VOID **) &PPIVariableServices\r
+              );\r
+  if (Status == EFI_SUCCESS) {\r
+    StrCpyS (CapsuleVarName, sizeof(CapsuleVarName)/sizeof(CHAR16), EFI_CAPSULE_VARIABLE_NAME);\r
+    TempVarName = CapsuleVarName + StrLen (CapsuleVarName);\r
+    Size = sizeof (CapsuleDataPtr64);\r
+    while (1) {\r
+      if (Index == 0) {\r
+        //\r
+        // For the first Capsule Image\r
+        //\r
+        Status = PPIVariableServices->GetVariable (\r
+                                        PPIVariableServices,\r
+                                        CapsuleVarName,\r
+                                        &gEfiCapsuleVendorGuid,\r
+                                        NULL,\r
+                                        &Size,\r
+                                        (VOID *) &CapsuleDataPtr64\r
+                                        );\r
+        if (EFI_ERROR (Status)) {\r
+          DEBUG ((DEBUG_INFO, "Capsule -- capsule variable not set\n"));\r
+          return EFI_NOT_FOUND;\r
+        }\r
         //\r
-        // Relocate the block\r
+        // We have a chicken/egg situation where the memory init code needs to\r
+        // know the boot mode prior to initializing memory. For this case, our\r
+        // validate function will fail. We can detect if this is the case if blocklist\r
+        // pointer is null. In that case, return success since we know that the\r
+        // variable is set.\r
         //\r
-        RelocBuffer = FindFreeMem (BlockList, MemBase, MemSize, (UINTN) TempBlockDesc->Length);\r
-        if (RelocBuffer == NULL) {\r
-          return NULL;\r
+        if (DescriptorBuffer == NULL) {\r
+          return EFI_SUCCESS;\r
+        }\r
+      } else {\r
+        UnicodeValueToStringS (\r
+          TempVarName,\r
+          sizeof (CapsuleVarName) - ((UINTN)TempVarName - (UINTN)CapsuleVarName),\r
+          0,\r
+          Index,\r
+          0\r
+          );\r
+        Status = PPIVariableServices->GetVariable (\r
+                                        PPIVariableServices,\r
+                                        CapsuleVarName,\r
+                                        &gEfiCapsuleVendorGuid,\r
+                                        NULL,\r
+                                        &Size,\r
+                                        (VOID *) &CapsuleDataPtr64\r
+                                        );\r
+        if (EFI_ERROR (Status)) {\r
+          break;\r
         }\r
 \r
-        CopyMem ((VOID *) RelocBuffer, (VOID *) (UINTN) TempBlockDesc->Union.DataBlock, (UINTN) 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
+        // If this BlockList has been linked before, skip this variable\r
+        //\r
+        Flag = FALSE;\r
+        for (TempIndex = 0; TempIndex < ValidIndex; TempIndex++) {\r
+          if (DescriptorBuffer[TempIndex] == CapsuleDataPtr64)  {\r
+            Flag = TRUE;\r
+            break;\r
+          }\r
+        }\r
+        if (Flag) {\r
+          Index ++;\r
+          continue;\r
+        }\r
       }\r
-    }\r
-    TempBlockDesc++;\r
-  }\r
-  //\r
-  // Now go through all the block descriptors to make sure that they're not\r
-  // in the memory region we want to copy them to.\r
-  //\r
-  CurrBlockDescHead = BlockList;\r
-  PrevBlockDescTail = NULL;\r
-  while ((CurrBlockDescHead != NULL) && (CurrBlockDescHead->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
-    //\r
-    // Get the size of this list then see if it overlaps our low region\r
-    //\r
-    TempBlockDesc = CurrBlockDescHead;\r
-    BlockListSize = sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
-    while (TempBlockDesc->Length != 0) {\r
-      BlockListSize += sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
-      TempBlockDesc++;\r
-    }\r
 \r
-    if (IsOverlapped (\r
-          (UINT8 *) NewBlockList,\r
-          BufferSize,\r
-          (UINT8 *) CurrBlockDescHead,\r
-          BlockListSize\r
-          )) {\r
       //\r
-      // Overlaps, so move it out of the way\r
-      //\r
-      RelocBuffer = FindFreeMem (BlockList, MemBase, MemSize, BlockListSize);\r
-      if (RelocBuffer == NULL) {\r
-        return NULL;\r
-      }\r
-      CopyMem ((VOID *) RelocBuffer, (VOID *) CurrBlockDescHead, BlockListSize);\r
-      DEBUG ((EFI_D_INFO, "Capsule reloc descriptor block #2\n"));\r
-      //\r
-      // Point the previous block's next point to this copied version. If\r
-      // the tail pointer is null, then this is the first descriptor block.\r
+      // Cache BlockList which has been processed\r
       //\r
-      if (PrevBlockDescTail == NULL) {\r
-        BlockList = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) RelocBuffer;\r
-      } else {\r
-        PrevBlockDescTail->Union.DataBlock = (EFI_PHYSICAL_ADDRESS) (UINTN) RelocBuffer;\r
-      }\r
-    }\r
-    //\r
-    // Save our new tail and jump to the next block list\r
-    //\r
-    PrevBlockDescTail = TempBlockDesc;\r
-    CurrBlockDescHead = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) TempBlockDesc->Union.ContinuationPointer;\r
-  }\r
-  //\r
-  // Cleared out low memory. Now copy the descriptors down there.\r
-  //\r
-  TempBlockDesc     = BlockList;\r
-  CurrBlockDescHead = NewBlockList;\r
-  while ((TempBlockDesc != NULL) && (TempBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
-    if (TempBlockDesc->Length != 0) {\r
-      CurrBlockDescHead->Union.DataBlock = TempBlockDesc->Union.DataBlock;\r
-      CurrBlockDescHead->Length = TempBlockDesc->Length;\r
-      CurrBlockDescHead++;\r
-      TempBlockDesc++;\r
-    } else {\r
-      TempBlockDesc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *) (UINTN) TempBlockDesc->Union.ContinuationPointer;\r
+      DescriptorBuffer[ValidIndex++] = CapsuleDataPtr64;\r
+      Index ++;\r
     }\r
   }\r
-  //\r
-  // Null terminate\r
-  //\r
-  CurrBlockDescHead->Union.ContinuationPointer   = (EFI_PHYSICAL_ADDRESS) (UINTN) NULL;\r
-  CurrBlockDescHead->Length = 0;\r
-  return NewBlockList;\r
+\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
   Capsule PPI service to coalesce a fragmented capsule in memory.\r
 \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
-                   |    PrivateData            |            |\r
-   DestPtr  ---->  +---------------------------+<-----------+\r
-                   |                           |            |\r
-                   |     FreeMem               |        FreeMemSize\r
-                   |                           |            |\r
-   FreeMemBase --->+---------------------------+<-----------+\r
-                   |    Terminator             |\r
-                   +---------------------------+\r
-                   |    BlockDescriptor n      |\r
-                   +---------------------------+\r
-                   |    .................      |\r
-                   +---------------------------+\r
-                   |    BlockDescriptor 1      |\r
-                   +---------------------------+\r
-                   |    BlockDescriptor 0      |\r
-                   +---------------------------+\r
-                   |    PrivateDataDesc 0      |\r
-      MemBase ---->+---------------------------+<----- BlockList\r
-\r
-\r
   @param PeiServices  General purpose services available to every PEIM.\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
@@ -768,54 +937,28 @@ CapsuleCoalesce (
   IN OUT UINTN                       *MemorySize\r
   )\r
 {\r
-  VOID                           *NewCapsuleBase;\r
-  VOID                           *DataPtr;\r
-  UINT8                          CapsuleIndex;\r
-  UINT8                          *FreeMemBase;\r
-  UINT8                          *DestPtr;\r
-  UINT8                          *RelocPtr;\r
-  UINT32                         CapsuleOffset[MAX_SUPPORT_CAPSULE_NUM]; \r
-  UINT32                         *AddDataPtr;\r
-  UINT32                         CapsuleTimes; \r
-  UINT64                         SizeLeft; \r
-  UINT64                         CapsuleImageSize; \r
-  UINTN                          CapsuleSize;\r
-  UINTN                          DescriptorsSize;\r
-  UINTN                          FreeMemSize;\r
-  UINTN                          NumDescriptors;\r
-  UINTN                          Index;\r
-  UINTN                          Size;\r
-  UINTN                          VariableCount;\r
-  CHAR16                         CapsuleVarName[30];\r
-  CHAR16                         *TempVarName;\r
-  EFI_PHYSICAL_ADDRESS           CapsuleDataPtr64;  \r
-  BOOLEAN                        IsCorrupted;\r
-  BOOLEAN                        CapsuleBeginFlag;\r
-  EFI_STATUS                     Status;\r
-  EFI_BOOT_MODE                  BootMode;\r
-  EFI_CAPSULE_HEADER             *CapsuleHeader;\r
-  EFI_CAPSULE_PEIM_PRIVATE_DATA  PrivateData;\r
-  EFI_CAPSULE_PEIM_PRIVATE_DATA  *PrivateDataPtr;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockList;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR  *CurrentBlockDesc;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR  *TempBlockDesc;\r
-  EFI_CAPSULE_BLOCK_DESCRIPTOR  PrivateDataDesc[2];\r
-  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;\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
-  Index            = 0;\r
-  VariableCount    = 0;\r
-  CapsuleVarName[0] = 0;  \r
+  UINTN                                Index;\r
+  UINTN                                Size;\r
+  UINTN                                VariableCount;\r
+  CHAR16                               CapsuleVarName[30];\r
+  CHAR16                               *TempVarName;\r
+  EFI_PHYSICAL_ADDRESS                 CapsuleDataPtr64;\r
+  EFI_STATUS                           Status;\r
+  EFI_BOOT_MODE                        BootMode;\r
+  EFI_PEI_READ_ONLY_VARIABLE2_PPI      *PPIVariableServices;\r
+  EFI_PHYSICAL_ADDRESS                 *VariableArrayAddress;\r
+  MEMORY_RESOURCE_DESCRIPTOR           *MemoryResource;\r
+#ifdef MDE_CPU_IA32\r
+  UINT16                               CoalesceImageMachineType;\r
+  EFI_PHYSICAL_ADDRESS                 CoalesceImageEntryPoint;\r
+  COALESCE_ENTRY                       CoalesceEntry;\r
+  EFI_CAPSULE_LONG_MODE_BUFFER         LongModeBuffer;\r
+#endif\r
+\r
+  Index                   = 0;\r
+  VariableCount           = 0;\r
+  CapsuleVarName[0]       = 0;\r
+  CapsuleDataPtr64        = 0;\r
 \r
   //\r
   // Someone should have already ascertained the boot mode. If it's not\r
@@ -823,9 +966,11 @@ CapsuleCoalesce (
   //\r
   Status = PeiServicesGetBootMode (&BootMode);\r
   if (EFI_ERROR (Status) || (BootMode != BOOT_ON_FLASH_UPDATE)) {\r
-    return EFI_NOT_FOUND;\r
+    DEBUG ((EFI_D_ERROR, "Boot mode is not correct for capsule update path.\n"));\r
+    Status = EFI_NOT_FOUND;\r
+    goto Done;\r
   }\r
-  \r
+\r
   //\r
   // User may set the same ScatterGatherList with several different variables,\r
   // so cache all ScatterGatherList for check later.\r
@@ -837,14 +982,20 @@ CapsuleCoalesce (
               (VOID **) &PPIVariableServices\r
               );\r
   if (EFI_ERROR (Status)) {\r
-    return Status;\r
+    goto Done;\r
   }\r
   Size = sizeof (CapsuleDataPtr64);\r
-  StrCpy (CapsuleVarName, EFI_CAPSULE_VARIABLE_NAME);\r
+  StrCpyS (CapsuleVarName, sizeof(CapsuleVarName)/sizeof(CHAR16), EFI_CAPSULE_VARIABLE_NAME);\r
   TempVarName = CapsuleVarName + StrLen (CapsuleVarName);\r
   while (TRUE) {\r
     if (Index > 0) {\r
-      UnicodeValueToString (TempVarName, 0, Index, 0);\r
+      UnicodeValueToStringS (\r
+        TempVarName,\r
+        sizeof (CapsuleVarName) - ((UINTN)TempVarName - (UINTN)CapsuleVarName),\r
+        0,\r
+        Index,\r
+        0\r
+        );\r
     }\r
     Status = PPIVariableServices->GetVariable (\r
                                     PPIVariableServices,\r
@@ -858,264 +1009,96 @@ CapsuleCoalesce (
       //\r
       // There is no capsule variables, quit\r
       //\r
-      DEBUG ((EFI_D_ERROR,"Capsule variable Index = %d\n", Index));\r
+      DEBUG ((EFI_D_INFO,"Capsule variable Index = %d\n", Index));\r
       break;\r
     }\r
     VariableCount++;\r
     Index++;\r
   }\r
-  \r
-  DEBUG ((EFI_D_ERROR,"Capsule variable count = %d\n", VariableCount));\r
-  \r
-  Status = PeiServicesAllocatePool (\r
-             VariableCount * sizeof (EFI_PHYSICAL_ADDRESS),\r
-             (VOID **)&mBufferAddress\r
-             );\r
-\r
-  if (Status != EFI_SUCCESS) {\r
-    DEBUG ((EFI_D_ERROR, "AllocatePages Failed!, Status = %x\n", Status));\r
-    return Status;\r
-  }\r
-  \r
-  //\r
-  // Find out if we actually have a capsule.\r
-  //\r
-  Status = GetCapsuleDescriptors (&BlockList);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
 \r
-  DEBUG_CODE (\r
-    CapsuleTestPatternPreCoalesce (PeiServices, BlockList);\r
-  );\r
-\r
-  //\r
-  // Get the size of our descriptors and the capsule size. GetCapsuleInfo()\r
-  // 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
-    return EFI_NOT_FOUND;\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   = CapsuleSize;\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
-  //\r
-  NumDescriptors  += 3;\r
-  CapsuleSize     += sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + sizeof(CapsuleOffset) + sizeof(UINT32);\r
-  BlockList        = PrivateDataDesc;\r
-  DescriptorsSize  = NumDescriptors * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);\r
+  DEBUG ((EFI_D_INFO,"Capsule variable count = %d\n", VariableCount));\r
 \r
   //\r
-  // Don't go below some min address. If the base is below it,\r
-  // then move it up and adjust the size accordingly.\r
+  // The last entry is the end flag.\r
   //\r
-  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
-      return EFI_BUFFER_TOO_SMALL;\r
-    } else {\r
-      *MemorySize = *MemorySize - ((UINTN) MIN_COALESCE_ADDR - (UINTN) *MemoryBase);\r
-      *MemoryBase = (VOID *) (UINTN) MIN_COALESCE_ADDR;\r
-    }\r
-  }\r
+  Status = PeiServicesAllocatePool (\r
+             (VariableCount + 1) * sizeof (EFI_PHYSICAL_ADDRESS),\r
+             (VOID **)&VariableArrayAddress\r
+             );\r
 \r
-  if (*MemorySize <= (CapsuleSize + DescriptorsSize)) {\r
-    return EFI_BUFFER_TOO_SMALL;\r
+  if (Status != EFI_SUCCESS) {\r
+    DEBUG ((EFI_D_ERROR, "AllocatePages Failed!, Status = %x\n", Status));\r
+    goto Done;\r
   }\r
 \r
-  FreeMemBase = *MemoryBase;\r
-  FreeMemSize = *MemorySize;\r
-  DEBUG ((EFI_D_INFO, "Capsule Free Memory from 0x%8X to 0x%8X\n", (UINTN) FreeMemBase, (UINTN) FreeMemBase + FreeMemSize));\r
+  ZeroMem (VariableArrayAddress, (VariableCount + 1) * sizeof (EFI_PHYSICAL_ADDRESS));\r
 \r
   //\r
-  // Relocate all the block descriptors to low memory to make further\r
-  // processing easier.\r
+  // Find out if we actually have a capsule.\r
+  // GetCapsuleDescriptors depends on variable PPI, so it should run in 32-bit environment.\r
   //\r
-  BlockList = RelocateBlockDescriptors (PeiServices, BlockList, FreeMemBase, FreeMemSize);\r
-  if (BlockList == NULL) {\r
-    //\r
-    // Not enough room to relocate the descriptors\r
-    //\r
-    return EFI_BUFFER_TOO_SMALL;\r
+  Status = GetCapsuleDescriptors (VariableArrayAddress);\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG ((EFI_D_ERROR, "Fail to find capsule variables.\n"));\r
+    goto Done;\r
   }\r
 \r
-  //\r
-  // Take the top of memory for the capsule. Naturally align.\r
-  //\r
-  DestPtr         = FreeMemBase + FreeMemSize - CapsuleSize;\r
-  DestPtr         = (UINT8 *) ((UINTN) DestPtr &~ (UINTN) (sizeof (UINTN) - 1));\r
-  FreeMemBase     = (UINT8 *) BlockList + DescriptorsSize;\r
-  FreeMemSize     = FreeMemSize - DescriptorsSize - CapsuleSize;\r
-  NewCapsuleBase  = (VOID *) DestPtr;\r
+  MemoryResource = BuildMemoryResourceDescriptor ();\r
 \r
-  //\r
-  // Move all the blocks to the top (high) of memory.\r
-  // Relocate all the obstructing blocks. Note that the block descriptors\r
-  // were coalesced when they were relocated, so we can just ++ the pointer.\r
-  //\r
-  CurrentBlockDesc = BlockList;\r
-  while ((CurrentBlockDesc->Length != 0) || (CurrentBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS) (UINTN) NULL)) {\r
+#ifdef MDE_CPU_IA32\r
+  if (FeaturePcdGet (PcdDxeIplSwitchToLongMode)) {\r
     //\r
-    // See if any of the remaining capsule blocks are in the way\r
+    // Switch to 64-bit mode to process capsule data when:\r
+    // 1. When DXE phase is 64-bit\r
+    // 2. When the buffer for 64-bit transition exists\r
+    // 3. When Capsule X64 image is built in BIOS image\r
+    // In 64-bit mode, we can process capsule data above 4GB.\r
     //\r
-    TempBlockDesc = CurrentBlockDesc;\r
-    while (TempBlockDesc->Length != 0) {\r
-      //\r
-      // Is this block in the way of where we want to copy the current descriptor to?\r
-      //\r
-      if (IsOverlapped (\r
-            (UINT8 *) DestPtr,\r
-            (UINTN) CurrentBlockDesc->Length,\r
-            (UINT8 *) (UINTN) TempBlockDesc->Union.DataBlock,\r
-            (UINTN) TempBlockDesc->Length\r
-            )) {\r
-        //\r
-        // Relocate the block\r
-        //\r
-        RelocPtr = FindFreeMem (BlockList, FreeMemBase, FreeMemSize, (UINTN) TempBlockDesc->Length);\r
-        if (RelocPtr == NULL) {\r
-          return EFI_BUFFER_TOO_SMALL;\r
-        }\r
-\r
-        CopyMem ((VOID *) RelocPtr, (VOID *) (UINTN) TempBlockDesc->Union.DataBlock, (UINTN) TempBlockDesc->Length);\r
-        DEBUG ((EFI_D_INFO, "Capsule reloc data block from 0x%8X to 0x%8X with size 0x%8X\n",\r
-                (UINTN) TempBlockDesc->Union.DataBlock, (UINTN) RelocPtr, (UINTN) TempBlockDesc->Length));\r
-\r
-        TempBlockDesc->Union.DataBlock = (EFI_PHYSICAL_ADDRESS) (UINTN) RelocPtr;\r
-      }\r
-      //\r
-      // Next descriptor\r
-      //\r
-      TempBlockDesc++;\r
+    CoalesceImageEntryPoint = 0;\r
+    Status = GetLongModeContext (&LongModeBuffer);\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((EFI_D_ERROR, "Fail to find the variable for long mode context!\n"));\r
+      Status = EFI_NOT_FOUND;\r
+      goto Done;\r
     }\r
-    //\r
-    // Ok, we made it through. Copy the block.\r
-    // we just support greping one capsule from the lists of block descs list.\r
-    //\r
-    CapsuleTimes ++;\r
-    //\r
-    //Skip the first block descriptor that filled with EFI_CAPSULE_PEIM_PRIVATE_DATA\r
-    //\r
-    if (CapsuleTimes > 1) {\r
-      //\r
-      //For every capsule entry point, check its header to determine whether to relocate it.\r
-      //If it is invalid, skip it and move on to the next capsule. If it is valid, relocate it.\r
-      //\r
-      if (CapsuleBeginFlag) {\r
-        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
-          //\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%8X to 0x%8X 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
-        //\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
-        //\r
-        //Go on relocating the current capule image.\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%8X to 0x%8X 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%8X to 0x%8X 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
-      }\r
-    } else {\r
-      //\r
-      //The first entry is the block descriptor for EFI_CAPSULE_PEIM_PRIVATE_DATA.\r
-      //\r
-      CopyMem ((VOID *) DestPtr, (VOID *) (UINTN) CurrentBlockDesc->Union.DataBlock, (UINTN) CurrentBlockDesc->Length);\r
-      DestPtr += CurrentBlockDesc->Length;\r
+    Status = FindCapsuleCoalesceImage (&CoalesceImageEntryPoint, &CoalesceImageMachineType);\r
+    if ((EFI_ERROR (Status)) || (CoalesceImageMachineType != EFI_IMAGE_MACHINE_X64)) {\r
+      DEBUG ((EFI_D_ERROR, "Fail to find CapsuleX64 module in FV!\n"));\r
+      Status = EFI_NOT_FOUND;\r
+      goto Done;\r
     }\r
+    ASSERT (CoalesceImageEntryPoint != 0);\r
+    CoalesceEntry = (COALESCE_ENTRY) (UINTN) CoalesceImageEntryPoint;\r
+    Status = ModeSwitch (&LongModeBuffer, CoalesceEntry, (EFI_PHYSICAL_ADDRESS)(UINTN)VariableArrayAddress, MemoryResource, MemoryBase, MemorySize);\r
+  } else {\r
     //\r
-    //Walk through the block descriptor list.\r
+    // Capsule is processed in IA32 mode.\r
     //\r
-    CurrentBlockDesc++;\r
+    Status = CapsuleDataCoalesce (PeiServices, (EFI_PHYSICAL_ADDRESS *)(UINTN)VariableArrayAddress, MemoryResource, MemoryBase, MemorySize);\r
   }\r
+#else\r
   //\r
-  // We return the base of memory we want reserved, and the size.\r
-  // The memory peim should handle it appropriately from there.\r
+  // Process capsule directly.\r
   //\r
-  *MemorySize = (UINTN) CapsuleImageSize;\r
-  *MemoryBase = (VOID *) NewCapsuleBase;\r
+  Status = CapsuleDataCoalesce (PeiServices, (EFI_PHYSICAL_ADDRESS *)(UINTN)VariableArrayAddress, MemoryResource, MemoryBase, MemorySize);\r
+#endif\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
+  DEBUG ((EFI_D_INFO, "Capsule Coalesce Status = %r!\n", Status));\r
 \r
-  *AddDataPtr++ = CapsuleIndex;\r
-\r
-  CopyMem (AddDataPtr, &CapsuleOffset[0], sizeof (UINT32) * CapsuleIndex);\r
+  if (Status == EFI_BUFFER_TOO_SMALL) {\r
+    DEBUG ((EFI_D_ERROR, "There is not enough memory to process capsule!\n"));\r
+  }\r
 \r
-  PrivateDataPtr = (EFI_CAPSULE_PEIM_PRIVATE_DATA *) NewCapsuleBase;\r
-  PrivateDataPtr->CapsuleSize = (UINTN)CapsuleImageSize;\r
+  if (Status == EFI_NOT_FOUND) {\r
+    DEBUG ((EFI_D_ERROR, "Fail to parse capsule descriptor in memory!\n"));\r
+    REPORT_STATUS_CODE (\r
+      EFI_ERROR_CODE | EFI_ERROR_MAJOR,\r
+      (EFI_SOFTWARE_PEI_MODULE | EFI_SW_PEI_EC_INVALID_CAPSULE_DESCRIPTOR)\r
+      );\r
+  }\r
 \r
+Done:\r
   return Status;\r
 }\r
 \r
@@ -1139,9 +1122,9 @@ CheckCapsuleUpdate (
   return Status;\r
 }\r
 /**\r
-  This function will look at a capsule and determine if it's a test pattern. \r
+  This function will look at a capsule and determine if it's a test pattern.\r
   If it is, then it will verify it and emit an error message if corruption is detected.\r
-  \r
+\r
   @param PeiServices   Standard pei services pointer\r
   @param CapsuleBase   Base address of coalesced capsule, which is preceeded\r
                        by private data. Very implementation specific.\r
@@ -1168,7 +1151,10 @@ CapsuleTestPattern (
   // is, then test it now.\r
   //\r
   TestPtr = (UINT32 *) CapsuleBase;\r
-  if (*TestPtr == CAPSULE_TEST_SIGNATURE) {\r
+  //\r
+  // 0x54534554 "TEST"\r
+  //\r
+  if (*TestPtr == 0x54534554) {\r
     RetValue = TRUE;\r
     DEBUG ((EFI_D_INFO, "Capsule test pattern mode activated...\n"));\r
     TestSize = TestPtr[1] / sizeof (UINT32);\r
@@ -1223,25 +1209,28 @@ CreateState (
   EFI_CAPSULE_PEIM_PRIVATE_DATA *PrivateData;\r
   UINTN                         Size;\r
   EFI_PHYSICAL_ADDRESS          NewBuffer;\r
-  UINT32                        *DataPtr;\r
-  UINT32                        CapsuleNumber;\r
+  UINTN                         CapsuleNumber;\r
   UINT32                        Index;\r
   EFI_PHYSICAL_ADDRESS          BaseAddress;\r
   UINT64                        Length;\r
\r
-  DataPtr        = NULL;\r
-  CapsuleNumber  = 0;\r
+\r
   PrivateData    = (EFI_CAPSULE_PEIM_PRIVATE_DATA *) CapsuleBase;\r
   if (PrivateData->Signature != EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE) {\r
     return EFI_VOLUME_CORRUPTED;\r
   }\r
+  if (PrivateData->CapsuleAllImageSize >= MAX_ADDRESS) {\r
+    DEBUG ((EFI_D_ERROR, "CapsuleAllImageSize too big - 0x%lx\n", PrivateData->CapsuleAllImageSize));\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+  if (PrivateData->CapsuleNumber >= MAX_ADDRESS) {\r
+    DEBUG ((EFI_D_ERROR, "CapsuleNumber too big - 0x%lx\n", PrivateData->CapsuleNumber));\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
   //\r
   // Capsule Number and Capsule Offset is in the tail of Capsule data.\r
   //\r
-  Size    = (UINTN) PrivateData->CapsuleSize;\r
-  DataPtr = (UINT32*)((UINTN)CapsuleBase + (UINTN)sizeof(EFI_CAPSULE_PEIM_PRIVATE_DATA)+ Size);\r
-  DataPtr = (UINT32*)(((UINTN) DataPtr + sizeof(UINT32) - 1) & ~(sizeof (UINT32) - 1));\r
-  CapsuleNumber = *DataPtr++;\r
+  Size          = (UINTN)PrivateData->CapsuleAllImageSize;\r
+  CapsuleNumber = (UINTN)PrivateData->CapsuleNumber;\r
   //\r
   // Allocate the memory so that it gets preserved into DXE\r
   //\r
@@ -1258,11 +1247,11 @@ CreateState (
   //\r
   // Copy to our new buffer for DXE\r
   //\r
-  DEBUG ((EFI_D_INFO, "Capsule copy from 0x%8X to 0x%8X with size 0x%8X\n", (UINTN) (PrivateData + 1), (UINTN) NewBuffer, Size));\r
-  CopyMem ((VOID *) (UINTN) NewBuffer, (VOID *) (UINTN) (PrivateData + 1), Size);\r
+  DEBUG ((EFI_D_INFO, "Capsule copy from 0x%8X to 0x%8X with size 0x%8X\n", (UINTN)((UINT8 *)PrivateData + sizeof(EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64)), (UINTN) NewBuffer, Size));\r
+  CopyMem ((VOID *) (UINTN) NewBuffer, (VOID *) (UINTN) ((UINT8 *)PrivateData + sizeof(EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof(UINT64)), Size);\r
   //\r
   // Check for test data pattern. If it is the test pattern, then we'll\r
-  // test it ans still create the HOB so that it can be used to verify\r
+  // test it and still create the HOB so that it can be used to verify\r
   // that capsules don't get corrupted all the way into BDS. BDS will\r
   // still try to turn it into a firmware volume, but will think it's\r
   // corrupted so nothing will happen.\r
@@ -1275,16 +1264,16 @@ CreateState (
   // Build the UEFI Capsule Hob for each capsule image.\r
   //\r
   for (Index = 0; Index < CapsuleNumber; Index ++) {\r
-    BaseAddress = NewBuffer + DataPtr[Index];\r
+    BaseAddress = NewBuffer + PrivateData->CapsuleOffset[Index];\r
     Length      = ((EFI_CAPSULE_HEADER *)((UINTN) BaseAddress))->CapsuleImageSize;\r
 \r
     BuildCvHob (BaseAddress, Length);\r
   }\r
-  \r
+\r
   return EFI_SUCCESS;\r
 }\r
 \r
-CONST PEI_CAPSULE_PPI        mCapsulePpi = {\r
+CONST EFI_PEI_CAPSULE_PPI        mCapsulePpi = {\r
   CapsuleCoalesce,\r
   CheckCapsuleUpdate,\r
   CreateState\r
@@ -1292,8 +1281,8 @@ CONST PEI_CAPSULE_PPI        mCapsulePpi = {
 \r
 CONST EFI_PEI_PPI_DESCRIPTOR mUefiPpiListCapsule = {\r
   (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),\r
-  &gPeiCapsulePpiGuid,\r
-  (PEI_CAPSULE_PPI *) &mCapsulePpi\r
+  &gEfiPeiCapsulePpiGuid,\r
+  (EFI_PEI_CAPSULE_PPI *) &mCapsulePpi\r
 };\r
 \r
 /**\r