Add capsule > 4GB support. When capsule data is put above 4GB, IA32 PEI transfers to long mode to get capsule data.

Signed-off-by: li-elvin
Reviewed-by: lgao4, mdkinney


git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@12264 6f19259b-4bc3-4df7-8a09-765794883524

MdeModulePkg/Include/Guid/CapsuleVendor.h

@@ -9,7 +9,7 @@
   @par Note: EDKII implementation of capsule updating has discarded this capsule GUID HOB data
              structure and used one UEFI Capsule HOB (defined in PI Specification 1.2) instead.
 
-Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials are licensed and made available under 
 the terms and conditions of the BSD License that accompanies this distribution.  
 The full text of the license may be found at
@@ -48,6 +48,18 @@ typedef struct {
   UINT32 Length;                     ///< Length of capsule data.
 } CAPSULE_HOB_INFO;
 
+//
+// The variable describes the long mode buffer used by IA32 Capsule PEIM
+// to call X64 CapsuleCoalesce code to handle >4GB capsule blocks.
+//
+#define EFI_CAPSULE_LONG_MODE_BUFFER_NAME L"CapsuleLongModeBuffer"
+
+typedef struct {
+  EFI_PHYSICAL_ADDRESS   PageTableAddress;
+  EFI_PHYSICAL_ADDRESS   StackBaseAddress;
+  UINT64                 StackSize;
+} EFI_CAPSULE_LONG_MODE_BUFFER;
+
 extern EFI_GUID gEfiCapsuleVendorGuid;
 
 #endif // #ifndef _EFI_CAPSULE_VENDOR_GUID_H_

MdeModulePkg/MdeModulePkg.dec

@@ -476,6 +476,9 @@
   ## FFS filename to find the ACPI tables
   gEfiMdeModulePkgTokenSpaceGuid.PcdAcpiTableStorageFile|{ 0x25, 0x4e, 0x37, 0x7e, 0x01, 0x8e, 0xee, 0x4f, 0x87, 0xf2, 0x39, 0xc, 0x23, 0xc6, 0x6, 0xcd }|VOID*|0x30000016
 
+  ## FFS filename to find the capsule coalesce image.
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCapsuleCoalesceFile|{ 0xA6, 0xE4, 0xFD, 0xF7, 0x4C, 0x29, 0x3c, 0x49, 0xB5, 0x0F, 0x97, 0x34, 0x55, 0x3B, 0xB7, 0x57 }|VOID*|0x30000017
+
   ## Single root I/O virtualization virtual function memory BAR alignment
   #  BITN set indicates 2 of n+12 power
   #  BIT0 set indicates 4KB alignment
@@ -729,6 +732,11 @@
   #  script node created in runtime phase.
   gEfiMdeModulePkgTokenSpaceGuid.PcdS3BootScriptRuntimeTableReservePageNumber|0x2|UINT16|0x0001005C
 
+  ## The PCD is used to specify the stack size when capsule IA32 PEI transfers to long mode in PEI phase.
+  #  The default size 32K. When changing the value of this PCD, the platform developer should
+  #  make sure the memory size is large enough to meet capsule PEI requiremnt in capsule update path.
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCapsulePeiLongModeStackSize|0x8000|UINT32|0x0001005C
+
 [PcdsPatchableInModule]
   ## Specify  memory size with page number for PEI code when 
   #  the feature of Loading Module at Fixed Address is enabled

MdeModulePkg/Universal/CapsulePei/Capsule.h

@@ -20,7 +20,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Uefi/UefiSpec.h>
 
 #include <Ppi/Capsule.h>
-
+#include <Ppi/LoadFile.h>
 #include <Ppi/ReadOnlyVariable2.h>
 #include <Guid/CapsuleVendor.h>
 
@@ -31,25 +31,71 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/HobLib.h>
 #include <Library/PeiServicesTablePointerLib.h>
 #include <Library/PrintLib.h>
