From 54aeed7e0086724820370e3a3cfac2c7438ebc8d Mon Sep 17 00:00:00 2001
From: Ray Ni <ray.ni@intel.com>
Date: Sat, 7 May 2022 17:10:49 +0800
Subject: [PATCH] MpInitLib: Allocate code buffer for PEI phase

Today's implementation assumes PEI phase runs at 32bit so
the execution-disable feature is not applicable.
It's not always TRUE.
The patch allocates 32bit&64bit code buffer for PEI phase as well.

Signed-off-by: Ray Ni <ray.ni@intel.com>
Reviewed-by: Eric Dong <eric.dong@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c |  2 +-
 UefiCpuPkg/Library/MpInitLib/MpLib.c    |  2 +-
 UefiCpuPkg/Library/MpInitLib/MpLib.h    |  2 +-
 UefiCpuPkg/Library/MpInitLib/PeiMpLib.c | 15 ++++++++++-----
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 60d14a5a0e..78cc3e2b93 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -162,7 +162,7 @@ GetWakeupBuffer (
   @retval 0       Cannot find free memory below 4GB.
 **/
 UINTN
-GetModeTransitionBuffer (
+AllocateCodeBuffer (
   IN UINTN  BufferSize
   )
 {
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 4a73787ee4..d761bdc487 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -1056,7 +1056,7 @@ AllocateResetVector (
                                    (CpuMpData->WakeupBuffer +
                                     CpuMpData->AddressMap.RendezvousFunnelSize +
                                     CpuMpData->AddressMap.SwitchToRealSize);
-    CpuMpData->WakeupBufferHigh = GetModeTransitionBuffer (
+    CpuMpData->WakeupBufferHigh = AllocateCodeBuffer (
                                     CpuMpData->AddressMap.RendezvousFunnelSize +
                                     CpuMpData->AddressMap.SwitchToRealSize -
                                     CpuMpData->AddressMap.ModeTransitionOffset
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index f8c52426dd..59ab960897 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -442,7 +442,7 @@ GetWakeupBuffer (
   @retval 0       Cannot find free memory below 4GB.
 **/
 UINTN
-GetModeTransitionBuffer (
+AllocateCodeBuffer (
   IN UINTN  BufferSize
   );
 
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
index efce574727..65400b95a2 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
@@ -299,14 +299,19 @@ GetWakeupBuffer (
   @retval 0       Cannot find free memory below 4GB.
 **/
 UINTN
-GetModeTransitionBuffer (
+AllocateCodeBuffer (
   IN UINTN  BufferSize
   )
 {
-  //
-  // PEI phase doesn't need to do such transition. So simply return 0.
-  //
-  return 0;
+  EFI_STATUS            Status;
+  EFI_PHYSICAL_ADDRESS  Address;
+
+  Status = PeiServicesAllocatePages (EfiBootServicesCode, EFI_SIZE_TO_PAGES (BufferSize), &Address);
+  if (EFI_ERROR (Status)) {
+    Address = 0;
+  }
+
+  return (UINTN)Address;
 }
 
 /**
-- 
2.39.2