]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/CpuDxe/CpuPageTable.c
BaseTools: Library hashing fix and optimization for --hash feature
[mirror_edk2.git] / UefiCpuPkg / CpuDxe / CpuPageTable.c
index 3ad55f65c58abd5d0ab47c907d87f9f5dfcb6f9a..c369b44f128e70a3e49e94adf5b484782fecdfb9 100644 (file)
@@ -4,29 +4,31 @@
   Copyright (c) 2017, 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/Cpuid.h>\r
+#include <Register/Msr.h>\r
 \r
 #include "CpuDxe.h"\r
 #include "CpuPageTable.h"\r
 \r
+///\r
+/// Paging registers\r
+///\r
+#define CR0_WP                      BIT16\r
+#define CR0_PG                      BIT31\r
+#define CR4_PSE                     BIT4\r
+#define CR4_PAE                     BIT5\r
+\r
 ///\r
 /// Page Table Entry\r
 ///\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
   Page4K,\r
@@ -87,70 +93,59 @@ PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {
   {Page1G,  SIZE_1GB, PAGING_1G_ADDRESS_MASK_64},\r
 };\r
 \r
-PAGE_TABLE_POOL   *mPageTablePool = NULL;\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
-/**\r
-  Enable write protection function for AP.\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
@@ -163,45 +158,60 @@ GetCurrentPagingContext (
   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
 \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
+    if ((AsmReadCr0 () & CR0_PG) != 0) {\r
+      mPagingContext.ContextData.X64.PageTableBase = (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);\r
+    } else {\r
+      mPagingContext.ContextData.X64.PageTableBase = 0;\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
+    if ((AsmReadCr4 () & CR4_PSE) != 0) {\r
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE;\r
+    }\r
+    if ((AsmReadCr4 () & CR4_PAE) != 0) {\r
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;\r
+    }\r
+    if ((AsmReadCr0 () & CR0_WP) != 0) {\r
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE;\r
+    }\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
+    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
+          mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED;\r
+        }\r
+      }\r
+\r
+      if (RegEdx.Bits.Page1GB != 0) {\r
+        mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;\r
       }\r
-    }\r
-    if ((RegEdx & BIT26) != 0) {\r
-      PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;\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
@@ -520,7 +530,7 @@ 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
@@ -528,7 +538,7 @@ SplitPage (
       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_PROGATE_BITS);\r
+      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | ((*PageEntry) & PAGE_ATTRIBUTE_BITS);\r
       return RETURN_SUCCESS;\r
     } else {\r
       return RETURN_UNSUPPORTED;\r
@@ -541,7 +551,7 @@ SplitPage (
     ASSERT (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
@@ -549,7 +559,7 @@ SplitPage (
       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_PROGATE_BITS);\r
+      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | ((*PageEntry) & PAGE_ATTRIBUTE_BITS);\r
       return RETURN_SUCCESS;\r
     } else {\r
       return RETURN_UNSUPPORTED;\r
@@ -571,21 +581,14 @@ IsReadOnlyPageWriteProtected (
   VOID\r
   )\r
 {\r
-  return ((AsmReadCr0 () & BIT16) != 0);\r
-}\r
-\r
-/**\r
-  Disable write protection function for AP.\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
+  //\r
+  // To avoid unforseen consequences, don't touch paging settings in SMM mode\r
+  // in this driver.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    return ((AsmReadCr0 () & CR0_WP) != 0);\r
+  }\r
+  return FALSE;\r
 }\r
 \r
 /**\r
@@ -596,8 +599,13 @@ DisableReadOnlyPageWriteProtect (
   VOID\r
   )\r
 {\r
-  AsmWriteCr0 (AsmReadCr0() & ~BIT16);\r
-  SyncMemoryPageAttributesAp (SyncCpuDisableWriteProtection);\r
+  //\r
+  // To avoid unforseen consequences, don't touch paging settings in SMM mode\r
+  // in this driver.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);\r
+  }\r
 }\r
 \r
 /**\r
@@ -608,8 +616,13 @@ EnableReadOnlyPageWriteProtect (
   VOID\r
   )\r
 {\r
-  AsmWriteCr0 (AsmReadCr0() | BIT16);\r
-  SyncMemoryPageAttributesAp (SyncCpuEnableWriteProtection);\r
+  //\r
+  // To avoid unforseen consequences, don't touch paging settings in SMM mode\r
+  // in this driver.\r
+  //\r
+  if (!IsInSmm ()) {\r
+    AsmWriteCr0 (AsmReadCr0 () | CR0_WP);\r
+  }\r
 }\r
 \r
 /**\r
@@ -837,10 +850,13 @@ AssignMemoryPageAttributes (
   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
     }\r
   }\r
 \r
@@ -983,7 +999,7 @@ RefreshGcdMemoryAttributesFromPaging (
                         );\r
         ASSERT_EFI_ERROR (Status);\r
         DEBUG ((\r
-          DEBUG_INFO,\r
+          DEBUG_VERBOSE,\r
           "Updated memory space attribute: [%lu] %016lx - %016lx (%016lx -> %016lx)\r\n",\r
           (UINT64)Index, BaseAddress, BaseAddress + Length - 1,\r
           MemorySpaceMap[Index].Attributes,\r
@@ -1025,6 +1041,16 @@ InitializePageTablePool (
   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
@@ -1035,9 +1061,15 @@ InitializePageTablePool (
   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
@@ -1071,7 +1103,9 @@ InitializePageTablePool (
     );\r
   ASSERT (IsModified == TRUE);\r
 \r
-  return TRUE;\r
+Done:\r
+  mPageTablePoolLock = FALSE;\r
+  return IsModified;\r
 }\r
 \r
 /**\r
@@ -1121,6 +1155,159 @@ AllocatePageTableMemory (
   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 (&PagingContext, PFAddress,\r
+                                               EFI_PAGE_SIZE, Attributes, NULL);\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
@@ -1144,6 +1331,15 @@ InitializePageTableLib (
     EnableReadOnlyPageWriteProtect ();\r
   }\r
 \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", CurrentPagingContext.MachineType));\r
   DEBUG ((DEBUG_INFO, "  MachineType   - 0x%x\n", CurrentPagingContext.MachineType));\r
   DEBUG ((DEBUG_INFO, "  PageTableBase - 0x%x\n", CurrentPagingContext.ContextData.X64.PageTableBase));\r