[PATCH] UefiCpuPkg/MpInitLib: fix wrong base address set as Stack Guard

[Jian J Wang <jian.j.wang@intel.com>]

The reason is that DXE part initialization will reuse the stack allocated
at PEI phase, if MP was initialized before. Some code added to check this
situation and use stack base address saved in HOB passed from PEI.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 40c1bf407a..05484c9ff3 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -295,6 +295,7 @@ InitMpGlobalData (
   UINTN                               Index;
   EFI_GCD_MEMORY_SPACE_DESCRIPTOR     MemDesc;
   UINTN                               StackBase;
+  CPU_INFO_IN_HOB                     *CpuInfoInHob;
 
   SaveCpuMpData (CpuMpData);
 
@@ -314,9 +315,18 @@ InitMpGlobalData (
       ASSERT (FALSE);
     }
 
-    for (Index = 0; Index < CpuMpData->CpuCount; ++Index) {
-      StackBase = CpuMpData->Buffer + Index * CpuMpData->CpuApStackSize;
+    //
+    // DXE will reuse stack allocated for APs at PEI phase if it's available.
+    // Let's check it here.
+    //
+    CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;
+    if (CpuInfoInHob != NULL && CpuInfoInHob->ApTopOfStack != 0) {
+      StackBase = CpuInfoInHob->ApTopOfStack;
+    } else {
+      StackBase = CpuMpData->Buffer;
+    }
 
+    for (Index = 0; Index < CpuMpData->CpuCount; ++Index) {
       Status = gDS->GetMemorySpaceDescriptor (StackBase, &MemDesc);
       ASSERT_EFI_ERROR (Status);
 
@@ -326,6 +336,9 @@ InitMpGlobalData (
                       MemDesc.Attributes | EFI_MEMORY_RP
                       );
       ASSERT_EFI_ERROR (Status);
+
+      DEBUG ((DEBUG_VERBOSE, "Stack Guard set at %x [cpu%d]!\n", StackBase, Index));
+      StackBase += CpuMpData->CpuApStackSize;
     }
   }
 
-- 
2.15.1.windows.2