]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/CpuDxe/CpuPageTable.c
UefiCpuPkg: Move AsmRelocateApLoopStart from Mpfuncs.nasm to AmdSev.nasm
[mirror_edk2.git] / UefiCpuPkg / CpuDxe / CpuPageTable.c
index 90d9823838138955f32613464d06de9395e507a5..5b1dec2d67603480d9019b26b60ffb0cf57eace9 100644 (file)
@@ -1,28 +1,22 @@
 /** @file\r
   Page table management support.\r
 \r
-  Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2017 - 2019, 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 of the BSD License\r
-  which accompanies this distribution.  The 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 <Base.h>\r
 #include <Uefi.h>\r
-#include <Library/BaseLib.h>\r
-#include <Library/CpuLib.h>\r
-#include <Library/BaseMemoryLib.h>\r
-#include <Library/MemoryAllocationLib.h>\r
-#include <Library/DebugLib.h>\r
-#include <Library/UefiBootServicesTableLib.h>\r
-#include <Protocol/MpService.h>\r
+#include <Library/PeCoffGetEntryPointLib.h>\r
+#include <Library/SerialPortLib.h>\r
+#include <Library/SynchronizationLib.h>\r
+#include <Library/PrintLib.h>\r
+#include <Protocol/SmmBase2.h>\r
+#include <Register/Intel/Cpuid.h>\r
+#include <Register/Intel/Msr.h>\r
 \r
 #include "CpuDxe.h"\r
 #include "CpuPageTable.h"\r
 ///\r
 /// Page Table Entry\r
 ///\r
-#define IA32_PG_P                   BIT0\r
-#define IA32_PG_RW                  BIT1\r
-#define IA32_PG_U                   BIT2\r
-#define IA32_PG_WT                  BIT3\r
-#define IA32_PG_CD                  BIT4\r
-#define IA32_PG_A                   BIT5\r
-#define IA32_PG_D                   BIT6\r
-#define IA32_PG_PS                  BIT7\r
-#define IA32_PG_PAT_2M              BIT12\r
-#define IA32_PG_PAT_4K              IA32_PG_PS\r
-#define IA32_PG_PMNT                BIT62\r
-#define IA32_PG_NX                  BIT63\r
-\r
-#define PAGE_ATTRIBUTE_BITS         (IA32_PG_D | IA32_PG_A | IA32_PG_U | IA32_PG_RW | IA32_PG_P)\r
+#define IA32_PG_P       BIT0\r
+#define IA32_PG_RW      BIT1\r
+#define IA32_PG_U       BIT2\r
+#define IA32_PG_WT      BIT3\r
+#define IA32_PG_CD      BIT4\r
+#define IA32_PG_A       BIT5\r
+#define IA32_PG_D       BIT6\r
+#define IA32_PG_PS      BIT7\r
+#define IA32_PG_PAT_2M  BIT12\r
+#define IA32_PG_PAT_4K  IA32_PG_PS\r
+#define IA32_PG_PMNT    BIT62\r
+#define IA32_PG_NX      BIT63\r
+\r
+#define PAGE_ATTRIBUTE_BITS             (IA32_PG_D | IA32_PG_A | IA32_PG_U | IA32_PG_RW | IA32_PG_P)\r
+#define PAGE_ATTRIBUTE_BITS_POST_SPLIT  (IA32_PG_RW | IA32_PG_P)\r
+\r
 //\r
 // Bits 1, 2, 5, 6 are reserved in the IA32 PAE PDPTE\r
 // X64 PAE PDPTE does not have such restriction\r
 //\r
-#define IA32_PAE_PDPTE_ATTRIBUTE_BITS    (IA32_PG_P)\r
+#define IA32_PAE_PDPTE_ATTRIBUTE_BITS  (IA32_PG_P)\r
 \r
-#define PAGE_PROGATE_BITS           (IA32_PG_NX | PAGE_ATTRIBUTE_BITS)\r
+#define PAGE_PROGATE_BITS  (IA32_PG_NX | PAGE_ATTRIBUTE_BITS)\r
 \r
 #define PAGING_4K_MASK  0xFFF\r
 #define PAGING_2M_MASK  0x1FFFFF\r
 \r
 #define PAGING_PAE_INDEX_MASK  0x1FF\r
 \r
-#define PAGING_4K_ADDRESS_MASK_64 0x000FFFFFFFFFF000ull\r
-#define PAGING_2M_ADDRESS_MASK_64 0x000FFFFFFFE00000ull\r
-#define PAGING_1G_ADDRESS_MASK_64 0x000FFFFFC0000000ull\r
+#define PAGING_4K_ADDRESS_MASK_64  0x000FFFFFFFFFF000ull\r
+#define PAGING_2M_ADDRESS_MASK_64  0x000FFFFFFFE00000ull\r
+#define PAGING_1G_ADDRESS_MASK_64  0x000FFFFFC0000000ull\r
+\r
+#define MAX_PF_ENTRY_COUNT        10\r
+#define MAX_DEBUG_MESSAGE_LENGTH  0x100\r
+#define IA32_PF_EC_ID             BIT4\r
 \r
 typedef enum {\r
   PageNone,\r
@@ -70,9 +70,9 @@ typedef enum {
 } PAGE_ATTRIBUTE;\r
 \r
 typedef struct {\r
-  PAGE_ATTRIBUTE   Attribute;\r
-  UINT64           Length;\r
-  UINT64           AddressMask;\r
+  PAGE_ATTRIBUTE    Attribute;\r
+  UINT64            Length;\r
+  UINT64            AddressMask;\r
 } PAGE_ATTRIBUTE_TABLE;\r
 \r
 typedef enum {\r
@@ -81,76 +81,65 @@ typedef enum {
   PageActionClear,\r
 } PAGE_ACTION;\r
 \r
-PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {\r
-  {Page4K,  SIZE_4KB, PAGING_4K_ADDRESS_MASK_64},\r
-  {Page2M,  SIZE_2MB, PAGING_2M_ADDRESS_MASK_64},\r
-  {Page1G,  SIZE_1GB, PAGING_1G_ADDRESS_MASK_64},\r
+PAGE_ATTRIBUTE_TABLE  mPageAttributeTable[] = {\r
+  { Page4K, SIZE_4KB, PAGING_4K_ADDRESS_MASK_64 },\r
+  { Page2M, SIZE_2MB, PAGING_2M_ADDRESS_MASK_64 },\r
+  { Page1G, SIZE_1GB, PAGING_1G_ADDRESS_MASK_64 },\r
 };\r
 \r
-PAGE_TABLE_POOL   *mPageTablePool = NULL;\r
-\r
-/**\r
-  Enable write protection function for AP.\r
+PAGE_TABLE_POOL                *mPageTablePool    = NULL;\r
+BOOLEAN                        mPageTablePoolLock = FALSE;\r
+PAGE_TABLE_LIB_PAGING_CONTEXT  mPagingContext;\r
+EFI_SMM_BASE2_PROTOCOL         *mSmmBase2 = NULL;\r
 \r
-  @param[in,out] Buffer  The pointer to private data buffer.\r
-**/\r
-VOID\r
-EFIAPI\r
-SyncCpuEnableWriteProtection (\r
-  IN OUT VOID *Buffer\r
-  )\r
-{\r
-  AsmWriteCr0 (AsmReadCr0 () | BIT16);\r
-}\r
+//\r
+// Record the page fault exception count for one instruction execution.\r
+//\r
+UINTN  *mPFEntryCount;\r
+UINT64                    *(*mLastPFEntryPointer)[MAX_PF_ENTRY_COUNT];\r
 \r
 /**\r
-  CpuFlushTlb function for AP.\r
-\r
-  @param[in,out] Buffer  The pointer to private data buffer.\r
+ Check if current execution environment is in SMM mode or not, via\r
+ EFI_SMM_BASE2_PROTOCOL.\r
+\r
+ This is necessary because of the fact that MdePkg\Library\SmmMemoryAllocationLib\r
+ supports to free memory outside SMRAM. The library will call gBS->FreePool() or\r
+ gBS->FreePages() and then SetMemorySpaceAttributes interface in turn to change\r
+ memory paging attributes during free operation, if some memory related features\r
+ are enabled (like Heap Guard).\r
+\r
+ This means that SetMemorySpaceAttributes() has chance to run in SMM mode. This\r
+ will cause incorrect result because SMM mode always loads its own page tables,\r
+ which are usually different from DXE. This function can be used to detect such\r
+ situation and help to avoid further misoperations.\r
+\r
+  @retval TRUE    In SMM mode.\r
+  @retval FALSE   Not in SMM mode.\r
 **/\r
