]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
UefiCpuPkg: PiSmmCpuDxeSmm skip MSR_IA32_MISC_ENABLE manipulation on AMD
[mirror_edk2.git] / UefiCpuPkg / PiSmmCpuDxeSmm / X64 / PageTbl.c
index f3791ce89792e92887ead38bd08e610d4de7fd38..810985df20aee6e5ac04b008241a378f5ebc2867 100644 (file)
@@ -1,16 +1,10 @@
 /** @file\r
 Page Fault (#PF) handler for X64 processors\r
 \r
-Copyright (c) 2009 - 2017, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2009 - 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
@@ -21,7 +15,27 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 \r
 LIST_ENTRY                          mPagePool = INITIALIZE_LIST_HEAD_VARIABLE (mPagePool);\r
 BOOLEAN                             m1GPageTableSupport = FALSE;\r
-BOOLEAN                             mCpuSmmStaticPageTable;\r
+BOOLEAN                             mCpuSmmRestrictedMemoryAccess;\r
+BOOLEAN                             m5LevelPagingNeeded;\r
+X86_ASSEMBLY_PATCH_LABEL            gPatch5LevelPagingNeeded;\r
+\r
+/**\r
+  Disable CET.\r
+**/\r
+VOID\r
+EFIAPI\r
+DisableCet (\r
+  VOID\r
+  );\r
+\r
+/**\r
+  Enable CET.\r
+**/\r
+VOID\r
+EFIAPI\r
+EnableCet (\r
+  VOID\r
+  );\r
 \r
 /**\r
   Check if 1-GByte pages is supported by processor or not.\r
@@ -48,6 +62,48 @@ Is1GPageSupport (
   return FALSE;\r
 }\r
 \r
+/**\r
+  The routine returns TRUE when CPU supports it (CPUID[7,0].ECX.BIT[16] is set) and\r
+  the max physical address bits is bigger than 48. Because 4-level paging can support\r
+  to address physical address up to 2^48 - 1, there is no need to enable 5-level paging\r
+  with max physical address bits <= 48.\r
+\r
+  @retval TRUE  5-level paging enabling is needed.\r
+  @retval FALSE 5-level paging enabling is not needed.\r
+**/\r
+BOOLEAN\r
+Is5LevelPagingNeeded (\r
+  VOID\r
+  )\r
+{\r
+  CPUID_VIR_PHY_ADDRESS_SIZE_EAX              VirPhyAddressSize;\r
+  CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS_ECX ExtFeatureEcx;\r
+  UINT32                                      MaxExtendedFunctionId;\r
+\r
+  AsmCpuid (CPUID_EXTENDED_FUNCTION, &MaxExtendedFunctionId, NULL, NULL, NULL);\r
+  if (MaxExtendedFunctionId >= CPUID_VIR_PHY_ADDRESS_SIZE) {\r
+    AsmCpuid (CPUID_VIR_PHY_ADDRESS_SIZE, &VirPhyAddressSize.Uint32, NULL, NULL, NULL);\r
+  } else {\r
+    VirPhyAddressSize.Bits.PhysicalAddressBits = 36;\r
+  }\r
+  AsmCpuidEx (\r
+    CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS,\r
+    CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS_SUB_LEAF_INFO,\r
+    NULL, NULL, &ExtFeatureEcx.Uint32, NULL\r
+    );\r
+  DEBUG ((\r
+    DEBUG_INFO, "PhysicalAddressBits = %d, 5LPageTable = %d.\n",\r
+    VirPhyAddressSize.Bits.PhysicalAddressBits, ExtFeatureEcx.Bits.FiveLevelPage\r
+    ));\r
+\r
+  if (VirPhyAddressSize.Bits.PhysicalAddressBits > 4 * 9 + 12) {\r
+    ASSERT (ExtFeatureEcx.Bits.FiveLevelPage == 1);\r
+    return TRUE;\r
+  } else {\r
+    return FALSE;\r
+  }\r
+}\r
+\r
 /**\r
   Set sub-entries number in entry.\r
 \r
@@ -118,14 +174,6 @@ CalculateMaximumSupportAddress (
       PhysicalAddressBits = 36;\r
     }\r
   }\r
-\r
-  //\r
-  // IA-32e paging translates 48-bit linear addresses to 52-bit physical addresses.\r
-  //\r
-  ASSERT (PhysicalAddressBits <= 52);\r
-  if (PhysicalAddressBits > 48) {\r
-    PhysicalAddressBits = 48;\r
-  }\r
   return PhysicalAddressBits;\r
 }\r
 \r
@@ -140,89 +188,137 @@ SetStaticPageTable (
   )\r
 {\r
   UINT64                                        PageAddress;\r
+  UINTN                                         NumberOfPml5EntriesNeeded;\r
   UINTN                                         NumberOfPml4EntriesNeeded;\r
   UINTN                                         NumberOfPdpEntriesNeeded;\r
+  UINTN                                         IndexOfPml5Entries;\r
   UINTN                                         IndexOfPml4Entries;\r
   UINTN                                         IndexOfPdpEntries;\r
   UINTN                                         IndexOfPageDirectoryEntries;\r
+  UINT64                                        *PageMapLevel5Entry;\r
   UINT64                                        *PageMapLevel4Entry;\r
   UINT64                                        *PageMap;\r
   UINT64                                        *PageDirectoryPointerEntry;\r
   UINT64                                        *PageDirectory1GEntry;\r
   UINT64                                        *PageDirectoryEntry;\r
 \r
-  if (mPhysicalAddressBits <= 39 ) {\r
-    NumberOfPml4EntriesNeeded = 1;\r
-    NumberOfPdpEntriesNeeded = (UINT32)LShiftU64 (1, (mPhysicalAddressBits - 30));\r
-  } else {\r
-    NumberOfPml4EntriesNeeded = (UINT32)LShiftU64 (1, (mPhysicalAddressBits - 39));\r
-    NumberOfPdpEntriesNeeded = 512;\r
+  //\r
+  // IA-32e paging translates 48-bit linear addresses to 52-bit physical addresses\r
+  //  when 5-Level Paging is disabled.\r
+  //\r
+  ASSERT (mPhysicalAddressBits <= 52);\r
+  if (!m5LevelPagingNeeded && mPhysicalAddressBits > 48) {\r
+    mPhysicalAddressBits = 48;\r
+  }\r
+\r
+  NumberOfPml5EntriesNeeded = 1;\r
+  if (mPhysicalAddressBits > 48) {\r
+    NumberOfPml5EntriesNeeded = (UINTN) LShiftU64 (1, mPhysicalAddressBits - 48);\r
+    mPhysicalAddressBits = 48;\r
   }\r
 \r
+  NumberOfPml4EntriesNeeded = 1;\r
+  if (mPhysicalAddressBits > 39) {\r
+    NumberOfPml4EntriesNeeded = (UINTN) LShiftU64 (1, mPhysicalAddressBits - 39);\r
+    mPhysicalAddressBits = 39;\r
+  }\r
+\r
+  NumberOfPdpEntriesNeeded = 1;\r
+  ASSERT (mPhysicalAddressBits > 30);\r
+  NumberOfPdpEntriesNeeded = (UINTN) LShiftU64 (1, mPhysicalAddressBits - 30);\r
+\r
   //\r
   // By architecture only one PageMapLevel4 exists - so lets allocate storage for it.\r
   //\r
   PageMap         = (VOID *) PageTable;\r
 \r
   PageMapLevel4Entry = PageMap;\r
-  PageAddress        = 0;\r
-  for (IndexOfPml4Entries = 0; IndexOfPml4Entries < NumberOfPml4EntriesNeeded; IndexOfPml4Entries++, PageMapLevel4Entry++) {\r
+  PageMapLevel5Entry = NULL;\r
+  if (m5LevelPagingNeeded) {\r
     //\r
-    // Each PML4 entry points to a page of Page Directory Pointer entries.\r
+    // By architecture only one PageMapLevel5 exists - so lets allocate storage for it.\r
     //\r
-    PageDirectoryPointerEntry = (UINT64 *) ((*PageMapLevel4Entry) & ~mAddressEncMask & gPhyMask);\r
-    if (PageDirectoryPointerEntry == NULL) {\r
-      PageDirectoryPointerEntry = AllocatePageTableMemory (1);\r
-      ASSERT(PageDirectoryPointerEntry != NULL);\r
-      ZeroMem (PageDirectoryPointerEntry, EFI_PAGES_TO_SIZE(1));\r
+    PageMapLevel5Entry = PageMap;\r
+  }\r
+  PageAddress        = 0;\r
 \r
-      *PageMapLevel4Entry = (UINT64)(UINTN)PageDirectoryPointerEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
+  for ( IndexOfPml5Entries = 0\r
+      ; IndexOfPml5Entries < NumberOfPml5EntriesNeeded\r
+      ; IndexOfPml5Entries++, PageMapLevel5Entry++) {\r
+    //\r
+    // Each PML5 entry points to a page of PML4 entires.\r
+    // So lets allocate space for them and fill them in in the IndexOfPml4Entries loop.\r
+    // When 5-Level Paging is disabled, below allocation happens only once.\r
+    //\r
+    if (m5LevelPagingNeeded) {\r
+      PageMapLevel4Entry = (UINT64 *) ((*PageMapLevel5Entry) & ~mAddressEncMask & gPhyMask);\r
+      if (PageMapLevel4Entry == NULL) {\r
+        PageMapLevel4Entry = AllocatePageTableMemory (1);\r
+        ASSERT(PageMapLevel4Entry != NULL);\r
+        ZeroMem (PageMapLevel4Entry, EFI_PAGES_TO_SIZE(1));\r
+\r
+        *PageMapLevel5Entry = (UINT64)(UINTN)PageMapLevel4Entry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
+      }\r
     }\r
 \r
-    if (m1GPageTableSupport) {\r
-      PageDirectory1GEntry = PageDirectoryPointerEntry;\r
-      for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectory1GEntry++, PageAddress += SIZE_1GB) {\r
-        if (IndexOfPml4Entries == 0 && IndexOfPageDirectoryEntries < 4) {\r
-          //\r
-          // Skip the < 4G entries\r
-          //\r
-          continue;\r
-        }\r
-        //\r
-        // Fill in the Page Directory entries\r
-        //\r
-        *PageDirectory1GEntry = PageAddress | mAddressEncMask | IA32_PG_PS | PAGE_ATTRIBUTE_BITS;\r
+    for (IndexOfPml4Entries = 0; IndexOfPml4Entries < (NumberOfPml5EntriesNeeded == 1 ? NumberOfPml4EntriesNeeded : 512); IndexOfPml4Entries++, PageMapLevel4Entry++) {\r
+      //\r
+      // Each PML4 entry points to a page of Page Directory Pointer entries.\r
+      //\r
+      PageDirectoryPointerEntry = (UINT64 *) ((*PageMapLevel4Entry) & ~mAddressEncMask & gPhyMask);\r
+      if (PageDirectoryPointerEntry == NULL) {\r
+        PageDirectoryPointerEntry = AllocatePageTableMemory (1);\r
+        ASSERT(PageDirectoryPointerEntry != NULL);\r
+        ZeroMem (PageDirectoryPointerEntry, EFI_PAGES_TO_SIZE(1));\r
+\r
+        *PageMapLevel4Entry = (UINT64)(UINTN)PageDirectoryPointerEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
       }\r
-    } else {\r
-      PageAddress = BASE_4GB;\r
-      for (IndexOfPdpEntries = 0; IndexOfPdpEntries < NumberOfPdpEntriesNeeded; IndexOfPdpEntries++, PageDirectoryPointerEntry++) {\r
-        if (IndexOfPml4Entries == 0 && IndexOfPdpEntries < 4) {\r
-          //\r
-          // Skip the < 4G entries\r
-          //\r
-          continue;\r
-        }\r
-        //\r
-        // Each Directory Pointer entries points to a page of Page Directory entires.\r
-        // So allocate space for them and fill them in in the IndexOfPageDirectoryEntries loop.\r
-        //\r
-        PageDirectoryEntry = (UINT64 *) ((*PageDirectoryPointerEntry) & ~mAddressEncMask & gPhyMask);\r
-        if (PageDirectoryEntry == NULL) {\r
-          PageDirectoryEntry = AllocatePageTableMemory (1);\r
-          ASSERT(PageDirectoryEntry != NULL);\r
-          ZeroMem (PageDirectoryEntry, EFI_PAGES_TO_SIZE(1));\r
 \r
+      if (m1GPageTableSupport) {\r
+        PageDirectory1GEntry = PageDirectoryPointerEntry;\r
+        for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectory1GEntry++, PageAddress += SIZE_1GB) {\r
+          if (IndexOfPml4Entries == 0 && IndexOfPageDirectoryEntries < 4) {\r
+            //\r
+            // Skip the < 4G entries\r
+            //\r
+            continue;\r
+          }\r
           //\r
-          // Fill in a Page Directory Pointer Entries\r
+          // Fill in the Page Directory entries\r
           //\r
-          *PageDirectoryPointerEntry = (UINT64)(UINTN)PageDirectoryEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
+          *PageDirectory1GEntry = PageAddress | mAddressEncMask | IA32_PG_PS | PAGE_ATTRIBUTE_BITS;\r
         }\r
-\r
-        for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectoryEntry++, PageAddress += SIZE_2MB) {\r
+      } else {\r
+        PageAddress = BASE_4GB;\r
+        for (IndexOfPdpEntries = 0; IndexOfPdpEntries < (NumberOfPml4EntriesNeeded == 1 ? NumberOfPdpEntriesNeeded : 512); IndexOfPdpEntries++, PageDirectoryPointerEntry++) {\r
+          if (IndexOfPml4Entries == 0 && IndexOfPdpEntries < 4) {\r
+            //\r
+            // Skip the < 4G entries\r
+            //\r
+            continue;\r
+          }\r
           //\r
-          // Fill in the Page Directory entries\r
+          // Each Directory Pointer entries points to a page of Page Directory entires.\r
+          // So allocate space for them and fill them in in the IndexOfPageDirectoryEntries loop.\r
           //\r
-          *PageDirectoryEntry = PageAddress | mAddressEncMask | IA32_PG_PS | PAGE_ATTRIBUTE_BITS;\r
+          PageDirectoryEntry = (UINT64 *) ((*PageDirectoryPointerEntry) & ~mAddressEncMask & gPhyMask);\r
+          if (PageDirectoryEntry == NULL) {\r
+            PageDirectoryEntry = AllocatePageTableMemory (1);\r
+            ASSERT(PageDirectoryEntry != NULL);\r
+            ZeroMem (PageDirectoryEntry, EFI_PAGES_TO_SIZE(1));\r
+\r
+            //\r
+            // Fill in a Page Directory Pointer Entries\r
+            //\r
+            *PageDirectoryPointerEntry = (UINT64)(UINTN)PageDirectoryEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
+          }\r
+\r
+          for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectoryEntry++, PageAddress += SIZE_2MB) {\r
+            //\r
+            // Fill in the Page Directory entries\r
+            //\r
+            *PageDirectoryEntry = PageAddress | mAddressEncMask | IA32_PG_PS | PAGE_ATTRIBUTE_BITS;\r
+          }\r
         }\r
       }\r
     }\r
@@ -247,19 +343,23 @@ SmmInitPageTable (
   UINTN                             PageFaultHandlerHookAddress;\r
   IA32_IDT_GATE_DESCRIPTOR          *IdtEntry;\r
   EFI_STATUS                        Status;\r
+  UINT64                            *Pml4Entry;\r
+  UINT64                            *Pml5Entry;\r
 \r
   //\r
   // Initialize spin lock\r
   //\r
   InitializeSpinLock (mPFLock);\r
 \r
-  mCpuSmmStaticPageTable = PcdGetBool (PcdCpuSmmStaticPageTable);\r
-  m1GPageTableSupport = Is1GPageSupport ();\r
-  DEBUG ((DEBUG_INFO, "1GPageTableSupport - 0x%x\n", m1GPageTableSupport));\r
-  DEBUG ((DEBUG_INFO, "PcdCpuSmmStaticPageTable - 0x%x\n", mCpuSmmStaticPageTable));\r
-\r
-  mPhysicalAddressBits = CalculateMaximumSupportAddress ();\r
-  DEBUG ((DEBUG_INFO, "PhysicalAddressBits - 0x%x\n", mPhysicalAddressBits));\r
+  mCpuSmmRestrictedMemoryAccess = PcdGetBool (PcdCpuSmmRestrictedMemoryAccess);\r
+  m1GPageTableSupport           = Is1GPageSupport ();\r
+  m5LevelPagingNeeded           = Is5LevelPagingNeeded ();\r
+  mPhysicalAddressBits          = CalculateMaximumSupportAddress ();\r
+  PatchInstructionX86 (gPatch5LevelPagingNeeded, m5LevelPagingNeeded, 1);\r
+  DEBUG ((DEBUG_INFO, "5LevelPaging Needed             - %d\n", m5LevelPagingNeeded));\r
+  DEBUG ((DEBUG_INFO, "1GPageTable Support             - %d\n", m1GPageTableSupport));\r
+  DEBUG ((DEBUG_INFO, "PcdCpuSmmRestrictedMemoryAccess - %d\n", mCpuSmmRestrictedMemoryAccess));\r
+  DEBUG ((DEBUG_INFO, "PhysicalAddressBits             - %d\n", mPhysicalAddressBits));\r
   //\r
   // Generate PAE page table for the first 4GB memory space\r
   //\r
@@ -276,17 +376,37 @@ SmmInitPageTable (
   //\r
   // Fill Page-Table-Level4 (PML4) entry\r
   //\r
-  PTEntry = (UINT64*)AllocatePageTableMemory (1);\r
-  ASSERT (PTEntry != NULL);\r
-  *PTEntry = Pages | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
-  ZeroMem (PTEntry + 1, EFI_PAGE_SIZE - sizeof (*PTEntry));\r
+  Pml4Entry = (UINT64*)AllocatePageTableMemory (1);\r
+  ASSERT (Pml4Entry != NULL);\r
+  *Pml4Entry = Pages | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
+  ZeroMem (Pml4Entry + 1, EFI_PAGE_SIZE - sizeof (*Pml4Entry));\r
 \r
   //\r
   // Set sub-entries number\r
   //\r
-  SetSubEntriesNum (PTEntry, 3);\r
+  SetSubEntriesNum (Pml4Entry, 3);\r
+  PTEntry = Pml4Entry;\r
 \r
-  if (mCpuSmmStaticPageTable) {\r
+  if (m5LevelPagingNeeded) {\r
+    //\r
+    // Fill PML5 entry\r
+    //\r
+    Pml5Entry = (UINT64*)AllocatePageTableMemory (1);\r
+    ASSERT (Pml5Entry != NULL);\r
+    *Pml5Entry = (UINTN) Pml4Entry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
+    ZeroMem (Pml5Entry + 1, EFI_PAGE_SIZE - sizeof (*Pml5Entry));\r
+    //\r
+    // Set sub-entries number\r
+    //\r
+    SetSubEntriesNum (Pml5Entry, 1);\r
+    PTEntry = Pml5Entry;\r
+  }\r
+\r
+  if (mCpuSmmRestrictedMemoryAccess) {\r
+    //\r
+    // When access to non-SMRAM memory is restricted, create page table\r
+    // that covers all memory space.\r
+    //\r
     SetStaticPageTable ((UINTN)PTEntry);\r
   } else {\r
     //\r
@@ -300,7 +420,9 @@ SmmInitPageTable (
     }\r
   }\r
 \r
-  if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {\r
+  if (FeaturePcdGet (PcdCpuSmmProfileEnable) ||\r
+      HEAP_GUARD_NONSTOP_MODE ||\r
+      NULL_DETECTION_NONSTOP_MODE) {\r
     //\r
     // Set own Page Fault entry instead of the default one, because SMM Profile\r
     // feature depends on IRET instruction to do Single Step\r
@@ -330,7 +452,7 @@ SmmInitPageTable (
   }\r
 \r
   //\r
-  // Return the address of PML4 (to set CR3)\r
+  // Return the address of PML4/PML5 (to set CR3)\r
   //\r
   return (UINT32)(UINTN)PTEntry;\r
 }\r
@@ -422,12 +544,16 @@ ReclaimPages (
   VOID\r
   )\r
 {\r
+  UINT64                       Pml5Entry;\r
+  UINT64                       *Pml5;\r
   UINT64                       *Pml4;\r
   UINT64                       *Pdpt;\r
   UINT64                       *Pdt;\r
+  UINTN                        Pml5Index;\r
   UINTN                        Pml4Index;\r
   UINTN                        PdptIndex;\r
   UINTN                        PdtIndex;\r
+  UINTN                        MinPml5;\r
   UINTN                        MinPml4;\r
   UINTN                        MinPdpt;\r
   UINTN                        MinPdt;\r
@@ -437,122 +563,167 @@ ReclaimPages (
   BOOLEAN                      PML4EIgnore;\r
   BOOLEAN                      PDPTEIgnore;\r
   UINT64                       *ReleasePageAddress;\r
+  IA32_CR4                     Cr4;\r
+  BOOLEAN                      Enable5LevelPaging;\r
+  UINT64                       PFAddress;\r
+  UINT64                       PFAddressPml5Index;\r
+  UINT64                       PFAddressPml4Index;\r
+  UINT64                       PFAddressPdptIndex;\r
+  UINT64                       PFAddressPdtIndex;\r
 \r
   Pml4 = NULL;\r
   Pdpt = NULL;\r
   Pdt  = NULL;\r
   MinAcc  = (UINT64)-1;\r
   MinPml4 = (UINTN)-1;\r
+  MinPml5 = (UINTN)-1;\r
   MinPdpt = (UINTN)-1;\r
   MinPdt  = (UINTN)-1;\r
   Acc     = 0;\r
   ReleasePageAddress = 0;\r
+  PFAddress = AsmReadCr2 ();\r
+  PFAddressPml5Index = BitFieldRead64 (PFAddress, 48, 48 + 8);\r
+  PFAddressPml4Index = BitFieldRead64 (PFAddress, 39, 39 + 8);\r
+  PFAddressPdptIndex = BitFieldRead64 (PFAddress, 30, 30 + 8);\r
+  PFAddressPdtIndex = BitFieldRead64 (PFAddress, 21, 21 + 8);\r
+\r
+  Cr4.UintN = AsmReadCr4 ();\r
+  Enable5LevelPaging = (BOOLEAN) (Cr4.Bits.LA57 == 1);\r
+  Pml5 = (UINT64*)(UINTN)(AsmReadCr3 () & gPhyMask);\r
+\r
+  if (!Enable5LevelPaging) {\r
+    //\r
+    // Create one fake PML5 entry for 4-Level Paging\r
+    // so that the page table parsing logic only handles 5-Level page structure.\r
+    //\r
+    Pml5Entry = (UINTN) Pml5 | IA32_PG_P;\r
+    Pml5 = &Pml5Entry;\r
+  }\r
 \r
   //\r
   // First, find the leaf entry has the smallest access record value\r
   //\r
-  Pml4 = (UINT64*)(UINTN)(AsmReadCr3 () & gPhyMask);\r
-  for (Pml4Index = 0; Pml4Index < EFI_PAGE_SIZE / sizeof (*Pml4); Pml4Index++) {\r
-    if ((Pml4[Pml4Index] & IA32_PG_P) == 0 || (Pml4[Pml4Index] & IA32_PG_PMNT) != 0) {\r
+  for (Pml5Index = 0; Pml5Index < (Enable5LevelPaging ? (EFI_PAGE_SIZE / sizeof (*Pml4)) : 1); Pml5Index++) {\r
+    if ((Pml5[Pml5Index] & IA32_PG_P) == 0 || (Pml5[Pml5Index] & IA32_PG_PMNT) != 0) {\r
       //\r
-      // If the PML4 entry is not present or is masked, skip it\r
+      // If the PML5 entry is not present or is masked, skip it\r
       //\r
       continue;\r
     }\r
-    Pdpt = (UINT64*)(UINTN)(Pml4[Pml4Index] & ~mAddressEncMask & gPhyMask);\r
-    PML4EIgnore = FALSE;\r
-    for (PdptIndex = 0; PdptIndex < EFI_PAGE_SIZE / sizeof (*Pdpt); PdptIndex++) {\r
-      if ((Pdpt[PdptIndex] & IA32_PG_P) == 0 || (Pdpt[PdptIndex] & IA32_PG_PMNT) != 0) {\r
+    Pml4 = (UINT64*)(UINTN)(Pml5[Pml5Index] & gPhyMask);\r
+    for (Pml4Index = 0; Pml4Index < EFI_PAGE_SIZE / sizeof (*Pml4); Pml4Index++) {\r
+      if ((Pml4[Pml4Index] & IA32_PG_P) == 0 || (Pml4[Pml4Index] & IA32_PG_PMNT) != 0) {\r
         //\r
-        // If the PDPT entry is not present or is masked, skip it\r
+        // If the PML4 entry is not present or is masked, skip it\r
         //\r
-        if ((Pdpt[PdptIndex] & IA32_PG_PMNT) != 0) {\r
-          //\r
-          // If the PDPT entry is masked, we will ignore checking the PML4 entry\r
-          //\r
-          PML4EIgnore = TRUE;\r
-        }\r
         continue;\r
       }\r
-      if ((Pdpt[PdptIndex] & IA32_PG_PS) == 0) {\r
-        //\r
-        // It's not 1-GByte pages entry, it should be a PDPT entry,\r
-        // we will not check PML4 entry more\r
-        //\r
-        PML4EIgnore = TRUE;\r
-        Pdt =  (UINT64*)(UINTN)(Pdpt[PdptIndex] & ~mAddressEncMask & gPhyMask);\r
-        PDPTEIgnore = FALSE;\r
-        for (PdtIndex = 0; PdtIndex < EFI_PAGE_SIZE / sizeof(*Pdt); PdtIndex++) {\r
-          if ((Pdt[PdtIndex] & IA32_PG_P) == 0 || (Pdt[PdtIndex] & IA32_PG_PMNT) != 0) {\r
+      Pdpt = (UINT64*)(UINTN)(Pml4[Pml4Index] & ~mAddressEncMask & gPhyMask);\r
+      PML4EIgnore = FALSE;\r
+      for (PdptIndex = 0; PdptIndex < EFI_PAGE_SIZE / sizeof (*Pdpt); PdptIndex++) {\r
+        if ((Pdpt[PdptIndex] & IA32_PG_P) == 0 || (Pdpt[PdptIndex] & IA32_PG_PMNT) != 0) {\r
+          //\r
+          // If the PDPT entry is not present or is masked, skip it\r
+          //\r
+          if ((Pdpt[PdptIndex] & IA32_PG_PMNT) != 0) {\r
             //\r
-            // If the PD entry is not present or is masked, skip it\r
+            // If the PDPT entry is masked, we will ignore checking the PML4 entry\r
             //\r
-            if ((Pdt[PdtIndex] & IA32_PG_PMNT) != 0) {\r
+            PML4EIgnore = TRUE;\r
+          }\r
+          continue;\r
+        }\r
+        if ((Pdpt[PdptIndex] & IA32_PG_PS) == 0) {\r
+          //\r
+          // It's not 1-GByte pages entry, it should be a PDPT entry,\r
+          // we will not check PML4 entry more\r
+          //\r
+          PML4EIgnore = TRUE;\r
+          Pdt = (UINT64*)(UINTN)(Pdpt[PdptIndex] & ~mAddressEncMask & gPhyMask);\r
+          PDPTEIgnore = FALSE;\r
+          for (PdtIndex = 0; PdtIndex < EFI_PAGE_SIZE / sizeof(*Pdt); PdtIndex++) {\r
+            if ((Pdt[PdtIndex] & IA32_PG_P) == 0 || (Pdt[PdtIndex] & IA32_PG_PMNT) != 0) {\r
+              //\r
+              // If the PD entry is not present or is masked, skip it\r
               //\r
-              // If the PD entry is masked, we will not PDPT entry more\r
+              if ((Pdt[PdtIndex] & IA32_PG_PMNT) != 0) {\r
+                //\r
+                // If the PD entry is masked, we will not PDPT entry more\r
+                //\r
+                PDPTEIgnore = TRUE;\r
+              }\r
+              continue;\r
+            }\r
+            if ((Pdt[PdtIndex] & IA32_PG_PS) == 0) {\r
+              //\r
+              // It's not 2 MByte page table entry, it should be PD entry\r
+              // we will find the entry has the smallest access record value\r
               //\r
               PDPTEIgnore = TRUE;\r
+              if (PdtIndex != PFAddressPdtIndex || PdptIndex != PFAddressPdptIndex ||\r
+                  Pml4Index != PFAddressPml4Index || Pml5Index != PFAddressPml5Index) {\r
+                Acc = GetAndUpdateAccNum (Pdt + PdtIndex);\r
+                if (Acc < MinAcc) {\r
+                  //\r
+                  // If the PD entry has the smallest access record value,\r
+                  // save the Page address to be released\r
+                  //\r
+                  MinAcc  = Acc;\r
+                  MinPml5 = Pml5Index;\r
+                  MinPml4 = Pml4Index;\r
+                  MinPdpt = PdptIndex;\r
+                  MinPdt  = PdtIndex;\r
+                  ReleasePageAddress = Pdt + PdtIndex;\r
+                }\r
+              }\r
             }\r
-            continue;\r
           }\r
-          if ((Pdt[PdtIndex] & IA32_PG_PS) == 0) {\r
+          if (!PDPTEIgnore) {\r
             //\r
-            // It's not 2 MByte page table entry, it should be PD entry\r
-            // we will find the entry has the smallest access record value\r
+            // If this PDPT entry has no PDT entries pointer to 4 KByte pages,\r
+            // it should only has the entries point to 2 MByte Pages\r
             //\r
-            PDPTEIgnore = TRUE;\r
-            Acc = GetAndUpdateAccNum (Pdt + PdtIndex);\r
-            if (Acc < MinAcc) {\r
-              //\r
-              // If the PD entry has the smallest access record value,\r
-              // save the Page address to be released\r
-              //\r
-              MinAcc  = Acc;\r
-              MinPml4 = Pml4Index;\r
-              MinPdpt = PdptIndex;\r
-              MinPdt  = PdtIndex;\r
-              ReleasePageAddress = Pdt + PdtIndex;\r
+            if (PdptIndex != PFAddressPdptIndex || Pml4Index != PFAddressPml4Index ||\r
+                Pml5Index != PFAddressPml5Index) {\r
+              Acc = GetAndUpdateAccNum (Pdpt + PdptIndex);\r
+              if (Acc < MinAcc) {\r
+                //\r
+                // If the PDPT entry has the smallest access record value,\r
+                // save the Page address to be released\r
+                //\r
+                MinAcc  = Acc;\r
+                MinPml5 = Pml5Index;\r
+                MinPml4 = Pml4Index;\r
+                MinPdpt = PdptIndex;\r
+                MinPdt  = (UINTN)-1;\r
+                ReleasePageAddress = Pdpt + PdptIndex;\r
+              }\r
             }\r
           }\r
         }\r
-        if (!PDPTEIgnore) {\r
-          //\r
-          // If this PDPT entry has no PDT entries pointer to 4 KByte pages,\r
-          // it should only has the entries point to 2 MByte Pages\r
-          //\r
-          Acc = GetAndUpdateAccNum (Pdpt + PdptIndex);\r
+      }\r
+      if (!PML4EIgnore) {\r
+        //\r
+        // If PML4 entry has no the PDPT entry pointer to 2 MByte pages,\r
+        // it should only has the entries point to 1 GByte Pages\r
+        //\r
+        if (Pml4Index != PFAddressPml4Index || Pml5Index != PFAddressPml5Index) {\r
+          Acc = GetAndUpdateAccNum (Pml4 + Pml4Index);\r
           if (Acc < MinAcc) {\r
             //\r
-            // If the PDPT entry has the smallest access record value,\r
+            // If the PML4 entry has the smallest access record value,\r
             // save the Page address to be released\r
             //\r
             MinAcc  = Acc;\r
+            MinPml5 = Pml5Index;\r
             MinPml4 = Pml4Index;\r
-            MinPdpt = PdptIndex;\r
+            MinPdpt = (UINTN)-1;\r
             MinPdt  = (UINTN)-1;\r
-            ReleasePageAddress = Pdpt + PdptIndex;\r
+            ReleasePageAddress = Pml4 + Pml4Index;\r
           }\r
         }\r
       }\r
     }\r
-    if (!PML4EIgnore) {\r
-      //\r
-      // If PML4 entry has no the PDPT entry pointer to 2 MByte pages,\r
-      // it should only has the entries point to 1 GByte Pages\r
-      //\r
-      Acc = GetAndUpdateAccNum (Pml4 + Pml4Index);\r
-      if (Acc < MinAcc) {\r
-        //\r
-        // If the PML4 entry has the smallest access record value,\r
-        // save the Page address to be released\r
-        //\r
-        MinAcc  = Acc;\r
-        MinPml4 = Pml4Index;\r
-        MinPdpt = (UINTN)-1;\r
-        MinPdt  = (UINTN)-1;\r
-        ReleasePageAddress = Pml4 + Pml4Index;\r
-      }\r
-    }\r
   }\r
   //\r
   // Make sure one PML4/PDPT/PD entry is selected\r
@@ -574,9 +745,11 @@ ReclaimPages (
       //\r
       // If 4 KByte Page Table is released, check the PDPT entry\r
       //\r
+      Pml4 = (UINT64 *) (UINTN) (Pml5[MinPml5] & gPhyMask);\r
       Pdpt = (UINT64*)(UINTN)(Pml4[MinPml4] & ~mAddressEncMask & gPhyMask);\r
       SubEntriesNum = GetSubEntriesNum(Pdpt + MinPdpt);\r
-      if (SubEntriesNum == 0) {\r
+      if (SubEntriesNum == 0 &&\r
+          (MinPdpt != PFAddressPdptIndex || MinPml4 != PFAddressPml4Index || MinPml5 != PFAddressPml5Index)) {\r
         //\r
         // Release the empty Page Directory table if there was no more 4 KByte Page Table entry\r
         // clear the Page directory entry\r
@@ -592,7 +765,7 @@ ReclaimPages (
       //\r
       // Update the sub-entries filed in PDPT entry and exit\r
       //\r
-      SetSubEntriesNum (Pdpt + MinPdpt, SubEntriesNum - 1);\r
+      SetSubEntriesNum (Pdpt + MinPdpt, (SubEntriesNum - 1) & 0x1FF);\r
       break;\r
     }\r
     if (MinPdpt != (UINTN)-1) {\r
@@ -600,7 +773,7 @@ ReclaimPages (
       // One 2MB Page Table is released or Page Directory table is released, check the PML4 entry\r
       //\r
       SubEntriesNum = GetSubEntriesNum (Pml4 + MinPml4);\r
-      if (SubEntriesNum == 0) {\r
+      if (SubEntriesNum == 0 && (MinPml4 != PFAddressPml4Index || MinPml5 != PFAddressPml5Index)) {\r
         //\r
         // Release the empty PML4 table if there was no more 1G KByte Page Table entry\r
         // clear the Page directory entry\r
@@ -613,7 +786,7 @@ ReclaimPages (
       //\r
       // Update the sub-entries filed in PML4 entry and exit\r
       //\r
-      SetSubEntriesNum (Pml4 + MinPml4, SubEntriesNum - 1);\r
+      SetSubEntriesNum (Pml4 + MinPml4, (SubEntriesNum - 1) & 0x1FF);\r
       break;\r
     }\r
     //\r
@@ -665,7 +838,7 @@ SmiDefaultPFHandler (
   )\r
 {\r
   UINT64                            *PageTable;\r
-  UINT64                            *Pml4;\r
+  UINT64                            *PageTableTop;\r
   UINT64                            PFAddress;\r
   UINTN                             StartBit;\r
   UINTN                             EndBit;\r
@@ -676,6 +849,8 @@ SmiDefaultPFHandler (
   UINTN                             PageAttribute;\r
   EFI_STATUS                        Status;\r
   UINT64                            *UpperEntry;\r
+  BOOLEAN                           Enable5LevelPaging;\r
+  IA32_CR4                          Cr4;\r
 \r
   //\r
   // Set default SMM page attribute\r
@@ -685,9 +860,12 @@ SmiDefaultPFHandler (
   PageAttribute = 0;\r
 \r
   EndBit = 0;\r
-  Pml4 = (UINT64*)(AsmReadCr3 () & gPhyMask);\r
+  PageTableTop = (UINT64*)(AsmReadCr3 () & gPhyMask);\r
   PFAddress = AsmReadCr2 ();\r
 \r
+  Cr4.UintN = AsmReadCr4 ();\r
+  Enable5LevelPaging = (BOOLEAN) (Cr4.Bits.LA57 != 0);\r
+\r
   Status = GetPlatformPageTableAttribute (PFAddress, &PageSize, &NumOfPages, &PageAttribute);\r
   //\r
   // If platform not support page table attribute, set default SMM page attribute\r
@@ -741,9 +919,9 @@ SmiDefaultPFHandler (
   }\r
 \r
   for (Index = 0; Index < NumOfPages; Index++) {\r
-    PageTable  = Pml4;\r
+    PageTable  = PageTableTop;\r
     UpperEntry = NULL;\r
-    for (StartBit = 39; StartBit > EndBit; StartBit -= 9) {\r
+    for (StartBit = Enable5LevelPaging ? 48 : 39; StartBit > EndBit; StartBit -= 9) {\r
       PTIndex = BitFieldRead64 (PFAddress, StartBit, StartBit + 8);\r
       if ((PageTable[PTIndex] & IA32_PG_P) == 0) {\r
         //\r
@@ -781,7 +959,7 @@ SmiDefaultPFHandler (
     PageTable[PTIndex] = ((PFAddress | mAddressEncMask) & gPhyMask & ~((1ull << EndBit) - 1)) |\r
                          PageAttribute | IA32_PG_A | PAGE_ATTRIBUTE_BITS;\r
     if (UpperEntry != NULL) {\r
-      SetSubEntriesNum (UpperEntry, GetSubEntriesNum (UpperEntry) + 1);\r
+      SetSubEntriesNum (UpperEntry, (GetSubEntriesNum (UpperEntry) + 1) & 0x1FF);\r
     }\r
     //\r
     // Get the next page address if we need to create more page tables\r
@@ -815,10 +993,11 @@ SmiPFHandler (
 \r
   PFAddress = AsmReadCr2 ();\r
 \r
-  if (mCpuSmmStaticPageTable && (PFAddress >= LShiftU64 (1, (mPhysicalAddressBits - 1)))) {\r
+  if (mCpuSmmRestrictedMemoryAccess && (PFAddress >= LShiftU64 (1, (mPhysicalAddressBits - 1)))) {\r
     DumpCpuContext (InterruptType, SystemContext);\r
     DEBUG ((DEBUG_ERROR, "Do not support address 0x%lx by processor!\n", PFAddress));\r
     CpuDeadLoop ();\r
+    goto Exit;\r
   }\r
 \r
   //\r
@@ -846,42 +1025,60 @@ SmiPFHandler (
           DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);\r
         );\r
       }\r
+\r
+      if (HEAP_GUARD_NONSTOP_MODE) {\r
+        GuardPagePFHandler (SystemContext.SystemContextX64->ExceptionData);\r
+        goto Exit;\r
+      }\r
     }\r
     CpuDeadLoop ();\r
+    goto Exit;\r
   }\r
 \r
   //\r
-  // If a page fault occurs in SMM range\r
+  // If a page fault occurs in non-SMRAM range.\r
   //\r
   if ((PFAddress < mCpuHotPlugData.SmrrBase) ||\r
       (PFAddress >= mCpuHotPlugData.SmrrBase + mCpuHotPlugData.SmrrSize)) {\r
-    DumpCpuContext (InterruptType, SystemContext);\r
     if ((SystemContext.SystemContextX64->ExceptionData & IA32_PF_EC_ID) != 0) {\r
+      DumpCpuContext (InterruptType, SystemContext);\r
       DEBUG ((DEBUG_ERROR, "Code executed on IP(0x%lx) out of SMM range after SMM is locked!\n", PFAddress));\r
       DEBUG_CODE (\r
         DumpModuleInfoByIp (*(UINTN *)(UINTN)SystemContext.SystemContextX64->Rsp);\r
       );\r
       CpuDeadLoop ();\r
+      goto Exit;\r
     }\r
-    if (IsSmmCommBufferForbiddenAddress (PFAddress)) {\r
-      DEBUG ((DEBUG_ERROR, "Access SMM communication forbidden address (0x%lx)!\n", PFAddress));\r
+\r
+    //\r
+    // If NULL pointer was just accessed\r
+    //\r
+    if ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0 &&\r
+        (PFAddress < EFI_PAGE_SIZE)) {\r
+      DumpCpuContext (InterruptType, SystemContext);\r
+      DEBUG ((DEBUG_ERROR, "!!! NULL pointer access !!!\n"));\r
       DEBUG_CODE (\r
         DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);\r
       );\r
+\r
+      if (NULL_DETECTION_NONSTOP_MODE) {\r
+        GuardPagePFHandler (SystemContext.SystemContextX64->ExceptionData);\r
+        goto Exit;\r
+      }\r
+\r
       CpuDeadLoop ();\r
+      goto Exit;\r
     }\r
-  }\r
 \r
-  //\r
-  // If NULL pointer was just accessed\r
-  //\r
-  if ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0 &&\r
-      (PFAddress < EFI_PAGE_SIZE)) {\r
-    DEBUG ((DEBUG_ERROR, "!!! NULL pointer access !!!\n"));\r
-    DEBUG_CODE (\r
-      DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);\r
-    );\r
-    CpuDeadLoop ();\r
+    if (mCpuSmmRestrictedMemoryAccess && IsSmmCommBufferForbiddenAddress (PFAddress)) {\r
+      DumpCpuContext (InterruptType, SystemContext);\r
+      DEBUG ((DEBUG_ERROR, "Access SMM communication forbidden address (0x%lx)!\n", PFAddress));\r
+      DEBUG_CODE (\r
+        DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);\r
+      );\r
+      CpuDeadLoop ();\r
+      goto Exit;\r
+    }\r
   }\r
 \r
   if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {\r
@@ -893,6 +1090,7 @@ SmiPFHandler (
     SmiDefaultPFHandler ();\r
   }\r
 \r
+Exit:\r
   ReleaseSpinLock (mPFLock);\r
 }\r
 \r
@@ -907,14 +1105,42 @@ SetPageTableAttributes (
   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
   BOOLEAN               IsSplitted;\r
   BOOLEAN               PageTableSplitted;\r
+  BOOLEAN               CetEnabled;\r
+  IA32_CR4              Cr4;\r
+  BOOLEAN               Enable5LevelPaging;\r
+\r
+  Cr4.UintN = AsmReadCr4 ();\r
+  Enable5LevelPaging = (BOOLEAN) (Cr4.Bits.LA57 == 1);\r
 \r
-  if (!mCpuSmmStaticPageTable) {\r
+  //\r
+  // Don't mark page table memory as read-only if\r
+  //  - no restriction on access to non-SMRAM memory; or\r
+  //  - SMM heap guard feature enabled; or\r
+  //      BIT2: SMM page guard enabled\r
+  //      BIT3: SMM pool guard enabled\r
+  //  - SMM profile feature enabled\r
+  //\r
+  if (!mCpuSmmRestrictedMemoryAccess ||\r
+      ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0) ||\r
+      FeaturePcdGet (PcdCpuSmmProfileEnable)) {\r
+    //\r
+    // Restriction on access to non-SMRAM memory and heap guard could not be enabled at the same time.\r
+    //\r
+    ASSERT (!(mCpuSmmRestrictedMemoryAccess &&\r
+              (PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0));\r
+\r
+    //\r
+    // Restriction on access to non-SMRAM memory and SMM profile could not be enabled at the same time.\r
+    //\r
+    ASSERT (!(mCpuSmmRestrictedMemoryAccess && FeaturePcdGet (PcdCpuSmmProfileEnable)));\r
     return ;\r
   }\r
 \r
@@ -924,49 +1150,71 @@ SetPageTableAttributes (
   // Disable write protection, because we need mark page table to be write protected.\r
   // We need *write* page table memory, to mark itself to be *read only*.\r
   //\r
+  CetEnabled = ((AsmReadCr4() & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;\r
+  if (CetEnabled) {\r
+    //\r
+    // CET must be disabled if WP is disabled.\r
+    //\r
+    DisableCet();\r
+  }\r
   AsmWriteCr0 (AsmReadCr0() & ~CR0_WP);\r
 \r
   do {\r
     DEBUG ((DEBUG_INFO, "Start...\n"));\r
     PageTableSplitted = FALSE;\r
-\r
-    L4PageTable = (UINT64 *)GetPageTableBase ();\r
-    SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L4PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
-    PageTableSplitted = (PageTableSplitted || IsSplitted);\r
-\r
-    for (Index4 = 0; Index4 < SIZE_4KB/sizeof(UINT64); Index4++) {\r
-      L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
-      if (L3PageTable == NULL) {\r
-        continue;\r
-      }\r
-\r
-      SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L3PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
+    L5PageTable = NULL;\r
+    if (Enable5LevelPaging) {\r
+      L5PageTable = (UINT64 *)GetPageTableBase ();\r
+      SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L5PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
       PageTableSplitted = (PageTableSplitted || IsSplitted);\r
+    }\r
 \r
-      for (Index3 = 0; Index3 < SIZE_4KB/sizeof(UINT64); Index3++) {\r
-        if ((L3PageTable[Index3] & IA32_PG_PS) != 0) {\r
-          // 1G\r
+    for (Index5 = 0; Index5 < (Enable5LevelPaging ? SIZE_4KB/sizeof(UINT64) : 1); Index5++) {\r
+      if (Enable5LevelPaging) {\r
+        L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
+        if (L4PageTable == NULL) {\r
           continue;\r
         }\r
-        L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
-        if (L2PageTable == NULL) {\r
+      } else {\r
+        L4PageTable = (UINT64 *)GetPageTableBase ();\r
+      }\r
+      SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L4PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
+      PageTableSplitted = (PageTableSplitted || IsSplitted);\r
+\r
+      for (Index4 = 0; Index4 < SIZE_4KB/sizeof(UINT64); Index4++) {\r
+        L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
+        if (L3PageTable == NULL) {\r
           continue;\r
         }\r
 \r
-        SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L2PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
+        SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L3PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
         PageTableSplitted = (PageTableSplitted || IsSplitted);\r
 \r
-        for (Index2 = 0; Index2 < SIZE_4KB/sizeof(UINT64); Index2++) {\r
-          if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {\r
-            // 2M\r
+        for (Index3 = 0; Index3 < SIZE_4KB/sizeof(UINT64); Index3++) {\r
+          if ((L3PageTable[Index3] & IA32_PG_PS) != 0) {\r
+            // 1G\r
             continue;\r
           }\r
-          L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
-          if (L1PageTable == NULL) {\r
+          L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
+          if (L2PageTable == NULL) {\r
             continue;\r
           }\r
-          SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L1PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
+\r
+          SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L2PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
           PageTableSplitted = (PageTableSplitted || IsSplitted);\r
+\r
+          for (Index2 = 0; Index2 < SIZE_4KB/sizeof(UINT64); Index2++) {\r
+            if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {\r
+              // 2M\r
+              continue;\r
+            }\r
+            L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
+            if (L1PageTable == NULL) {\r
+              continue;\r
+            }\r
+            SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L1PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);\r
+            PageTableSplitted = (PageTableSplitted || IsSplitted);\r
+          }\r
         }\r
       }\r
     }\r
@@ -976,6 +1224,62 @@ SetPageTableAttributes (
   // Enable write protection, after page table updated.\r
   //\r
   AsmWriteCr0 (AsmReadCr0() | CR0_WP);\r
+  if (CetEnabled) {\r
+    //\r
+    // re-enable CET.\r
+    //\r
+    EnableCet();\r
+  }\r
 \r
   return ;\r
 }\r
+\r
+/**\r
+  This function reads CR2 register when on-demand paging is enabled.\r
+\r
+  @param[out]  *Cr2  Pointer to variable to hold CR2 register value.\r
+**/\r
+VOID\r
+SaveCr2 (\r
+  OUT UINTN  *Cr2\r
+  )\r
+{\r
+  if (!mCpuSmmRestrictedMemoryAccess) {\r
+    //\r
+    // On-demand paging is enabled when access to non-SMRAM is not restricted.\r
+    //\r
+    *Cr2 = AsmReadCr2 ();\r
+  }\r
+}\r
+\r
+/**\r
+  This function restores CR2 register when on-demand paging is enabled.\r
+\r
+  @param[in]  Cr2  Value to write into CR2 register.\r
+**/\r
+VOID\r
+RestoreCr2 (\r
+  IN UINTN  Cr2\r
+  )\r
+{\r
+  if (!mCpuSmmRestrictedMemoryAccess) {\r
+    //\r
+    // On-demand paging is enabled when access to non-SMRAM is not restricted.\r
+    //\r
+    AsmWriteCr2 (Cr2);\r
+  }\r
+}\r
+\r
+/**\r
+  Return whether access to non-SMRAM is restricted.\r
+\r
+  @retval TRUE  Access to non-SMRAM is restricted.\r
+  @retval FALSE Access to non-SMRAM is not restricted.\r
+**/\r
+BOOLEAN\r
+IsRestrictedMemoryAccess (\r
+  VOID\r
+  )\r
+{\r
+  return mCpuSmmRestrictedMemoryAccess;\r
+}\r