[PATCH v5] UefiCpuPkg\CpuSmm: Save & restore CR2 on-demand paging in SMM

[nkvangup <narendra.k.vanguput@intel.com>]

BZ: https://bugzilla.tianocore.org/show_bug.cgi?id=1593

For every SMI occurrence, save and restore CR2 register only when SMM
on-demand paging support is enabled in 64 bit operation mode.
This is not a bug but to have better improvement of code.

Patch5 is updated with separate functions for Save and Restore of CR2
based on review feedback.

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Vanguput Narendra K <narendra.k.vanguput@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Yao Jiewen <jiewen.yao@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c   | 22 ++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c      |  9 +++++----
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h | 16 ++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c    | 28 ++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
index b734a1ea8c..3750332ca8 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
@@ -316,3 +316,25 @@ SetPageTableAttributes (
 
   return ;
 }
+
+/**
+  This function returns with no action for 32 bit.
+**/
+VOID
+SaveCr2 (
+  VOID
+  )
+{
+// Do Nothing
+}
+
+/**
+  This function returns with no action for 32 bit.
+**/
+VOID
+RestoreCr2 (
+  VOID
+  )
+{
+// Do Nothing
+}
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 3b0b3b52ac..6a5736a3eb 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -1107,14 +1107,14 @@ SmiRendezvous (
   BOOLEAN                        IsBsp;
   BOOLEAN                        BspInProgress;
   UINTN                          Index;
-  UINTN                          Cr2;
 
   ASSERT(CpuIndex < mMaxNumberOfCpus);
 
   //
-  // Save Cr2 because Page Fault exception in SMM may override its value
+  // Save Cr2 because Page Fault exception in SMM may override its value,
+  // when using on-demand paging for above 4G memory.
   //
-  Cr2 = AsmReadCr2 ();
+  SaveCr2 ();
 
   //
   // Perform CPU specific entry hooks
@@ -1253,10 +1253,11 @@ SmiRendezvous (
 
 Exit:
   SmmCpuFeaturesRendezvousExit (CpuIndex);
+
   //
   // Restore Cr2
   //
-  AsmWriteCr2 (Cr2);
+  RestoreCr2 ();
 }
 
 /**
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 84efb22981..71a8c13960 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -1243,4 +1243,20 @@ EFIAPI
 PiSmmCpuSmiEntryFixupAddress (
  );
 
+/**
+  This function saves CR2 register for 64 bit and no action for 32 bit.
+**/
+VOID
+SaveCr2 (
+  VOID
+  );
+
+/**
+  This function restores CR2 register for 64 bit and no action for 32 bit.
+**/
+VOID
+RestoreCr2 (
+  VOID
+  );
+
 #endif
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
index 2c77cb47a4..76a30de171 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
@@ -22,6 +22,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 LIST_ENTRY                          mPagePool = INITIALIZE_LIST_HEAD_VARIABLE (mPagePool);
 BOOLEAN                             m1GPageTableSupport = FALSE;
 BOOLEAN                             mCpuSmmStaticPageTable;
+UINTN                               Cr2 = 0;
 
 /**
   Disable CET.
@@ -1053,3 +1054,30 @@ SetPageTableAttributes (
 
   return ;
 }
+
+/**
+  This function saves CR2 register.
+**/
+VOID
+SaveCr2 (
+  VOID
+  )
+{
+  if (!mCpuSmmStaticPageTable) {
+    Cr2 = AsmReadCr2 ();
+  }
+}
+
+/**
+  This function restores CR2 register.
+**/
+VOID
+RestoreCr2 (
+  VOID
+  )
+{
+  if ((!mCpuSmmStaticPageTable) && (Cr2 != 0)) {
+    AsmWriteCr2 (Cr2);
+    Cr2 = 0;
+  }
+}
-- 
2.16.2.windows.1