/** @file
  Capsule update PEIM for UEFI2.0

Copyright (c) 2006 - 2019, Intel Corporation. All rights reserved.<BR>
Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>

SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "Capsule.h"

#define DEFAULT_SG_LIST_HEADS  (20)

#ifdef MDE_CPU_IA32
//
// Global Descriptor Table (GDT)
//
GLOBAL_REMOVE_IF_UNREFERENCED IA32_SEGMENT_DESCRIPTOR  mGdtEntries[] = {
  /* selector { Global Segment Descriptor                              } */
  /* 0x00 */ {
    { 0,      0, 0, 0,   0, 0, 0, 0,   0, 0, 0, 0, 0 }
  },                                                                      // null descriptor
  /* 0x08 */ {
    { 0xffff, 0, 0, 0x3, 1, 0, 1, 0xf, 0, 0, 1, 1, 0 }
  },                                                                      // linear data segment descriptor
  /* 0x10 */ {
    { 0xffff, 0, 0, 0xf, 1, 0, 1, 0xf, 0, 0, 1, 1, 0 }
  },                                                                      // linear code segment descriptor
  /* 0x18 */ {
    { 0xffff, 0, 0, 0x3, 1, 0, 1, 0xf, 0, 0, 1, 1, 0 }
  },                                                                      // system data segment descriptor
  /* 0x20 */ {
    { 0xffff, 0, 0, 0xb, 1, 0, 1, 0xf, 0, 0, 1, 1, 0 }
  },                                                                      // system code segment descriptor
  /* 0x28 */ {
    { 0,      0, 0, 0,   0, 0, 0, 0,   0, 0, 0, 0, 0 }
  },                                                                      // spare segment descriptor
  /* 0x30 */ {
    { 0xffff, 0, 0, 0x3, 1, 0, 1, 0xf, 0, 0, 1, 1, 0 }
  },                                                                      // system data segment descriptor
  /* 0x38 */ {
    { 0xffff, 0, 0, 0xb, 1, 0, 1, 0xf, 0, 1, 0, 1, 0 }
  },                                                                      // system code segment descriptor
  /* 0x40 */ {
    { 0,      0, 0, 0,   0, 0, 0, 0,   0, 0, 0, 0, 0 }
  },                                                                      // spare segment descriptor
};

//
// IA32 Gdt register
//
GLOBAL_REMOVE_IF_UNREFERENCED CONST IA32_DESCRIPTOR  mGdt = {
  sizeof (mGdtEntries) - 1,
  (UINTN)mGdtEntries
};

/**
  The function will check if 1G page is supported.

  @retval TRUE   1G page is supported.
  @retval FALSE  1G page is not supported.

**/
BOOLEAN
IsPage1GSupport (
  VOID
  )
{
  UINT32   RegEax;
  UINT32   RegEdx;
  BOOLEAN  Page1GSupport;

  Page1GSupport = FALSE;
  if (PcdGetBool (PcdUse1GPageTable)) {
    AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);
    if (RegEax >= 0x80000001) {
      AsmCpuid (0x80000001, NULL, NULL, NULL, &RegEdx);
      if ((RegEdx & BIT26) != 0) {
        Page1GSupport = TRUE;
      }
    }
  }

  return Page1GSupport;
}

/**
  Calculate the total size of page table.

  @param[in] Page1GSupport      1G page support or not.

  @return The size of page table.

**/
UINTN
CalculatePageTableSize (
  IN BOOLEAN  Page1GSupport
  )
{
  UINTN   ExtraPageTablePages;
  UINTN   TotalPagesNum;
  UINT8   PhysicalAddressBits;
  UINT32  NumberOfPml4EntriesNeeded;
  UINT32  NumberOfPdpEntriesNeeded;

  //
  // Create 4G page table by default,
  // and let PF handler to handle > 4G request.
  //
  PhysicalAddressBits = 32;
  ExtraPageTablePages = EXTRA_PAGE_TABLE_PAGES;

  //
  // Calculate the table entries needed.
  //
  if (PhysicalAddressBits <= 39 ) {
    NumberOfPml4EntriesNeeded = 1;
    NumberOfPdpEntriesNeeded  = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 30));
  } else {
    NumberOfPml4EntriesNeeded = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 39));
    NumberOfPdpEntriesNeeded  = 512;
  }

  if (!Page1GSupport) {
    TotalPagesNum = (NumberOfPdpEntriesNeeded + 1) * NumberOfPml4EntriesNeeded + 1;
  } else {
    TotalPagesNum = NumberOfPml4EntriesNeeded + 1;
  }

  TotalPagesNum += ExtraPageTablePages;

  return EFI_PAGES_TO_SIZE (TotalPagesNum);
}

