[Patch V4 07/15] UefiCpuPkg/PiSmmCpuDxeSmm: Add 2 function to disable/enable CR0.WP

["duntan" <dun.tan@intel.com>]

Add two functions to disable/enable CR0.WP. These two unctions
will also be used in later commits. This commit doesn't change any
functionality.

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>
Cc: Gerd Hoffmann <kraxel@redhat.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h         |  24 ++++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 115 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-------------------------------------------------
 2 files changed, 90 insertions(+), 49 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index ba341cadc6..e0c4ca76dc 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -1565,4 +1565,28 @@ SmmWaitForApArrival (
   VOID
   );
 
+/**
+  Disable Write Protect on pages marked as read-only if Cr0.Bits.WP is 1.
+
+  @param[out]  WpEnabled      If Cr0.WP is enabled.
+  @param[out]  CetEnabled     If CET is enabled.
+**/
+VOID
+DisableReadOnlyPageWriteProtect (
+  OUT BOOLEAN  *WpEnabled,
+  OUT BOOLEAN  *CetEnabled
+  );
+
+/**
+  Enable Write Protect on pages marked as read-only.
+
+  @param[out]  WpEnabled      If Cr0.WP should be enabled.
+  @param[out]  CetEnabled     If CET should be enabled.
+**/
+VOID
+EnableReadOnlyPageWriteProtect (
+  BOOLEAN  WpEnabled,
+  BOOLEAN  CetEnabled
+  );
+
 #endif
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 2faee8f859..4b512edf68 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -40,6 +40,64 @@ PAGE_TABLE_POOL  *mPageTablePool = NULL;
 //
 BOOLEAN  mIsReadOnlyPageTable = FALSE;
 
+/**
+  Disable Write Protect on pages marked as read-only if Cr0.Bits.WP is 1.
+
+  @param[out]  WpEnabled      If Cr0.WP is enabled.
+  @param[out]  CetEnabled     If CET is enabled.
+**/
+VOID
+DisableReadOnlyPageWriteProtect (
+  OUT BOOLEAN  *WpEnabled,
+  OUT BOOLEAN  *CetEnabled
+  )
+{
+  IA32_CR0  Cr0;
+
+  *CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;
+  Cr0.UintN   = AsmReadCr0 ();
+  *WpEnabled  = (Cr0.Bits.WP != 0) ? TRUE : FALSE;
+  if (*WpEnabled) {
+    if (*CetEnabled) {
+      //
+      // CET must be disabled if WP is disabled. Disable CET before clearing CR0.WP.
+      //
+      DisableCet ();
+    }
+
+    Cr0.Bits.WP = 0;
+    AsmWriteCr0 (Cr0.UintN);
+  }
+}
+
+/**
+  Enable Write Protect on pages marked as read-only.
+
+  @param[out]  WpEnabled      If Cr0.WP should be enabled.
+  @param[out]  CetEnabled     If CET should be enabled.
+**/
+VOID
+EnableReadOnlyPageWriteProtect (
+  BOOLEAN  WpEnabled,
+  BOOLEAN  CetEnabled
+  )
+{
+  IA32_CR0  Cr0;
+
+  if (WpEnabled) {
+    Cr0.UintN   = AsmReadCr0 ();
+    Cr0.Bits.WP = 1;
+    AsmWriteCr0 (Cr0.UintN);
+
+    if (CetEnabled) {
+      //
+      // re-enable CET.
+      //
+      EnableCet ();
+    }
+  }
+}
+
 /**
   Initialize a buffer pool for page table use only.
 
@@ -62,10 +120,9 @@ InitializePageTablePool (
   IN UINTN  PoolPages
   )
 {
-  VOID      *Buffer;
-  BOOLEAN   CetEnabled;
-  BOOLEAN   WpEnabled;
-  IA32_CR0  Cr0;
+  VOID     *Buffer;
+  BOOLEAN  WpEnabled;
+  BOOLEAN  CetEnabled;
 
   //
   // Always reserve at least PAGE_TABLE_POOL_UNIT_PAGES, including one page for
@@ -102,34 +159,9 @@ InitializePageTablePool (
   // If page table memory has been marked as RO, mark the new pool pages as read-only.
   //
   if (mIsReadOnlyPageTable) {
-    CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;
-    Cr0.UintN  = AsmReadCr0 ();
-    WpEnabled  = (Cr0.Bits.WP != 0) ? TRUE : FALSE;
-    if (WpEnabled) {
-      if (CetEnabled) {
-        //
-        // CET must be disabled if WP is disabled. Disable CET before clearing CR0.WP.
-        //
-        DisableCet ();
-      }
-
-      Cr0.Bits.WP = 0;
-      AsmWriteCr0 (Cr0.UintN);
-    }
-
+    DisableReadOnlyPageWriteProtect (&WpEnabled, &CetEnabled);
     SmmSetMemoryAttributes ((EFI_PHYSICAL_ADDRESS)(UINTN)Buffer, EFI_PAGES_TO_SIZE (PoolPages), EFI_MEMORY_RO);
-    if (WpEnabled) {
-      Cr0.UintN   = AsmReadCr0 ();
-      Cr0.Bits.WP = 1;
-      AsmWriteCr0 (Cr0.UintN);
-
-      if (CetEnabled) {
-        //
-        // re-enable CET.
-        //
-        EnableCet ();
-      }
-    }
+    EnableReadOnlyPageWriteProtect (WpEnabled, CetEnabled);
   }
 
   return TRUE;
@@ -1782,6 +1814,7 @@ SetPageTableAttributes (
   VOID
   )
 {
+  BOOLEAN  WpEnabled;
   BOOLEAN  CetEnabled;
 
   if (!IfReadOnlyPageTableNeeded ()) {
@@ -1794,15 +1827,7 @@ SetPageTableAttributes (
   // 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);
+  DisableReadOnlyPageWriteProtect (&WpEnabled, &CetEnabled);
 
   // Set memory used by page table as Read Only.
   DEBUG ((DEBUG_INFO, "Start...\n"));
@@ -1811,20 +1836,12 @@ SetPageTableAttributes (
   //
   // Enable write protection, after page table attribute updated.
   //
-  AsmWriteCr0 (AsmReadCr0 () | CR0_WP);
+  EnableReadOnlyPageWriteProtect (TRUE, CetEnabled);
   mIsReadOnlyPageTable = TRUE;
 
   //
   // Flush TLB after mark all page table pool as read only.
   //
   FlushTlbForAll ();
-
-  if (CetEnabled) {
-    //
-    // re-enable CET.
-    //
-    EnableCet ();
-  }
-
   return;
 }
-- 
2.31.1.windows.1