[PATCH 3/4] UefiCpuPkg/CpuDxe: implement non-stop mode for uefi

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

Same as SMM profile feature, a special #PF is used to set page attribute
to 'present' and a special #DB handler to reset it back to 'not-present',
right after the instruction causing #PF got executed.

Since the new #PF handler won't enter into dead-loop, the instruction
which caused the #PF will get chance to re-execute with accessible pages.

The exception message will still be printed out on debug console so that
the developer/QA can find that there's potential heap overflow or null
pointer access occurred.

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/CpuDxe/CpuDxe.h       |  39 ++++++
 UefiCpuPkg/CpuDxe/CpuDxe.inf     |   3 +
 UefiCpuPkg/CpuDxe/CpuMp.c        |  34 ++++-
 UefiCpuPkg/CpuDxe/CpuPageTable.c | 271 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 341 insertions(+), 6 deletions(-)

diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.h b/UefiCpuPkg/CpuDxe/CpuDxe.h
index 540f5f2dbf..7d65e39e90 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.h
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.h
@@ -57,6 +57,12 @@
                                        EFI_MEMORY_RO    \
                                        )
 
+#define HEAP_GUARD_NONSTOP_MODE       \
+        ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT6|BIT1|BIT0)) > BIT6)
+
+#define NULL_DETECTION_NONSTOP_MODE   \
+        ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & (BIT6|BIT0)) > BIT6)
+
 /**
   Flush CPU data cache. If the instruction cache is fully coherent
   with all DMA operations then function can just return EFI_SUCCESS.
@@ -273,7 +279,40 @@ RefreshGcdMemoryAttributesFromPaging (
   VOID
   );
 
+/**
+  Special handler for #DB exception, which will restore the page attributes
+  (not-present). It should work with #PF handler which will set pages to
+  'present'.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+**/
+VOID
+EFIAPI
+DebugExceptionHandler (
+  IN EFI_EXCEPTION_TYPE   InterruptType,
+  IN EFI_SYSTEM_CONTEXT   SystemContext
+  );
+
+/**
+  Special handler for #PF exception, which will set the pages which caused
+  #PF to be 'present'. The attribute of those pages should be restored in
+  the subsequent #DB handler.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+**/
+VOID
+EFIAPI
+PageFaultExceptionHandler (
+  IN EFI_EXCEPTION_TYPE   InterruptType,
+  IN EFI_SYSTEM_CONTEXT   SystemContext
+  );
+
 extern BOOLEAN mIsAllocatingPageTable;
+extern UINTN   mNumberOfProcessors;
 
 #endif
 
diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.inf b/UefiCpuPkg/CpuDxe/CpuDxe.inf
index 6a199b72f7..97a381b046 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.inf
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.inf
@@ -46,6 +46,7 @@
   ReportStatusCodeLib
   MpInitLib
   TimerLib
+  PeCoffGetEntryPointLib
 
 [Sources]
   CpuDxe.c
@@ -79,6 +80,8 @@
 [Pcd]
   gEfiMdeModulePkgTokenSpaceGuid.PcdPteMemoryEncryptionAddressOrMask    ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard                       ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPropertyMask               ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdNullPointerDetectionPropertyMask    ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdCpuStackSwitchExceptionList              ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdCpuKnownGoodStackSize                    ## CONSUMES
 
diff --git a/UefiCpuPkg/CpuDxe/CpuMp.c b/UefiCpuPkg/CpuDxe/CpuMp.c
index 82145e7624..f8489eb1c3 100644
--- a/UefiCpuPkg/CpuDxe/CpuMp.c
+++ b/UefiCpuPkg/CpuDxe/CpuMp.c
@@ -673,10 +673,6 @@ InitializeMpExceptionStackSwitchHandlers (
   UINT8                           *GdtBuffer;
   UINT8                           *StackTop;
 
-  if (!PcdGetBool (PcdCpuStackGuard)) {
-    return;
-  }
-
   ExceptionNumber = FixedPcdGetSize (PcdCpuStackSwitchExceptionList);
   NewStackSize = FixedPcdGet32 (PcdCpuKnownGoodStackSize) * ExceptionNumber;
 
@@ -790,6 +786,32 @@ InitializeMpExceptionStackSwitchHandlers (
   }
 }
 