/**
  Allocates and fills in the Page Directory and Page Table Entries to
  establish a 4G page table.

  @param[in] PageTablesAddress  The base address of page table.
  @param[in] Page1GSupport      1G page support or not.

**/
VOID
Create4GPageTables (
  IN EFI_PHYSICAL_ADDRESS  PageTablesAddress,
  IN BOOLEAN               Page1GSupport
  )
{
  UINT8                           PhysicalAddressBits;
  EFI_PHYSICAL_ADDRESS            PageAddress;
  UINTN                           IndexOfPml4Entries;
  UINTN                           IndexOfPdpEntries;
  UINTN                           IndexOfPageDirectoryEntries;
  UINT32                          NumberOfPml4EntriesNeeded;
  UINT32                          NumberOfPdpEntriesNeeded;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PageMapLevel4Entry;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PageMap;
  PAGE_MAP_AND_DIRECTORY_POINTER  *PageDirectoryPointerEntry;
  PAGE_TABLE_ENTRY                *PageDirectoryEntry;
  UINTN                           BigPageAddress;
  PAGE_TABLE_1G_ENTRY             *PageDirectory1GEntry;
  UINT64                          AddressEncMask;

  //
  // Make sure AddressEncMask is contained to smallest supported address field.
  //
  AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;

  //
  // Create 4G page table by default,
  // and let PF handler to handle > 4G request.
  //
  PhysicalAddressBits = 32;

  //
  // Calculate the table entries needed.
  //
  if (PhysicalAddressBits <= 39 ) {
    NumberOfPml4EntriesNeeded = 1;
    NumberOfPdpEntriesNeeded  = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 30));
  } else {
    NumberOfPml4EntriesNeeded = (UINT32)LShiftU64 (1, (PhysicalAddressBits - 39));
    NumberOfPdpEntriesNeeded  = 512;
  }

  //
  // Pre-allocate big pages to avoid later allocations.
  //
  BigPageAddress = (UINTN)PageTablesAddress;

  //
  // By architecture only one PageMapLevel4 exists - so lets allocate storage for it.
  //
  PageMap         = (VOID *)BigPageAddress;
  BigPageAddress += SIZE_4KB;

  PageMapLevel4Entry = PageMap;
  PageAddress        = 0;
  for (IndexOfPml4Entries = 0; IndexOfPml4Entries < NumberOfPml4EntriesNeeded; IndexOfPml4Entries++, PageMapLevel4Entry++) {
    //
    // Each PML4 entry points to a page of Page Directory Pointer entires.
    // So lets allocate space for them and fill them in in the IndexOfPdpEntries loop.
    //
    PageDirectoryPointerEntry = (VOID *)BigPageAddress;
    BigPageAddress           += SIZE_4KB;

    //
    // Make a PML4 Entry
    //
    PageMapLevel4Entry->Uint64         = (UINT64)(UINTN)PageDirectoryPointerEntry | AddressEncMask;
    PageMapLevel4Entry->Bits.ReadWrite = 1;
    PageMapLevel4Entry->Bits.Present   = 1;

    if (Page1GSupport) {
      PageDirectory1GEntry = (VOID *)PageDirectoryPointerEntry;

      for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectory1GEntry++, PageAddress += SIZE_1GB) {
        //
        // Fill in the Page Directory entries
        //
        PageDirectory1GEntry->Uint64         = (UINT64)PageAddress | AddressEncMask;
        PageDirectory1GEntry->Bits.ReadWrite = 1;
        PageDirectory1GEntry->Bits.Present   = 1;
        PageDirectory1GEntry->Bits.MustBe1   = 1;
      }
    } else {
      for (IndexOfPdpEntries = 0; IndexOfPdpEntries < NumberOfPdpEntriesNeeded; IndexOfPdpEntries++, PageDirectoryPointerEntry++) {
        //
        // Each Directory Pointer entries points to a page of Page Directory entires.
        // So allocate space for them and fill them in in the IndexOfPageDirectoryEntries loop.
        //
        PageDirectoryEntry = (VOID *)BigPageAddress;
        BigPageAddress    += SIZE_4KB;

        //
        // Fill in a Page Directory Pointer Entries
        //
        PageDirectoryPointerEntry->Uint64         = (UINT64)(UINTN)PageDirectoryEntry | AddressEncMask;
        PageDirectoryPointerEntry->Bits.ReadWrite = 1;
        PageDirectoryPointerEntry->Bits.Present   = 1;

        for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectoryEntry++, PageAddress += SIZE_2MB) {
          //
          // Fill in the Page Directory entries
          //
          PageDirectoryEntry->Uint64         = (UINT64)PageAddress | AddressEncMask;
          PageDirectoryEntry->Bits.ReadWrite = 1;
          PageDirectoryEntry->Bits.Present   = 1;
          PageDirectoryEntry->Bits.MustBe1   = 1;
        }
      }

      for ( ; IndexOfPdpEntries < 512; IndexOfPdpEntries++, PageDirectoryPointerEntry++) {
        ZeroMem (
          PageDirectoryPointerEntry,
          sizeof (PAGE_MAP_AND_DIRECTORY_POINTER)
          );
      }
    }
  }

  //
  // For the PML4 entries we are not using fill in a null entry.
  //
  for ( ; IndexOfPml4Entries < 512; IndexOfPml4Entries++, PageMapLevel4Entry++) {
    ZeroMem (
      PageMapLevel4Entry,
      sizeof (PAGE_MAP_AND_DIRECTORY_POINTER)
      );
  }
}

/**
  Return function from long mode to 32-bit mode.

  @param  EntrypointContext  Context for mode switching
  @param  ReturnContext      Context for mode switching

**/
VOID
ReturnFunction (
  SWITCH_32_TO_64_CONTEXT  *EntrypointContext,
  SWITCH_64_TO_32_CONTEXT  *ReturnContext
  )
{
  //
  // Restore original GDT
  //
  AsmWriteGdtr (&ReturnContext->Gdtr);

  //
  // return to original caller
  //
  LongJump ((BASE_LIBRARY_JUMP_BUFFER  *)(UINTN)EntrypointContext->JumpBuffer, 1);

  //
  // never be here
  //
  ASSERT (FALSE);
}

/**
  Thunk function from 32-bit protection mode to long mode.

  @param  PageTableAddress  Page table base address
  @param  Context           Context for mode switching
  @param  ReturnContext     Context for mode switching

  @retval EFI_SUCCESS  Function successfully executed.

**/
EFI_STATUS
Thunk32To64 (
  EFI_PHYSICAL_ADDRESS     PageTableAddress,
  SWITCH_32_TO_64_CONTEXT  *Context,
  SWITCH_64_TO_32_CONTEXT  *ReturnContext
  )
{
  UINTN       SetJumpFlag;
  EFI_STATUS  Status;

  //
  // Save return address, LongJump will return here then
  //
  SetJumpFlag = SetJump ((BASE_LIBRARY_JUMP_BUFFER  *)(UINTN)Context->JumpBuffer);

  if (SetJumpFlag == 0) {
    //
    // Build 4G Page Tables.
    //
    Create4GPageTables (PageTableAddress, Context->Page1GSupport);

    //
    // Create 64-bit GDT
    //
    AsmWriteGdtr (&mGdt);

    //
    // Write CR3
    //
    AsmWriteCr3 ((UINTN)PageTableAddress);

    DEBUG ((
      DEBUG_INFO,
      "%a() Stack Base: 0x%lx, Stack Size: 0x%lx\n",
      __FUNCTION__,
      Context->StackBufferBase,
      Context->StackBufferLength
      ));

    //
    // Disable interrupt of Debug timer, since the IDT table cannot work in long mode
    //
    SaveAndSetDebugTimerInterrupt (FALSE);
    //
    // Transfer to long mode
    //
    AsmEnablePaging64 (
      0x38,
      (UINT64)Context->EntryPoint,
      (UINT64)(UINTN)Context,
      (UINT64)(UINTN)ReturnContext,
      Context->StackBufferBase + Context->StackBufferLength
      );
  }

  //
  // Convert to 32-bit Status and return
  //
  Status = EFI_SUCCESS;
  if ((UINTN)ReturnContext->ReturnStatus != 0) {
    Status = ENCODE_ERROR ((UINTN)ReturnContext->ReturnStatus);
  }

  return Status;
}

