[PATCH 0/3] Introduce page table pool mechanism in SMM page table

[PATCH 0/3] Introduce page table pool mechanism in SMM page table

[duntan]

1.Introduce page table pool mechanism in SMM page table
2.Remove unused API in SmmCpuFeaturesLib.h
3.Simplify the code to set smm page table as RO

Dun Tan (3):
  UefiCpuPkg/PiSmmCpuDxeSmm: Introduce page table pool mechanism
  UefiCpuPkg: Remove unused API in SmmCpuFeaturesLib.h
  UefiCpuPkg: Simplify the code to set smm page table as RO

 UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h     |  25 -------------------------
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c           | 122 --------------------------------------------------------------------------------------------------------------------------
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c         |  30 ------------------------------
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h         |  30 ++++++++++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 256 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c            | 169 -------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 6 files changed, 286 insertions(+), 346 deletions(-)

-- 
2.31.1.windows.1

[PATCH 1/3] UefiCpuPkg/PiSmmCpuDxeSmm: Introduce page table pool mechanism

[duntan]

Introduce page table pool mechanism for smm page table to simplify
page table memory management and protection. This mechanism has been
used in DxeIpl. The basic idea is to allocate a bunch of continuous
pages of memory in advance, and all future page tables consumption
will happen in those pool instead of system memory.
Since we have centralized page tables, we only need to mark all page
table pools as RO, instead of searching page table memory layer by
layer in smm page table. Once current page table pool has been used
up, another memory pool will be allocated and the new pool will also
be set as RO if current page table memory has been marked as RO.

Signed-off-by: Dun Tan <dun.tan@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c           |   1 +
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c         |  30 ------------------------------
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h         |  13 +++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 122 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c            |   1 +
 5 files changed, 137 insertions(+), 30 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
index 97058a2810..26efa71eff 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
@@ -319,6 +319,7 @@ SetPageTableAttributes (
     //
     EnableCet ();
   }
