3 files changed, 69 insertions, 18 deletions
diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 7839c24976..93fc63bf93 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -29,6 +29,11 @@ VOID             *mReservedApLoopFunc = NULL;
 UINTN            mReservedTopOfApStack;
 volatile UINT32  mNumberToFinish = 0;
 
+//
+// Begin wakeup buffer allocation below 0x88000
+//
+STATIC EFI_PHYSICAL_ADDRESS mSevEsDxeWakeupBuffer = 0x88000;
+
 /**
   Enable Debug Agent to support source debugging on AP function.
 
@@ -102,7 +107,14 @@ GetWakeupBuffer (
   // LagacyBios driver depends on CPU Arch protocol which guarantees below
   // allocation runs earlier than LegacyBios driver.
   //
-  StartAddress = 0x88000;
+  if (PcdGetBool (PcdSevEsIsEnabled)) {
+    //
+    // SEV-ES Wakeup buffer should be under 0x88000 and under any previous one
+    //
+    StartAddress = mSevEsDxeWakeupBuffer;
+  } else {
+    StartAddress = 0x88000;
+  }
   Status = gBS->AllocatePages (
                   AllocateMaxAddress,
                   MemoryType,
@@ -112,6 +124,11 @@ GetWakeupBuffer (
   ASSERT_EFI_ERROR (Status);
   if (EFI_ERROR (Status)) {
     StartAddress = (EFI_PHYSICAL_ADDRESS) -1;
+  } else if (PcdGetBool (PcdSevEsIsEnabled)) {
+    //
+    // Next SEV-ES wakeup buffer allocation must be below this allocation
+    //
+    mSevEsDxeWakeupBuffer = StartAddress;
   }
 
   DEBUG ((DEBUG_INFO, "WakeupBufferStart = %x, WakeupBufferSize = %x\n",
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index dc2a54aa31..b9a06747ed 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -1164,20 +1164,6 @@ GetApResetVectorSize (
            AddressMap->SwitchToRealSize +
            sizeof (MP_CPU_EXCHANGE_INFO);
 
-  //
-  // The AP reset stack is only used by SEV-ES guests. Do not add to the
-  // allocation if SEV-ES is not enabled.
-  //
-  if (PcdGetBool (PcdSevEsIsEnabled)) {
-    //
-    // Stack location is based on APIC ID, so use the total number of
-    // processors for calculating the total stack area.
-    //
-    Size += AP_RESET_STACK_SIZE * PcdGet32 (PcdCpuMaxLogicalProcessorNumber);
-
-    Size = ALIGN_VALUE (Size, CPU_STACK_ALIGNMENT);
-  }
-
   return Size;
 }
 
@@ -1192,6 +1178,7 @@ AllocateResetVector (
   )
 {
   UINTN           ApResetVectorSize;
+  UINTN           ApResetStackSize;
 
   if (CpuMpData->WakeupBuffer == (UINTN) -1) {
     ApResetVectorSize = GetApResetVectorSize (&CpuMpData->AddressMap);
@@ -1207,9 +1194,39 @@ AllocateResetVector (
                                     CpuMpData->AddressMap.ModeTransitionOffset
                                     );
     //
-    // The reset stack starts at the end of the buffer.
+    // The AP reset stack is only used by SEV-ES guests. Do not allocate it
+    // if SEV-ES is not enabled.
     //
-    CpuMpData->SevEsAPResetStackStart = CpuMpData->WakeupBuffer + ApResetVectorSize;
+    if (PcdGetBool (PcdSevEsIsEnabled)) {
+      //
+      // Stack location is based on ProcessorNumber, so use the total number
+      // of processors for calculating the total stack area.
+      //
+      ApResetStackSize = (AP_RESET_STACK_SIZE *
+                          PcdGet32 (PcdCpuMaxLogicalProcessorNumber));
+
+      //
+      // Invoke GetWakeupBuffer a second time to allocate the stack area
+      // below 1MB. The returned buffer will be page aligned and sized and
+      // below the previously allocated buffer.
+      //
+      CpuMpData->SevEsAPResetStackStart = GetWakeupBuffer (ApResetStackSize);
+
+      //
+      // Check to be sure that the "allocate below" behavior hasn't changed.
+      // This will also catch a failed allocation, as "-1" is returned on
+      // failure.
+      //
+      if (CpuMpData->SevEsAPResetStackStart >= CpuMpData->WakeupBuffer) {
+        DEBUG ((
+          DEBUG_ERROR,
+          "SEV-ES AP reset stack is not below wakeup buffer\n"
+          ));
+
+        ASSERT (FALSE);
+        CpuDeadLoop ();
+      }
+    }
   }
   BackupAndPrepareWakeupBuffer (CpuMpData);
 }
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
index 3989bd6a7a..90015c650c 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
@@ -11,6 +11,8 @@
 #include <Guid/S3SmmInitDone.h>
 #include <Ppi/ShadowMicrocode.h>
 
+STATIC UINT64 mSevEsPeiWakeupBuffer = BASE_1MB;
+
 /**
   S3 SMM Init Done notification function.
 
@@ -220,7 +222,13 @@ GetWakeupBuffer (
         // Need memory under 1MB to be collected here
         //
         WakeupBufferEnd = Hob.ResourceDescriptor->PhysicalStart + Hob.ResourceDescriptor->ResourceLength;
-        if (WakeupBufferEnd > BASE_1MB) {
+        if (PcdGetBool (PcdSevEsIsEnabled) &&
+            WakeupBufferEnd > mSevEsPeiWakeupBuffer) {
+          //
+          // SEV-ES Wakeup buffer should be under 1MB and under any previous one
+          //
+          WakeupBufferEnd = mSevEsPeiWakeupBuffer;
+        } else if (WakeupBufferEnd > BASE_1MB) {
           //
           // Wakeup buffer should be under 1MB
           //
@@ -244,6 +252,15 @@ GetWakeupBuffer (
           }
           DEBUG ((DEBUG_INFO, "WakeupBufferStart = %x, WakeupBufferSize = %x\n",
                                WakeupBufferStart, WakeupBufferSize));
+
+          if (PcdGetBool (PcdSevEsIsEnabled)) {
+            //
+            // Next SEV-ES wakeup buffer allocation must be below this
+            // allocation
+            //
+            mSevEsPeiWakeupBuffer = WakeupBufferStart;
+          }
+
           return (UINTN)WakeupBufferStart;
         }
       }