/**
  If in 32 bit protection mode, and coalesce image is of X64, switch to long mode.

  @param  LongModeBuffer            The context of long mode.
  @param  CoalesceEntry             Entry of coalesce image.
  @param  BlockListAddr             Address of block list.
  @param  MemoryResource            Pointer to the buffer of memory resource descriptor.
  @param  MemoryBase                Base of memory range.
  @param  MemorySize                Size of memory range.

  @retval EFI_SUCCESS               Successfully switched to long mode and execute coalesce.
  @retval Others                    Failed to execute coalesce in long mode.

**/
EFI_STATUS
ModeSwitch (
  IN EFI_CAPSULE_LONG_MODE_BUFFER  *LongModeBuffer,
  IN COALESCE_ENTRY                CoalesceEntry,
  IN EFI_PHYSICAL_ADDRESS          BlockListAddr,
  IN MEMORY_RESOURCE_DESCRIPTOR    *MemoryResource,
  IN OUT VOID                      **MemoryBase,
  IN OUT UINTN                     *MemorySize
  )
{
  EFI_STATUS                Status;
  EFI_PHYSICAL_ADDRESS      MemoryBase64;
  UINT64                    MemorySize64;
  EFI_PHYSICAL_ADDRESS      MemoryEnd64;
  SWITCH_32_TO_64_CONTEXT   Context;
  SWITCH_64_TO_32_CONTEXT   ReturnContext;
  BASE_LIBRARY_JUMP_BUFFER  JumpBuffer;
  EFI_PHYSICAL_ADDRESS      ReservedRangeBase;
  EFI_PHYSICAL_ADDRESS      ReservedRangeEnd;
  BOOLEAN                   Page1GSupport;

  ZeroMem (&Context, sizeof (SWITCH_32_TO_64_CONTEXT));
  ZeroMem (&ReturnContext, sizeof (SWITCH_64_TO_32_CONTEXT));

  MemoryBase64 = (UINT64)(UINTN)*MemoryBase;
  MemorySize64 = (UINT64)(UINTN)*MemorySize;
  MemoryEnd64  = MemoryBase64 + MemorySize64;

  Page1GSupport = IsPage1GSupport ();

  //
  // Merge memory range reserved for stack and page table
  //
  if (LongModeBuffer->StackBaseAddress < LongModeBuffer->PageTableAddress) {
    ReservedRangeBase = LongModeBuffer->StackBaseAddress;
    ReservedRangeEnd  = LongModeBuffer->PageTableAddress + CalculatePageTableSize (Page1GSupport);
  } else {
    ReservedRangeBase = LongModeBuffer->PageTableAddress;
    ReservedRangeEnd  = LongModeBuffer->StackBaseAddress + LongModeBuffer->StackSize;
  }

  //
  // Check if memory range reserved is overlap with MemoryBase ~ MemoryBase + MemorySize.
  // If they are overlapped, get a larger range to process capsule data.
  //
  if (ReservedRangeBase <= MemoryBase64) {
    if (ReservedRangeEnd < MemoryEnd64) {
      MemoryBase64 = ReservedRangeEnd;
    } else {
      DEBUG ((DEBUG_ERROR, "Memory is not enough to process capsule!\n"));
      return EFI_OUT_OF_RESOURCES;
    }
  } else if (ReservedRangeBase < MemoryEnd64) {
    if ((ReservedRangeEnd < MemoryEnd64) &&
        (ReservedRangeBase - MemoryBase64 < MemoryEnd64 - ReservedRangeEnd))
    {
      MemoryBase64 = ReservedRangeEnd;
    } else {
      MemorySize64 = (UINT64)(UINTN)(ReservedRangeBase - MemoryBase64);
    }
  }

  //
  // Initialize context jumping to 64-bit enviroment
  //
  Context.JumpBuffer        = (EFI_PHYSICAL_ADDRESS)(UINTN)&JumpBuffer;
  Context.StackBufferBase   = LongModeBuffer->StackBaseAddress;
  Context.StackBufferLength = LongModeBuffer->StackSize;
  Context.EntryPoint        = (EFI_PHYSICAL_ADDRESS)(UINTN)CoalesceEntry;
  Context.BlockListAddr     = BlockListAddr;
  Context.MemoryResource    = (EFI_PHYSICAL_ADDRESS)(UINTN)MemoryResource;
  Context.MemoryBase64Ptr   = (EFI_PHYSICAL_ADDRESS)(UINTN)&MemoryBase64;
  Context.MemorySize64Ptr   = (EFI_PHYSICAL_ADDRESS)(UINTN)&MemorySize64;
  Context.Page1GSupport     = Page1GSupport;
  Context.AddressEncMask    = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;

  //
  // Prepare data for return back
  //
  ReturnContext.ReturnCs         = 0x10;
  ReturnContext.ReturnEntryPoint = (EFI_PHYSICAL_ADDRESS)(UINTN)ReturnFunction;
  //
  // Will save the return status of processing capsule
  //
  ReturnContext.ReturnStatus = 0;

  //
  // Save original GDT
  //
  AsmReadGdtr ((IA32_DESCRIPTOR *)&ReturnContext.Gdtr);

  Status = Thunk32To64 (LongModeBuffer->PageTableAddress, &Context, &ReturnContext);

  if (!EFI_ERROR (Status)) {
    *MemoryBase = (VOID *)(UINTN)MemoryBase64;
    *MemorySize = (UINTN)MemorySize64;
  }

  return Status;
}

