[PATCH] UefiCpuPkg/MpInitLib: AP uses memory preceding IDT to store CpuMpData

[Ruiyu Ni <ruiyu.ni@intel.com>]

Today's MpInitLib PEI implementation directly calls
PeiServices->GetHobList() from AP which may cause racing issue.

This patch fixes this issue by storing the CpuMpData to memory
preceding IDT. Pointer to PeiServices pointer is stored there,
so after AP procedure returns, the PeiServices pointer should be
restored.

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Jeff Fan <vanjeff_919@hotmail.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Fish Andrew <afish@apple.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c | 33 ++++++++++++++++++-
 UefiCpuPkg/Library/MpInitLib/MpLib.c    |  8 +++++
 UefiCpuPkg/Library/MpInitLib/MpLib.h    | 27 +++++++++++++++-
 UefiCpuPkg/Library/MpInitLib/PeiMpLib.c | 56 +++++++++++++++++++++++++++++++--
 4 files changed, 119 insertions(+), 5 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index e7ed21c6cd..26fead2c66 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -1,7 +1,7 @@
 /** @file
   MP initialize support functions for DXE phase.
 
-  Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2016 - 2018, Intel Corporation. All rights reserved.<BR>
   This program and the accompanying materials
   are licensed and made available under the terms and conditions of the BSD License
   which accompanies this distribution.  The full text of the license may be found at
@@ -75,6 +75,37 @@ SaveCpuMpData (
   mCpuMpData = CpuMpData;
 }
 
+/**
+  Push the CpuMpData for AP to use.
+
+  @param[in]  The pointer to CPU MP Data structure will be pushed.
+  @param[out] The pointer to the context which will be passed to PopCpuMpData().
+
+  @return  The pointer value which was stored in where the CPU MP Data is pushed.
+**/
+VOID *
+PushCpuMpData (
+  IN  CPU_MP_DATA    *CpuMpData,
+  OUT VOID           **Context
+  )
+{
+  return NULL;
+}
+
+/**
+  Pop the CpuMpData.
+
+  @param[in] Pointer  The pointer value which was stored in where the CPU MP Data is pushed.
+  @param[in] Context  The context of push/pop operation.
+**/
+VOID
+PopCpuMpData (
+  IN VOID           *Pointer,
+  IN VOID           *Context
+  )
+{  
+}
+
 /**
   Get available system memory below 1MB by specified size.
 
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index f2ff40417a..786a7825d5 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -580,6 +580,8 @@ ApWakeupFunction (
   CPU_INFO_IN_HOB            *CpuInfoInHob;
   UINT64                     ApTopOfStack;
   UINTN                      CurrentApicMode;
+  VOID                       *BackupPtr;
+  VOID                       *Context;
 
   //
   // AP finished assembly code and begin to execute C code
@@ -659,8 +661,14 @@ ApWakeupFunction (
           EnableDebugAgent ();
           //
           // Invoke AP function here
+          // Use a BSP owned area (PeiServices Pointer storage) to store the CpuMpData.
+          // It's required in PEI phase because CpuMpData cannot be cached in global variable as in DXE phase.
+          // DXE version of Pushxxx andPopxxx is dummy implementation.
           //
+          BackupPtr = PushCpuMpData (CpuMpData, &Context);
           Procedure (Parameter);
+          PopCpuMpData (BackupPtr, Context);
+
           CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;
           if (CpuMpData->SwitchBspFlag) {
             //
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index e7f9a4de0a..270d62ff20 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -1,7 +1,7 @@
 /** @file
   Common header file for MP Initialize Library.
 
-  Copyright (c) 2016 - 2017, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2016 - 2018, Intel Corporation. All rights reserved.<BR>
   This program and the accompanying materials
   are licensed and made available under the terms and conditions of the BSD License
   which accompanies this distribution.  The full text of the license may be found at
@@ -321,6 +321,31 @@ SaveCpuMpData (
   IN CPU_MP_DATA   *CpuMpData
   );
 
+/**
+  Push the CpuMpData for AP to use.
+
+  @param[in]  The pointer to CPU MP Data structure will be pushed.
+  @param[out] The pointer to the context which will be passed to PopCpuMpData().
+
+  @return  The pointer value which was stored in where the CPU MP Data is pushed.
+**/
+VOID *
+PushCpuMpData (
+  IN  CPU_MP_DATA    *CpuMpData,
+  OUT VOID           **Context
+  );
+
+/**
+  Pop the CpuMpData.
+
+  @param[in] Pointer  The pointer value which was stored in where the CPU MP Data is pushed.
+  @param[in] Context  The context of push/pop operation.
+**/
+VOID
+PopCpuMpData (
+  IN VOID           *Pointer,
+  IN VOID           *Context
+  );
 
 /**
   Get available system memory below 1MB by specified size.
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
index 791ae9db6e..5c9c4b3b1e 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
@@ -27,6 +27,9 @@ EnableDebugAgent (
 
 /**
   Get pointer to CPU MP Data structure.
+  For BSP, the pointer is retrieved from HOB.
+  For AP, the pointer is retrieved from the location which stores the PeiServices pointer.
+  It's safe because BSP is blocking and has no chance to use PeiServices pointer when AP is executing.
 
   @return  The pointer to CPU MP Data structure.
 **/
