From f7cf7d5e848f401b49ab6d70b8ba933b96d73405 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 BSP and AP finish the
sync and before the MMI handlers.

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

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 210b23af21d8d..6a62be8a637df 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -520,16 +520,18 @@ 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);
 
-    if (SmmCpuFeaturesNeedConfigureMtrrs ()) {
-      //
-      // Signal all APs it's time for backup MTRRs
-      //
-      ReleaseAllAPs ();
+    //
+    // Signal all APs it's time for:
+    // 1. Backup MTRRs if needed.
+    // 2. Perform SMM CPU Platform Hook before executing MMI Handler.
+    //
+    ReleaseAllAPs ();
 
+    if (SmmCpuFeaturesNeedConfigureMtrrs ()) {
       //
       // 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 +566,11 @@ BSPHandler (
     }
   }
 
+  //
+  // Perform SMM CPU Platform Hook before executing MMI Handler
+  //
+  SmmCpuPlatformHookBeforeMmiHandler ();
+
   //
   // The BUSY lock is initialized to Acquired state
   //
@@ -806,14 +813,16 @@ APHandler (
     // Notify BSP of arrival at this point
     //
     SmmCpuSyncReleaseBsp (mSmmMpSyncData->SyncContext, CpuIndex, BspIndex);
-  }
 
-  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);
+  }
 
+  if (SmmCpuFeaturesNeedConfigureMtrrs ()) {
     //
     // Backup OS MTRRs
     //
@@ -840,6 +849,11 @@ APHandler (
     SmmCpuSyncReleaseBsp (mSmmMpSyncData->SyncContext, CpuIndex, BspIndex);
   }
 
+  //
+  // Perform SMM CPU Platform Hook before executing MMI Handler
+  //
+  SmmCpuPlatformHookBeforeMmiHandler ();
+
   while (TRUE) {
     //
     // Wait for something to happen