1 files changed, 35 insertions, 2 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
index 670a5cf663..fb4655128f 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
@@ -970,9 +970,9 @@ PiCpuSmmEntry (
   //
   BufferPages = EFI_SIZE_TO_PAGES (SIZE_32KB + TileSize * (mMaxNumberOfCpus - 1));
   if ((FamilyId == 4) || (FamilyId == 5)) {
-    Buffer = AllocateAlignedCodePages (BufferPages, SIZE_32KB);
+    Buffer = AllocateAlignedPages (BufferPages, SIZE_32KB);
   } else {
-    Buffer = AllocateAlignedCodePages (BufferPages, SIZE_4KB);
+    Buffer = AllocateAlignedPages (BufferPages, SIZE_4KB);
   }
   ASSERT (Buffer != NULL);
   DEBUG ((EFI_D_INFO, "SMRAM SaveState Buffer (0x%08x, 0x%08x)\n", Buffer, EFI_PAGES_TO_SIZE(BufferPages)));
@@ -1416,6 +1416,35 @@ ConfigSmmCodeAccessCheck (
 }
 
 /**
+  This API provides a way to allocate memory for page table.
+
+  This API can be called more once to allocate memory for page tables.
+
+  Allocates the number of 4KB pages of type EfiRuntimeServicesData and returns a pointer to the
+  allocated buffer.  The buffer returned is aligned on a 4KB boundary.  If Pages is 0, then NULL
+  is returned.  If there is not enough memory remaining to satisfy the request, then NULL is
+  returned.
+
+  @param  Pages                 The number of 4 KB pages to allocate.
+
+  @return A pointer to the allocated buffer or NULL if allocation fails.
+
+**/
+VOID *
+AllocatePageTableMemory (
+  IN UINTN           Pages
+  )
+{
+  VOID  *Buffer;
+
+  Buffer = SmmCpuFeaturesAllocatePageTableMemory (Pages);
+  if (Buffer != NULL) {
+    return Buffer;
+  }
+  return AllocatePages (Pages);
+}
+
+/**
   Perform the remaining tasks.
 
 **/
@@ -1440,6 +1469,8 @@ PerformRemainingTasks (
     //
     ConfigSmmCodeAccessCheck ();
 
+    SmmCpuFeaturesCompleteSmmReadyToLock ();
+
     //
     // Clean SMM ready to lock flag
     //
@@ -1465,6 +1496,8 @@ PerformPreTasks (
     //
     ConfigSmmCodeAccessCheck ();
 
+    SmmCpuFeaturesCompleteSmmReadyToLock ();
+
     mRestoreSmmConfigurationInS3 = FALSE;
   }
 }