[PATCH 4/4] UefiCpuPkg/PiSmmCpuDxeSmm: implement non-stop mode for SMM

[Jian J Wang <jian.j.wang@intel.com>]

Since SMM profile feature has already implemented non-stop mode if #PF
occurred, this patch just makes use of the existing implementation to
accommodate heap guard and NULL pointer detection feature.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c         | 43 +++++++++++------
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiException.nasm |  3 +-
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c           | 60 +++++++++++++++++++++++-
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.h           | 15 ++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h   |  6 +++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c          | 43 +++++++++++------
 6 files changed, 139 insertions(+), 31 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
index 9300a232e4..a32b736089 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
@@ -38,7 +38,9 @@ SmmInitPageTable (
 
   mPhysicalAddressBits = 32;
 
-  if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
+  if (FeaturePcdGet (PcdCpuSmmProfileEnable) ||
+      HEAP_GUARD_NONSTOP_MODE ||
+      NULL_DETECTION_NONSTOP_MODE) {
     //
     // Set own Page Fault entry instead of the default one, because SMM Profile
     // feature depends on IRET instruction to do Single Step
@@ -129,6 +131,11 @@ SmiPFHandler (
           DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextIa32->Eip);
         );
       }
+
+      if (HEAP_GUARD_NONSTOP_MODE) {
+        GuardPagePFHandler (SystemContext.SystemContextIa32->ExceptionData);
+        goto Exit;
+      }
     }
     CpuDeadLoop ();
   }
@@ -146,6 +153,26 @@ SmiPFHandler (
       );
       CpuDeadLoop ();
     }
+
+    //
+    // If NULL pointer was just accessed
+    //
+    if ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0 &&
+        (PFAddress < EFI_PAGE_SIZE)) {
+      DumpCpuContext (InterruptType, SystemContext);
+      DEBUG ((DEBUG_ERROR, "!!! NULL pointer access !!!\n"));
+      DEBUG_CODE (
+        DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextIa32->Eip);
+      );
+
+      if (NULL_DETECTION_NONSTOP_MODE) {
+        GuardPagePFHandler (SystemContext.SystemContextIa32->ExceptionData);
+        goto Exit;
+      }
+
+      CpuDeadLoop ();
+    }
+
     if (IsSmmCommBufferForbiddenAddress (PFAddress)) {
       DumpCpuContext (InterruptType, SystemContext);
       DEBUG ((DEBUG_ERROR, "Access SMM communication forbidden address (0x%x)!\n", PFAddress));
@@ -156,19 +183,6 @@ SmiPFHandler (
     }
   }
 
-  //
-  // If NULL pointer was just accessed
-  //
-  if ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0 &&
-      (PFAddress < EFI_PAGE_SIZE)) {
-    DumpCpuContext (InterruptType, SystemContext);
-    DEBUG ((DEBUG_ERROR, "!!! NULL pointer access !!!\n"));
-    DEBUG_CODE (
-      DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextIa32->Eip);
-    );
-    CpuDeadLoop ();
-  }
-
   if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
     SmmProfilePFHandler (
       SystemContext.SystemContextIa32->Eip,
@@ -179,6 +193,7 @@ SmiPFHandler (
     SmiDefaultPFHandler ();
   }
 
+Exit:
   ReleaseSpinLock (mPFLock);
 }
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiException.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiException.nasm
index fa02c1016c..879fa0ba63 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiException.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiException.nasm
@@ -20,6 +20,7 @@
 
 extern  ASM_PFX(FeaturePcdGet (PcdCpuSmmProfileEnable))
 extern  ASM_PFX(SmiPFHandler)
+extern  ASM_PFX(mSetupDebugTrap)
 
 global  ASM_PFX(gcSmiIdtr)
 global  ASM_PFX(gcSmiGdtr)
