[PATCH 2/3] UefiCpuPkg/DxeMpLib: Allocate new safe stack < 4GB

[Jeff Fan <jeff.fan@intel.com>]

For long mode DXE, we will disable paging on AP to protected mode to execute AP
safe loop code in ACPI NVS range under 4GB. But we forget to allocate stack for
AP under 4GB and AP still are using original AP stack. If original AP stack is
larger than 4GB, it cannot be used after AP is transferred to protected mode.

Moreover, even though AP stack is always under 4GB on protected mode DXE, AP
stack(in BootServiceData) maybe crashed by OS after Exit Boot Service event.

This fix is to allocate ACPI NVS range under 4GB together with AP safe loop
code. APs will switch to new stack at beginning of safe loop code.

Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Feng Tian <feng.tian@intel.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c        | 19 +++++++++++++++++--
 UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm |  9 +++++++--
 UefiCpuPkg/Library/MpInitLib/MpLib.h           |  3 ++-
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm  |  3 +++
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index a0d5eeb..d0f9f7e 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -18,6 +18,7 @@
 #include <Library/UefiBootServicesTableLib.h>
 
 #define  AP_CHECK_INTERVAL     (EFI_TIMER_PERIOD_MILLISECONDS (100))
+#define  AP_SAFT_STACK_SIZE    128
 
 CPU_MP_DATA      *mCpuMpData = NULL;
 EFI_EVENT        mCheckAllApsEvent = NULL;
@@ -25,6 +26,7 @@ EFI_EVENT        mMpInitExitBootServicesEvent = NULL;
 EFI_EVENT        mLegacyBootEvent = NULL;
 volatile BOOLEAN mStopCheckAllApsStatus = TRUE;
 VOID             *mReservedApLoopFunc = NULL;
+UINTN            mReservedTopOfApStack;
 
 /**
   Get the pointer to CPU MP Data structure.
@@ -241,11 +243,18 @@ RelocateApLoop (
   CPU_MP_DATA            *CpuMpData;
   BOOLEAN                MwaitSupport;
   ASM_RELOCATE_AP_LOOP   AsmRelocateApLoopFunc;
+  UINTN                  ProcessorNumber;
 
+  MpInitLibWhoAmI (&ProcessorNumber); 
   CpuMpData    = GetCpuMpData ();
   MwaitSupport = IsMwaitSupport ();
   AsmRelocateApLoopFunc = (ASM_RELOCATE_AP_LOOP) (UINTN) mReservedApLoopFunc;
-  AsmRelocateApLoopFunc (MwaitSupport, CpuMpData->ApTargetCState, CpuMpData->PmCodeSegment);
+  AsmRelocateApLoopFunc (
+    MwaitSupport,
+    CpuMpData->ApTargetCState,
+    CpuMpData->PmCodeSegment,
+    mReservedTopOfApStack - ProcessorNumber * AP_SAFT_STACK_SIZE
+    );
   //
   // It should never reach here
   //
@@ -289,6 +298,7 @@ InitMpGlobalData (
 {
   EFI_STATUS                 Status;
   EFI_PHYSICAL_ADDRESS       Address;
+  UINTN                      ApSafeBufferSize;
 
   SaveCpuMpData (CpuMpData);
 
@@ -307,16 +317,21 @@ InitMpGlobalData (
   // Allocating it in advance since memory services are not available in
   // Exit Boot Services callback function.
   //
+  ApSafeBufferSize  = CpuMpData->AddressMap.RelocateApLoopFuncSize;
+  ApSafeBufferSize += CpuMpData->CpuCount * AP_SAFT_STACK_SIZE;
+
   Address = BASE_4GB - 1;
   Status  = gBS->AllocatePages (
                    AllocateMaxAddress,
                    EfiReservedMemoryType,
-                   EFI_SIZE_TO_PAGES (sizeof (CpuMpData->AddressMap.RelocateApLoopFuncSize)),
+                   EFI_SIZE_TO_PAGES (ApSafeBufferSize),
                    &Address
                    );
   ASSERT_EFI_ERROR (Status);
   mReservedApLoopFunc = (VOID *) (UINTN) Address;
   ASSERT (mReservedApLoopFunc != NULL);
+  mReservedTopOfApStack = (UINTN) Address + EFI_SIZE_TO_PAGES (ApSafeBufferSize) * SIZE_4KB;
+  ASSERT ((mReservedTopOfApStack & (UINTN)(CPU_STACK_ALIGNMENT - 1)) == 0);
   CopyMem (
     mReservedApLoopFunc,
     CpuMpData->AddressMap.RelocateApLoopFuncAddress,
diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
index 64e51d8..4e55760 100644
--- a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
@@ -217,14 +217,19 @@ RendezvousFunnelProcEnd:
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
 AsmRelocateApLoopStart:
-    cmp        byte [esp + 4], 1
+    push       ebp
+    mov        ebp, esp
+    mov        ecx, [ebp + 8]      ; MwaitSupport
+    mov        ebx, [ebp + 12]     ; ApTargetCState
+    mov        esp, [ebp + 20]     ; TopOfApStack
+    cmp        cl,  1              ; Check mwait-monitor support
     jnz        HltLoop
 MwaitLoop:
     mov        eax, esp
     xor        ecx, ecx
     xor        edx, edx
     monitor
-    mov        eax, [esp + 8]    ; Mwait Cx, Target C-State per eax[7:4]
+    mov        eax, ebx            ; Mwait Cx, Target C-State per eax[7:4]
     shl        eax, 4
     mwait
     jmp        MwaitLoop
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index f73a469..e6dea18 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -250,7 +250,8 @@ VOID
 (EFIAPI * ASM_RELOCATE_AP_LOOP) (
   IN BOOLEAN                 MwaitSupport,
   IN UINTN                   ApTargetCState,
-  IN UINTN                   PmCodeSegment
+  IN UINTN                   PmCodeSegment,
+  IN UINTN                   TopOfApStack
   );
 
 /**
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index aaabb50..7c8fa45 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
@@ -224,6 +224,9 @@ RendezvousFunnelProcEnd:
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
 AsmRelocateApLoopStart:
+    push       rbp
+    mov        rbp, rsp
+    mov        rsp, r9
     push       rcx
     push       rdx
 
-- 
2.9.3.windows.2