/**
  Locates the coalesce image entry point, and detects its machine type.

  @param CoalesceImageEntryPoint   Pointer to coalesce image entry point for output.
  @param CoalesceImageMachineType  Pointer to machine type of coalesce image.

  @retval EFI_SUCCESS     Coalesce image successfully located.
  @retval Others          Failed to locate the coalesce image.

**/
EFI_STATUS
FindCapsuleCoalesceImage (
  OUT EFI_PHYSICAL_ADDRESS  *CoalesceImageEntryPoint,
  OUT UINT16                *CoalesceImageMachineType
  )
{
  EFI_STATUS             Status;
  UINTN                  Instance;
  EFI_PEI_LOAD_FILE_PPI  *LoadFile;
  EFI_PEI_FV_HANDLE      VolumeHandle;
  EFI_PEI_FILE_HANDLE    FileHandle;
  EFI_PHYSICAL_ADDRESS   CoalesceImageAddress;
  UINT64                 CoalesceImageSize;
  UINT32                 AuthenticationState;

  Instance = 0;

  while (TRUE) {
    Status = PeiServicesFfsFindNextVolume (Instance++, &VolumeHandle);
    if (EFI_ERROR (Status)) {
      return Status;
    }

    Status = PeiServicesFfsFindFileByName (PcdGetPtr (PcdCapsuleCoalesceFile), VolumeHandle, &FileHandle);
    if (!EFI_ERROR (Status)) {
      Status = PeiServicesLocatePpi (&gEfiPeiLoadFilePpiGuid, 0, NULL, (VOID **)&LoadFile);
      ASSERT_EFI_ERROR (Status);

      Status = LoadFile->LoadFile (
                           LoadFile,
                           FileHandle,
                           &CoalesceImageAddress,
                           &CoalesceImageSize,
                           CoalesceImageEntryPoint,
                           &AuthenticationState
                           );
      if (EFI_ERROR (Status)) {
        DEBUG ((DEBUG_ERROR, "Unable to find PE32 section in CapsuleX64 image ffs %r!\n", Status));
        return Status;
      }

      *CoalesceImageMachineType = PeCoffLoaderGetMachineType ((VOID *)(UINTN)CoalesceImageAddress);
      break;
    } else {
      continue;
    }
  }

  return Status;
}

/**
  Gets the reserved long mode buffer.

  @param  LongModeBuffer  Pointer to the long mode buffer for output.

  @retval EFI_SUCCESS     Long mode buffer successfully retrieved.
  @retval Others          Variable storing long mode buffer not found.

**/
EFI_STATUS
GetLongModeContext (
  OUT EFI_CAPSULE_LONG_MODE_BUFFER  *LongModeBuffer
  )
{
  EFI_STATUS                       Status;
  UINTN                            Size;
  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;

  Status = PeiServicesLocatePpi (
             &gEfiPeiReadOnlyVariable2PpiGuid,
             0,
             NULL,
             (VOID **)&PPIVariableServices
             );
  ASSERT_EFI_ERROR (Status);

  Size   = sizeof (EFI_CAPSULE_LONG_MODE_BUFFER);
  Status = PPIVariableServices->GetVariable (
                                  PPIVariableServices,
                                  EFI_CAPSULE_LONG_MODE_BUFFER_NAME,
                                  &gEfiCapsuleVendorGuid,
                                  NULL,
                                  &Size,
                                  LongModeBuffer
                                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Error Get LongModeBuffer variable %r!\n", Status));
  }

  return Status;
}

#endif

#if defined (MDE_CPU_IA32) || defined (MDE_CPU_X64)

/**
  Get physical address bits.

  @return Physical address bits.

**/
UINT8
GetPhysicalAddressBits (
  VOID
  )
{
  UINT32  RegEax;
  UINT8   PhysicalAddressBits;
  VOID    *Hob;

  //
  // Get physical address bits supported.
  //
  Hob = GetFirstHob (EFI_HOB_TYPE_CPU);
  if (Hob != NULL) {
    PhysicalAddressBits = ((EFI_HOB_CPU *)Hob)->SizeOfMemorySpace;
  } else {
    AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);
    if (RegEax >= 0x80000008) {
      AsmCpuid (0x80000008, &RegEax, NULL, NULL, NULL);
      PhysicalAddressBits = (UINT8)RegEax;
    } else {
      PhysicalAddressBits = 36;
    }
  }

  //
  // IA-32e paging translates 48-bit linear addresses to 52-bit physical addresses.
  //
  ASSERT (PhysicalAddressBits <= 52);
  if (PhysicalAddressBits > 48) {
    PhysicalAddressBits = 48;
  }

  return PhysicalAddressBits;
}

#endif

/**
  Sort memory resource entries based upon PhysicalStart, from low to high.

  @param[in, out] MemoryResource    A pointer to the memory resource entry buffer.

**/
VOID
SortMemoryResourceDescriptor (
  IN OUT MEMORY_RESOURCE_DESCRIPTOR  *MemoryResource
  )
{
  MEMORY_RESOURCE_DESCRIPTOR  *MemoryResourceEntry;
  MEMORY_RESOURCE_DESCRIPTOR  *NextMemoryResourceEntry;
  MEMORY_RESOURCE_DESCRIPTOR  TempMemoryResource;

  MemoryResourceEntry     = MemoryResource;
  NextMemoryResourceEntry = MemoryResource + 1;
  while (MemoryResourceEntry->ResourceLength != 0) {
    while (NextMemoryResourceEntry->ResourceLength != 0) {
      if (MemoryResourceEntry->PhysicalStart > NextMemoryResourceEntry->PhysicalStart) {
        CopyMem (&TempMemoryResource, MemoryResourceEntry, sizeof (MEMORY_RESOURCE_DESCRIPTOR));
        CopyMem (MemoryResourceEntry, NextMemoryResourceEntry, sizeof (MEMORY_RESOURCE_DESCRIPTOR));
        CopyMem (NextMemoryResourceEntry, &TempMemoryResource, sizeof (MEMORY_RESOURCE_DESCRIPTOR));
      }

      NextMemoryResourceEntry = NextMemoryResourceEntry + 1;
    }

    MemoryResourceEntry     = MemoryResourceEntry + 1;
    NextMemoryResourceEntry = MemoryResourceEntry + 1;
  }
}