@@ -673,7 +674,7 @@ o16 mov     [ecx + IA32_TSS._SS], ax
     mov     esp, ebp
 
 ; Set single step DB# if SMM profile is enabled and page fault exception happens
-    cmp     byte [dword ASM_PFX(FeaturePcdGet (PcdCpuSmmProfileEnable))], 0
+    cmp     byte [dword ASM_PFX(mSetupDebugTrap)], 0
     jz      @Done2
 
 ; Create return context for iretd in stub function
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
index b4fe0bc23b..2cc46c8c61 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
@@ -51,6 +51,11 @@ BOOLEAN                   mBtsSupported     = TRUE;
 //
 BOOLEAN                   mSmmProfileStart = FALSE;
 
+//
+// The flag indicates if #DB will be setup in #PF handler.
+//
+BOOLEAN                   mSetupDebugTrap = FALSE;
+
 //
 // Record the page fault exception count for one instruction execution.
 //
@@ -229,7 +234,9 @@ DebugExceptionHandler (
   UINTN  CpuIndex;
   UINTN  PFEntry;
 
-  if (!mSmmProfileStart) {
+  if (!mSmmProfileStart &&
+      !HEAP_GUARD_NONSTOP_MODE &&
+      !NULL_DETECTION_NONSTOP_MODE) {
     return;
   }
   CpuIndex = GetCpuIndex ();
@@ -1174,7 +1181,9 @@ InitSmmProfile (
   //
   // Skip SMM profile initialization if feature is disabled
   //
-  if (!FeaturePcdGet (PcdCpuSmmProfileEnable)) {
+  if (!FeaturePcdGet (PcdCpuSmmProfileEnable) &&
+      !HEAP_GUARD_NONSTOP_MODE &&
+      !NULL_DETECTION_NONSTOP_MODE) {
     return;
   }
 
@@ -1187,6 +1196,11 @@ InitSmmProfile (
   // Initialize profile IDT.
   //
   InitIdtr ();
+
+  //
+  // Tell #PF handler to prepare a #DB subsequently.
+  //
+  mSetupDebugTrap = TRUE;
 }
 
 /**
@@ -1294,6 +1308,48 @@ RestorePageTableBelow4G (
   }
 }
 
+/**
+  Handler for Page Fault triggered by Guard page.
+
+  @param  ErrorCode  The Error code of exception.
+
+**/
+VOID
+GuardPagePFHandler (
+  UINTN ErrorCode
+  )
+{
+  UINT64                *PageTable;
+  UINT64                PFAddress;
+  UINT64                RestoreAddress;
+  UINTN                 RestorePageNumber;
+  UINTN                 CpuIndex;
+  BOOLEAN               IsValidPFAddress;
+
+  IsValidPFAddress  = FALSE;
+  PageTable         = (UINT64 *)AsmReadCr3 ();
+  PFAddress         = AsmReadCr2 ();
+  CpuIndex          = GetCpuIndex ();
+
+  //
+  // Memory operation cross pages, like "rep mov" instruction, will cause
+  // infinite loop between this and Debug Trap handler. We have to make sure
+  // that current page and the page followed are both in PRESENT state.
+  //
+  RestorePageNumber = 2;
+  RestoreAddress = PFAddress;
+  while (RestorePageNumber > 0) {
+    RestorePageTableBelow4G (PageTable, RestoreAddress, CpuIndex, ErrorCode);
+    RestoreAddress += EFI_PAGE_SIZE;
+    RestorePageNumber--;
+  }
+
+  //
+  // Flush TLB
+  //
+  CpuFlushTlb ();
+}
+
 /**
   The Page fault handler to save SMM profile data.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.h b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.h
index 04a3dfb2e8..c2a48235ab 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.h
@@ -114,6 +114,17 @@ GetCpuIndex (
   VOID
   );
 
+/**
+  Handler for Page Fault triggered by Guard page.
+
+  @param  ErrorCode  The Error code of exception.
+
+**/
+VOID
+GuardPagePFHandler (
+  UINTN ErrorCode
+  );
+
 //
 // The flag indicates if execute-disable is supported by processor.
 //
@@ -122,5 +133,9 @@ extern BOOLEAN    mXdSupported;
 // The flag indicates if execute-disable is enabled on processor.
 //
 extern BOOLEAN    mXdEnabled;
+//
+// The flag indicates if #DB will be setup in #PF handler.
+//
+extern BOOLEAN    mSetupDebugTrap;
 
 #endif // _SMM_PROFILE_H_
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h
index 1613e9cd5c..bacb2f8ad3 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h
@@ -64,6 +64,12 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #define   MSR_DEBUG_CTL_BTINT        0x100
 #define MSR_DS_AREA                  0x600
 
+#define HEAP_GUARD_NONSTOP_MODE      \
+        ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT6|BIT3|BIT2)) > BIT6)
+
+#define NULL_DETECTION_NONSTOP_MODE  \
+        ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & (BIT6|BIT1)) > BIT6)
+
 typedef struct {
   EFI_PHYSICAL_ADDRESS   Base;
   EFI_PHYSICAL_ADDRESS   Top;
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
index 0fe944fc18..5bb7d57238 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
@@ -300,7 +300,9 @@ SmmInitPageTable (
     }
   }
 
