1 files changed, 24 insertions, 2 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
index 2483f2ea84..9c8e2d15ac 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
@@ -29,6 +29,26 @@ EnableCet (
   );
 
 /**
+  Get page table base address and the depth of the page table.
+
+  @param[out] Base        Page table base address.
+  @param[out] FiveLevels  TRUE means 5 level paging. FALSE means 4 level paging.
+**/
+VOID
+GetPageTable (
+  OUT UINTN   *Base,
+  OUT BOOLEAN *FiveLevels OPTIONAL
+  )
+{
+  *Base = ((mInternalCr3 == 0) ?
+            (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64) :
+            mInternalCr3);
+  if (FiveLevels != NULL) {
+    *FiveLevels = FALSE;
+  }
+}
+
+/**
   Create PageTable for SMM use.
 
   @return     PageTable Address
@@ -226,6 +246,7 @@ SetPageTableAttributes (
   UINT64                *L1PageTable;
   UINT64                *L2PageTable;
   UINT64                *L3PageTable;
+  UINTN                 PageTableBase;
   BOOLEAN               IsSplitted;
   BOOLEAN               PageTableSplitted;
   BOOLEAN               CetEnabled;
@@ -268,9 +289,10 @@ SetPageTableAttributes (
     DEBUG ((DEBUG_INFO, "Start...\n"));
     PageTableSplitted = FALSE;
 
-    L3PageTable = (UINT64 *)GetPageTableBase ();
+    GetPageTable (&PageTableBase, NULL);
+    L3PageTable = (UINT64 *)PageTableBase;
 
-    SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)(UINTN)L3PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
+    SmmSetMemoryAttributesEx ((EFI_PHYSICAL_ADDRESS)PageTableBase, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
     PageTableSplitted = (PageTableSplitted || IsSplitted);
 
     for (Index3 = 0; Index3 < 4; Index3++) {