From 7319b5990bc0e7b7d9c26e928461b8a420d14145 Mon Sep 17 00:00:00 2001
From: Jiaxin Wu <jiaxin.wu@intel.com>
Date: Tue, 10 Sep 2024 11:57:09 +0800
Subject: [PATCH] UefiCpuPkg/PiSmmCpuDxeSmm: Consume
 SmmCpuPlatformHookBeforeMmiHandler func

This patch is for PiSmmCpuDxeSmm driver to add one round wait/release sync
for BSP and AP to perform the SMM CPU Platform Hook before executing MMI
Handler: SmmCpuPlatformHookBeforeMmiHandler (). With the function, SMM CPU
driver can perform the platform specific items after one round BSP and AP
sync (to make sure all APs in SMI) and before the MMI handlers.

After the change, steps #1 and #2 are additional requirements if the
MmCpuSyncModeTradition mode is selected.

Signed-off-by: Jiaxin Wu <jiaxin.wu@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c | 31 ++++++++++++++++++++-------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 2df2cb9f22dbc..e3986fdc15a3d 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -520,16 +520,19 @@ BSPHandler (
     ApCount = CpuCount - 1;
 
     //
-    // Wait for all APs to get ready for programming MTRRs
+    // Wait for all APs of arrival at this point
     //
     SmmCpuSyncWaitForAPs (mSmmMpSyncData->SyncContext, ApCount, CpuIndex); /// #1: Wait APs
 
+    //
+    // Signal all APs it's time for:
+    // 1. Backup MTRRs if needed.
+    // 2. Perform SMM CPU Platform Hook before executing MMI Handler.
+    //
+    ReleaseAllAPs (); /// #2: Signal APs
+
     if (SmmCpuFeaturesNeedConfigureMtrrs ()) {
       //
-      // Signal all APs it's time for backup MTRRs
-      //
-      ReleaseAllAPs (); /// #2: Signal APs
-
       //
       // SmmCpuSyncWaitForAPs() may wait for ever if an AP happens to enter SMM at
       // exactly this point. Please make sure PcdCpuSmmMaxSyncLoops has been set
@@ -564,6 +567,11 @@ BSPHandler (
     }
   }
 
+  //
+  // Perform SMM CPU Platform Hook before executing MMI Handler
+  //
+  SmmCpuPlatformHookBeforeMmiHandler ();
+
   //
   // The BUSY lock is initialized to Acquired state
   //
@@ -806,14 +814,16 @@ APHandler (
     // Notify BSP of arrival at this point
     //
     SmmCpuSyncReleaseBsp (mSmmMpSyncData->SyncContext, CpuIndex, BspIndex); /// #1: Signal BSP
-  }
 
-  if (SmmCpuFeaturesNeedConfigureMtrrs ()) {
     //
-    // Wait for the signal from BSP to backup MTRRs
+    // Wait for the signal from BSP to:
+    // 1. Backup MTRRs if needed.
+    // 2. Perform SMM CPU Platform Hook before executing MMI Handler.
     //
     SmmCpuSyncWaitForBsp (mSmmMpSyncData->SyncContext, CpuIndex, BspIndex); /// #2: Wait BSP
+  }
 
+  if (SmmCpuFeaturesNeedConfigureMtrrs ()) {
     //
     // Backup OS MTRRs
     //
@@ -840,6 +850,11 @@ APHandler (
     SmmCpuSyncReleaseBsp (mSmmMpSyncData->SyncContext, CpuIndex, BspIndex); /// #5: Signal BSP
   }
 
+  //
+  // Perform SMM CPU Platform Hook before executing MMI Handler
+  //
+  SmmCpuPlatformHookBeforeMmiHandler ();
+
   while (TRUE) {
     //
     // Wait for something to happen