/**
  Merge continous memory resource entries.

  @param[in, out] MemoryResource    A pointer to the memory resource entry buffer.

**/
VOID
MergeMemoryResourceDescriptor (
  IN OUT MEMORY_RESOURCE_DESCRIPTOR  *MemoryResource
  )
{
  MEMORY_RESOURCE_DESCRIPTOR  *MemoryResourceEntry;
  MEMORY_RESOURCE_DESCRIPTOR  *NewMemoryResourceEntry;
  MEMORY_RESOURCE_DESCRIPTOR  *NextMemoryResourceEntry;
  MEMORY_RESOURCE_DESCRIPTOR  *MemoryResourceEnd;

  MemoryResourceEntry    = MemoryResource;
  NewMemoryResourceEntry = MemoryResource;
  while (MemoryResourceEntry->ResourceLength != 0) {
    CopyMem (NewMemoryResourceEntry, MemoryResourceEntry, sizeof (MEMORY_RESOURCE_DESCRIPTOR));
    NextMemoryResourceEntry = MemoryResourceEntry + 1;

    while ((NextMemoryResourceEntry->ResourceLength != 0) &&
           (NextMemoryResourceEntry->PhysicalStart == (MemoryResourceEntry->PhysicalStart + MemoryResourceEntry->ResourceLength)))
    {
      MemoryResourceEntry->ResourceLength += NextMemoryResourceEntry->ResourceLength;
      if (NewMemoryResourceEntry != MemoryResourceEntry) {
        NewMemoryResourceEntry->ResourceLength += NextMemoryResourceEntry->ResourceLength;
      }

      NextMemoryResourceEntry = NextMemoryResourceEntry + 1;
    }

    MemoryResourceEntry    = NextMemoryResourceEntry;
    NewMemoryResourceEntry = NewMemoryResourceEntry + 1;
  }

  //
  // Set NULL terminate memory resource descriptor after merging.
  //
  MemoryResourceEnd = NewMemoryResourceEntry;
  ZeroMem (MemoryResourceEnd, sizeof (MEMORY_RESOURCE_DESCRIPTOR));
}

/**
  Build memory resource descriptor from resource descriptor in HOB list.

  @return Pointer to the buffer of memory resource descriptor.
          NULL if no memory resource descriptor reported in HOB list
          before capsule Coalesce.

**/
MEMORY_RESOURCE_DESCRIPTOR *
BuildMemoryResourceDescriptor (
  VOID
  )
{
  EFI_PEI_HOB_POINTERS         Hob;
  UINTN                        Index;
  EFI_HOB_RESOURCE_DESCRIPTOR  *ResourceDescriptor;
  MEMORY_RESOURCE_DESCRIPTOR   *MemoryResource;
  EFI_STATUS                   Status;

  //
  // Get the count of memory resource descriptor.
  //
  Index   = 0;
  Hob.Raw = GetFirstHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR);
  while (Hob.Raw != NULL) {
    ResourceDescriptor = (EFI_HOB_RESOURCE_DESCRIPTOR *)Hob.Raw;
    if (ResourceDescriptor->ResourceType == EFI_RESOURCE_SYSTEM_MEMORY) {
      Index++;
    }

    Hob.Raw = GET_NEXT_HOB (Hob);
    Hob.Raw = GetNextHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR, Hob.Raw);
  }

  if (Index == 0) {
    DEBUG ((DEBUG_INFO | DEBUG_WARN, "No memory resource descriptor reported in HOB list before capsule Coalesce\n"));
 #if defined (MDE_CPU_IA32) || defined (MDE_CPU_X64)
    //
    // Allocate memory to hold memory resource descriptor,
    // include extra one NULL terminate memory resource descriptor.
    //
    Status = PeiServicesAllocatePool ((1 + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR), (VOID **)&MemoryResource);
    ASSERT_EFI_ERROR (Status);
    ZeroMem (MemoryResource, (1 + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR));

    MemoryResource[0].PhysicalStart  = 0;
    MemoryResource[0].ResourceLength = LShiftU64 (1, GetPhysicalAddressBits ());
    DEBUG ((
      DEBUG_INFO,
      "MemoryResource[0x0] - Start(0x%0lx) Length(0x%0lx)\n",
      MemoryResource[0x0].PhysicalStart,
      MemoryResource[0x0].ResourceLength
      ));
    return MemoryResource;
 #else
    return NULL;
 #endif
  }

  //
  // Allocate memory to hold memory resource descriptor,
  // include extra one NULL terminate memory resource descriptor.
  //
  Status = PeiServicesAllocatePool ((Index + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR), (VOID **)&MemoryResource);
  ASSERT_EFI_ERROR (Status);
  ZeroMem (MemoryResource, (Index + 1) * sizeof (MEMORY_RESOURCE_DESCRIPTOR));

  //
  // Get the content of memory resource descriptor.
  //
  Index   = 0;
  Hob.Raw = GetFirstHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR);
  while (Hob.Raw != NULL) {
    ResourceDescriptor = (EFI_HOB_RESOURCE_DESCRIPTOR *)Hob.Raw;
    if (ResourceDescriptor->ResourceType == EFI_RESOURCE_SYSTEM_MEMORY) {
      DEBUG ((
        DEBUG_INFO,
        "MemoryResource[0x%x] - Start(0x%0lx) Length(0x%0lx)\n",
        Index,
        ResourceDescriptor->PhysicalStart,
        ResourceDescriptor->ResourceLength
        ));
      MemoryResource[Index].PhysicalStart  = ResourceDescriptor->PhysicalStart;
      MemoryResource[Index].ResourceLength = ResourceDescriptor->ResourceLength;
      Index++;
    }

    Hob.Raw = GET_NEXT_HOB (Hob);
    Hob.Raw = GetNextHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR, Hob.Raw);
  }

  SortMemoryResourceDescriptor (MemoryResource);
  MergeMemoryResourceDescriptor (MemoryResource);

  DEBUG ((DEBUG_INFO, "Dump MemoryResource[] after sorted and merged\n"));
  for (Index = 0; MemoryResource[Index].ResourceLength != 0; Index++) {
    DEBUG ((
      DEBUG_INFO,
      "  MemoryResource[0x%x] - Start(0x%0lx) Length(0x%0lx)\n",
      Index,
      MemoryResource[Index].PhysicalStart,
      MemoryResource[Index].ResourceLength
      ));
  }

  return MemoryResource;
}

