[PATCH] UefiCpuPkg:Fixed AsmRelocateApLoopStart and ensure allocated memory <4GB

["Yuanhao Xie" <yuanhao.xie@intel.com>]

Kept 4GB allocation limit for the case that APs are still transferred to
32-bit protected mode on long mode DXE.
Fixed AsmRelocateApLoopStart stack offset in Ia32.

Bugzilla: https://bugzilla.tianocore.org/show_bug.cgi?id=4234

Cc: Eric Dong eric.dong@intel.com
Cc: Ray Ni ray.ni@intel.com
Cc: Rahul Kumar rahul1.kumar@intel.com
Signed-off-by: Yuanhao Xie <yuanhao.xie@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c       | 35 ++++++++++++-------
 .../Library/MpInitLib/Ia32/MpFuncs.nasm       |  9 ++---
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index beab06a5b1..1994ee44c2 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -389,7 +389,7 @@ RelocateApLoop (
   MpInitLibWhoAmI (&ProcessorNumber);
   CpuMpData    = GetCpuMpData ();
   MwaitSupport = IsMwaitSupport ();
-  if (StandardSignatureIsAuthenticAMD ()) {
+  if (StandardSignatureIsAuthenticAMD () && (sizeof (UINTN) == sizeof (UINT64))) {
     StackStart               = CpuMpData->UseSevEsAPMethod ? CpuMpData->SevEsAPResetStackStart : mReservedTopOfApStack;
     AsmRelocateApLoopFuncAmd = (ASM_RELOCATE_AP_LOOP_AMD)(UINTN)mReservedApLoopFunc;
     AsmRelocateApLoopFuncAmd (
@@ -480,6 +480,7 @@ InitMpGlobalData (
   EFI_GCD_MEMORY_SPACE_DESCRIPTOR  MemDesc;
   UINTN                            StackBase;
   CPU_INFO_IN_HOB                  *CpuInfoInHob;
+  EFI_PHYSICAL_ADDRESS             Address;
 
   SaveCpuMpData (CpuMpData);
 
@@ -536,9 +537,9 @@ InitMpGlobalData (
 
   //
   // Avoid APs access invalid buffer data which allocated by BootServices,
-  // so we will allocate reserved data for AP loop code. We also need to
-  // allocate this buffer below 4GB due to APs may be transferred to 32bit
-  // protected mode on long mode DXE.
+  // so we will allocate reserved data for AP loop code. We need to
+  // allocate this buffer below 4GB for the case that APs are transferred
+  // to 32-bit protected mode on long mode DXE.
   // Allocating it in advance since memory services are not available in
   // Exit Boot Services callback function.
   //
@@ -555,19 +556,25 @@ InitMpGlobalData (
                          )
                        );
 
-  mReservedTopOfApStack = (UINTN)AllocateReservedPages (EFI_SIZE_TO_PAGES (ApSafeBufferSize));
-  ASSERT (mReservedTopOfApStack != 0);
-  ASSERT ((mReservedTopOfApStack & (UINTN)(CPU_STACK_ALIGNMENT - 1)) == 0);
-  ASSERT ((AP_SAFE_STACK_SIZE & (CPU_STACK_ALIGNMENT - 1)) == 0);
-
-  mReservedApLoopFunc = (VOID *)(mReservedTopOfApStack + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
-  if (StandardSignatureIsAuthenticAMD ()) {
+  if (StandardSignatureIsAuthenticAMD () && (sizeof (UINTN) == sizeof (UINT64))) {
+    Address = BASE_4GB - 1;
+    Status  = gBS->AllocatePages (
+                     AllocateMaxAddress,
+                     EfiReservedMemoryType,
+                     EFI_SIZE_TO_PAGES (ApSafeBufferSize),
+                     &Address
+                     );
+    ASSERT_EFI_ERROR (Status);
+    mReservedApLoopFunc = (VOID *)((UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
     CopyMem (
       mReservedApLoopFunc,
       CpuMpData->AddressMap.RelocateApLoopFuncAddressAmd,
       CpuMpData->AddressMap.RelocateApLoopFuncSizeAmd
       );
   } else {
+    Address = (UINTN)AllocateReservedPages (EFI_SIZE_TO_PAGES (ApSafeBufferSize));
+    ASSERT (Address != 0);
+    mReservedApLoopFunc = (VOID *)((UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
     CopyMem (
       mReservedApLoopFunc,
       CpuMpData->AddressMap.RelocateApLoopFuncAddress,
@@ -575,12 +582,14 @@ InitMpGlobalData (
       );
 
     mApPageTable = CreatePageTable (
-                     mReservedTopOfApStack,
+                     (UINTN)Address,
                      ApSafeBufferSize
                      );
   }
 
-  mReservedTopOfApStack += CpuMpData->CpuCount * AP_SAFE_STACK_SIZE;
+  mReservedTopOfApStack = (UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE;
+  ASSERT ((mReservedTopOfApStack & (UINTN)(CPU_STACK_ALIGNMENT - 1)) == 0);
+  ASSERT ((AP_SAFE_STACK_SIZE & (CPU_STACK_ALIGNMENT - 1)) == 0);
 
   Status = gBS->CreateEvent (
                   EVT_TIMER | EVT_NOTIFY_SIGNAL,
diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
index bfcdbd31c1..5cffa632ab 100644
--- a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
@@ -219,20 +219,17 @@ SwitchToRealProcEnd:
 RendezvousFunnelProcEnd:
 
 ;-------------------------------------------------------------------------------------
-;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack, CountTofinish, Pm16CodeSegment, SevEsAPJumpTable, WakeupBuffer);
-;
-;  The last three parameters (Pm16CodeSegment, SevEsAPJumpTable and WakeupBuffer) are
-;  specific to SEV-ES support and are not applicable on IA32.
+;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, TopOfApStack, CountTofinish, Cr3);
 ;-------------------------------------------------------------------------------------
 AsmRelocateApLoopStart:
     mov        eax, esp
-    mov        esp, [eax + 16]     ; TopOfApStack
+    mov        esp, [eax + 12]     ; 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
-    mov        eax, [eax + 20]     ; CountTofinish
+    mov        eax, [eax + 16]     ; CountTofinish
     lock dec   dword [eax]         ; (*CountTofinish)--
     cmp        cl,  1              ; Check mwait-monitor support
     jnz        HltLoop
-- 
2.36.1.windows.1