+#include <Library/PeCoffLib.h>
+#include <Library/PeCoffGetEntryPointLib.h>
+#include <Library/PcdLib.h>
+#include <Library/ReportStatusCodeLib.h>
+#include <IndustryStandard/PeImage.h>
+#include "Common/CommonHeader.h"
+
+#pragma pack(1)
 
 //
-// We want to avoid using memory at 0 for coalescing, so set a
+// Page-Map Level-4 Offset (PML4) and
-// min address.
+// Page-Directory-Pointer Offset (PDPE) entries 4K & 2MB
 //
-#define MIN_COALESCE_ADDR 0x100000
+
-#define MAX_SUPPORT_CAPSULE_NUM 50
+typedef union {
+  struct {
+    UINT64  Present:1;                // 0 = Not present in memory, 1 = Present in memory
+    UINT64  ReadWrite:1;              // 0 = Read-Only, 1= Read/Write
+    UINT64  UserSupervisor:1;         // 0 = Supervisor, 1=User
+    UINT64  WriteThrough:1;           // 0 = Write-Back caching, 1=Write-Through caching
+    UINT64  CacheDisabled:1;          // 0 = Cached, 1=Non-Cached
+    UINT64  Accessed:1;               // 0 = Not accessed, 1 = Accessed (set by CPU)
+    UINT64  Reserved:1;               // Reserved
+    UINT64  MustBeZero:2;             // Must Be Zero
+    UINT64  Available:3;              // Available for use by system software
+    UINT64  PageTableBaseAddress:40;  // Page Table Base Address
+    UINT64  AvabilableHigh:11;        // Available for use by system software
+    UINT64  Nx:1;                     // No Execute bit
+  } Bits;
+  UINT64    Uint64;
+} PAGE_MAP_AND_DIRECTORY_POINTER;
 
 //
-// This capsule PEIM puts its private data at the start of the
+// Page Table Entry 2MB
-// coalesced capsule. Here's the structure definition.
 //
-#define EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE SIGNATURE_32 ('C', 'a', 'p', 'D')
+typedef union {
+  struct {
+    UINT64  Present:1;                // 0 = Not present in memory, 1 = Present in memory
+    UINT64  ReadWrite:1;              // 0 = Read-Only, 1= Read/Write
+    UINT64  UserSupervisor:1;         // 0 = Supervisor, 1=User
+    UINT64  WriteThrough:1;           // 0 = Write-Back caching, 1=Write-Through caching
+    UINT64  CacheDisabled:1;          // 0 = Cached, 1=Non-Cached
+    UINT64  Accessed:1;               // 0 = Not accessed, 1 = Accessed (set by CPU)
+    UINT64  Dirty:1;                  // 0 = Not Dirty, 1 = written by processor on access to page
+    UINT64  MustBe1:1;                // Must be 1 
+    UINT64  Global:1;                 // 0 = Not global page, 1 = global page TLB not cleared on CR3 write
+    UINT64  Available:3;              // Available for use by system software
+    UINT64  PAT:1;                    //
+    UINT64  MustBeZero:8;             // Must be zero;
+    UINT64  PageTableBaseAddress:31;  // Page Table Base Address
+    UINT64  AvabilableHigh:11;        // Available for use by system software
+    UINT64  Nx:1;                     // 0 = Execute Code, 1 = No Code Execution
+  } Bits;
+  UINT64    Uint64;
+} PAGE_TABLE_ENTRY;
 
-typedef struct {
+#pragma pack()
-  UINT32  Signature;
-  UINTN   CapsuleSize;
-} EFI_CAPSULE_PEIM_PRIVATE_DATA;
 
-#define CAPSULE_TEST_SIGNATURE SIGNATURE_32('T', 'E', 'S', 'T')
+typedef
+EFI_STATUS
+(*COALESCE_ENTRY) (
+  IN EFI_PEI_SERVICES                **PeiServices,
+  IN EFI_CAPSULE_BLOCK_DESCRIPTOR    *BlockList,
+  IN OUT VOID                        **MemoryBase,
+  IN OUT UINTN                       *MemorySize
+  );
 
 #endif