/**
  Check if the capsules are staged.

  @retval TRUE              The capsules are staged.
  @retval FALSE             The capsules are not staged.

**/
BOOLEAN
AreCapsulesStaged (
  VOID
  )
{
  EFI_STATUS                       Status;
  UINTN                            Size;
  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;
  EFI_PHYSICAL_ADDRESS             CapsuleDataPtr64;

  CapsuleDataPtr64 = 0;

  Status = PeiServicesLocatePpi (
             &gEfiPeiReadOnlyVariable2PpiGuid,
             0,
             NULL,
             (VOID **)&PPIVariableServices
             );

  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Failed to find ReadOnlyVariable2PPI\n"));
    return FALSE;
  }

  //
  // Check for Update capsule
  //
  Size   = sizeof (CapsuleDataPtr64);
  Status = PPIVariableServices->GetVariable (
                                  PPIVariableServices,
                                  EFI_CAPSULE_VARIABLE_NAME,
                                  &gEfiCapsuleVendorGuid,
                                  NULL,
                                  &Size,
                                  (VOID *)&CapsuleDataPtr64
                                  );

  if (!EFI_ERROR (Status)) {
    return TRUE;
  }

  return FALSE;
}

/**
  Check all the variables for SG list heads and get the count and addresses.

  @param ListLength               A pointer would return the SG list length.
  @param HeadList                 A ponter to the capsule SG list.

  @retval EFI_SUCCESS             a valid capsule is present
  @retval EFI_NOT_FOUND           if a valid capsule is not present
  @retval EFI_INVALID_PARAMETER   the input parameter is invalid
  @retval EFI_OUT_OF_RESOURCES    fail to allocate memory

**/
EFI_STATUS
GetScatterGatherHeadEntries (
  OUT UINTN                 *ListLength,
  OUT EFI_PHYSICAL_ADDRESS  **HeadList
  )
{
  EFI_STATUS                       Status;
  UINTN                            Size;
  UINTN                            Index;
  UINTN                            TempIndex;
  UINTN                            ValidIndex;
  BOOLEAN                          Flag;
  CHAR16                           CapsuleVarName[30];
  CHAR16                           *TempVarName;
  EFI_PHYSICAL_ADDRESS             CapsuleDataPtr64;
  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;
  EFI_PHYSICAL_ADDRESS             *TempList;
  EFI_PHYSICAL_ADDRESS             *EnlargedTempList;
  UINTN                            TempListLength;

  Index             = 0;
  TempVarName       = NULL;
  CapsuleVarName[0] = 0;
  ValidIndex        = 0;
  CapsuleDataPtr64  = 0;

  if ((ListLength == NULL) || (HeadList == NULL)) {
    DEBUG ((DEBUG_ERROR, "%a Invalid parameters.  Inputs can't be NULL\n", __FUNCTION__));
    ASSERT (ListLength != NULL);
    ASSERT (HeadList != NULL);
    return EFI_INVALID_PARAMETER;
  }

  *ListLength = 0;
  *HeadList   = NULL;

  Status = PeiServicesLocatePpi (
             &gEfiPeiReadOnlyVariable2PpiGuid,
             0,
             NULL,
             (VOID **)&PPIVariableServices
             );

  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Failed to find ReadOnlyVariable2PPI\n"));
    return Status;
  }

  //
  // Allocate memory for sg list head
  //
  TempListLength = DEFAULT_SG_LIST_HEADS * sizeof (EFI_PHYSICAL_ADDRESS);
  TempList       = AllocateZeroPool (TempListLength);
  if (TempList == NULL) {
    DEBUG ((DEBUG_ERROR, "Failed to allocate memory\n"));
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // setup var name buffer for update capsules
  //
  StrCpyS (CapsuleVarName, sizeof (CapsuleVarName) / sizeof (CHAR16), EFI_CAPSULE_VARIABLE_NAME);
  TempVarName = CapsuleVarName + StrLen (CapsuleVarName);
  while (TRUE) {
    if (Index != 0) {
      UnicodeValueToStringS (
        TempVarName,
        (sizeof (CapsuleVarName) - ((UINTN)TempVarName - (UINTN)CapsuleVarName)),
        0,
        Index,
        0
        );
    }

    Size   = sizeof (CapsuleDataPtr64);
    Status = PPIVariableServices->GetVariable (
                                    PPIVariableServices,
                                    CapsuleVarName,
                                    &gEfiCapsuleVendorGuid,
                                    NULL,
                                    &Size,
                                    (VOID *)&CapsuleDataPtr64
                                    );

    if (EFI_ERROR (Status)) {
      if (Status != EFI_NOT_FOUND) {
        DEBUG ((DEBUG_ERROR, "Unexpected error getting Capsule Update variable.  Status = %r\n", Status));
      }

      break;
    }

    //
    // If this BlockList has been linked before, skip this variable
    //
    Flag = FALSE;
    for (TempIndex = 0; TempIndex < ValidIndex; TempIndex++) {
      if (TempList[TempIndex] == CapsuleDataPtr64) {
        Flag = TRUE;
        break;
      }
    }

    if (Flag) {
      Index++;
      continue;
    }

    //
    // The TempList is full, enlarge it
    //
    if ((ValidIndex + 1) >= TempListLength) {
      EnlargedTempList = AllocateZeroPool (TempListLength * 2);
      if (EnlargedTempList == NULL) {
        DEBUG ((DEBUG_ERROR, "Fail to allocate memory!\n"));
        return EFI_OUT_OF_RESOURCES;
      }

      CopyMem (EnlargedTempList, TempList, TempListLength);
      FreePool (TempList);
      TempList        = EnlargedTempList;
      TempListLength *= 2;
    }

    //
    // add it to the cached list
    //
    TempList[ValidIndex++] = CapsuleDataPtr64;
    Index++;
  }

  if (ValidIndex == 0) {
    DEBUG ((DEBUG_ERROR, "%a didn't find any SG lists in variables\n", __FUNCTION__));
    return EFI_NOT_FOUND;
  }

  *HeadList = AllocateZeroPool ((ValidIndex + 1) * sizeof (EFI_PHYSICAL_ADDRESS));
  if (*HeadList == NULL) {
    DEBUG ((DEBUG_ERROR, "Failed to allocate memory\n"));
    return EFI_OUT_OF_RESOURCES;
  }

  CopyMem (*HeadList, TempList, (ValidIndex) * sizeof (EFI_PHYSICAL_ADDRESS));
  *ListLength = ValidIndex;

  return EFI_SUCCESS;
}