@@ -35,9 +38,17 @@ GetCpuMpData (
   VOID
   )
 {
-  CPU_MP_DATA      *CpuMpData;
-
-  CpuMpData = GetCpuMpDataFromGuidedHob ();
+  CPU_MP_DATA                  *CpuMpData;
+  MSR_IA32_APIC_BASE_REGISTER  ApicBaseMsr;
+  IA32_DESCRIPTOR              Idtr;
+
+  ApicBaseMsr.Uint64 = AsmReadMsr64 (MSR_IA32_APIC_BASE);
+  if (ApicBaseMsr.Bits.BSP == 1) {
+    CpuMpData = GetCpuMpDataFromGuidedHob ();
+  } else {
+    AsmReadIdtr (&Idtr);
+    CpuMpData = (CPU_MP_DATA *)(*(UINTN *) (Idtr.Base - sizeof (UINTN)));
+  }
   ASSERT (CpuMpData != NULL);
   return CpuMpData;
 }
@@ -64,6 +75,45 @@ SaveCpuMpData (
     );
 }
 
+/**
+  Push the CpuMpData for AP to use.
+
+  @param[in]  The pointer to CPU MP Data structure will be pushed.
+  @param[out] The pointer to the context which will be passed to PopCpuMpData().
+
+  @return  The pointer value which was stored in where the CPU MP Data is pushed.
+**/
+VOID *
+PushCpuMpData (
+  IN  CPU_MP_DATA    *CpuMpData,
+  OUT VOID           **Context
+  )
+{
+  EFI_PEI_SERVICES  **PeiServices;
+  IA32_DESCRIPTOR   Idtr;
+
+  AsmReadIdtr (&Idtr);
+  *Context = (VOID *) (Idtr.Base - sizeof (UINTN));
+  PeiServices = (EFI_PEI_SERVICES **)(*(UINTN *)(*Context));
+  *(UINTN *)(*Context) = (UINTN)CpuMpData;
+  return PeiServices;
+}
+
+/**
+  Pop the CpuMpData.
+
+  @param[in] Pointer  The pointer value which was stored in where the CPU MP Data is pushed.
+  @param[in] Context  The context of push/pop operation.
+**/
+VOID
+PopCpuMpData (
+  IN VOID           *Pointer,
+  IN VOID           *Context
+  )
+{
+  *(UINTN *)Context = (UINTN)Pointer;
+}
+
 /**
   Check if AP wakeup buffer is overlapped with existing allocated buffer.
 
-- 
2.16.1.windows.1