+/**
+  Initializes MP exceptions handlers for special features, such as Heap Guard
+  and Stack Guard.
+**/
+VOID
+InitializeMpExceptionHandlers (
+  VOID
+  )
+{
+  //
+  // Enable non-stop mode for #PF triggered by Heap Guard or NULL Pointer
+  // Detection.
+  //
+  if (HEAP_GUARD_NONSTOP_MODE || NULL_DETECTION_NONSTOP_MODE) {
+    RegisterCpuInterruptHandler(EXCEPT_IA32_DEBUG, DebugExceptionHandler);
+    RegisterCpuInterruptHandler(EXCEPT_IA32_PAGE_FAULT, PageFaultExceptionHandler);
+  }
+
+  //
+  // Setup stack switch for Stack Guard feature.
+  //
+  if (PcdGetBool (PcdCpuStackGuard)) {
+    InitializeMpExceptionStackSwitchHandlers();
+  }
+}
+
 /**
   Initialize Multi-processor support.
 
@@ -814,9 +836,9 @@ InitializeMpSupport (
   DEBUG ((DEBUG_INFO, "Detect CPU count: %d\n", mNumberOfProcessors));
 
   //
-  // Initialize exception stack switch handlers for each logic processor.
+  // Initialize special exception handlers for each logic processor.
   //
-  InitializeMpExceptionStackSwitchHandlers ();
+  InitializeMpExceptionHandlers ();
 
   //
   // Update CPU healthy information from Guided HOB
diff --git a/UefiCpuPkg/CpuDxe/CpuPageTable.c b/UefiCpuPkg/CpuDxe/CpuPageTable.c
index df021798c0..b65f74bd72 100644
--- a/UefiCpuPkg/CpuDxe/CpuPageTable.c
+++ b/UefiCpuPkg/CpuDxe/CpuPageTable.c
@@ -22,6 +22,10 @@
 #include <Library/MemoryAllocationLib.h>
 #include <Library/DebugLib.h>
 #include <Library/UefiBootServicesTableLib.h>
+#include <Library/PeCoffGetEntryPointLib.h>
+#include <Library/SerialPortLib.h>
+#include <Library/SynchronizationLib.h>
+#include <Library/PrintLib.h>
 #include <Protocol/MpService.h>
 #include <Protocol/SmmBase2.h>
 #include <Register/Cpuid.h>
@@ -73,6 +77,10 @@
 #define PAGING_2M_ADDRESS_MASK_64 0x000FFFFFFFE00000ull
 #define PAGING_1G_ADDRESS_MASK_64 0x000FFFFFC0000000ull
 
+#define MAX_PF_ENTRY_COUNT        10
+#define MAX_DEBUG_MESSAGE_LENGTH  0x100
+#define IA32_PF_EC_ID             BIT4
+
 typedef enum {
   PageNone,
   Page4K,
@@ -102,6 +110,13 @@ PAGE_TABLE_POOL                   *mPageTablePool = NULL;
 PAGE_TABLE_LIB_PAGING_CONTEXT     mPagingContext;
 EFI_SMM_BASE2_PROTOCOL            *mSmmBase2 = NULL;
 
+//
+// Record the page fault exception count for one instruction execution.
+//
+UINTN                     *mPFEntryCount;
+UINT64                    (*mLastPFEntryValue)[MAX_PF_ENTRY_COUNT];
+UINT64                    *(*mLastPFEntryPointer)[MAX_PF_ENTRY_COUNT];
+
 /**
  Check if current execution environment is in SMM mode or not, via
  EFI_SMM_BASE2_PROTOCOL.
@@ -1135,6 +1150,253 @@ AllocatePageTableMemory (
   return Buffer;
 }
 
+/**
+  Prints a message to the serial port.
+
+  @param  Format      Format string for the message to print.
+  @param  ...         Variable argument list whose contents are accessed
+                      based on the format string specified by Format.
+
+**/
+STATIC
+VOID
+EFIAPI
+InternalPrintMessage (
+  IN  CONST CHAR8  *Format,
+  ...
+  )
+{
+  CHAR8    Buffer[MAX_DEBUG_MESSAGE_LENGTH];
+  VA_LIST  Marker;
+
+  //
+  // Convert the message to an ASCII String
+  //
+  VA_START (Marker, Format);
+  AsciiVSPrint (Buffer, sizeof (Buffer), Format, Marker);
+  VA_END (Marker);
+
+  //
+  // Send the print string to a Serial Port
+  //
+  SerialPortWrite ((UINT8 *)Buffer, AsciiStrLen (Buffer));
+}
+
+/**
+  Find and display image base address and return image base and its entry point.
+
+  @param CurrentIp      Current instruction pointer.
+
+**/
+STATIC
+VOID
+DumpModuleImageInfo (
+  IN  UINTN              CurrentIp
+  )
+{
+  EFI_STATUS                           Status;
+  UINTN                                Pe32Data;
+  VOID                                 *PdbPointer;
+  VOID                                 *EntryPoint;
+
+  Pe32Data = PeCoffSearchImageBase (CurrentIp);
+  if (Pe32Data == 0) {
+    InternalPrintMessage ("!!!! Can't find image information. !!!!\n");
+  } else {
+    //
+    // Find Image Base entry point
+    //
+    Status = PeCoffLoaderGetEntryPoint ((VOID *) Pe32Data, &EntryPoint);
+    if (EFI_ERROR (Status)) {
+      EntryPoint = NULL;
+    }
+    InternalPrintMessage ("!!!! Find image based on IP(0x%x) ", CurrentIp);
+    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *) Pe32Data);
+    if (PdbPointer != NULL) {
+      InternalPrintMessage ("%a", PdbPointer);
+    } else {
+      InternalPrintMessage ("(No PDB) " );
+    }
+    InternalPrintMessage (
+      " (ImageBase=%016lp, EntryPoint=%016p) !!!!\n",
+      (VOID *) Pe32Data,
+      EntryPoint
+      );
+  }
+}
+
+/**
+  Special handler for #DB exception, which will restore the page attributes
+  (not-present). It should work with #PF handler which will set pages to
+  'present'.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+**/
+VOID
+EFIAPI
+DebugExceptionHandler (
+  IN EFI_EXCEPTION_TYPE   ExceptionType,
+  IN EFI_SYSTEM_CONTEXT   SystemContext
+  )
+{
+  UINTN     CpuIndex;
+  UINTN     PFEntry;
+  BOOLEAN   IsWpEnabled;
+
+  MpInitLibWhoAmI (&CpuIndex);
+
+  //
+  // Clear last PF entries
+  //
+  IsWpEnabled = IsReadOnlyPageWriteProtected ();
+  if (IsWpEnabled) {
+    DisableReadOnlyPageWriteProtect ();
+  }
+
+  for (PFEntry = 0; PFEntry < mPFEntryCount[CpuIndex]; PFEntry++) {
+    if (mLastPFEntryPointer[CpuIndex][PFEntry] != NULL) {
+      *mLastPFEntryPointer[CpuIndex][PFEntry] &= ~IA32_PG_P;
+    }
+  }
+
+  if (IsWpEnabled) {
+    EnableReadOnlyPageWriteProtect ();
+  }
+
+  //
+  // Reset page fault exception count for next page fault.
+  //
+  mPFEntryCount[CpuIndex] = 0;
+
+  //
+  // Flush TLB
+  //
+  CpuFlushTlb ();
+
+  //
+  // Clear TF in EFLAGS
+  //
+  if (mPagingContext.MachineType == IMAGE_FILE_MACHINE_I386) {
+    SystemContext.SystemContextIa32->Eflags &= (UINT32)~BIT8;
+  } else {
+    SystemContext.SystemContextX64->Rflags &= (UINT64)~BIT8;
+  }
+}
+
+/**
+  Special handler for #PF exception, which will set the pages which caused
+  #PF to be 'present'. The attribute of those pages should be restored in
+  the subsequent #DB handler.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+**/
+VOID
+EFIAPI
+PageFaultExceptionHandler (
+  IN EFI_EXCEPTION_TYPE   ExceptionType,
+  IN EFI_SYSTEM_CONTEXT   SystemContext
+  )
+{
+  EFI_STATUS                      Status;
+  UINT64                          PFAddress;
+  PAGE_TABLE_LIB_PAGING_CONTEXT   PagingContext;
+  PAGE_ATTRIBUTE                  PageAttribute;
+  UINT64                          Attributes;
+  UINT64                          *PageEntry;
+  UINTN                           Index;
+  UINTN                           CpuIndex;
+  UINTN                           PageNumber;
+  UINTN                           CurrentIp;
+  BOOLEAN                         NonStopMode;
+
+  PFAddress = AsmReadCr2 () & ~EFI_PAGE_MASK;
+  if (PFAddress < BASE_4KB) {
+    NonStopMode = NULL_DETECTION_NONSTOP_MODE ? TRUE : FALSE;
+  } else {
+    NonStopMode = HEAP_GUARD_NONSTOP_MODE ? TRUE : FALSE;
+  }
+
+  if (NonStopMode) {
+    MpInitLibWhoAmI(&CpuIndex);
+    GetCurrentPagingContext (&PagingContext);
+    //
+    // Memory operation cross page boundary, 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.
+    //
+    PageNumber = 2;
+    while (PageNumber > 0) {
+      PageEntry = GetPageTableEntry(&PagingContext, PFAddress, &PageAttribute);
+      ASSERT(PageEntry != NULL);
+
+      if (PageEntry != NULL) {
+        Attributes = GetAttributesFromPageEntry(PageEntry);
+        if ((Attributes & EFI_MEMORY_RP) != 0) {
+          Attributes &= ~EFI_MEMORY_RP;
+          Status = AssignMemoryPageAttributes (&PagingContext, PFAddress,
+                                               EFI_PAGE_SIZE, Attributes, NULL);
+          if (!EFI_ERROR(Status)) {
+            Index = mPFEntryCount[CpuIndex];
+            //
+            // Re-retrieve page entry because above calling might update page
+            // table due to table split.
+            //
+            PageEntry = GetPageTableEntry(&PagingContext, PFAddress, &PageAttribute);
+            mLastPFEntryPointer[CpuIndex][Index++] = PageEntry;
+            mPFEntryCount[CpuIndex] = Index;
+          }
+        }
+      }
+
+      PFAddress += EFI_PAGE_SIZE;
+      --PageNumber;
+    }
+  }
+
+  //
+  // Initialize the serial port before dumping.
+  //
+  SerialPortInitialize ();
+  //
+  // Display ExceptionType, CPU information and Image information
+  //
+  DumpCpuContext (ExceptionType, SystemContext);
+  //
+  // Dump module image base and module entry point by RIP
+  //
+  if (mPagingContext.MachineType == IMAGE_FILE_MACHINE_I386) {
+    if ((SystemContext.SystemContextIa32->ExceptionData & IA32_PF_EC_ID) != 0) {
+      //
+      // The EIP in SystemContext could not be used
+      // if it is page fault with I/D set.
+      //
+      CurrentIp = *(UINTN *)(UINTN)SystemContext.SystemContextIa32->Esp;
+    } else {
+      CurrentIp = (UINTN)SystemContext.SystemContextIa32->Eip;
+    }
+  } else {
+    if ((SystemContext.SystemContextX64->ExceptionData & IA32_PF_EC_ID) != 0) {
+      //
+      // The RIP in SystemContext could not be used
+      // if it is page fault with I/D set.
+      //
+      CurrentIp = *(UINTN *)(UINTN)SystemContext.SystemContextX64->Rsp;
+    } else {
+      CurrentIp = (UINTN)SystemContext.SystemContextX64->Rip;
+    }
+  }
+
+  DumpModuleImageInfo (CurrentIp);
+
+  if (!NonStopMode) {
+    CpuDeadLoop();
+  }
+}
+
 /**
   Initialize the Page Table lib.
 **/
@@ -1158,6 +1420,15 @@ InitializePageTableLib (
     EnableReadOnlyPageWriteProtect ();
   }
 
+  if (HEAP_GUARD_NONSTOP_MODE || NULL_DETECTION_NONSTOP_MODE) {
+    mPFEntryCount = (UINTN *)AllocateZeroPool (sizeof (UINTN) * mNumberOfProcessors);
+    ASSERT (mPFEntryCount != NULL);
+
+    mLastPFEntryPointer = (UINT64 *(*)[MAX_PF_ENTRY_COUNT])
+                          AllocateZeroPool (sizeof (mLastPFEntryPointer[0]) * mNumberOfProcessors);
+    ASSERT (mLastPFEntryPointer != NULL);
+  }
+
   DEBUG ((DEBUG_INFO, "CurrentPagingContext:\n", CurrentPagingContext.MachineType));
   DEBUG ((DEBUG_INFO, "  MachineType   - 0x%x\n", CurrentPagingContext.MachineType));
   DEBUG ((DEBUG_INFO, "  PageTableBase - 0x%x\n", CurrentPagingContext.ContextData.X64.PageTableBase));
-- 
2.16.2.windows.1