/**
  Capsule PPI service to coalesce a fragmented capsule in memory.

  @param PeiServices  General purpose services available to every PEIM.
  @param MemoryBase   Pointer to the base of a block of memory that we can walk
                      all over while trying to coalesce our buffers.
                      On output, this variable will hold the base address of
                      a coalesced capsule.
  @param MemorySize   Size of the memory region pointed to by MemoryBase.
                      On output, this variable will contain the size of the
                      coalesced capsule.

  @retval EFI_NOT_FOUND   if we can't determine the boot mode
                          if the boot mode is not flash-update
                          if we could not find the capsule descriptors

  @retval EFI_BUFFER_TOO_SMALL
                          if we could not coalesce the capsule in the memory
                          region provided to us

  @retval EFI_SUCCESS     if there's no capsule, or if we processed the
                          capsule successfully.
**/
EFI_STATUS
EFIAPI
CapsuleCoalesce (
  IN     EFI_PEI_SERVICES  **PeiServices,
  IN OUT VOID              **MemoryBase,
  IN OUT UINTN             *MemorySize
  )
{
  EFI_STATUS                  Status;
  EFI_BOOT_MODE               BootMode;
  UINTN                       ListLength;
  EFI_PHYSICAL_ADDRESS        *VariableArrayAddress;
  MEMORY_RESOURCE_DESCRIPTOR  *MemoryResource;

 #ifdef MDE_CPU_IA32
  UINT16                        CoalesceImageMachineType;
  EFI_PHYSICAL_ADDRESS          CoalesceImageEntryPoint;
  COALESCE_ENTRY                CoalesceEntry;
  EFI_CAPSULE_LONG_MODE_BUFFER  LongModeBuffer;
 #endif

  ListLength           = 0;
  VariableArrayAddress = NULL;

  //
  // Someone should have already ascertained the boot mode. If it's not
  // capsule update, then return normally.
  //
  Status = PeiServicesGetBootMode (&BootMode);
  if (EFI_ERROR (Status) || (BootMode != BOOT_ON_FLASH_UPDATE)) {
    DEBUG ((DEBUG_ERROR, "Boot mode is not correct for capsule update path.\n"));
    Status = EFI_NOT_FOUND;
    goto Done;
  }

  //
  // Get SG list entries
  //
  Status = GetScatterGatherHeadEntries (&ListLength, &VariableArrayAddress);
  if (EFI_ERROR (Status) || (VariableArrayAddress == NULL)) {
    DEBUG ((DEBUG_ERROR, "%a failed to get Scatter Gather List Head Entries.  Status = %r\n", __FUNCTION__, Status));
    goto Done;
  }

  MemoryResource = BuildMemoryResourceDescriptor ();

 #ifdef MDE_CPU_IA32
  if (FeaturePcdGet (PcdDxeIplSwitchToLongMode)) {
    //
    // Switch to 64-bit mode to process capsule data when:
    // 1. When DXE phase is 64-bit
    // 2. When the buffer for 64-bit transition exists
    // 3. When Capsule X64 image is built in BIOS image
    // In 64-bit mode, we can process capsule data above 4GB.
    //
    CoalesceImageEntryPoint = 0;
    Status                  = GetLongModeContext (&LongModeBuffer);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "Fail to find the variable for long mode context!\n"));
      Status = EFI_NOT_FOUND;
      goto Done;
    }

    Status = FindCapsuleCoalesceImage (&CoalesceImageEntryPoint, &CoalesceImageMachineType);
    if ((EFI_ERROR (Status)) || (CoalesceImageMachineType != EFI_IMAGE_MACHINE_X64)) {
      DEBUG ((DEBUG_ERROR, "Fail to find CapsuleX64 module in FV!\n"));
      Status = EFI_NOT_FOUND;
      goto Done;
    }

    ASSERT (CoalesceImageEntryPoint != 0);
    CoalesceEntry = (COALESCE_ENTRY)(UINTN)CoalesceImageEntryPoint;
    Status        = ModeSwitch (&LongModeBuffer, CoalesceEntry, (EFI_PHYSICAL_ADDRESS)(UINTN)VariableArrayAddress, MemoryResource, MemoryBase, MemorySize);
  } else {
    //
    // Capsule is processed in IA32 mode.
    //
    Status = CapsuleDataCoalesce (PeiServices, (EFI_PHYSICAL_ADDRESS *)(UINTN)VariableArrayAddress, MemoryResource, MemoryBase, MemorySize);
  }

 #else
  //
  // Process capsule directly.
  //
  Status = CapsuleDataCoalesce (PeiServices, (EFI_PHYSICAL_ADDRESS *)(UINTN)VariableArrayAddress, MemoryResource, MemoryBase, MemorySize);
 #endif

  DEBUG ((DEBUG_INFO, "Capsule Coalesce Status = %r!\n", Status));

  if (Status == EFI_BUFFER_TOO_SMALL) {
    DEBUG ((DEBUG_ERROR, "There is not enough memory to process capsule!\n"));
  }

  if (Status == EFI_NOT_FOUND) {
    DEBUG ((DEBUG_ERROR, "Fail to parse capsule descriptor in memory!\n"));
    REPORT_STATUS_CODE (
      EFI_ERROR_CODE | EFI_ERROR_MAJOR,
      (EFI_SOFTWARE_PEI_MODULE | EFI_SW_PEI_EC_INVALID_CAPSULE_DESCRIPTOR)
      );
  }

Done:
  return Status;
}

/**
  Determine if we're in capsule update boot mode.

  @param PeiServices  PEI services table

  @retval EFI_SUCCESS   if we have a capsule available
  @retval EFI_NOT_FOUND no capsule detected

**/
EFI_STATUS
EFIAPI
CheckCapsuleUpdate (
  IN EFI_PEI_SERVICES  **PeiServices
  )
{
  if (AreCapsulesStaged ()) {
    return EFI_SUCCESS;
  } else {
    return EFI_NOT_FOUND;
  }
}