MdeModulePkg/Universal/CapsulePei/CapsulePei.inf

@@ -3,7 +3,7 @@
 #
 # Capsule update module supports EFI and UEFI.
 #
-# Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+# Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>
 #
 # This program and the accompanying materials
 # are licensed and made available under the terms and conditions
@@ -34,6 +34,7 @@
 [Sources]
   UefiCapsule.c
   Capsule.h
+  Common/CapsuleCoalesce.c  
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -48,6 +49,10 @@
   DebugLib
   PeiServicesTablePointerLib
   PrintLib
+  PeCoffLib
+  PeCoffGetEntryPointLib
+  PcdLib
+  ReportStatusCodeLib
 
 [Guids]
   gEfiCapsuleVendorGuid                         # ALWAYS_CONSUMED
@@ -56,7 +61,11 @@
 [Ppis]
   gEfiPeiReadOnlyVariable2PpiGuid               # PPI ALWAYS_CONSUMED
   gPeiCapsulePpiGuid                            # PPI ALWAYS_CONSUMED
+  gEfiPeiLoadFilePpiGuid                        # PPI ALWAYS_CONSUMED
 
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCapsuleCoalesceFile
+  gEfiMdeModulePkgTokenSpaceGuid.PcdDxeIplSwitchToLongMode
 
 [Depex]
   gEfiPeiReadOnlyVariable2PpiGuid

MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf

@@ -0,0 +1,48 @@
+## @file
+# Component description file for CapsuleX64 module.
+#
+# The X64 entrypoint to process capsule in long mode.
+# This module is built as X64.
+#
+# Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+#
+# This program and the accompanying materials
+# are licensed and made available under the terms and conditions
+# of the BSD License which accompanies this distribution.  The
+# full text of the license may be found at
+# http://opensource.org/licenses/bsd-license.php
+#
+# THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+# WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+#
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = CapsuleX64
+  FILE_GUID                      = F7FDE4A6-294C-493c-B50F-9734553BB757
+  MODULE_TYPE                    = PEIM
+  VERSION_STRING                 = 1.0
+
+#
+# The following information is for reference only and not required by the build tools.
+#
+#  VALID_ARCHITECTURES           = X64
+#
+
+[Sources]
+  X64/X64Entry.c
+  Common/CapsuleCoalesce.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  MdeModulePkg/MdeModulePkg.dec
+
+[LibraryClasses]
+  BaseLib
+  DebugLib
+
+[Depex]
+  FALSE
+
+

MdeModulePkg/Universal/CapsulePei/Common/CommonHeader.h