+  mIsReadOnlyPageTable = TRUE;
 
   return;
 }
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
index 37e3cfc449..655175a2c6 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
@@ -1322,36 +1322,6 @@ 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);
-}
-
 /**
   Allocate pages for code.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 0bfba7e359..a0daaa1900 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -260,12 +260,25 @@ extern UINTN                 mNumberOfCpus;
 extern EFI_SMM_CPU_PROTOCOL  mSmmCpu;
 extern EFI_MM_MP_PROTOCOL    mSmmMp;
 extern BOOLEAN               m5LevelPagingNeeded;
+extern BOOLEAN               mIsReadOnlyPageTable;
 
 ///
 /// The mode of the CPU at the time an SMI occurs
 ///
 extern UINT8  mSmmSaveStateRegisterLma;
 
+#define PAGE_TABLE_POOL_ALIGNMENT   BASE_128KB
+#define PAGE_TABLE_POOL_UNIT_SIZE   BASE_128KB
+#define PAGE_TABLE_POOL_UNIT_PAGES  EFI_SIZE_TO_PAGES (PAGE_TABLE_POOL_UNIT_SIZE)
+#define PAGE_TABLE_POOL_ALIGN_MASK  \
+  (~(EFI_PHYSICAL_ADDRESS)(PAGE_TABLE_POOL_ALIGNMENT - 1))
+
+typedef struct {
+  VOID     *NextPool;
+  UINTN    Offset;
+  UINTN    FreePages;
+} PAGE_TABLE_POOL;
+
 //
 // SMM CPU Protocol function prototypes.
 //
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 773ab927e6..8f0c6410e6 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -35,6 +35,128 @@ PAGE_ATTRIBUTE_TABLE  mPageAttributeTable[] = {
 BOOLEAN  mIsShadowStack      = FALSE;
 BOOLEAN  m5LevelPagingNeeded = FALSE;
 
+//
+// Global variable to keep track current available memory used as page table.
+//
+PAGE_TABLE_POOL  *mPageTablePool = NULL;
+
+//
+// If memory used by SMM page table has been mareked as ReadOnly.
+//
+BOOLEAN  mIsReadOnlyPageTable = FALSE;
+
+/**
+  Initialize a buffer pool for page table use only.
+
+  To reduce the potential split operation on page table, the pages reserved for
+  page table should be allocated in the times of PAGE_TABLE_POOL_UNIT_PAGES and
+  at the boundary of PAGE_TABLE_POOL_ALIGNMENT. So the page pool is always
+  initialized with number of pages greater than or equal to the given PoolPages.
+
+  Once the pages in the pool are used up, this method should be called again to
+  reserve at least another PAGE_TABLE_POOL_UNIT_PAGES. But usually this won't
+  happen in practice.
+
+  @param PoolPages  The least page number of the pool to be created.
+
+  @retval TRUE    The pool is initialized successfully.
+  @retval FALSE   The memory is out of resource.
+**/
+BOOLEAN
+InitializePageTablePool (
+  IN UINTN  PoolPages
+  )
+{
+  VOID  *Buffer;
+
+  //
+  // Always reserve at least PAGE_TABLE_POOL_UNIT_PAGES, including one page for
+  // header.
+  //
+  PoolPages += 1;   // Add one page for header.
+  PoolPages  = ((PoolPages - 1) / PAGE_TABLE_POOL_UNIT_PAGES + 1) *
+               PAGE_TABLE_POOL_UNIT_PAGES;
+  Buffer = AllocateAlignedPages (PoolPages, PAGE_TABLE_POOL_ALIGNMENT);
+  if (Buffer == NULL) {
+    DEBUG ((DEBUG_ERROR, "ERROR: Out of aligned pages\r\n"));
+    return FALSE;
+  }
+
+  //
+  // Link all pools into a list for easier track later.
+  //
+  if (mPageTablePool == NULL) {
+    mPageTablePool           = Buffer;
+    mPageTablePool->NextPool = mPageTablePool;
+  } else {
+    ((PAGE_TABLE_POOL *)Buffer)->NextPool = mPageTablePool->NextPool;
+    mPageTablePool->NextPool              = Buffer;
+    mPageTablePool                        = Buffer;
+  }
+
+  //
+  // Reserve one page for pool header.
+  //
+  mPageTablePool->FreePages = PoolPages - 1;
+  mPageTablePool->Offset    = EFI_PAGES_TO_SIZE (1);
+
+  //
+  // If page table memory has been marked as RO, mark the new pool pages as read-only.
+  //
+  if (mIsReadOnlyPageTable) {
+    AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);
+    SmmSetMemoryAttributes ((EFI_PHYSICAL_ADDRESS)(UINTN)Buffer, EFI_PAGES_TO_SIZE (PoolPages), EFI_MEMORY_RO);
+    AsmWriteCr0 (AsmReadCr0 () | CR0_WP);
+  }
+
+  return TRUE;
+}
+
+/**
+  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;
+
+  if (Pages == 0) {
+    return NULL;
+  }
+
+  //
+  // Renew the pool if necessary.
+  //
+  if ((mPageTablePool == NULL) ||
+      (Pages > mPageTablePool->FreePages))
+  {
+    if (!InitializePageTablePool (Pages)) {
+      return NULL;
+    }
+  }
+
+  Buffer = (UINT8 *)mPageTablePool + mPageTablePool->Offset;
+
+  mPageTablePool->Offset    += EFI_PAGES_TO_SIZE (Pages);
+  mPageTablePool->FreePages -= Pages;
+
+  return Buffer;
+}
+
 /**
   Return length according to page attributes.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
index bf90050503..d714ca5b5a 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
@@ -1304,6 +1304,7 @@ SetPageTableAttributes (
     //
     EnableCet ();
   }
+  mIsReadOnlyPageTable = TRUE;
 
   return;
 }
-- 
2.31.1.windows.1

[PATCH 2/3] UefiCpuPkg: Remove unused API in SmmCpuFeaturesLib.h

[duntan]

Remove SmmCpuFeaturesAllocatePageTableMemory in this headfile.
This API is not used by PiSmmCpuDxeSmm driver any more. Also
no other files use this API.

Signed-off-by: Dun Tan <dun.tan@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
---
 UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h | 25 -------------------------
 1 file changed, 25 deletions(-)

diff --git a/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h b/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h
index 54cae865a2..52160c7145 100644
--- a/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h
+++ b/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h
@@ -386,29 +386,4 @@ SmmCpuFeaturesCompleteSmmReadyToLock (
   VOID
   );
 
-/**
-  This API provides a method for a CPU to allocate a specific region for storing page tables.
-
-  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.
-
-  This function can also return NULL if there is no preference on where the page tables are allocated in SMRAM.
-
-  @param  Pages                 The number of 4 KB pages to allocate.
-
-  @return A pointer to the allocated buffer for page tables.
-  @retval NULL      Fail to allocate a specific region for storing page tables,
-                    Or there is no preference on where the page tables are allocated in SMRAM.
-
-**/
-VOID *
-EFIAPI
-SmmCpuFeaturesAllocatePageTableMemory (
-  IN UINTN  Pages
-  );
-
 #endif
