[Patch 1/1] UefiCpuPkg/Library/MpInitLib: Fix AP VolatileRegisters race condition

["Michael D Kinney" <michael.d.kinney@intel.com>]

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=3182

Fix the order of operations in ApWakeupFunction() when PcdCpuApLoopMode
is set to HLT mode that uses INIT-SIPI-SIPI to wake APs.  In this mode,
volatile state is restored and saved each time a INIT-SIPI-SIPI is sent
to an AP to request a function to be executed on the AP.  When the
function is completed the volatile state of the AP is saved.  However,
the counters NumApsExecuting and FinishedCount are updated before
the volatile state is saved.  This allows for a race condition window
for the BSP that is waiting on these counters to request a new
INIT-SIPI-SIPI before all the APs have completely saved their volatile
state.  The fix is to save the AP volatile state before updating the
NumApsExecuting and FinishedCount counters.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Signed-off-by: Michael D Kinney <michael.d.kinney@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/MpLib.c | 31 ++++++++++++++++------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 681fa79b4cff..8b1f7f84bad6 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -769,15 +769,6 @@ ApWakeupFunction (
       RestoreVolatileRegisters (&CpuMpData->CpuData[0].VolatileRegisters, FALSE);
       InitializeApData (CpuMpData, ProcessorNumber, BistData, ApTopOfStack);
       ApStartupSignalBuffer = CpuMpData->CpuData[ProcessorNumber].StartupApSignal;
-
-      //
-      // Delay decrementing the APs executing count when SEV-ES is enabled
-      // to allow the APs to issue an AP_RESET_HOLD before the BSP possibly
-      // performs another INIT-SIPI-SIPI sequence.
-      //
-      if (!CpuMpData->SevEsIsEnabled) {
-        InterlockedDecrement ((UINT32 *) &CpuMpData->MpCpuExchangeInfo->NumApsExecuting);
-      }
     } else {
       //
       // Execute AP function if AP is ready
@@ -866,19 +857,33 @@ ApWakeupFunction (
       }
     }
 
+    if (CpuMpData->ApLoopMode == ApInHltLoop) {
+      //
+      // Save AP volatile registers
+      //
+      SaveVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters);
+    }
+
     //
     // AP finished executing C code
     //
     InterlockedIncrement ((UINT32 *) &CpuMpData->FinishedCount);
 
+    if (CpuMpData->InitFlag == ApInitConfig) {
+      //
+      // Delay decrementing the APs executing count when SEV-ES is enabled
+      // to allow the APs to issue an AP_RESET_HOLD before the BSP possibly
+      // performs another INIT-SIPI-SIPI sequence.
+      //
+      if (!CpuMpData->SevEsIsEnabled) {
+        InterlockedDecrement ((UINT32 *) &CpuMpData->MpCpuExchangeInfo->NumApsExecuting);
+      }
+    }
+
     //
     // Place AP is specified loop mode
     //
     if (CpuMpData->ApLoopMode == ApInHltLoop) {
-      //
-      // Save AP volatile registers
-      //
-      SaveVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters);
       //
       // Place AP in HLT-loop
       //
-- 
2.29.2.windows.2