UefiCpuPkg/PiSmmCpuDxeSmm: Enable single step after SmmProfile start

There is a bug in the existing code: the single step is always enabled
once the Page Fault (#PF) occurs, but it is only disabled when the SMM
Profile feature actually starts (see DebugExceptionHandler).
If the SMM Profile feature has not been started, this will result in
the single-step mode remaining enabled if a Page Fault occurs.

This patch is to enable the single-step debugging mode by setting the
Trap Flag only after SmmProfile feature starts.

Signed-off-by: Jiaxin Wu <jiaxin.wu@intel.com>

UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c

@@ -759,6 +759,11 @@ SmmProfileStart (
   // The flag indicates SMM profile starts to work.
   //
   mSmmProfileStart = TRUE;
+
+  //
+  // Tell #PF handler to prepare a #DB subsequently.
+  //
+  mSetupDebugTrap = TRUE;
 }
 
 /**
@@ -1146,11 +1151,6 @@ InitSmmProfile (
   // Initialize profile IDT.
   //
   InitIdtr ();
-
-  //
-  // Tell #PF handler to prepare a #DB subsequently.
-  //
-  mSetupDebugTrap = TRUE;
 }
 
 /**

UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiException.nasm

@@ -13,6 +13,7 @@
 ;-------------------------------------------------------------------------------
 
 extern  ASM_PFX(SmiPFHandler)
+extern  ASM_PFX(mSetupDebugTrap)
 
 global  ASM_PFX(gcSmiIdtr)
 global  ASM_PFX(gcSmiGdtr)
@@ -369,9 +370,14 @@ ASM_PFX(PageFaultIdtHandlerSmmProfile):
 
     mov     rsp, rbp
 
+; Check if mSetupDebugTrap is TRUE (non-zero)
+    cmp     byte [dword ASM_PFX(mSetupDebugTrap)], 0
+    jz      SkipSettingTF
+
 ; Enable TF bit after page fault handler runs
     bts     dword [rsp + 40], 8  ;RFLAGS
 
+SkipSettingTF:
     pop     rbp
     add     rsp, 16           ; skip INT# & ErrCode
     iretq