/**
  This function will look at a capsule and determine if it's a test pattern.
  If it is, then it will verify it and emit an error message if corruption is detected.

  @param PeiServices   Standard pei services pointer
  @param CapsuleBase   Base address of coalesced capsule, which is preceeded
                       by private data. Very implementation specific.

  @retval TRUE    Capsule image is the test image
  @retval FALSE   Capsule image is not the test image.

**/
BOOLEAN
CapsuleTestPattern (
  IN EFI_PEI_SERVICES  **PeiServices,
  IN VOID              *CapsuleBase
  )
{
  UINT32   *TestPtr;
  UINT32   TestCounter;
  UINT32   TestSize;
  BOOLEAN  RetValue;

  RetValue = FALSE;

  //
  // Look at the capsule data and determine if it's a test pattern. If it
  // is, then test it now.
  //
  TestPtr = (UINT32 *)CapsuleBase;
  //
  // 0x54534554 "TEST"
  //
  if (*TestPtr == 0x54534554) {
    RetValue = TRUE;
    DEBUG ((DEBUG_INFO, "Capsule test pattern mode activated...\n"));
    TestSize = TestPtr[1] / sizeof (UINT32);
    //
    // Skip over the signature and the size fields in the pattern data header
    //
    TestPtr    += 2;
    TestCounter = 0;
    while (TestSize > 0) {
      if (*TestPtr != TestCounter) {
        DEBUG ((DEBUG_INFO, "Capsule test pattern mode FAILED: BaseAddr/FailAddr 0x%X 0x%X\n", (UINT32)(UINTN)(EFI_CAPSULE_PEIM_PRIVATE_DATA *)CapsuleBase, (UINT32)(UINTN)TestPtr));
        return TRUE;
      }

      TestPtr++;
      TestCounter++;
      TestSize--;
    }

    DEBUG ((DEBUG_INFO, "Capsule test pattern mode SUCCESS\n"));
  }

  return RetValue;
}

/**
  Capsule PPI service that gets called after memory is available. The
  capsule coalesce function, which must be called first, returns a base
  address and size, which can be anything actually. Once the memory init
  PEIM has discovered memory, then it should call this function and pass in
  the base address and size returned by the coalesce function. Then this
  function can create a capsule HOB and return.

  @param PeiServices   standard pei services pointer
  @param CapsuleBase   address returned by the capsule coalesce function. Most
                       likely this will actually be a pointer to private data.
  @param CapsuleSize   value returned by the capsule coalesce function.

  @retval EFI_VOLUME_CORRUPTED  CapsuleBase does not appear to point to a
                                coalesced capsule
  @retval EFI_SUCCESS           if all goes well.
**/
EFI_STATUS
EFIAPI
CreateState (
  IN EFI_PEI_SERVICES  **PeiServices,
  IN VOID              *CapsuleBase,
  IN UINTN             CapsuleSize
  )
{
  EFI_STATUS                     Status;
  EFI_CAPSULE_PEIM_PRIVATE_DATA  *PrivateData;
  UINTN                          Size;
  EFI_PHYSICAL_ADDRESS           NewBuffer;
  UINTN                          CapsuleNumber;
  UINT32                         Index;
  EFI_PHYSICAL_ADDRESS           BaseAddress;
  UINT64                         Length;

  PrivateData = (EFI_CAPSULE_PEIM_PRIVATE_DATA *)CapsuleBase;
  if (PrivateData->Signature != EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE) {
    return EFI_VOLUME_CORRUPTED;
  }

  if (PrivateData->CapsuleAllImageSize >= MAX_ADDRESS) {
    DEBUG ((DEBUG_ERROR, "CapsuleAllImageSize too big - 0x%lx\n", PrivateData->CapsuleAllImageSize));
    return EFI_OUT_OF_RESOURCES;
  }

  if (PrivateData->CapsuleNumber >= MAX_ADDRESS) {
    DEBUG ((DEBUG_ERROR, "CapsuleNumber too big - 0x%lx\n", PrivateData->CapsuleNumber));
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // Capsule Number and Capsule Offset is in the tail of Capsule data.
  //
  Size          = (UINTN)PrivateData->CapsuleAllImageSize;
  CapsuleNumber = (UINTN)PrivateData->CapsuleNumber;
  //
  // Allocate the memory so that it gets preserved into DXE
  //
  Status = PeiServicesAllocatePages (
             EfiRuntimeServicesData,
             EFI_SIZE_TO_PAGES (Size),
             &NewBuffer
             );

  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "AllocatePages Failed!\n"));
    return Status;
  }

  //
  // Copy to our new buffer for DXE
  //
  DEBUG ((DEBUG_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));
  CopyMem ((VOID *)(UINTN)NewBuffer, (VOID *)(UINTN)((UINT8 *)PrivateData + sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof (UINT64)), Size);
  //
  // Check for test data pattern. If it is the test pattern, then we'll
  // test it and still create the HOB so that it can be used to verify
  // that capsules don't get corrupted all the way into BDS. BDS will
  // still try to turn it into a firmware volume, but will think it's
  // corrupted so nothing will happen.
  //
  DEBUG_CODE (
    CapsuleTestPattern (PeiServices, (VOID *)(UINTN)NewBuffer);
    );

  //
  // Build the UEFI Capsule Hob for each capsule image.
  //
  for (Index = 0; Index < CapsuleNumber; Index++) {
    BaseAddress = NewBuffer + PrivateData->CapsuleOffset[Index];
    Length      = ((EFI_CAPSULE_HEADER *)((UINTN)BaseAddress))->CapsuleImageSize;

    BuildCvHob (BaseAddress, Length);
  }

  return EFI_SUCCESS;
}

CONST EFI_PEI_CAPSULE_PPI  mCapsulePpi = {
  CapsuleCoalesce,
  CheckCapsuleUpdate,
  CreateState
};

CONST EFI_PEI_PPI_DESCRIPTOR  mUefiPpiListCapsule = {
  (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiPeiCapsulePpiGuid,
  (EFI_PEI_CAPSULE_PPI *)&mCapsulePpi
};

/**
  Entry point function for the PEIM

  @param FileHandle      Handle of the file being invoked.
  @param PeiServices     Describes the list of possible PEI Services.

  @return EFI_SUCCESS    If we installed our PPI

**/
EFI_STATUS
EFIAPI
CapsuleMain (
  IN       EFI_PEI_FILE_HANDLE  FileHandle,
  IN CONST EFI_PEI_SERVICES     **PeiServices
  )
{
  //
  // Just produce our PPI
  //
  return PeiServicesInstallPpi (&mUefiPpiListCapsule);
}