@@ -0,0 +1,84 @@
+/** @file
+  Common header file.
+
+Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+This program and the accompanying materials
+are licensed and made available under the terms and conditions of the BSD License
+which accompanies this distribution.  The full text of the license may be found at
+http://opensource.org/licenses/bsd-license.php
+
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+
+#ifndef _CAPSULE_THUNK_32_TO_64_
+#define _CAPSULE_THUNK_32_TO_64_
+
+#include <Uefi.h>
+#include "PiPei.h"
+
+//
+// This capsule PEIM puts its private data at the start of the
+// coalesced capsule. Here's the structure definition.
+//
+#define EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE SIGNATURE_32 ('C', 'a', 'p', 'D')
+
+typedef struct {
+  UINT32  Signature;
+  UINTN   CapsuleSize;
+} EFI_CAPSULE_PEIM_PRIVATE_DATA;
+
+#define CAPSULE_TEST_SIGNATURE SIGNATURE_32('T', 'E', 'S', 'T')
+
+typedef struct {
+  EFI_PHYSICAL_ADDRESS  EntryPoint;
+  EFI_PHYSICAL_ADDRESS  StackBufferBase;
+  UINT64                StackBufferLength;
+  EFI_PHYSICAL_ADDRESS  JumpBuffer;
+  EFI_PHYSICAL_ADDRESS  BlockListAddr;
+  EFI_PHYSICAL_ADDRESS  MemoryBase64Ptr;
+  EFI_PHYSICAL_ADDRESS  MemorySize64Ptr;
+} SWITCH_32_TO_64_CONTEXT;
+
+typedef struct {
+  UINT16                ReturnCs;
+  EFI_PHYSICAL_ADDRESS  ReturnEntryPoint;
+  UINT64                ReturnStatus;
+  IA32_DESCRIPTOR       Gdtr;
+} SWITCH_64_TO_32_CONTEXT;
+
+/**
+  The function to coalesce a fragmented capsule in memory.
+
+  @param PeiServices        General purpose services available to every PEIM.
+  @param BlockListBuffer    Point to the buffer of Capsule Descriptor Variables.
+  @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
+CapsuleDataCoalesce (
+  IN EFI_PEI_SERVICES                **PeiServices,
+  IN IN EFI_PHYSICAL_ADDRESS         *BlockListBuffer,
+  IN OUT VOID                        **MemoryBase,
+  IN OUT UINTN                       *MemorySize
+  );
+
+#endif

MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c

@@ -0,0 +1,66 @@
+/** @file
+  The X64 entrypoint is used to process capsule in long mode.
+
+Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+This program and the accompanying materials
+are licensed and made available under the terms and conditions of the BSD License
+which accompanies this distribution.  The full text of the license may be found at
+http://opensource.org/licenses/bsd-license.php
+
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include "CommonHeader.h"
+
+/**
+  The X64 entrypoint is used to process capsule in long mode then
+  return to 32-bit protected mode.
+
+  @param  EntrypointContext   Pointer to the context of long mode.
+  @param  ReturnContext       Pointer to the context of 32-bit protected mode.
+
+  @retval This function should never return actually.
+
+**/
+EFI_STATUS
+EFIAPI
+_ModuleEntryPoint (
+  SWITCH_32_TO_64_CONTEXT       *EntrypointContext,
+  SWITCH_64_TO_32_CONTEXT       *ReturnContext
+)
+{
+  EFI_STATUS    Status;
+
+  //
+  // Call CapsuleDataCoalesce to process capsule.
+  //
+  Status = CapsuleDataCoalesce (
+             NULL,
+             (EFI_PHYSICAL_ADDRESS *) (UINTN) EntrypointContext->BlockListAddr,
+             (VOID **) (UINTN) EntrypointContext->MemoryBase64Ptr,
+             (UINTN *) (UINTN) EntrypointContext->MemorySize64Ptr
+             );
+  
+  ReturnContext->ReturnStatus = Status;
+
+  //
+  // Finish to coalesce capsule, and return to 32-bit mode.
+  //
+  AsmDisablePaging64 (
+    ReturnContext->ReturnCs,
+    (UINT32) ReturnContext->ReturnEntryPoint,
+    (UINT32) (UINTN) EntrypointContext,
+    (UINT32) (UINTN) ReturnContext,
+    (UINT32) (EntrypointContext->StackBufferBase + EntrypointContext->StackBufferLength)
+    );  
+  
+  //
+  // Should never be here.
+  //
+  ASSERT (FALSE);
+  return EFI_SUCCESS;
+}

MdeModulePkg/Universal/CapsuleRuntimeDxe/CapsuleRuntimeDxe.inf

@@ -33,6 +33,12 @@
 [Sources]
   CapsuleService.c
 
+[Sources.Ia32, Sources.IPF, Sources.EBC, Sources.ARM]
+  SaveLongModeContext.c
+
+[Sources.X64]
+  X64/SaveLongModeContext.c
+
 [Packages]
   MdePkg/MdePkg.dec
   MdeModulePkg/MdeModulePkg.dec
@@ -47,19 +53,38 @@
   UefiRuntimeLib
   BaseLib
   PrintLib