-  if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
+  if (FeaturePcdGet (PcdCpuSmmProfileEnable) ||
+      HEAP_GUARD_NONSTOP_MODE ||
+      NULL_DETECTION_NONSTOP_MODE) {
     //
     // Set own Page Fault entry instead of the default one, because SMM Profile
     // feature depends on IRET instruction to do Single Step
@@ -846,6 +848,11 @@ SmiPFHandler (
           DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);
         );
       }
+
+      if (HEAP_GUARD_NONSTOP_MODE) {
+        GuardPagePFHandler (SystemContext.SystemContextX64->ExceptionData);
+        goto Exit;
+      }
     }
     CpuDeadLoop ();
   }
@@ -863,6 +870,26 @@ SmiPFHandler (
       );
       CpuDeadLoop ();
     }
+
+    //
+    // If NULL pointer was just accessed
+    //
+    if ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0 &&
+        (PFAddress < EFI_PAGE_SIZE)) {
+      DumpCpuContext (InterruptType, SystemContext);
+      DEBUG ((DEBUG_ERROR, "!!! NULL pointer access !!!\n"));
+      DEBUG_CODE (
+        DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);
+      );
+
+      if (NULL_DETECTION_NONSTOP_MODE) {
+        GuardPagePFHandler (SystemContext.SystemContextX64->ExceptionData);
+        goto Exit;
+      }
+
+      CpuDeadLoop ();
+    }
+
     if (IsSmmCommBufferForbiddenAddress (PFAddress)) {
       DumpCpuContext (InterruptType, SystemContext);
       DEBUG ((DEBUG_ERROR, "Access SMM communication forbidden address (0x%lx)!\n", PFAddress));
@@ -873,19 +900,6 @@ SmiPFHandler (
     }
   }
 
-  //
-  // If NULL pointer was just accessed
-  //
-  if ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0 &&
-      (PFAddress < EFI_PAGE_SIZE)) {
-    DumpCpuContext (InterruptType, SystemContext);
-    DEBUG ((DEBUG_ERROR, "!!! NULL pointer access !!!\n"));
-    DEBUG_CODE (
-      DumpModuleInfoByIp ((UINTN)SystemContext.SystemContextX64->Rip);
-    );
-    CpuDeadLoop ();
-  }
-
   if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
     SmmProfilePFHandler (
       SystemContext.SystemContextX64->Rip,
@@ -895,6 +909,7 @@ SmiPFHandler (
     SmiDefaultPFHandler ();
   }
 
+Exit:
   ReleaseSpinLock (mPFLock);
 }
 
-- 
2.16.2.windows.1