[Patch V2 1/2] UefiCpuPkg/PiSmmCpuDxeSmm: Add a new mIsShadowStack flag

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

This patch is code refactoring and doesn't change any functionality.
Add a new IsShadowStack flag to identify whether current memory is
shadow stack. Previous smm code logic regards a RO range as shadow
stack and set the dirty bit in corresponding page table entry if
mInternalCr3 is not 0, which may be confusing.

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/SmmCpuMemoryManagement.c | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 1f7cc15727..237742d7e6 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -33,6 +33,7 @@ PAGE_ATTRIBUTE_TABLE  mPageAttributeTable[] = {
 };
 
 UINTN  mInternalCr3;
+UINTN  mIsShadowStack = FALSE;
 
 /**
   Set the internal page table base address.
@@ -249,7 +250,7 @@ ConvertPageEntryAttribute (
   if ((Attributes & EFI_MEMORY_RO) != 0) {
     if (IsSet) {
       NewPageEntry &= ~(UINT64)IA32_PG_RW;
-      if (mInternalCr3 != 0) {
+      if (mIsShadowStack) {
         // Environment setup
         // ReadOnly page need set Dirty bit for shadow stack
         NewPageEntry |= IA32_PG_D;
@@ -734,10 +735,11 @@ SetShadowStack (
   EFI_STATUS  Status;
 
   SetPageTableBase (Cr3);
-
-  Status = SmmSetMemoryAttributes (BaseAddress, Length, EFI_MEMORY_RO);
+  mIsShadowStack = TRUE;
+  Status         = SmmSetMemoryAttributes (BaseAddress, Length, EFI_MEMORY_RO);
 
   SetPageTableBase (0);
+  mIsShadowStack = FALSE;
 
   return Status;
 }
-- 
2.31.1.windows.1