+  
+[LibraryClasses.X64]
+  LockBoxLib
+  UefiLib
+  BaseMemoryLib
+  HobLib  
 
 [Guids]
   gEfiCapsuleVendorGuid                         ## SOMETIMES_PRODUCED (Process across reset capsule image) ## Variable:L"CapsuleUpdateData" for capsule updated data
 
+[Guids.X64]
+  gEfiAcpiVariableGuid                          # ALWAYS_CONSUMED
+  gEfiAcpiS3ContextGuid                         # ALWAYS_CONSUMED
+
 [Protocols]
   gEfiCapsuleArchProtocolGuid                   ## PRODUCED
 
+[Protocols.X64]
+  gEfiDxeSmmReadyToLockProtocolGuid             # ALWAYS_CONSUMED
+
 [FeaturePcd]
   gEfiMdeModulePkgTokenSpaceGuid.PcdSupportUpdateCapsuleReset
 
+[FeaturePcd.X64]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdDxeIplSwitchToLongMode
+
 [Pcd]
   gEfiMdeModulePkgTokenSpaceGuid.PcdMaxSizeNonPopulateCapsule
   gEfiMdeModulePkgTokenSpaceGuid.PcdMaxSizePopulateCapsule || gEfiMdeModulePkgTokenSpaceGuid.PcdSupportUpdateCapsuleReset  ## Populate Image requires reset support.
 
+[Pcd.X64]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCapsulePeiLongModeStackSize
+
 [Depex]
   gEfiVariableWriteArchProtocolGuid             ## Depends on variable write functionality to produce capsule data variable

MdeModulePkg/Universal/CapsuleRuntimeDxe/CapsuleService.c

@@ -4,7 +4,7 @@
   It installs the Capsule Architectural Protocol defined in PI1.0a to signify 
   the capsule runtime services are ready.
 
-Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -40,6 +40,15 @@ EFI_HANDLE  mNewHandle = NULL;
 //
 UINTN       mTimes      = 0;
 
+/**
+  Create the variable to save the base address of page table and stack
+  for transferring into long mode in IA32 PEI.
+**/
+VOID
+SaveLongModeContext (
+  VOID
+  );
+
 /**
   Passes capsules to the firmware with both virtual and physical mapping. Depending on the intended
   consumption, the firmware may process the capsule immediately. If the payload should persist
@@ -333,6 +342,15 @@ CapsuleServiceInitialize (
 {
   EFI_STATUS  Status;
   
+  //
+  // When PEI phase is IA32, DXE phase is X64, it is possible that capsule data are 
+  // put above 4GB, so capsule PEI will transfer to long mode to get capsule data.
+  // The page table and stack is used to transfer processor mode from IA32 to long mode.
+  // Create the base address of page table and stack, and save them into variable.
+  // This is not needed when capsule with reset type is not supported.
+  //
+  SaveLongModeContext ();
+    
   //
   // Install capsule runtime services into UEFI runtime service tables.
   //

MdeModulePkg/Universal/CapsuleRuntimeDxe/SaveLongModeContext.c

@@ -0,0 +1,27 @@
+/** @file
+  Create the NULL function to pass build in IA32/IPF/ARM/EBC.
+
+Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+This program and the accompanying materials
+are licensed and made available under the terms and conditions of the BSD License
+which accompanies this distribution.  The full text of the license may be found at
+http://opensource.org/licenses/bsd-license.php
+
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+
+#include <Uefi.h>
+
+/**
+  Only when PEI is IA32 and DXE is X64, we need transfer to long mode in PEI
+  in order to process capsule data above 4GB. So create a NULL function here for
+  other cases.
+**/
+VOID
+SaveLongModeContext (
+  VOID
+  )
+{
+}

MdeModulePkg/Universal/CapsuleRuntimeDxe/X64/SaveLongModeContext.c

