From 83d5871184d1e09332565bfc939e5fc8354b5b79 Mon Sep 17 00:00:00 2001
From: Dun Tan <dun.tan@intel.com>
Date: Tue, 9 Aug 2022 15:22:07 +0800
Subject: [PATCH] UefiCpuPkg/PiSmmCpuDxeSmm: Add a new mIsShadowStack flag

This patch is code refactoring and doesn't change any functionality.
Add a new mIsShadowStack 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: Rahul Kumar <rahul1.kumar@intel.com>
Reviewed-by: Ray Ni <ray.ni@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 1f7cc15727..10de8dd8e4 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -32,7 +32,8 @@ PAGE_ATTRIBUTE_TABLE  mPageAttributeTable[] = {
   { Page1G, SIZE_1GB, PAGING_1G_ADDRESS_MASK_64 },
 };
 
-UINTN  mInternalCr3;
+UINTN    mInternalCr3;
+BOOLEAN  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;
 }