UefiCpuPkg/DxeMpLib: Allocate new safe stack < 4GB

For long mode DXE, we will disable paging on AP to protected mode to execute AP
safe loop code in reserved memory 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. Besides MwaitSupport == TRUE, AP stack is still required during phase of
disabling paging in long mode DXE.

Moreover, even though AP stack is always under 4GB (a) in Ia32 DXE and (b) with
this patch, after transferring to protected mode from X64 DXE, AP stack
(in BootServiceData) maybe crashed by OS after Exit Boot Service event.

This fix is to allocate reserved memory range under 4GB together with AP safe
loop code. APs will switch to new stack in 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>
Reviewed-by: Laszlo Ersek <lersek@redhat.com>
Tested-by: Laszlo Ersek <lersek@redhat.com>

UefiCpuPkg/Library/MpInitLib/DxeMpLib.c

@@ -18,6 +18,7 @@
 #include <Library/UefiBootServicesTableLib.h>
 
 #define  AP_CHECK_INTERVAL     (EFI_TIMER_PERIOD_MILLISECONDS (100))
+#define  AP_SAFE_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_SAFE_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_SAFE_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_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES (ApSafeBufferSize));
+  ASSERT ((mReservedTopOfApStack & (UINTN)(CPU_STACK_ALIGNMENT - 1)) == 0);
   CopyMem (
     mReservedApLoopFunc,
     CpuMpData->AddressMap.RelocateApLoopFuncAddress,

UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm

@@ -215,19 +215,26 @@ CProcedureInvoke:
 RendezvousFunnelProcEnd:
 
 ;-------------------------------------------------------------------------------------
-;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment);
+;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack);
 ;-------------------------------------------------------------------------------------
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
 AsmRelocateApLoopStart:
-    cmp        byte [esp + 4], 1
+    mov        eax, esp
+    mov        esp, [eax + 16]     ; TopOfApStack
+    push       dword [eax]         ; push return address for stack trace
+    push       ebp
+    mov        ebp, esp
+    mov        ebx, [eax + 8]      ; ApTargetCState
+    mov        ecx, [eax + 4]      ; MwaitSupport
+    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

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
   );
 
 /**

UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm

@@ -222,11 +222,12 @@ CProcedureInvoke:
 RendezvousFunnelProcEnd:
 
 ;-------------------------------------------------------------------------------------
-;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment);
+;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack);
 ;-------------------------------------------------------------------------------------
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
 AsmRelocateApLoopStart:
+    mov        rsp, r9
     push       rcx
     push       rdx