@@ -0,0 +1,224 @@
+/** @file
+  Create the variable to save the base address of page table and stack
+  for transferring into long mode in IA32 capsule PEI.
+
+Copyright (c) 2011, Intel Corporation. All rights reserved.<BR>
+This program and the accompanying materials
+are licensed and made available under the terms and conditions of the BSD License
+which accompanies this distribution.  The full text of the license may be found at
+http://opensource.org/licenses/bsd-license.php
+
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+
+#include <Uefi.h>
+
+#include <Protocol/Capsule.h>
+#include <Protocol/DxeSmmReadyToLock.h>
+
+#include <Guid/CapsuleVendor.h>
+#include <Guid/AcpiS3Context.h>
+
+#include <Library/DebugLib.h>
+#include <Library/PcdLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiRuntimeServicesTableLib.h>
+#include <Library/UefiRuntimeLib.h>
+#include <Library/BaseLib.h>
+#include <Library/LockBoxLib.h>
+#include <Library/UefiLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/HobLib.h>
+
+/**
+  Allocate EfiACPIMemoryNVS below 4G memory address.
+
+  This function allocates EfiACPIMemoryNVS below 4G memory address.
+
+  @param  Size         Size of memory to allocate.
+  
+  @return Allocated address for output.
+
+**/
+VOID*
+AllocateAcpiNvsMemoryBelow4G (
+  IN   UINTN   Size
+  )
+{
+  UINTN                 Pages;
+  EFI_PHYSICAL_ADDRESS  Address;
+  EFI_STATUS            Status;
+  VOID*                 Buffer;
+
+  Pages = EFI_SIZE_TO_PAGES (Size);
+  Address = 0xffffffff;
+
+  Status  = gBS->AllocatePages (
+                   AllocateMaxAddress,
+                   EfiACPIMemoryNVS,
+                   Pages,
+                   &Address
+                   );
+  ASSERT_EFI_ERROR (Status);
+
+  Buffer = (VOID *) (UINTN) Address;
+  ZeroMem (Buffer, Size);
+
+  return Buffer;
+}
+
+/**
+  DxeSmmReadyToLock Protocol notification event handler.
+  We reuse S3 ACPI NVS reserved memory to do capsule process
+  after reset.
+
+  @param[in] Event    Event whose notification function is being invoked.
+  @param[in] Context  Pointer to the notification function's context.
+**/
+VOID
+EFIAPI
+DxeSmmReadyToLockNotification (
+  IN EFI_EVENT  Event,
+  IN VOID       *Context
+  )
+{
+  EFI_STATUS                    Status;
+  VOID                          *DxeSmmReadyToLock;
+  UINTN                         VarSize;
+  EFI_PHYSICAL_ADDRESS          TempAcpiS3Context;
+  ACPI_S3_CONTEXT               *AcpiS3Context;
+  EFI_CAPSULE_LONG_MODE_BUFFER  LongModeBuffer;
+  UINTN                         TotalPagesNum;
+  UINT8                         PhysicalAddressBits;
+  VOID                          *Hob;
+  UINT32                        NumberOfPml4EntriesNeeded;
+  UINT32                        NumberOfPdpEntriesNeeded;
+  BOOLEAN                       LockBoxFound;
+
+  Status = gBS->LocateProtocol (
+                  &gEfiDxeSmmReadyToLockProtocolGuid,
+                  NULL,
+                  &DxeSmmReadyToLock
+                  );
+  if (EFI_ERROR (Status)) {
+    return ;
+  }
+
+  //
+  // Get the ACPI NVS pages reserved by AcpiS3Save
+  //
+  LockBoxFound = FALSE;
+  VarSize = sizeof (EFI_PHYSICAL_ADDRESS);
+  Status = RestoreLockBox (
+             &gEfiAcpiVariableGuid,
+             &TempAcpiS3Context,
+             &VarSize
+             );
+  if (!EFI_ERROR (Status)) {
+    AcpiS3Context = (ACPI_S3_CONTEXT *)(UINTN)TempAcpiS3Context;
+    ASSERT (AcpiS3Context != NULL);
+    
+    Status = RestoreLockBox (
+               &gEfiAcpiS3ContextGuid,
+               NULL,
+               NULL
+               );
+    if (!EFI_ERROR (Status)) {
+      LongModeBuffer.PageTableAddress = AcpiS3Context->S3NvsPageTableAddress;
+      LongModeBuffer.StackBaseAddress = AcpiS3Context->BootScriptStackBase;
+      LongModeBuffer.StackSize        = AcpiS3Context->BootScriptStackSize;
+      LockBoxFound                    = TRUE;
+    }
+  }
+  
+  if (!LockBoxFound) {
+    //
+    // Page table base address and stack base address can not be found in lock box,
+    // allocate both here. 
+    //
+
+    //
+    // Get physical address bits supported from CPU HOB.
+    //
+    PhysicalAddressBits = 36;
+    
+    Hob = GetFirstHob (EFI_HOB_TYPE_CPU);
+    if (Hob != NULL) {
+      PhysicalAddressBits = ((EFI_HOB_CPU *) Hob)->SizeOfMemorySpace;
+    }
+    
+    //
+    // IA-32e paging translates 48-bit linear addresses to 52-bit physical addresses.
+    //
+    ASSERT (PhysicalAddressBits <= 52);
+    if (PhysicalAddressBits > 48) {
+      PhysicalAddressBits = 48;
+    }
+    
+    //
+    // Calculate page table size and allocate memory for it.
+    //
+    if (PhysicalAddressBits <= 39 ) {
+      NumberOfPml4EntriesNeeded = 1;
+      NumberOfPdpEntriesNeeded =  1 << (PhysicalAddressBits - 30);
+    } else {
+      NumberOfPml4EntriesNeeded = 1 << (PhysicalAddressBits - 39);
+      NumberOfPdpEntriesNeeded = 512;
+    }
+    
+    TotalPagesNum = (NumberOfPdpEntriesNeeded + 1) * NumberOfPml4EntriesNeeded + 1;
+    LongModeBuffer.PageTableAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)AllocateAcpiNvsMemoryBelow4G (EFI_PAGES_TO_SIZE (TotalPagesNum));
+    ASSERT (LongModeBuffer.PageTableAddress != 0);
+    
+    //
+    // Allocate stack
+    //
+    LongModeBuffer.StackSize        = PcdGet32 (PcdCapsulePeiLongModeStackSize);
+    LongModeBuffer.StackBaseAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)AllocateAcpiNvsMemoryBelow4G (PcdGet32 (PcdCapsulePeiLongModeStackSize));
+    ASSERT (LongModeBuffer.StackBaseAddress != 0);
+  }
+
+  Status = gRT->SetVariable (
+                  EFI_CAPSULE_LONG_MODE_BUFFER_NAME,
+                  &gEfiCapsuleVendorGuid,
+                  EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
+                  sizeof (EFI_CAPSULE_LONG_MODE_BUFFER),
+                  &LongModeBuffer
+                  );
+  ASSERT_EFI_ERROR (Status);
+
+  //
+  // Close event, so it will not be invoked again.
+  //
+  gBS->CloseEvent (Event);
+
+  return ;
+}
+
+/**
+  Create the variable to save the base address of page table and stack
+  for transferring into long mode in IA32 capsule PEI.
+**/
+VOID
+SaveLongModeContext (
+  VOID
+  )
+{
+  VOID        *Registration;
+  
+  if ((FeaturePcdGet(PcdSupportUpdateCapsuleReset)) && (FeaturePcdGet (PcdDxeIplSwitchToLongMode))) {
+    //
+    // Register event to get ACPI NVS pages reserved from lock box, these pages will be used by
+    // Capsule IA32 PEI to transfer to long mode to access capsule above 4GB.
+    //
+    EfiCreateProtocolNotifyEvent (
+      &gEfiDxeSmmReadyToLockProtocolGuid,
+      TPL_CALLBACK,
+      DxeSmmReadyToLockNotification,
+      NULL,
+      &Registration
+      );
+  }
+}