-VOID\r
-EFIAPI\r
-SyncCpuFlushTlb (\r
-  IN OUT VOID *Buffer\r
+BOOLEAN\r
+IsInSmm (\r
+  VOID\r
   )\r
 {\r
-  CpuFlushTlb();\r
-}\r
+  BOOLEAN  InSmm;\r
 \r
-/**\r
-  Sync memory page attributes for AP.\r
+  InSmm = FALSE;\r
+  if (mSmmBase2 == NULL) {\r
+    gBS->LocateProtocol (&gEfiSmmBase2ProtocolGuid, NULL, (VOID **)&mSmmBase2);\r
+  }\r
+\r
+  if (mSmmBase2 != NULL) {\r
+    mSmmBase2->InSmm (mSmmBase2, &InSmm);\r
+  }\r
 \r
-  @param[in] Procedure            A pointer to the function to be run on enabled APs of\r
-                                  the system.\r
-**/\r
-VOID\r
-SyncMemoryPageAttributesAp (\r
-  IN EFI_AP_PROCEDURE            Procedure\r
-  )\r
-{\r
-  EFI_STATUS                Status;\r
-  EFI_MP_SERVICES_PROTOCOL  *MpService;\r
-\r
-  Status = gBS->LocateProtocol (\r
-                  &gEfiMpServiceProtocolGuid,\r
-                  NULL,\r
-                  (VOID **)&MpService\r
-                  );\r
   //\r
-  // Synchronize the update with all APs\r
+  // mSmmBase2->InSmm() can only detect if the caller is running in SMRAM\r
+  // or from SMM driver. It cannot tell if the caller is running in SMM mode.\r
+  // Check page table base address to guarantee that because SMM mode willl\r
+  // load its own page table.\r
   //\r
-  if (!EFI_ERROR (Status)) {\r
-    Status = MpService->StartupAllAPs (\r
-                          MpService,          // This\r
-                          Procedure,          // Procedure\r
-                          FALSE,              // SingleThread\r
-                          NULL,               // WaitEvent\r
-                          0,                  // TimeoutInMicrosecsond\r
-                          NULL,               // ProcedureArgument\r
-                          NULL                // FailedCpuList\r
-                          );\r
-    ASSERT (Status == EFI_SUCCESS || Status == EFI_NOT_STARTED || Status == EFI_NOT_READY);\r
-  }\r
+  return (InSmm &&\r
+          mPagingContext.ContextData.X64.PageTableBase != (UINT64)AsmReadCr3 ());\r
 }\r
 \r
 /**\r
@@ -160,48 +149,79 @@ SyncMemoryPageAttributesAp (
 **/\r
 VOID\r
 GetCurrentPagingContext (\r
-  IN OUT PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext\r
+  IN OUT PAGE_TABLE_LIB_PAGING_CONTEXT  *PagingContext\r
   )\r
 {\r
-  UINT32                         RegEax;\r
-  UINT32                         RegEdx;\r
+  UINT32                      RegEax;\r
+  CPUID_EXTENDED_CPU_SIG_EDX  RegEdx;\r
+  MSR_IA32_EFER_REGISTER      MsrEfer;\r
+  IA32_CR4                    Cr4;\r
+  IA32_CR0                    Cr0;\r
+  UINT32                      *Attributes;\r
+  UINTN                       *PageTableBase;\r
 \r
-  ZeroMem(PagingContext, sizeof(*PagingContext));\r
-  if (sizeof(UINTN) == sizeof(UINT64)) {\r
-    PagingContext->MachineType = IMAGE_FILE_MACHINE_X64;\r
-  } else {\r
-    PagingContext->MachineType = IMAGE_FILE_MACHINE_I386;\r
-  }\r
-  if ((AsmReadCr0 () & BIT31) != 0) {\r
-    PagingContext->ContextData.X64.PageTableBase = (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);\r
-  } else {\r
-    PagingContext->ContextData.X64.PageTableBase = 0;\r
-  }\r
+  //\r
+  // Don't retrieve current paging context from processor if in SMM mode.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    ZeroMem (&mPagingContext, sizeof (mPagingContext));\r
+    if (sizeof (UINTN) == sizeof (UINT64)) {\r
+      mPagingContext.MachineType = IMAGE_FILE_MACHINE_X64;\r
+    } else {\r
+      mPagingContext.MachineType = IMAGE_FILE_MACHINE_I386;\r
+    }\r
 \r
-  if ((AsmReadCr4 () & BIT4) != 0) {\r
-    PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE;\r
-  }\r
-  if ((AsmReadCr4 () & BIT5) != 0) {\r
-    PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;\r
-  }\r
-  if ((AsmReadCr0 () & BIT16) != 0) {\r
-    PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE;\r
-  }\r
+    GetPagingDetails (&mPagingContext.ContextData, &PageTableBase, &Attributes);\r
 \r
-  AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);\r
-  if (RegEax > 0x80000000) {\r
-    AsmCpuid (0x80000001, NULL, NULL, NULL, &RegEdx);\r
-    if ((RegEdx & BIT20) != 0) {\r
-      // XD supported\r
-      if ((AsmReadMsr64 (0xC0000080) & BIT11) != 0) {\r
-        // XD activated\r
-        PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED;\r
-      }\r
+    Cr0.UintN = AsmReadCr0 ();\r
+    Cr4.UintN = AsmReadCr4 ();\r
+\r
+    if (Cr0.Bits.PG != 0) {\r
+      *PageTableBase = (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);\r
+    } else {\r
+      *PageTableBase = 0;\r
+    }\r
+\r
+    if (Cr0.Bits.WP  != 0) {\r
+      *Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE;\r
+    }\r
+\r
+    if (Cr4.Bits.PSE != 0) {\r
+      *Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE;\r
+    }\r
+\r
+    if (Cr4.Bits.PAE != 0) {\r
+      *Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;\r
+    }\r
+\r
+    if (Cr4.Bits.LA57 != 0) {\r
+      *Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL;\r
     }\r
-    if ((RegEdx & BIT26) != 0) {\r
-      PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;\r
+\r
+    AsmCpuid (CPUID_EXTENDED_FUNCTION, &RegEax, NULL, NULL, NULL);\r
+    if (RegEax >= CPUID_EXTENDED_CPU_SIG) {\r
+      AsmCpuid (CPUID_EXTENDED_CPU_SIG, NULL, NULL, NULL, &RegEdx.Uint32);\r
+\r
+      if (RegEdx.Bits.NX != 0) {\r
+        // XD supported\r
+        MsrEfer.Uint64 = AsmReadMsr64 (MSR_CORE_IA32_EFER);\r
+        if (MsrEfer.Bits.NXE != 0) {\r
+          // XD activated\r
+          *Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED;\r
+        }\r
+      }\r
+\r
+      if (RegEdx.Bits.Page1GB != 0) {\r
+        *Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;\r
+      }\r
     }\r
   }\r
+\r
+  //\r
+  // This can avoid getting SMM paging context if in SMM mode. We cannot assume\r
+  // SMM mode shares the same paging context as DXE.\r
+  //\r
+  CopyMem (PagingContext, &mPagingContext, sizeof (mPagingContext));\r
 }\r
 \r
 /**\r
@@ -217,11 +237,13 @@ PageAttributeToLength (
   )\r
 {\r
   UINTN  Index;\r
-  for (Index = 0; Index < sizeof(mPageAttributeTable)/sizeof(mPageAttributeTable[0]); Index++) {\r
+\r
+  for (Index = 0; Index < sizeof (mPageAttributeTable)/sizeof (mPageAttributeTable[0]); Index++) {\r
     if (PageAttribute == mPageAttributeTable[Index].Attribute) {\r
       return (UINTN)mPageAttributeTable[Index].Length;\r
     }\r
   }\r
+\r
   return 0;\r
 }\r
 \r
@@ -238,11 +260,13 @@ PageAttributeToMask (
   )\r
 {\r
   UINTN  Index;\r
-  for (Index = 0; Index < sizeof(mPageAttributeTable)/sizeof(mPageAttributeTable[0]); Index++) {\r
+\r
+  for (Index = 0; Index < sizeof (mPageAttributeTable)/sizeof (mPageAttributeTable[0]); Index++) {\r
     if (PageAttribute == mPageAttributeTable[Index].Attribute) {\r
       return (UINTN)mPageAttributeTable[Index].AddressMask;\r
     }\r
   }\r
+\r
   return 0;\r
 }\r
 \r
@@ -257,23 +281,26 @@ PageAttributeToMask (
 **/\r
 VOID *\r
 GetPageTableEntry (\r
-  IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext,\r
-  IN  PHYSICAL_ADDRESS                  Address,\r
-  OUT PAGE_ATTRIBUTE                    *PageAttribute\r
+  IN  PAGE_TABLE_LIB_PAGING_CONTEXT  *PagingContext,\r
+  IN  PHYSICAL_ADDRESS               Address,\r
+  OUT PAGE_ATTRIBUTE                 *PageAttribute\r
   )\r
 {\r
-  UINTN                 Index1;\r
-  UINTN                 Index2;\r
-  UINTN                 Index3;\r
-  UINTN                 Index4;\r
-  UINT64                *L1PageTable;\r
-  UINT64                *L2PageTable;\r
-  UINT64                *L3PageTable;\r
-  UINT64                *L4PageTable;\r
-  UINT64                AddressEncMask;\r
+  UINTN   Index1;\r
+  UINTN   Index2;\r
+  UINTN   Index3;\r
+  UINTN   Index4;\r
+  UINTN   Index5;\r
+  UINT64  *L1PageTable;\r
+  UINT64  *L2PageTable;\r
+  UINT64  *L3PageTable;\r
+  UINT64  *L4PageTable;\r
+  UINT64  *L5PageTable;\r
+  UINT64  AddressEncMask;\r
 \r
   ASSERT (PagingContext != NULL);\r
 \r
+  Index5 = ((UINTN)RShiftU64 (Address, 48)) & PAGING_PAE_INDEX_MASK;\r
   Index4 = ((UINTN)RShiftU64 (Address, 39)) & PAGING_PAE_INDEX_MASK;\r
   Index3 = ((UINTN)Address >> 30) & PAGING_PAE_INDEX_MASK;\r
   Index2 = ((UINTN)Address >> 21) & PAGING_PAE_INDEX_MASK;\r
@@ -282,9 +309,23 @@ GetPageTableEntry (
   // Make sure AddressEncMask is contained to smallest supported address field.\r
   //\r
   AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;\r
+  if (AddressEncMask == 0) {\r
+    AddressEncMask = PcdGet64 (PcdTdxSharedBitMask) & PAGING_1G_ADDRESS_MASK_64;\r
+  }\r
 \r
   if (PagingContext->MachineType == IMAGE_FILE_MACHINE_X64) {\r
-    L4PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;\r
+    if ((PagingContext->ContextData.X64.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL) != 0) {\r
+      L5PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;\r
+      if (L5PageTable[Index5] == 0) {\r
+        *PageAttribute = PageNone;\r
+        return NULL;\r
+      }\r
+\r
+      L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
+    } else {\r
+      L4PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;\r
+    }\r
+\r
     if (L4PageTable[Index4] == 0) {\r
       *PageAttribute = PageNone;\r
       return NULL;\r
@@ -292,13 +333,15 @@ GetPageTableEntry (
 \r
     L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
   } else {\r
-    ASSERT((PagingContext->ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) != 0);\r
+    ASSERT ((PagingContext->ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) != 0);\r
     L3PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.Ia32.PageTableBase;\r
   }\r
+\r
   if (L3PageTable[Index3] == 0) {\r
     *PageAttribute = PageNone;\r
     return NULL;\r
   }\r
+\r
   if ((L3PageTable[Index3] & IA32_PG_PS) != 0) {\r
     // 1G\r
     *PageAttribute = Page1G;\r
@@ -310,6 +353,7 @@ GetPageTableEntry (
     *PageAttribute = PageNone;\r
     return NULL;\r
   }\r
+\r
   if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {\r
     // 2M\r
     *PageAttribute = Page2M;\r
@@ -322,6 +366,7 @@ GetPageTableEntry (
     *PageAttribute = PageNone;\r
     return NULL;\r
   }\r
+\r
   *PageAttribute = Page4K;\r
   return &L1PageTable[Index1];\r
 }\r
@@ -335,20 +380,24 @@ GetPageTableEntry (
 **/\r
 UINT64\r
 GetAttributesFromPageEntry (\r
-  IN  UINT64                            *PageEntry\r
+  IN  UINT64  *PageEntry\r
   )\r
 {\r
   UINT64  Attributes;\r
+\r
   Attributes = 0;\r
   if ((*PageEntry & IA32_PG_P) == 0) {\r
     Attributes |= EFI_MEMORY_RP;\r
   }\r
+\r
   if ((*PageEntry & IA32_PG_RW) == 0) {\r
     Attributes |= EFI_MEMORY_RO;\r
   }\r
+\r
   if ((*PageEntry & IA32_PG_NX) != 0) {\r
     Attributes |= EFI_MEMORY_XP;\r
   }\r
+\r
   return Attributes;\r
 }\r
 \r
@@ -363,80 +412,86 @@ GetAttributesFromPageEntry (
 **/\r
 VOID\r
 ConvertPageEntryAttribute (\r
-  IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext,\r
-  IN  UINT64                            *PageEntry,\r
-  IN  UINT64                            Attributes,\r
-  IN  PAGE_ACTION                       PageAction,\r
-  OUT BOOLEAN                           *IsModified\r
+  IN  PAGE_TABLE_LIB_PAGING_CONTEXT  *PagingContext,\r
+  IN  UINT64                         *PageEntry,\r
+  IN  UINT64                         Attributes,\r
+  IN  PAGE_ACTION                    PageAction,\r
+  OUT BOOLEAN                        *IsModified\r
   )\r
 {\r
   UINT64  CurrentPageEntry;\r
   UINT64  NewPageEntry;\r
+  UINT32  *PageAttributes;\r
 \r
   CurrentPageEntry = *PageEntry;\r
-  NewPageEntry = CurrentPageEntry;\r
+  NewPageEntry     = CurrentPageEntry;\r
   if ((Attributes & EFI_MEMORY_RP) != 0) {\r
     switch (PageAction) {\r
-    case PageActionAssign:\r
-    case PageActionSet:\r
-      NewPageEntry &= ~(UINT64)IA32_PG_P;\r
-      break;\r
-    case PageActionClear:\r
-      NewPageEntry |= IA32_PG_P;\r
-      break;\r
+      case PageActionAssign:\r
+      case PageActionSet:\r
+        NewPageEntry &= ~(UINT64)IA32_PG_P;\r
+        break;\r
+      case PageActionClear:\r
+        NewPageEntry |= IA32_PG_P;\r
+        break;\r
     }\r
   } else {\r
     switch (PageAction) {\r
-    case PageActionAssign:\r
-      NewPageEntry |= IA32_PG_P;\r
-      break;\r
-    case PageActionSet:\r
-    case PageActionClear:\r
-      break;\r
+      case PageActionAssign:\r
+        NewPageEntry |= IA32_PG_P;\r
+        break;\r
+      case PageActionSet:\r
+      case PageActionClear:\r
+        break;\r
     }\r
   }\r
+\r
   if ((Attributes & EFI_MEMORY_RO) != 0) {\r
     switch (PageAction) {\r
-    case PageActionAssign:\r
-    case PageActionSet:\r
-      NewPageEntry &= ~(UINT64)IA32_PG_RW;\r
-      break;\r
-    case PageActionClear:\r
-      NewPageEntry |= IA32_PG_RW;\r
-      break;\r
-    }\r
-  } else {\r
-    switch (PageAction) {\r
-    case PageActionAssign:\r
-      NewPageEntry |= IA32_PG_RW;\r
-      break;\r
-    case PageActionSet:\r
-    case PageActionClear:\r
-      break;\r
-    }\r
-  }\r
-  if ((PagingContext->ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED) != 0) {\r
-    if ((Attributes & EFI_MEMORY_XP) != 0) {\r
-      switch (PageAction) {\r
       case PageActionAssign:\r
       case PageActionSet:\r
-        NewPageEntry |= IA32_PG_NX;\r
+        NewPageEntry &= ~(UINT64)IA32_PG_RW;\r
         break;\r
       case PageActionClear:\r
-        NewPageEntry &= ~IA32_PG_NX;\r
+        NewPageEntry |= IA32_PG_RW;\r
         break;\r
-      }\r
-    } else {\r
-      switch (PageAction) {\r
+    }\r
+  } else {\r
+    switch (PageAction) {\r
       case PageActionAssign:\r
-        NewPageEntry &= ~IA32_PG_NX;\r
+        NewPageEntry |= IA32_PG_RW;\r
         break;\r
       case PageActionSet:\r
       case PageActionClear:\r
         break;\r
+    }\r
+  }\r
+\r
+  GetPagingDetails (&PagingContext->ContextData, NULL, &PageAttributes);\r
+\r
+  if ((*PageAttributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED) != 0) {\r
+    if ((Attributes & EFI_MEMORY_XP) != 0) {\r
+      switch (PageAction) {\r
+        case PageActionAssign:\r
+        case PageActionSet:\r
+          NewPageEntry |= IA32_PG_NX;\r
+          break;\r
+        case PageActionClear:\r
+          NewPageEntry &= ~IA32_PG_NX;\r
+          break;\r
+      }\r
+    } else {\r
+      switch (PageAction) {\r
+        case PageActionAssign:\r
+          NewPageEntry &= ~IA32_PG_NX;\r
+          break;\r
+        case PageActionSet:\r
+        case PageActionClear:\r
+          break;\r
       }\r
     }\r
   }\r
+\r
   *PageEntry = NewPageEntry;\r
   if (CurrentPageEntry != NewPageEntry) {\r
     *IsModified = TRUE;\r
@@ -459,13 +514,13 @@ ConvertPageEntryAttribute (
 **/\r
 PAGE_ATTRIBUTE\r
 NeedSplitPage (\r
-  IN  PHYSICAL_ADDRESS                  BaseAddress,\r
-  IN  UINT64                            Length,\r
-  IN  UINT64                            *PageEntry,\r
-  IN  PAGE_ATTRIBUTE                    PageAttribute\r
+  IN  PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64            Length,\r
+  IN  UINT64            *PageEntry,\r
+  IN  PAGE_ATTRIBUTE    PageAttribute\r
   )\r
 {\r
-  UINT64                PageEntryLength;\r
+  UINT64  PageEntryLength;\r
 \r
   PageEntryLength = PageAttributeToLength (PageAttribute);\r
 \r
@@ -494,16 +549,16 @@ NeedSplitPage (
 **/\r
 RETURN_STATUS\r
 SplitPage (\r
-  IN  UINT64                            *PageEntry,\r
-  IN  PAGE_ATTRIBUTE                    PageAttribute,\r
-  IN  PAGE_ATTRIBUTE                    SplitAttribute,\r
-  IN  PAGE_TABLE_LIB_ALLOCATE_PAGES     AllocatePagesFunc\r
+  IN  UINT64                         *PageEntry,\r
+  IN  PAGE_ATTRIBUTE                 PageAttribute,\r
+  IN  PAGE_ATTRIBUTE                 SplitAttribute,\r
+  IN  PAGE_TABLE_LIB_ALLOCATE_PAGES  AllocatePagesFunc\r
   )\r
 {\r
-  UINT64   BaseAddress;\r
-  UINT64   *NewPageEntry;\r
-  UINTN    Index;\r
-  UINT64   AddressEncMask;\r
+  UINT64  BaseAddress;\r
+  UINT64  *NewPageEntry;\r
+  UINTN   Index;\r
+  UINT64  AddressEncMask;\r
 \r
   ASSERT (PageAttribute == Page2M || PageAttribute == Page1G);\r
 \r
@@ -520,15 +575,17 @@ SplitPage (
     ASSERT (SplitAttribute == Page4K);\r
     if (SplitAttribute == Page4K) {\r
       NewPageEntry = AllocatePagesFunc (1);\r
-      DEBUG ((DEBUG_INFO, "Split - 0x%x\n", NewPageEntry));\r
+      DEBUG ((DEBUG_VERBOSE, "Split - 0x%x\n", NewPageEntry));\r
       if (NewPageEntry == NULL) {\r
         return RETURN_OUT_OF_RESOURCES;\r
       }\r
+\r
       BaseAddress = *PageEntry & ~AddressEncMask & PAGING_2M_ADDRESS_MASK_64;\r
-      for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {\r
+      for (Index = 0; Index < SIZE_4KB / sizeof (UINT64); Index++) {\r
         NewPageEntry[Index] = (BaseAddress + SIZE_4KB * Index) | AddressEncMask | ((*PageEntry) & PAGE_PROGATE_BITS);\r
       }\r
-      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | ((*PageEntry) & PAGE_ATTRIBUTE_BITS);\r
+\r
+      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | PAGE_ATTRIBUTE_BITS_POST_SPLIT;\r
       return RETURN_SUCCESS;\r
     } else {\r
       return RETURN_UNSUPPORTED;\r
@@ -539,17 +596,19 @@ SplitPage (
     // No need support 1G->4K directly, we should use 1G->2M, then 2M->4K to get more compact page table.\r
     //\r
     ASSERT (SplitAttribute == Page2M || SplitAttribute == Page4K);\r
-    if ((SplitAttribute == Page2M || SplitAttribute == Page4K)) {\r
+    if (((SplitAttribute == Page2M) || (SplitAttribute == Page4K))) {\r
       NewPageEntry = AllocatePagesFunc (1);\r
-      DEBUG ((DEBUG_INFO, "Split - 0x%x\n", NewPageEntry));\r
+      DEBUG ((DEBUG_VERBOSE, "Split - 0x%x\n", NewPageEntry));\r
       if (NewPageEntry == NULL) {\r
         return RETURN_OUT_OF_RESOURCES;\r
       }\r
+\r
       BaseAddress = *PageEntry & ~AddressEncMask  & PAGING_1G_ADDRESS_MASK_64;\r
-      for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {\r
+      for (Index = 0; Index < SIZE_4KB / sizeof (UINT64); Index++) {\r
         NewPageEntry[Index] = (BaseAddress + SIZE_2MB * Index) | AddressEncMask | IA32_PG_PS | ((*PageEntry) & PAGE_PROGATE_BITS);\r
       }\r
-      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | ((*PageEntry) & PAGE_ATTRIBUTE_BITS);\r
+\r
+      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | PAGE_ATTRIBUTE_BITS_POST_SPLIT;\r
       return RETURN_SUCCESS;\r
     } else {\r
       return RETURN_UNSUPPORTED;\r
@@ -571,21 +630,18 @@ IsReadOnlyPageWriteProtected (
   VOID\r
   )\r
 {\r
-  return ((AsmReadCr0 () & BIT16) != 0);\r
-}\r
+  IA32_CR0  Cr0;\r
 \r
-/**\r
-  Disable write protection function for AP.\r
+  //\r
+  // To avoid unforseen consequences, don't touch paging settings in SMM mode\r
+  // in this driver.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    Cr0.UintN = AsmReadCr0 ();\r
+    return (BOOLEAN)(Cr0.Bits.WP != 0);\r
+  }\r
 \r
-  @param[in,out] Buffer  The pointer to private data buffer.\r
-**/\r
-VOID\r
-EFIAPI\r
-SyncCpuDisableWriteProtection (\r
-  IN OUT VOID *Buffer\r
-  )\r
-{\r
-  AsmWriteCr0 (AsmReadCr0() & ~BIT16);\r
+  return FALSE;\r
 }\r
 \r
 /**\r
@@ -596,8 +652,17 @@ DisableReadOnlyPageWriteProtect (
   VOID\r
   )\r
 {\r
-  AsmWriteCr0 (AsmReadCr0() & ~BIT16);\r
-  SyncMemoryPageAttributesAp (SyncCpuDisableWriteProtection);\r
+  IA32_CR0  Cr0;\r
+\r
+  //\r
+  // To avoid unforseen consequences, don't touch paging settings in SMM mode\r
+  // in this driver.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    Cr0.UintN   = AsmReadCr0 ();\r
+    Cr0.Bits.WP = 0;\r
+    AsmWriteCr0 (Cr0.UintN);\r
+  }\r
 }\r
 \r
 /**\r
@@ -608,8 +673,17 @@ EnableReadOnlyPageWriteProtect (
   VOID\r
   )\r
 {\r
-  AsmWriteCr0 (AsmReadCr0() | BIT16);\r
-  SyncMemoryPageAttributesAp (SyncCpuEnableWriteProtection);\r
+  IA32_CR0  Cr0;\r
+\r
+  //\r
+  // To avoid unforseen consequences, don't touch paging settings in SMM mode\r
+  // in this driver.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    Cr0.UintN   = AsmReadCr0 ();\r
+    Cr0.Bits.WP = 1;\r
+    AsmWriteCr0 (Cr0.UintN);\r
+  }\r
 }\r
 \r
 /**\r
@@ -643,39 +717,41 @@ EnableReadOnlyPageWriteProtect (
 **/\r
 RETURN_STATUS\r
 ConvertMemoryPageAttributes (\r
-  IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext OPTIONAL,\r
-  IN  PHYSICAL_ADDRESS                  BaseAddress,\r
-  IN  UINT64                            Length,\r
-  IN  UINT64                            Attributes,\r
-  IN  PAGE_ACTION                       PageAction,\r
-  IN  PAGE_TABLE_LIB_ALLOCATE_PAGES     AllocatePagesFunc OPTIONAL,\r
-  OUT BOOLEAN                           *IsSplitted,  OPTIONAL\r
-  OUT BOOLEAN                           *IsModified   OPTIONAL\r
+  IN  PAGE_TABLE_LIB_PAGING_CONTEXT  *PagingContext OPTIONAL,\r
+  IN  PHYSICAL_ADDRESS               BaseAddress,\r
+  IN  UINT64                         Length,\r
+  IN  UINT64                         Attributes,\r
+  IN  PAGE_ACTION                    PageAction,\r
+  IN  PAGE_TABLE_LIB_ALLOCATE_PAGES  AllocatePagesFunc OPTIONAL,\r
+  OUT BOOLEAN                        *IsSplitted   OPTIONAL,\r
+  OUT BOOLEAN                        *IsModified   OPTIONAL\r
   )\r
 {\r
-  PAGE_TABLE_LIB_PAGING_CONTEXT     CurrentPagingContext;\r
-  UINT64                            *PageEntry;\r
-  PAGE_ATTRIBUTE                    PageAttribute;\r
-  UINTN                             PageEntryLength;\r
-  PAGE_ATTRIBUTE                    SplitAttribute;\r
-  RETURN_STATUS                     Status;\r
-  BOOLEAN                           IsEntryModified;\r
-  BOOLEAN                           IsWpEnabled;\r
+  PAGE_TABLE_LIB_PAGING_CONTEXT  CurrentPagingContext;\r
+  UINT64                         *PageEntry;\r
+  PAGE_ATTRIBUTE                 PageAttribute;\r
+  UINTN                          PageEntryLength;\r
+  PAGE_ATTRIBUTE                 SplitAttribute;\r
+  RETURN_STATUS                  Status;\r
+  BOOLEAN                        IsEntryModified;\r
+  BOOLEAN                        IsWpEnabled;\r
 \r
   if ((BaseAddress & (SIZE_4KB - 1)) != 0) {\r
     DEBUG ((DEBUG_ERROR, "BaseAddress(0x%lx) is not aligned!\n", BaseAddress));\r
     return EFI_UNSUPPORTED;\r
   }\r
+\r
   if ((Length & (SIZE_4KB - 1)) != 0) {\r
     DEBUG ((DEBUG_ERROR, "Length(0x%lx) is not aligned!\n", Length));\r
     return EFI_UNSUPPORTED;\r
   }\r
+\r
   if (Length == 0) {\r
     DEBUG ((DEBUG_ERROR, "Length is 0!\n"));\r
     return RETURN_INVALID_PARAMETER;\r
   }\r
 \r
-  if ((Attributes & ~(EFI_MEMORY_RP | EFI_MEMORY_RO | EFI_MEMORY_XP)) != 0) {\r
+  if ((Attributes & ~EFI_MEMORY_ATTRIBUTE_MASK) != 0) {\r
     DEBUG ((DEBUG_ERROR, "Attributes(0x%lx) has unsupported bit\n", Attributes));\r
     return EFI_UNSUPPORTED;\r
   }\r
@@ -683,44 +759,50 @@ ConvertMemoryPageAttributes (
   if (PagingContext == NULL) {\r
     GetCurrentPagingContext (&CurrentPagingContext);\r
   } else {\r
-    CopyMem (&CurrentPagingContext, PagingContext, sizeof(CurrentPagingContext));\r
+    CopyMem (&CurrentPagingContext, PagingContext, sizeof (CurrentPagingContext));\r
   }\r
-  switch(CurrentPagingContext.MachineType) {\r
-  case IMAGE_FILE_MACHINE_I386:\r
-    if (CurrentPagingContext.ContextData.Ia32.PageTableBase == 0) {\r
-      if (Attributes == 0) {\r
-        return EFI_SUCCESS;\r
-      } else {\r
-        DEBUG ((DEBUG_ERROR, "PageTable is 0!\n"));\r
+\r
+  switch (CurrentPagingContext.MachineType) {\r
+    case IMAGE_FILE_MACHINE_I386:\r
+      if (CurrentPagingContext.ContextData.Ia32.PageTableBase == 0) {\r
+        if (Attributes == 0) {\r
+          return EFI_SUCCESS;\r
+        } else {\r
+          DEBUG ((DEBUG_ERROR, "PageTable is 0!\n"));\r
+          return EFI_UNSUPPORTED;\r
+        }\r
+      }\r
+\r
+      if ((CurrentPagingContext.ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) == 0) {\r
+        DEBUG ((DEBUG_ERROR, "Non-PAE Paging!\n"));\r
         return EFI_UNSUPPORTED;\r
       }\r
-    }\r
-    if ((CurrentPagingContext.ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) == 0) {\r
-      DEBUG ((DEBUG_ERROR, "Non-PAE Paging!\n"));\r
-      return EFI_UNSUPPORTED;\r
-    }\r
-    if ((BaseAddress + Length) > BASE_4GB) {\r
-      DEBUG ((DEBUG_ERROR, "Beyond 4GB memory in 32-bit mode!\n"));\r
+\r
+      if ((BaseAddress + Length) > BASE_4GB) {\r
+        DEBUG ((DEBUG_ERROR, "Beyond 4GB memory in 32-bit mode!\n"));\r
+        return EFI_UNSUPPORTED;\r
+      }\r
+\r
+      break;\r
+    case IMAGE_FILE_MACHINE_X64:\r
+      ASSERT (CurrentPagingContext.ContextData.X64.PageTableBase != 0);\r
+      break;\r
+    default:\r
+      ASSERT (FALSE);\r
       return EFI_UNSUPPORTED;\r
-    }\r
-    break;\r
-  case IMAGE_FILE_MACHINE_X64:\r
-    ASSERT (CurrentPagingContext.ContextData.X64.PageTableBase != 0);\r
-    break;\r
-  default:\r
-    ASSERT(FALSE);\r
-    return EFI_UNSUPPORTED;\r
-    break;\r
+      break;\r
   }\r
 \r
-//  DEBUG ((DEBUG_ERROR, "ConvertMemoryPageAttributes(%x) - %016lx, %016lx, %02lx\n", IsSet, BaseAddress, Length, Attributes));\r
+  //  DEBUG ((DEBUG_ERROR, "ConvertMemoryPageAttributes(%x) - %016lx, %016lx, %02lx\n", IsSet, BaseAddress, Length, Attributes));\r
 \r
   if (IsSplitted != NULL) {\r
     *IsSplitted = FALSE;\r
   }\r
+\r
   if (IsModified != NULL) {\r
     *IsModified = FALSE;\r
   }\r
+\r
   if (AllocatePagesFunc == NULL) {\r
     AllocatePagesFunc = AllocatePageTableMemory;\r
   }\r
@@ -734,7 +816,7 @@ ConvertMemoryPageAttributes (
   }\r
 \r
   //\r
-  // Below logic is to check 2M/4K page to make sure we donot waist memory.\r
+  // Below logic is to check 2M/4K page to make sure we do not waste memory.\r
   //\r
   Status = EFI_SUCCESS;\r
   while (Length != 0) {\r
@@ -743,8 +825,9 @@ ConvertMemoryPageAttributes (
       Status = RETURN_UNSUPPORTED;\r
       goto Done;\r
     }\r
+\r
     PageEntryLength = PageAttributeToLength (PageAttribute);\r
-    SplitAttribute = NeedSplitPage (BaseAddress, Length, PageEntry, PageAttribute);\r
+    SplitAttribute  = NeedSplitPage (BaseAddress, Length, PageEntry, PageAttribute);\r
     if (SplitAttribute == PageNone) {\r
       ConvertPageEntryAttribute (&CurrentPagingContext, PageEntry, Attributes, PageAction, &IsEntryModified);\r
       if (IsEntryModified) {\r
@@ -752,27 +835,32 @@ ConvertMemoryPageAttributes (
           *IsModified = TRUE;\r
         }\r
       }\r
+\r
       //\r
       // Convert success, move to next\r
       //\r
       BaseAddress += PageEntryLength;\r
-      Length -= PageEntryLength;\r
+      Length      -= PageEntryLength;\r
     } else {\r
       if (AllocatePagesFunc == NULL) {\r
         Status = RETURN_UNSUPPORTED;\r
         goto Done;\r
       }\r
+\r
       Status = SplitPage (PageEntry, PageAttribute, SplitAttribute, AllocatePagesFunc);\r
       if (RETURN_ERROR (Status)) {\r
         Status = RETURN_UNSUPPORTED;\r
         goto Done;\r
       }\r
+\r
       if (IsSplitted != NULL) {\r
         *IsSplitted = TRUE;\r
       }\r
+\r
       if (IsModified != NULL) {\r
         *IsModified = TRUE;\r
       }\r
+\r
       //\r
       // Just split current page\r
       // Convert success in next around\r
@@ -787,6 +875,7 @@ Done:
   if (IsWpEnabled) {\r
     EnableReadOnlyPageWriteProtect ();\r
   }\r
+\r
   return Status;\r
 }\r
 \r
@@ -796,7 +885,7 @@ Done:
 \r
   Caller should make sure BaseAddress and Length is at page boundary.\r
 \r
-  Caller need guarentee the TPL <= TPL_NOTIFY, if there is split page request.\r
+  Caller need guarantee the TPL <= TPL_NOTIFY, if there is split page request.\r
 \r
   @param[in]  PagingContext     The paging context. NULL means get page table from current CPU context.\r
   @param[in]  BaseAddress       The physical address that is the start address of a memory region.\r
@@ -821,26 +910,29 @@ Done:
 RETURN_STATUS\r
 EFIAPI\r
 AssignMemoryPageAttributes (\r
-  IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext OPTIONAL,\r
-  IN  PHYSICAL_ADDRESS                  BaseAddress,\r
-  IN  UINT64                            Length,\r
-  IN  UINT64                            Attributes,\r
-  IN  PAGE_TABLE_LIB_ALLOCATE_PAGES     AllocatePagesFunc OPTIONAL\r
+  IN  PAGE_TABLE_LIB_PAGING_CONTEXT  *PagingContext OPTIONAL,\r
+  IN  PHYSICAL_ADDRESS               BaseAddress,\r
+  IN  UINT64                         Length,\r
+  IN  UINT64                         Attributes,\r
+  IN  PAGE_TABLE_LIB_ALLOCATE_PAGES  AllocatePagesFunc OPTIONAL\r
   )\r
 {\r
   RETURN_STATUS  Status;\r
   BOOLEAN        IsModified;\r
   BOOLEAN        IsSplitted;\r
 \r
-//  DEBUG((DEBUG_INFO, "AssignMemoryPageAttributes: 0x%lx - 0x%lx (0x%lx)\n", BaseAddress, Length, Attributes));\r
+  //  DEBUG((DEBUG_INFO, "AssignMemoryPageAttributes: 0x%lx - 0x%lx (0x%lx)\n", BaseAddress, Length, Attributes));\r
   Status = ConvertMemoryPageAttributes (PagingContext, BaseAddress, Length, Attributes, PageActionAssign, AllocatePagesFunc, &IsSplitted, &IsModified);\r
-  if (!EFI_ERROR(Status)) {\r
+  if (!EFI_ERROR (Status)) {\r
     if ((PagingContext == NULL) && IsModified) {\r
       //\r
-      // Flush TLB as last step\r
+      // Flush TLB as last step.\r
+      //\r
+      // Note: Since APs will always init CR3 register in HLT loop mode or do\r
+      // TLB flush in MWAIT loop mode, there's no need to flush TLB for them\r
+      // here.\r
       //\r
-      CpuFlushTlb();\r
-      SyncMemoryPageAttributesAp (SyncCpuFlushTlb);\r
+      CpuFlushTlb ();\r
     }\r
   }\r
 \r
@@ -855,7 +947,7 @@ IsExecuteDisableEnabled (
   VOID\r
   )\r
 {\r
-  MSR_CORE_IA32_EFER_REGISTER    MsrEfer;\r
+  MSR_CORE_IA32_EFER_REGISTER  MsrEfer;\r
 \r
   MsrEfer.Uint64 = AsmReadMsr64 (MSR_IA32_EFER);\r
   return (MsrEfer.Bits.NXE == 1);\r
@@ -869,21 +961,21 @@ RefreshGcdMemoryAttributesFromPaging (
   VOID\r
   )\r
 {\r
-  EFI_STATUS                          Status;\r
-  UINTN                               NumberOfDescriptors;\r
-  EFI_GCD_MEMORY_SPACE_DESCRIPTOR     *MemorySpaceMap;\r
-  PAGE_TABLE_LIB_PAGING_CONTEXT       PagingContext;\r
-  PAGE_ATTRIBUTE                      PageAttribute;\r
-  UINT64                              *PageEntry;\r
-  UINT64                              PageLength;\r
-  UINT64                              MemorySpaceLength;\r
-  UINT64                              Length;\r
-  UINT64                              BaseAddress;\r
-  UINT64                              PageStartAddress;\r
-  UINT64                              Attributes;\r
-  UINT64                              Capabilities;\r
-  UINT64                              NewAttributes;\r
-  UINTN                               Index;\r
+  EFI_STATUS                       Status;\r
+  UINTN                            NumberOfDescriptors;\r
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *MemorySpaceMap;\r
+  PAGE_TABLE_LIB_PAGING_CONTEXT    PagingContext;\r
+  PAGE_ATTRIBUTE                   PageAttribute;\r
+  UINT64                           *PageEntry;\r
+  UINT64                           PageLength;\r
+  UINT64                           MemorySpaceLength;\r
+  UINT64                           Length;\r
+  UINT64                           BaseAddress;\r
+  UINT64                           PageStartAddress;\r
+  UINT64                           Attributes;\r
+  UINT64                           Capabilities;\r
+  UINT64                           NewAttributes;\r
+  UINTN                            Index;\r
 \r
   //\r
   // Assuming that memory space map returned is sorted already; otherwise sort\r
@@ -894,10 +986,10 @@ RefreshGcdMemoryAttributesFromPaging (
 \r
   GetCurrentPagingContext (&PagingContext);\r
 \r
-  Attributes      = 0;\r
-  NewAttributes   = 0;\r
-  BaseAddress     = 0;\r
-  PageLength      = 0;\r
+  Attributes    = 0;\r
+  NewAttributes = 0;\r
+  BaseAddress   = 0;\r
+  PageLength    = 0;\r
 \r
   if (IsExecuteDisableEnabled ()) {\r
     Capabilities = EFI_MEMORY_RO | EFI_MEMORY_RP | EFI_MEMORY_XP;\r
@@ -924,13 +1016,14 @@ RefreshGcdMemoryAttributesFromPaging (
                     );\r
     if (EFI_ERROR (Status)) {\r
       //\r
-      // If we cannot udpate the capabilities, we cannot update its\r
+      // If we cannot update the capabilities, we cannot update its\r
       // attributes either. So just simply skip current block of memory.\r
       //\r
       DEBUG ((\r
         DEBUG_WARN,\r
         "Failed to update capability: [%lu] %016lx - %016lx (%016lx -> %016lx)\r\n",\r
-        (UINT64)Index, MemorySpaceMap[Index].BaseAddress,\r
+        (UINT64)Index,\r
+        MemorySpaceMap[Index].BaseAddress,\r
         MemorySpaceMap[Index].BaseAddress + MemorySpaceMap[Index].Length - 1,\r
         MemorySpaceMap[Index].Capabilities,\r
         MemorySpaceMap[Index].Capabilities | Capabilities\r
@@ -966,16 +1059,17 @@ RefreshGcdMemoryAttributesFromPaging (
         //\r
         // Note current memory space might start in the middle of a page\r
         //\r
-        PageStartAddress  = (*PageEntry) & (UINT64)PageAttributeToMask(PageAttribute);\r
-        PageLength        = PageAttributeToLength (PageAttribute) - (BaseAddress - PageStartAddress);\r
-        Attributes        = GetAttributesFromPageEntry (PageEntry);\r
+        PageStartAddress = (*PageEntry) & (UINT64)PageAttributeToMask (PageAttribute);\r
+        PageLength       = PageAttributeToLength (PageAttribute) - (BaseAddress - PageStartAddress);\r
+        Attributes       = GetAttributesFromPageEntry (PageEntry);\r
       }\r
 \r
       Length = MIN (PageLength, MemorySpaceLength);\r
       if (Attributes != (MemorySpaceMap[Index].Attributes &\r
-                         EFI_MEMORY_PAGETYPE_MASK)) {\r
+                         EFI_MEMORY_ATTRIBUTE_MASK))\r
+      {\r
         NewAttributes = (MemorySpaceMap[Index].Attributes &\r
-                         ~EFI_MEMORY_PAGETYPE_MASK) | Attributes;\r
+                         ~EFI_MEMORY_ATTRIBUTE_MASK) | Attributes;\r
         Status = gDS->SetMemorySpaceAttributes (\r
                         BaseAddress,\r
                         Length,\r
@@ -985,7 +1079,9 @@ RefreshGcdMemoryAttributesFromPaging (
         DEBUG ((\r
           DEBUG_VERBOSE,\r
           "Updated memory space attribute: [%lu] %016lx - %016lx (%016lx -> %016lx)\r\n",\r
-          (UINT64)Index, BaseAddress, BaseAddress + Length - 1,\r
+          (UINT64)Index,\r
+          BaseAddress,\r
+          BaseAddress + Length - 1,\r
           MemorySpaceMap[Index].Attributes,\r
           NewAttributes\r
           ));\r
@@ -1019,42 +1115,58 @@ RefreshGcdMemoryAttributesFromPaging (
 **/\r
 BOOLEAN\r
 InitializePageTablePool (\r
-  IN  UINTN                           PoolPages\r
+  IN  UINTN  PoolPages\r
   )\r
 {\r
-  VOID                      *Buffer;\r
-  BOOLEAN                   IsModified;\r
+  VOID     *Buffer;\r
+  BOOLEAN  IsModified;\r
+\r
+  //\r
+  // Do not allow re-entrance.\r
+  //\r
+  if (mPageTablePoolLock) {\r
+    return FALSE;\r
+  }\r
+\r
+  mPageTablePoolLock = TRUE;\r
+  IsModified         = FALSE;\r
 \r
   //\r
   // Always reserve at least PAGE_TABLE_POOL_UNIT_PAGES, including one page for\r
   // header.\r
   //\r
   PoolPages += 1;   // Add one page for header.\r
-  PoolPages = ((PoolPages - 1) / PAGE_TABLE_POOL_UNIT_PAGES + 1) *\r
-              PAGE_TABLE_POOL_UNIT_PAGES;\r
+  PoolPages  = ((PoolPages - 1) / PAGE_TABLE_POOL_UNIT_PAGES + 1) *\r
+               PAGE_TABLE_POOL_UNIT_PAGES;\r
   Buffer = AllocateAlignedPages (PoolPages, PAGE_TABLE_POOL_ALIGNMENT);\r
   if (Buffer == NULL) {\r
     DEBUG ((DEBUG_ERROR, "ERROR: Out of aligned pages\r\n"));\r
-    return FALSE;\r
+    goto Done;\r
   }\r
 \r
+  DEBUG ((\r
+    DEBUG_INFO,\r
+    "Paging: added %lu pages to page table pool\r\n",\r
+    (UINT64)PoolPages\r
+    ));\r
+\r
   //\r
   // Link all pools into a list for easier track later.\r
   //\r
   if (mPageTablePool == NULL) {\r
-    mPageTablePool = Buffer;\r
+    mPageTablePool           = Buffer;\r
     mPageTablePool->NextPool = mPageTablePool;\r
   } else {\r
     ((PAGE_TABLE_POOL *)Buffer)->NextPool = mPageTablePool->NextPool;\r
-    mPageTablePool->NextPool = Buffer;\r
-    mPageTablePool = Buffer;\r
+    mPageTablePool->NextPool              = Buffer;\r
+    mPageTablePool                        = Buffer;\r
   }\r
 \r
   //\r
   // Reserve one page for pool header.\r
   //\r
-  mPageTablePool->FreePages  = PoolPages - 1;\r
-  mPageTablePool->Offset = EFI_PAGES_TO_SIZE (1);\r
+  mPageTablePool->FreePages = PoolPages - 1;\r
+  mPageTablePool->Offset    = EFI_PAGES_TO_SIZE (1);\r
 \r
   //\r
   // Mark the whole pool pages as read-only.\r
@@ -1071,7 +1183,9 @@ InitializePageTablePool (
     );\r
   ASSERT (IsModified == TRUE);\r
 \r
-  return TRUE;\r
+Done:\r
+  mPageTablePoolLock = FALSE;\r
+  return IsModified;\r
 }\r
 \r
 /**\r
@@ -1094,10 +1208,10 @@ InitializePageTablePool (
 VOID *\r
 EFIAPI\r
 AllocatePageTableMemory (\r
-  IN UINTN           Pages\r
+  IN UINTN  Pages\r
   )\r
 {\r
-  VOID                            *Buffer;\r
+  VOID  *Buffer;\r
 \r
   if (Pages == 0) {\r
     return NULL;\r
@@ -1106,8 +1220,9 @@ AllocatePageTableMemory (
   //\r
   // Renew the pool if necessary.\r
   //\r
-  if (mPageTablePool == NULL ||\r
-      Pages > mPageTablePool->FreePages) {\r
+  if ((mPageTablePool == NULL) ||\r
+      (Pages > mPageTablePool->FreePages))\r
+  {\r
     if (!InitializePageTablePool (Pages)) {\r
       return NULL;\r
     }\r
@@ -1115,12 +1230,170 @@ AllocatePageTableMemory (
 \r
   Buffer = (UINT8 *)mPageTablePool + mPageTablePool->Offset;\r
 \r
-  mPageTablePool->Offset     += EFI_PAGES_TO_SIZE (Pages);\r
-  mPageTablePool->FreePages  -= Pages;\r
+  mPageTablePool->Offset    += EFI_PAGES_TO_SIZE (Pages);\r
+  mPageTablePool->FreePages -= Pages;\r
 \r
   return Buffer;\r
 }\r
 \r
+/**\r
+  Special handler for #DB exception, which will restore the page attributes\r
+  (not-present). It should work with #PF handler which will set pages to\r
+  'present'.\r
+\r
+  @param ExceptionType  Exception type.\r
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.\r
+\r
+**/\r
+VOID\r
+EFIAPI\r
+DebugExceptionHandler (\r
+  IN EFI_EXCEPTION_TYPE  ExceptionType,\r
+  IN EFI_SYSTEM_CONTEXT  SystemContext\r
+  )\r
+{\r
+  UINTN    CpuIndex;\r
+  UINTN    PFEntry;\r
+  BOOLEAN  IsWpEnabled;\r
+\r
+  MpInitLibWhoAmI (&CpuIndex);\r
+\r
+  //\r
+  // Clear last PF entries\r
+  //\r
+  IsWpEnabled = IsReadOnlyPageWriteProtected ();\r
+  if (IsWpEnabled) {\r
+    DisableReadOnlyPageWriteProtect ();\r
+  }\r
+\r
+  for (PFEntry = 0; PFEntry < mPFEntryCount[CpuIndex]; PFEntry++) {\r
+    if (mLastPFEntryPointer[CpuIndex][PFEntry] != NULL) {\r
+      *mLastPFEntryPointer[CpuIndex][PFEntry] &= ~(UINT64)IA32_PG_P;\r
+    }\r
+  }\r
+\r
+  if (IsWpEnabled) {\r
+    EnableReadOnlyPageWriteProtect ();\r
+  }\r
+\r
+  //\r
+  // Reset page fault exception count for next page fault.\r
+  //\r
+  mPFEntryCount[CpuIndex] = 0;\r
+\r
+  //\r
+  // Flush TLB\r
+  //\r
+  CpuFlushTlb ();\r
+\r
+  //\r
+  // Clear TF in EFLAGS\r
+  //\r
+  if (mPagingContext.MachineType == IMAGE_FILE_MACHINE_I386) {\r
+    SystemContext.SystemContextIa32->Eflags &= (UINT32) ~BIT8;\r
+  } else {\r
+    SystemContext.SystemContextX64->Rflags &= (UINT64) ~BIT8;\r
+  }\r
+}\r
+\r
+/**\r
+  Special handler for #PF exception, which will set the pages which caused\r
+  #PF to be 'present'. The attribute of those pages should be restored in\r
+  the subsequent #DB handler.\r
+\r
+  @param ExceptionType  Exception type.\r
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.\r
+\r
+**/\r
+VOID\r
+EFIAPI\r
+PageFaultExceptionHandler (\r
+  IN EFI_EXCEPTION_TYPE  ExceptionType,\r
+  IN EFI_SYSTEM_CONTEXT  SystemContext\r
+  )\r
+{\r
+  EFI_STATUS                     Status;\r
+  UINT64                         PFAddress;\r
+  PAGE_TABLE_LIB_PAGING_CONTEXT  PagingContext;\r
+  PAGE_ATTRIBUTE                 PageAttribute;\r
+  UINT64                         Attributes;\r
+  UINT64                         *PageEntry;\r
+  UINTN                          Index;\r
+  UINTN                          CpuIndex;\r
+  UINTN                          PageNumber;\r
+  BOOLEAN                        NonStopMode;\r
+\r
+  PFAddress = AsmReadCr2 () & ~EFI_PAGE_MASK;\r
+  if (PFAddress < BASE_4KB) {\r
+    NonStopMode = NULL_DETECTION_NONSTOP_MODE ? TRUE : FALSE;\r
+  } else {\r
+    NonStopMode = HEAP_GUARD_NONSTOP_MODE ? TRUE : FALSE;\r
+  }\r
+\r
+  if (NonStopMode) {\r
+    MpInitLibWhoAmI (&CpuIndex);\r
+    GetCurrentPagingContext (&PagingContext);\r
+    //\r
+    // Memory operation cross page boundary, like "rep mov" instruction, will\r
+    // cause infinite loop between this and Debug Trap handler. We have to make\r
+    // sure that current page and the page followed are both in PRESENT state.\r
+    //\r
+    PageNumber = 2;\r
+    while (PageNumber > 0) {\r
+      PageEntry = GetPageTableEntry (&PagingContext, PFAddress, &PageAttribute);\r
+      ASSERT (PageEntry != NULL);\r
+\r
+      if (PageEntry != NULL) {\r
+        Attributes = GetAttributesFromPageEntry (PageEntry);\r
+        if ((Attributes & EFI_MEMORY_RP) != 0) {\r
+          Attributes &= ~EFI_MEMORY_RP;\r
+          Status      = AssignMemoryPageAttributes (\r
+                          &PagingContext,\r
+                          PFAddress,\r
+                          EFI_PAGE_SIZE,\r
+                          Attributes,\r
+                          NULL\r
+                          );\r
+          if (!EFI_ERROR (Status)) {\r
+            Index = mPFEntryCount[CpuIndex];\r
+            //\r
+            // Re-retrieve page entry because above calling might update page\r
+            // table due to table split.\r
+            //\r
+            PageEntry                              = GetPageTableEntry (&PagingContext, PFAddress, &PageAttribute);\r
+            mLastPFEntryPointer[CpuIndex][Index++] = PageEntry;\r
+            mPFEntryCount[CpuIndex]                = Index;\r
+          }\r
+        }\r
+      }\r
+\r
+      PFAddress += EFI_PAGE_SIZE;\r
+      --PageNumber;\r
+    }\r
+  }\r
+\r
+  //\r
+  // Initialize the serial port before dumping.\r
+  //\r
+  SerialPortInitialize ();\r
+  //\r
+  // Display ExceptionType, CPU information and Image information\r
+  //\r
+  DumpCpuContext (ExceptionType, SystemContext);\r
+  if (NonStopMode) {\r
+    //\r
+    // Set TF in EFLAGS\r
+    //\r
+    if (mPagingContext.MachineType == IMAGE_FILE_MACHINE_I386) {\r
+      SystemContext.SystemContextIa32->Eflags |= (UINT32)BIT8;\r
+    } else {\r
+      SystemContext.SystemContextX64->Rflags |= (UINT64)BIT8;\r
+    }\r
+  } else {\r
+    CpuDeadLoop ();\r
+  }\r
+}\r
+\r
 /**\r
   Initialize the Page Table lib.\r
 **/\r
@@ -1129,26 +1402,38 @@ InitializePageTableLib (
   VOID\r
   )\r
 {\r
-  PAGE_TABLE_LIB_PAGING_CONTEXT     CurrentPagingContext;\r
+  PAGE_TABLE_LIB_PAGING_CONTEXT  CurrentPagingContext;\r
+  UINT32                         *Attributes;\r
+  UINTN                          *PageTableBase;\r
 \r
   GetCurrentPagingContext (&CurrentPagingContext);\r
 \r
+  GetPagingDetails (&CurrentPagingContext.ContextData, &PageTableBase, &Attributes);\r
+\r
   //\r
   // Reserve memory of page tables for future uses, if paging is enabled.\r
   //\r
-  if (CurrentPagingContext.ContextData.X64.PageTableBase != 0 &&\r
-      (CurrentPagingContext.ContextData.Ia32.Attributes &\r
-       PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) != 0) {\r
+  if ((*PageTableBase != 0) &&\r
+      ((*Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) != 0))\r
+  {\r
     DisableReadOnlyPageWriteProtect ();\r
     InitializePageTablePool (1);\r
     EnableReadOnlyPageWriteProtect ();\r
   }\r
 \r
-  DEBUG ((DEBUG_INFO, "CurrentPagingContext:\n", CurrentPagingContext.MachineType));\r
+  if (HEAP_GUARD_NONSTOP_MODE || NULL_DETECTION_NONSTOP_MODE) {\r
+    mPFEntryCount = (UINTN *)AllocateZeroPool (sizeof (UINTN) * mNumberOfProcessors);\r
+    ASSERT (mPFEntryCount != NULL);\r
+\r
+    mLastPFEntryPointer = (UINT64 *(*)[MAX_PF_ENTRY_COUNT])\r
+                          AllocateZeroPool (sizeof (mLastPFEntryPointer[0]) * mNumberOfProcessors);\r
+    ASSERT (mLastPFEntryPointer != NULL);\r
+  }\r
+\r
+  DEBUG ((DEBUG_INFO, "CurrentPagingContext:\n"));\r
   DEBUG ((DEBUG_INFO, "  MachineType   - 0x%x\n", CurrentPagingContext.MachineType));\r
-  DEBUG ((DEBUG_INFO, "  PageTableBase - 0x%x\n", CurrentPagingContext.ContextData.X64.PageTableBase));\r
-  DEBUG ((DEBUG_INFO, "  Attributes    - 0x%x\n", CurrentPagingContext.ContextData.X64.Attributes));\r
+  DEBUG ((DEBUG_INFO, "  PageTableBase - 0x%Lx\n", (UINT64)*PageTableBase));\r
+  DEBUG ((DEBUG_INFO, "  Attributes    - 0x%x\n", *Attributes));\r
 \r
-  return ;\r
+  return;\r
 }\r
-\r