-- 
2.31.1.windows.1

[PATCH 3/3] UefiCpuPkg: Simplify the code to set smm page table as RO

[duntan]

Simplify the code to set memory used by smm page table as RO.
Since memory used by smm page table are in PageTablePool list,
we only need to set all PageTablePool as ReadOnly in smm page
table itself. Also, we only need to flush tlb once after
setting all page table pool as Read Only.

Signed-off-by: Dun Tan <dun.tan@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c           | 123 ---------------------------------------------------------------------------------------------------------------------------
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h         |  19 ++++++++++++++++++-
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 134 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c            | 170 --------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 4 files changed, 152 insertions(+), 294 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
index 26efa71eff..26bbba77b0 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
@@ -10,24 +10,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "PiSmmCpuDxeSmm.h"
 
-/**
-  Disable CET.
-**/
-VOID
-EFIAPI
-DisableCet (
-  VOID
-  );
-
-/**
-  Enable CET.
-**/
-VOID
-EFIAPI
-EnableCet (
-  VOID
-  );
-
 /**
   Create PageTable for SMM use.
 
@@ -221,111 +203,6 @@ Exit:
 }
 
 /**
-  This function sets memory attribute for page table.
-**/
-VOID
-SetPageTableAttributes (
-  VOID
-  )
-{
-  UINTN    Index2;
-  UINTN    Index3;
-  UINT64   *L1PageTable;
-  UINT64   *L2PageTable;
-  UINT64   *L3PageTable;
-  UINTN    PageTableBase;
-  BOOLEAN  IsSplitted;
-  BOOLEAN  PageTableSplitted;
-  BOOLEAN  CetEnabled;
-
-  //
-  // Don't mark page table to read-only if heap guard is enabled.
-  //
-  //      BIT2: SMM page guard enabled
-  //      BIT3: SMM pool guard enabled
-  //
-  if ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0) {
-    DEBUG ((DEBUG_INFO, "Don't mark page table to read-only as heap guard is enabled\n"));
-    return;
-  }
-
-  //
-  // Don't mark page table to read-only if SMM profile is enabled.
-  //
-  if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
-    DEBUG ((DEBUG_INFO, "Don't mark page table to read-only as SMM profile is enabled\n"));
-    return;
-  }
-
-  DEBUG ((DEBUG_INFO, "SetPageTableAttributes\n"));
-
-  //
-  // Disable write protection, because we need mark page table to be write protected.
-  // We need *write* page table memory, to mark itself to be *read only*.
-  //
-  CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;
-  if (CetEnabled) {
-    //
-    // CET must be disabled if WP is disabled.
-    //
-    DisableCet ();
-  }
-
-  AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);
-
-  do {
-    DEBUG ((DEBUG_INFO, "Start...\n"));
-    PageTableSplitted = FALSE;
-
-    PageTableBase = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;
-    L3PageTable   = (UINT64 *)PageTableBase;
-
-    SmmSetMemoryAttributesEx (PageTableBase, FALSE, (EFI_PHYSICAL_ADDRESS)PageTableBase, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-    PageTableSplitted = (PageTableSplitted || IsSplitted);
-
-    for (Index3 = 0; Index3 < 4; Index3++) {
-      L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
-      if (L2PageTable == NULL) {
-        continue;
-      }
-
-      SmmSetMemoryAttributesEx (PageTableBase, FALSE, (EFI_PHYSICAL_ADDRESS)(UINTN)L2PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-      PageTableSplitted = (PageTableSplitted || IsSplitted);
-
-      for (Index2 = 0; Index2 < SIZE_4KB/sizeof (UINT64); Index2++) {
-        if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {
-          // 2M
-          continue;
-        }
-
-        L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
-        if (L1PageTable == NULL) {
-          continue;
-        }
-
-        SmmSetMemoryAttributesEx (PageTableBase, FALSE, (EFI_PHYSICAL_ADDRESS)(UINTN)L1PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-        PageTableSplitted = (PageTableSplitted || IsSplitted);
-      }
-    }
-  } while (PageTableSplitted);
-
-  //
-  // Enable write protection, after page table updated.
-  //
-  AsmWriteCr0 (AsmReadCr0 () | CR0_WP);
-  if (CetEnabled) {
-    //
-    // re-enable CET.
-    //
-    EnableCet ();
-  }
-  mIsReadOnlyPageTable = TRUE;
-
-  return;
-}
-
-/**
-  This function returns with no action for 32 bit.
 
   @param[out]  *Cr2  Pointer to variable to hold CR2 register value.
 **/
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index a0daaa1900..5f0a38e400 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -260,7 +260,6 @@ extern UINTN                 mNumberOfCpus;
 extern EFI_SMM_CPU_PROTOCOL  mSmmCpu;
 extern EFI_MM_MP_PROTOCOL    mSmmMp;
 extern BOOLEAN               m5LevelPagingNeeded;
-extern BOOLEAN               mIsReadOnlyPageTable;
 
 ///
 /// The mode of the CPU at the time an SMI occurs
@@ -279,6 +278,24 @@ typedef struct {
   UINTN    FreePages;
 } PAGE_TABLE_POOL;
 
+/**
+  Disable CET.
+**/
+VOID
+EFIAPI
+DisableCet (
+  VOID
+  );
+
+/**
+  Enable CET.
+**/
+VOID
+EFIAPI
+EnableCet (
+  VOID
+  );
+
 //
 // SMM CPU Protocol function prototypes.
 //
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 8f0c6410e6..9f091c6485 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -1738,3 +1738,137 @@ EdkiiSmmGetMemoryAttributes (
 
   return EFI_SUCCESS;
 }
+
+/**
+  Prevent the memory pages used for SMM page table from been overwritten.
+**/
+VOID
+EnablePageTableProtection (
+  VOID
+  )
+{
+  PAGE_TABLE_POOL       *HeadPool;
+  PAGE_TABLE_POOL       *Pool;
+  UINT64                PoolSize;
+  EFI_PHYSICAL_ADDRESS  Address;
+  UINTN                 PageTableBase;
+
+  if (mPageTablePool == NULL) {
+    return;
+  }
+
+  PageTableBase = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;
+
+  //
+  // ConvertMemoryPageAttributes might update mPageTablePool. It's safer to
+  // remember original one in advance.
+  //
+  HeadPool = mPageTablePool;
+  Pool     = HeadPool;
+  do {
+    Address  = (EFI_PHYSICAL_ADDRESS)(UINTN)Pool & PAGE_TABLE_POOL_ALIGN_MASK;
+    PoolSize = Pool->Offset + EFI_PAGES_TO_SIZE (Pool->FreePages);
+
+    ConvertMemoryPageAttributes (PageTableBase, m5LevelPagingNeeded, Address, PoolSize, EFI_MEMORY_RO, TRUE, NULL, NULL);
+    Pool = Pool->NextPool;
+  } while (Pool != HeadPool);
+}
+
+/**
+  Return whether memory used by SMM page table need to be set as Read Only.
+
+  @retval TRUE  Need to set SMM page table as Read Only.
+  @retval FALSE Do not set SMM page table as Read Only.
+**/
+BOOLEAN
+IfReadOnlyPageTableNeeded (
+  VOID
+  )
+{
+  //
+  // Don't mark page table memory as read-only if
+  //  - no restriction on access to non-SMRAM memory; or
+  //  - SMM heap guard feature enabled; or
+  //      BIT2: SMM page guard enabled
+  //      BIT3: SMM pool guard enabled
+  //  - SMM profile feature enabled
+  //
+  if (!IsRestrictedMemoryAccess () ||
+      ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0) ||
+      FeaturePcdGet (PcdCpuSmmProfileEnable))
+  {
+    if (sizeof (UINTN) == sizeof (UINT64)) {
+      //
+      // Restriction on access to non-SMRAM memory and heap guard could not be enabled at the same time.
+      //
+      ASSERT (
+        !(IsRestrictedMemoryAccess () &&
+          (PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0)
+        );
+
+      //
+      // Restriction on access to non-SMRAM memory and SMM profile could not be enabled at the same time.
+      //
+      ASSERT (!(IsRestrictedMemoryAccess () && FeaturePcdGet (PcdCpuSmmProfileEnable)));
+    }
+
+    return FALSE;
+  }
+
+  return TRUE;
+}
+
+/**
+  This function sets memory attribute for page table.
+**/
+VOID
+SetPageTableAttributes (
+  VOID
+  )
+{
+  BOOLEAN  CetEnabled;
+
+  if (!IfReadOnlyPageTableNeeded ()) {
+    return;
+  }
+
+  DEBUG ((DEBUG_INFO, "SetPageTableAttributes\n"));
+
+  //
+  // Disable write protection, because we need mark page table to be write protected.
+  // We need *write* page table memory, to mark itself to be *read only*.
+  //
+  CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;
+  if (CetEnabled) {
+    //
+    // CET must be disabled if WP is disabled.
+    //
+    DisableCet ();
+  }
+
+  AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);
+
+  // Set memory used by page table as Read Only.
+  DEBUG ((DEBUG_INFO, "Start...\n"));
+  EnablePageTableProtection ();
+
+  //
+  // Enable write protection, after page table attribute updated.
+  //
+  AsmWriteCr0 (AsmReadCr0 () | CR0_WP);
+  mIsReadOnlyPageTable = TRUE;
+
+  //
+  // Flush TLB after mark all page table pool as read only.
+  //
+  FlushTlbForAll ();
+
+  if (CetEnabled) {
+    //
+    // re-enable CET.
+    //
+    EnableCet ();
+  }
+
+  return;
+}
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
index d714ca5b5a..3deb1ffd67 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
@@ -20,24 +20,6 @@ BOOLEAN                   m1GPageTableSupport = FALSE;
 BOOLEAN                   mCpuSmmRestrictedMemoryAccess;
 X86_ASSEMBLY_PATCH_LABEL  gPatch5LevelPagingNeeded;
 
-/**
-  Disable CET.
-**/
-VOID
-EFIAPI
-DisableCet (
-  VOID
-  );
-
-/**
-  Enable CET.
-**/
-VOID
-EFIAPI
-EnableCet (
-  VOID
-  );
-
 /**
   Check if 1-GByte pages is supported by processor or not.
 
@@ -1157,158 +1139,6 @@ Exit:
   ReleaseSpinLock (mPFLock);
 }
 
-/**
-  This function sets memory attribute for page table.
-**/
-VOID
-SetPageTableAttributes (
-  VOID
-  )
-{
-  UINTN     Index2;
-  UINTN     Index3;
-  UINTN     Index4;
-  UINTN     Index5;
-  UINT64    *L1PageTable;
-  UINT64    *L2PageTable;
-  UINT64    *L3PageTable;
-  UINT64    *L4PageTable;
-  UINT64    *L5PageTable;
-  UINTN     PageTableBase;
-  BOOLEAN   IsSplitted;
-  BOOLEAN   PageTableSplitted;
-  BOOLEAN   CetEnabled;
-  BOOLEAN   Enable5LevelPaging;
-  IA32_CR4  Cr4;
-
-  //
-  // Don't mark page table memory as read-only if
-  //  - no restriction on access to non-SMRAM memory; or
-  //  - SMM heap guard feature enabled; or
-  //      BIT2: SMM page guard enabled
-  //      BIT3: SMM pool guard enabled
-  //  - SMM profile feature enabled
-  //
-  if (!mCpuSmmRestrictedMemoryAccess ||
-      ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0) ||
-      FeaturePcdGet (PcdCpuSmmProfileEnable))
-  {
-    //
-    // Restriction on access to non-SMRAM memory and heap guard could not be enabled at the same time.
-    //
-    ASSERT (
-      !(mCpuSmmRestrictedMemoryAccess &&
-        (PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0)
-      );
-
-    //
-    // Restriction on access to non-SMRAM memory and SMM profile could not be enabled at the same time.
-    //
-    ASSERT (!(mCpuSmmRestrictedMemoryAccess && FeaturePcdGet (PcdCpuSmmProfileEnable)));
-    return;
-  }
-
-  DEBUG ((DEBUG_INFO, "SetPageTableAttributes\n"));
-
-  //
-  // Disable write protection, because we need mark page table to be write protected.
-  // We need *write* page table memory, to mark itself to be *read only*.
-  //
-  CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;
-  if (CetEnabled) {
-    //
-    // CET must be disabled if WP is disabled.
-    //
-    DisableCet ();
-  }
-
-  AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);
-
-  do {
-    DEBUG ((DEBUG_INFO, "Start...\n"));
-    PageTableSplitted = FALSE;
-    L5PageTable       = NULL;
-
-    PageTableBase      = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;
-    Cr4.UintN          = AsmReadCr4 ();
-    Enable5LevelPaging = (BOOLEAN)(Cr4.Bits.LA57 == 1);
-
-    if (Enable5LevelPaging) {
-      L5PageTable = (UINT64 *)PageTableBase;
-      SmmSetMemoryAttributesEx (PageTableBase, Enable5LevelPaging, (EFI_PHYSICAL_ADDRESS)PageTableBase, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-      PageTableSplitted = (PageTableSplitted || IsSplitted);
-    }
-
-    for (Index5 = 0; Index5 < (Enable5LevelPaging ? SIZE_4KB/sizeof (UINT64) : 1); Index5++) {
-      if (Enable5LevelPaging) {
-        L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
-        if (L4PageTable == NULL) {
-          continue;
-        }
-      } else {
-        L4PageTable = (UINT64 *)PageTableBase;
-      }
-
-      SmmSetMemoryAttributesEx (PageTableBase, Enable5LevelPaging, (EFI_PHYSICAL_ADDRESS)(UINTN)L4PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-      PageTableSplitted = (PageTableSplitted || IsSplitted);
-
-      for (Index4 = 0; Index4 < SIZE_4KB/sizeof (UINT64); Index4++) {
-        L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
-        if (L3PageTable == NULL) {
-          continue;
-        }
-
-        SmmSetMemoryAttributesEx (PageTableBase, Enable5LevelPaging, (EFI_PHYSICAL_ADDRESS)(UINTN)L3PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-        PageTableSplitted = (PageTableSplitted || IsSplitted);
-
-        for (Index3 = 0; Index3 < SIZE_4KB/sizeof (UINT64); Index3++) {
-          if ((L3PageTable[Index3] & IA32_PG_PS) != 0) {
-            // 1G
-            continue;
-          }
-
-          L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
-          if (L2PageTable == NULL) {
-            continue;
-          }
-
-          SmmSetMemoryAttributesEx (PageTableBase, Enable5LevelPaging, (EFI_PHYSICAL_ADDRESS)(UINTN)L2PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-          PageTableSplitted = (PageTableSplitted || IsSplitted);
-
-          for (Index2 = 0; Index2 < SIZE_4KB/sizeof (UINT64); Index2++) {
-            if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {
-              // 2M
-              continue;
-            }
-
-            L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
-            if (L1PageTable == NULL) {
-              continue;
-            }
-
-            SmmSetMemoryAttributesEx (PageTableBase, Enable5LevelPaging, (EFI_PHYSICAL_ADDRESS)(UINTN)L1PageTable, SIZE_4KB, EFI_MEMORY_RO, &IsSplitted);
-            PageTableSplitted = (PageTableSplitted || IsSplitted);
-          }
-        }
-      }
-    }
-  } while (PageTableSplitted);
-
-  //
-  // Enable write protection, after page table updated.
-  //
-  AsmWriteCr0 (AsmReadCr0 () | CR0_WP);
-  if (CetEnabled) {
-    //
-    // re-enable CET.
-    //
-    EnableCet ();
-  }
-  mIsReadOnlyPageTable = TRUE;
-
-  return;
-}
-
 /**
   This function reads CR2 register when on-demand paging is enabled.
 
-- 
2.31.1.windows.1

Re: [PATCH 1/3] UefiCpuPkg/PiSmmCpuDxeSmm: Introduce page table pool mechanism

[Ni, Ray]

> +  //
> +  // If page table memory has been marked as RO, mark the new pool pages as read-only.
> +  //
> +  if (mIsReadOnlyPageTable) {
> +    AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);
> +    SmmSetMemoryAttributes ((EFI_PHYSICAL_ADDRESS)(UINTN)Buffer, EFI_PAGES_TO_SIZE (PoolPages),
> EFI_MEMORY_RO);
> +    AsmWriteCr0 (AsmReadCr0 () | CR0_WP);
> +  }
> +

If you check other logic that clears WP bit in CR0, CET is disabled before clearing WP.
You need to follow the same pattern here.

Other looks good to me.

Re: [PATCH 2/3] UefiCpuPkg: Remove unused API in SmmCpuFeaturesLib.h

[Ni, Ray]

Reviewed-by: Ray Ni <ray.ni@intel.com>

> -----Original Message-----
> From: Tan, Dun <dun.tan@intel.com>
> Sent: Friday, December 16, 2022 11:01 AM
> To: devel@edk2.groups.io
> Cc: Dong, Eric <eric.dong@intel.com>; Ni, Ray <ray.ni@intel.com>; Kumar, Rahul R <rahul.r.kumar@intel.com>
> Subject: [PATCH 2/3] UefiCpuPkg: Remove unused API in SmmCpuFeaturesLib.h
> 
> Remove SmmCpuFeaturesAllocatePageTableMemory in this headfile.
> This API is not used by PiSmmCpuDxeSmm driver any more. Also
> no other files use this API.
> 
> Signed-off-by: Dun Tan <dun.tan@intel.com>
> Cc: Eric Dong <eric.dong@intel.com>
> Cc: Ray Ni <ray.ni@intel.com>
> Cc: Rahul Kumar <rahul1.kumar@intel.com>
> ---
>  UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h | 25 -------------------------
>  1 file changed, 25 deletions(-)
> 
> diff --git a/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h b/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h
> index 54cae865a2..52160c7145 100644
> --- a/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h
> +++ b/UefiCpuPkg/Include/Library/SmmCpuFeaturesLib.h
> @@ -386,29 +386,4 @@ SmmCpuFeaturesCompleteSmmReadyToLock (
>    VOID
>    );
> 
> -/**
> -  This API provides a method for a CPU to allocate a specific region for storing page tables.
> -
> -  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.
> -
> -  This function can also return NULL if there is no preference on where the page tables are allocated in SMRAM.
> -
> -  @param  Pages                 The number of 4 KB pages to allocate.
> -
> -  @return A pointer to the allocated buffer for page tables.
> -  @retval NULL      Fail to allocate a specific region for storing page tables,
> -                    Or there is no preference on where the page tables are allocated in SMRAM.
> -
> -**/
> -VOID *
> -EFIAPI
> -SmmCpuFeaturesAllocatePageTableMemory (
> -  IN UINTN  Pages
> -  );
> -
>  #endif
> --
> 2.31.1.windows.1

Re: [PATCH 3/3] UefiCpuPkg: Simplify the code to set smm page table as RO

[Ni, Ray]

> +  PageTableBase = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;
> +
> +  //
> +  // ConvertMemoryPageAttributes might update mPageTablePool. It's safer to
> +  // remember original one in advance.
> +  //
> +  HeadPool = mPageTablePool;
> +  Pool     = HeadPool;
> +  do {
> +    Address  = (EFI_PHYSICAL_ADDRESS)(UINTN)Pool & PAGE_TABLE_POOL_ALIGN_MASK;

1. When is the Pool not aligned on 128KB boundary? If it's guaranteed, can we remove the "& PAGE_TABLE_POOL_ALIGN_MASK"?


> +    PoolSize = Pool->Offset + EFI_PAGES_TO_SIZE (Pool->FreePages);
> +
> +    ConvertMemoryPageAttributes (PageTableBase, m5LevelPagingNeeded, Address, PoolSize, EFI_MEMORY_RO, TRUE,
> NULL, NULL);

2. Can you please explain in comments that above call is to make the entire pool including header, used-memory, free-memory
    as read-only?

3. It's better to use LinkedList library APIs from BaseLib. The comments apply to the first patch as well. But I am fine if you decide not to do it in this patch.

> +  {
> +    if (sizeof (UINTN) == sizeof (UINT64)) {
> +      //
> +      // Restriction on access to non-SMRAM memory and heap guard could not be enabled at the same time.
> +      //
> +      ASSERT (
> +        !(IsRestrictedMemoryAccess () &&
> +          (PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0)
> +        );
> +
> +      //
> +      // Restriction on access to non-SMRAM memory and SMM profile could not be enabled at the same time.
> +      //
> +      ASSERT (!(IsRestrictedMemoryAccess () && FeaturePcdGet (PcdCpuSmmProfileEnable)));
> +    }

4. I don't think we still need the above two assertions. But let's not clean up the code in your patch. @Wang, Jian J

Re: [PATCH 3/3] UefiCpuPkg: Simplify the code to set smm page table as RO

[duntan]

Thanks Ray for the comments. I'll modify the code in V2 patch set.
For comments 3, I'll replace the code by LinkedList library APIs from BaseLib in future patches.

Thanks,
Dun

-----Original Message-----
From: Ni, Ray <ray.ni@intel.com> 
Sent: Monday, December 19, 2022 3:53 PM
To: Tan, Dun <dun.tan@intel.com>; devel@edk2.groups.io; Wang, Jian J <jian.j.wang@intel.com>
Cc: Dong, Eric <eric.dong@intel.com>; Kumar, Rahul R <rahul.r.kumar@intel.com>
Subject: RE: [PATCH 3/3] UefiCpuPkg: Simplify the code to set smm page table as RO

> +  PageTableBase = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;
> +
> +  //
> +  // ConvertMemoryPageAttributes might update mPageTablePool. It's 
> + safer to  // remember original one in advance.
> +  //
> +  HeadPool = mPageTablePool;
> +  Pool     = HeadPool;
> +  do {
> +    Address  = (EFI_PHYSICAL_ADDRESS)(UINTN)Pool & 
> + PAGE_TABLE_POOL_ALIGN_MASK;

1. When is the Pool not aligned on 128KB boundary? If it's guaranteed, can we remove the "& PAGE_TABLE_POOL_ALIGN_MASK"?


> +    PoolSize = Pool->Offset + EFI_PAGES_TO_SIZE (Pool->FreePages);
> +
> +    ConvertMemoryPageAttributes (PageTableBase, m5LevelPagingNeeded, 
> + Address, PoolSize, EFI_MEMORY_RO, TRUE,
> NULL, NULL);

2. Can you please explain in comments that above call is to make the entire pool including header, used-memory, free-memory
    as read-only?

3. It's better to use LinkedList library APIs from BaseLib. The comments apply to the first patch as well. But I am fine if you decide not to do it in this patch.

> +  {
> +    if (sizeof (UINTN) == sizeof (UINT64)) {
> +      //
> +      // Restriction on access to non-SMRAM memory and heap guard could not be enabled at the same time.
> +      //
> +      ASSERT (
> +        !(IsRestrictedMemoryAccess () &&
> +          (PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0)
> +        );
> +
> +      //
> +      // Restriction on access to non-SMRAM memory and SMM profile could not be enabled at the same time.
> +      //
> +      ASSERT (!(IsRestrictedMemoryAccess () && FeaturePcdGet (PcdCpuSmmProfileEnable)));
> +    }

4. I don't think we still need the above two assertions. But let's not clean up the code in your patch. @Wang, Jian J