[Patch V3 2/4] UefiCpuPkg: Add Unit tests for PeiCpuExceptionHandlerLib

["duntan" <dun.tan@intel.com>]

The previous change adds unit test for DxeCpuExeptionHandlerLib
in 64bit mode. This change create a PEIM to add unit test for
PeiCpuExceptionHandlerLib based on previous change.It can run
in both 32bit and 64bit modes.

Signed-off-by: Dun Tan <dun.tan@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
---
 UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h             |   9 +++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c       | 135 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTestAsm.nasm | 208 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerLibUnitTest.inf |  61 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerUnitTest.c      | 204 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 617 insertions(+)

diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
index 936098fde8..bad3387db5 100644
--- a/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
@@ -93,6 +93,15 @@ typedef struct {
   UINT64    R15;
 } GENERAL_REGISTER;
 
+typedef struct {
+  UINT32    Edi;
+  UINT32    Esi;
+  UINT32    Ebx;
+  UINT32    Edx;
+  UINT32    Ecx;
+  UINT32    Eax;
+} GENERAL_REGISTER_IA32;
+
 extern UINTN               mFaultInstructionLength;
 extern EFI_EXCEPTION_TYPE  mExceptionType;
 extern UINTN               mRspAddress[];
diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c
new file mode 100644
index 0000000000..8bb27249dc
--- /dev/null
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c
@@ -0,0 +1,135 @@
+/** @file
+  Unit tests of the CpuExceptionHandlerLib.
+
+  Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "CpuExceptionHandlerTest.h"
+
+GENERAL_REGISTER_IA32  mActualContextInHandler;
+GENERAL_REGISTER_IA32  mActualContextAfterException;
+
+//
+// In TestCpuContextConsistency, Cpu registers will be set to mExpectedContextInHandler/mExpectedContextAfterException.
+// Ecx in mExpectedContextInHandler is set runtime since Ecx is needed in assembly code.
+// For GP and PF, Ecx is set to FaultParameter. For other exception triggered by INTn, Ecx is set to ExceptionType.
+//
+GENERAL_REGISTER_IA32  mExpectedContextInHandler      = { 1, 2, 3, 4, 5, 0 };
+GENERAL_REGISTER_IA32  mExpectedContextAfterException = { 11, 12, 13, 14, 15, 16 };
+
+/**
+  Special handler for fault exception.
+  Rip/Eip in SystemContext will be modified to the instruction after the exception instruction.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+**/
+VOID
+EFIAPI
+AdjustRipForFaultHandler (
+  IN EFI_EXCEPTION_TYPE  ExceptionType,
+  IN EFI_SYSTEM_CONTEXT  SystemContext
+  )
+{
+  mExceptionType                        = ExceptionType;
+  SystemContext.SystemContextIa32->Eip += mFaultInstructionLength;
+}
+
+/**
+  Special handler for ConsistencyOfCpuContext test case.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+**/
+VOID
+EFIAPI
+AdjustCpuContextHandler (
+  IN EFI_EXCEPTION_TYPE  ExceptionType,
+  IN EFI_SYSTEM_CONTEXT  SystemContext
+  )
+{
+  //
+  // Store SystemContext in exception handler.
+  //
+  mActualContextInHandler.Edi = SystemContext.SystemContextIa32->Edi;
+  mActualContextInHandler.Esi = SystemContext.SystemContextIa32->Esi;
+  mActualContextInHandler.Ebx = SystemContext.SystemContextIa32->Ebx;
+  mActualContextInHandler.Edx = SystemContext.SystemContextIa32->Edx;
+  mActualContextInHandler.Ecx = SystemContext.SystemContextIa32->Ecx;
+  mActualContextInHandler.Eax = SystemContext.SystemContextIa32->Eax;
+
+  //
+  // Modify cpu context. These registers will be stored in mActualContextAfterException.
+  // Do not handle Esp and Ebp in SystemContext. CpuExceptionHandlerLib doesn't set Esp and
+  // Esp register to the value in SystemContext.
+  //
+  SystemContext.SystemContextIa32->Edi = mExpectedContextAfterException.Edi;
+  SystemContext.SystemContextIa32->Esi = mExpectedContextAfterException.Esi;
+  SystemContext.SystemContextIa32->Ebx = mExpectedContextAfterException.Ebx;
+  SystemContext.SystemContextIa32->Edx = mExpectedContextAfterException.Edx;
+  SystemContext.SystemContextIa32->Ecx = mExpectedContextAfterException.Ecx;
+  SystemContext.SystemContextIa32->Eax = mExpectedContextAfterException.Eax;
+
+  //
+  // When fault exception happens, eip/rip points to the faulting instruction.
+  // For now, olny GP and PF are tested in fault exception.
+  //
+  if ((ExceptionType == EXCEPT_IA32_PAGE_FAULT) || (ExceptionType == EXCEPT_IA32_GP_FAULT)) {
+    AdjustRipForFaultHandler (ExceptionType, SystemContext);
+  }
+}
+
+/**
+  Compare cpu context in ConsistencyOfCpuContext test case.
+  1.Compare mActualContextInHandler with mExpectedContextInHandler.
+  2.Compare mActualContextAfterException with mExpectedContextAfterException.
+
+  @retval  UNIT_TEST_PASSED             The Unit test has completed and it was successful.
+  @retval  UNIT_TEST_ERROR_TEST_FAILED  A test case assertion has failed.
+**/
+UNIT_TEST_STATUS
+CompareCpuContext (
+  VOID
+  )
+{
+  UT_ASSERT_EQUAL (mActualContextInHandler.Edi, mExpectedContextInHandler.Edi);
+  UT_ASSERT_EQUAL (mActualContextInHandler.Esi, mExpectedContextInHandler.Esi);
+  UT_ASSERT_EQUAL (mActualContextInHandler.Ebx, mExpectedContextInHandler.Ebx);
+  UT_ASSERT_EQUAL (mActualContextInHandler.Edx, mExpectedContextInHandler.Edx);
+  UT_ASSERT_EQUAL (mActualContextInHandler.Ecx, mExpectedContextInHandler.Ecx);
+  UT_ASSERT_EQUAL (mActualContextInHandler.Eax, mExpectedContextInHandler.Eax);
+
+  UT_ASSERT_EQUAL (mActualContextAfterException.Edi, mExpectedContextAfterException.Edi);
+  UT_ASSERT_EQUAL (mActualContextAfterException.Esi, mExpectedContextAfterException.Esi);
+  UT_ASSERT_EQUAL (mActualContextAfterException.Ebx, mExpectedContextAfterException.Ebx);
+  UT_ASSERT_EQUAL (mActualContextAfterException.Edx, mExpectedContextAfterException.Edx);
+  UT_ASSERT_EQUAL (mActualContextAfterException.Ecx, mExpectedContextAfterException.Ecx);
+  UT_ASSERT_EQUAL (mActualContextAfterException.Eax, mExpectedContextAfterException.Eax);
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  Special handler for CpuStackGuard test case.
+
+  @param ExceptionType  Exception type.
+  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+**/
+VOID
+EFIAPI
+CpuStackGuardExceptionHandler (
+  IN EFI_EXCEPTION_TYPE  ExceptionType,
+  IN EFI_SYSTEM_CONTEXT  SystemContext
+  )
+{
+  UINTN  LocalVariable;
+
+  AdjustRipForFaultHandler (ExceptionType, SystemContext);
+  mRspAddress[0] = (UINTN)SystemContext.SystemContextIa32->Esp;
+  mRspAddress[1] = (UINTN)(&LocalVariable);
+
+  return;
+}
diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTestAsm.nasm b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTestAsm.nasm
new file mode 100644
index 0000000000..48031a5109
--- /dev/null
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTestAsm.nasm
@@ -0,0 +1,208 @@
+;------------------------------------------------------------------------------
+;
+; Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
+; SPDX-License-Identifier: BSD-2-Clause-Patent
+;
+; Module Name:
+;
+;   ArchExceptionHandlerTestAsm.nasm
+;
+; Abstract:
+;
+;   ia32 CPU Exception Handler Lib Unit test
+;
+;------------------------------------------------------------------------------
+
+    SECTION .text
+
+struc GENERAL_REGISTER_IA32
+  .Edi:    resd    1
+  .Esi:    resd    1
+  .Ebx:    resd    1
+  .Edx:    resd    1
+  .Ecx:    resd    1
+  .Eax:    resd    1
+
+endstruc
+
+extern ASM_PFX(mExpectedContextInHandler)
+extern ASM_PFX(mActualContextAfterException)
+extern ASM_PFX(mFaultInstructionLength)
+
+;------------------------------------------------------------------------------
+; VOID
+; EFIAPI
+; TriggerGPException (
+;  UINTN  Cr4ReservedBit
+;  );
+;------------------------------------------------------------------------------
+global ASM_PFX(TriggerGPException)
+ASM_PFX(TriggerGPException):
+    ;
+    ; Set reserved bit 15 of cr4 to 1
+    ;
+    lea  ecx, [ASM_PFX(mFaultInstructionLength)]
+    mov  dword[ecx], TriggerGPExceptionAfter - TriggerGPExceptionBefore
+    mov  ecx, dword [esp + 0x4]
+TriggerGPExceptionBefore:
+    mov  cr4, ecx
+TriggerGPExceptionAfter:
+    ret
+
+;------------------------------------------------------------------------------
+; VOID
+; EFIAPI
+; TriggerPFException (
+;  UINTN  PfAddress
+;  );
+;------------------------------------------------------------------------------
+global ASM_PFX(TriggerPFException)
+ASM_PFX(TriggerPFException):
+    lea  ecx, [ASM_PFX(mFaultInstructionLength)]
+    mov  dword[ecx], TriggerPFExceptionAfter - TriggerPFExceptionBefore
+    mov  ecx, dword [esp + 0x4]
+TriggerPFExceptionBefore:
+    mov  dword[ecx], 0x1
+TriggerPFExceptionAfter:
+    ret
+
+;------------------------------------------------------------------------------
+; ModifyEcxInGlobalBeforeException;
+; This function is writed by assebly code because it's only called in this file.
+; It's used to set Ecx in mExpectedContextInHandler for different exception.
+;------------------------------------------------------------------------------
+global ASM_PFX(ModifyEcxInGlobalBeforeException)
+ASM_PFX(ModifyEcxInGlobalBeforeException):
+    push eax
+    lea  eax, [ASM_PFX(mExpectedContextInHandler)]
+    mov  [eax + GENERAL_REGISTER_IA32.Ecx], ecx
+    pop  eax
+    ret
+
+;------------------------------------------------------------------------------
+;VOID
+;EFIAPI
+;AsmTestConsistencyOfCpuContext (
+;  IN  EFI_EXCEPTION_TYPE ExceptionType
+;  IN  UINTN              FaultParameter   OPTIONAL
+;  );
+;------------------------------------------------------------------------------
+global ASM_PFX(AsmTestConsistencyOfCpuContext)
+ASM_PFX(AsmTestConsistencyOfCpuContext):
+    ;
+    ; push 7 general register plus 4 bytes
+    ;
+    pushad
+
+    ;
+    ; Modify register to mExpectedContextInHandler. Do not handle Esp and Ebp.
+    ; CpuExceptionHandlerLib doesn't set Esp and Esp register to the value in SystemContext.
+    ;
+    lea eax, [ASM_PFX(mExpectedContextInHandler)]
+    mov edi, [eax + GENERAL_REGISTER_IA32.Edi]
+    mov esi, [eax + GENERAL_REGISTER_IA32.Esi]
+    mov ebx, [eax + GENERAL_REGISTER_IA32.Ebx]
+    mov edx, [eax + GENERAL_REGISTER_IA32.Edx]
+    ;
+    ; Set ecx to ExceptionType
+    ;
+    mov ecx, dword [esp + 0x24]
+    mov eax, [eax + GENERAL_REGISTER_IA32.Eax]
+
+    cmp  ecx, 0xd
+    jz   GPException
+    cmp  ecx, 0xe
+    jz   PFException
+    jmp  INTnException
+
+PFException:
+    mov  ecx, dword [esp + 0x28]                    ; Set ecx to PFAddress.
+    call ASM_PFX(ModifyEcxInGlobalBeforeException)  ; Set mExpectedContextInHandler.Ecx to PFAddress.
+    push ecx                                        ; Push PfAddress into stack.
+    call ASM_PFX(TriggerPFException)
+    jmp  AfterException
+
+GPException:
+    mov  ecx, dword [esp + 0x28]                    ; Set ecx to CR4_RESERVED_BIT.
+    call ASM_PFX(ModifyEcxInGlobalBeforeException)  ; Set mExpectedContextInHandler.Ecx to CR4_RESERVED_BIT.
+    push ecx                                        ; Push CR4_RESERVED_BIT into stack.
+    call ASM_PFX(TriggerGPException)
+    jmp  AfterException
+
+INTnException:
+    call ASM_PFX(ModifyEcxInGlobalBeforeException)  ; Set mExpectedContextInHandler.Ecx to ExceptionType.
+    push ecx                                        ; Push ExceptionType into stack.
+    call ASM_PFX(TriggerINTnException)
+
+AfterException:
+    ;
+    ; Save register in mActualContextAfterException.
+    ;
+    push eax
+    lea  eax, [ASM_PFX(mActualContextAfterException)]
+    mov  [eax + GENERAL_REGISTER_IA32.Edi], edi
+    mov  [eax + GENERAL_REGISTER_IA32.Esi], esi
+    mov  [eax + GENERAL_REGISTER_IA32.Ebx], ebx
+    mov  [eax + GENERAL_REGISTER_IA32.Edx], edx
+    mov  [eax + GENERAL_REGISTER_IA32.Ecx], ecx
+    pop  ecx
+    mov  [eax + GENERAL_REGISTER_IA32.Eax], ecx
+    add  esp, 4
+
+    ;
+    ; restore original register
+    ;
+    popad
+    ret
+
+;------------------------------------------------------------------------------
+; VOID
+; EFIAPI
+; TriggerStackOverflow (
+;  VOID
+;  );
+;------------------------------------------------------------------------------
+global ASM_PFX(TriggerStackOverflow)
+ASM_PFX(TriggerStackOverflow):
+    lea  ecx, [ASM_PFX(mFaultInstructionLength)]
+    mov  dword[ecx], TriggerCpuStackGuardAfter - TriggerCpuStackGuardBefore
+TriggerCpuStackGuardBefore:
+    ;
+    ; Clear CR0.TS since it is set after return from a nested DF
+    ;
+    call TriggerCpuStackGuardBefore
+    clts
+TriggerCpuStackGuardAfter:
+    ret
+
+;------------------------------------------------------------------------------
+; VOID
+; EFIAPI
+; TriggerINTnException (
+;  IN  EFI_EXCEPTION_TYPE ExceptionType
+;  );
+;------------------------------------------------------------------------------
+global ASM_PFX(TriggerINTnException)
+ASM_PFX(TriggerINTnException):
+    push eax
+    push edx
+    lea  eax, [AsmTriggerException1 - AsmTriggerException0]
+    mov  ecx, dword [esp + 0xc]
+    push ecx
+    mul  ecx
+    mov  ecx, AsmTriggerException0
+    add  eax, ecx
+    pop  ecx
+    pop  edx
+    jmp  eax
+    ;
+    ; eax = AsmTriggerException0 + (AsmTriggerException1 - AsmTriggerException0) * ecx
+    ;
+%assign Vector 0
+%rep  22
+AsmTriggerException %+ Vector:
+    pop eax
+    INT Vector
+    ret
+%assign Vector Vector+1
+%endrep
diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerLibUnitTest.inf b/UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerLibUnitTest.inf
new file mode 100644
index 0000000000..25f8f8dbe0
--- /dev/null
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerLibUnitTest.inf
@@ -0,0 +1,61 @@
+## @file
+# Unit tests of the PeiCpuExceptionHandlerLib instance.
+#
+# Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010006
+  BASE_NAME                      = CpuExceptionHandlerPeiTest
+  FILE_GUID                      = 39A96CF7-F369-4357-9234-4B52F98A007F
+  MODULE_TYPE                    = PEIM
+  VERSION_STRING                 = 1.0
+  ENTRY_POINT                    = PeiEntryPoint
+
+#
+# The following information is for reference only and not required by the build tools.
+#
+#  VALID_ARCHITECTURES           = IA32 X64
+#
+[Sources.Ia32]
+  Ia32/ArchExceptionHandlerTestAsm.nasm
+  Ia32/ArchExceptionHandlerTest.c
+
+[Sources.X64]
+  X64/ArchExceptionHandlerTestAsm.nasm
+  X64/ArchExceptionHandlerTest.c
+
+[Sources.common]
+  CpuExceptionHandlerTest.h
+  CpuExceptionHandlerTestCommon.c
+  PeiCpuExceptionHandlerUnitTest.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  MdeModulePkg/MdeModulePkg.dec
+  UefiCpuPkg/UefiCpuPkg.dec
+
+[LibraryClasses]
+  BaseLib
+  BaseMemoryLib
+  DebugLib
+  UnitTestLib
+  MemoryAllocationLib
+  CpuExceptionHandlerLib
+  PeimEntryPoint
+  HobLib
+  PeiServicesLib
+  CpuPageTableLib
+  PeiServicesTablePointerLib
+
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard   ## CONSUMES
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuApStackSize       ## CONSUMES
+
+[Ppis]
+  gEdkiiPeiMpServices2PpiGuid                       ## CONSUMES
+
+[Depex]
+  gEdkiiPeiMpServices2PpiGuid AND
+  gEfiPeiMemoryDiscoveredPpiGuid
diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerUnitTest.c b/UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerUnitTest.c
new file mode 100644
index 0000000000..d9408d2f5e
--- /dev/null
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerUnitTest.c
@@ -0,0 +1,204 @@
+/** @file
+  Unit tests of the CpuExceptionHandlerLib.
+
+  Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "CpuExceptionHandlerTest.h"
+#include <Library/PeimEntryPoint.h>
+#include <Library/PeiServicesLib.h>
+#include <Library/PeiServicesTablePointerLib.h>
+
+/**
+  Initialize Bsp Idt with a new Idt table and return the IA32_DESCRIPTOR buffer.
+  In PEIM, store original PeiServicePointer before new Idt table.
+
+  @return Pointer to the allocated IA32_DESCRIPTOR buffer.
+**/
+VOID *
+InitializeBspIdt (
+  VOID
+  )
+{
+  UINTN            *NewIdtTable;
+  IA32_DESCRIPTOR  *Idtr;
+
+  Idtr = AllocateZeroPool (sizeof (IA32_DESCRIPTOR));
+  ASSERT (Idtr != NULL);
+  NewIdtTable = AllocateZeroPool (sizeof (IA32_IDT_GATE_DESCRIPTOR) * CPU_INTERRUPT_NUM + sizeof (UINTN));
+  ASSERT (NewIdtTable != NULL);
+  //
+  // Store original PeiServicePointer before new Idt table
+  //
+  *NewIdtTable = (UINTN)GetPeiServicesTablePointer ();
+  NewIdtTable  = (UINTN *)((UINTN)NewIdtTable + sizeof (UINTN));
+
+  Idtr->Base  = (UINTN)NewIdtTable;
+  Idtr->Limit = (UINT16)(sizeof (IA32_IDT_GATE_DESCRIPTOR) * CPU_INTERRUPT_NUM - 1);
+
+  AsmWriteIdtr (Idtr);
+  return Idtr;
+}
+
+/**
+  Retrieve the number of logical processor in the platform and the number of those logical processors that
+  are enabled on this boot.
+
+  @param[in]  MpServices          MP_SERVICES structure.
+  @param[out] NumberOfProcessors  Pointer to the total number of logical processors in the system, including
+                                  the BSP and disabled APs.
+  @param[out] NumberOfEnabledProcessors Pointer to the number of processors in the system that are enabled.
+
+  @retval EFI_SUCCESS       Retrieve the number of logical processor successfully
+  @retval Others            Retrieve the number of logical processor unsuccessfully
+**/
+EFI_STATUS
+MpServicesUnitTestGetNumberOfProcessors (
+  IN MP_SERVICES  MpServices,
+  OUT UINTN       *NumberOfProcessors,
+  OUT UINTN       *NumberOfEnabledProcessors
+  )
+{
+  return MpServices.Ppi->GetNumberOfProcessors (MpServices.Ppi, NumberOfProcessors, NumberOfEnabledProcessors);
+}
+
+/**
+  Caller gets one enabled AP to execute a caller-provided function.
+
+  @param[in]  MpServices    MP_SERVICES structure.
+  @param[in]  Procedure     Pointer to the function to be run on enabled APs of the system.
+  @param[in]  ProcessorNumber       The handle number of the AP.
+  @param[in]  TimeoutInMicroSeconds Indicates the time limit in microseconds for APs to return from Procedure,
+                                    for blocking mode only. Zero means infinity.
+  @param[in]  ProcedureArgument     The parameter passed into Procedure for all APs.
+
+
+  @retval EFI_SUCCESS       Caller gets one enabled AP to execute a caller-provided function successfully
+  @retval Others            Caller gets one enabled AP to execute a caller-provided function unsuccessfully
+**/
+EFI_STATUS
+MpServicesUnitTestStartupThisAP (
+  IN MP_SERVICES       MpServices,
+  IN EFI_AP_PROCEDURE  Procedure,
+  IN UINTN             ProcessorNumber,
+  IN UINTN             TimeoutInMicroSeconds,
+  IN VOID              *ProcedureArgument
+  )
+{
+  return MpServices.Ppi->StartupThisAP (MpServices.Ppi, Procedure, ProcessorNumber, TimeoutInMicroSeconds, ProcedureArgument);
+}
+
+/**
+  Execute a caller provided function on all enabled APs.
+
+  @param[in]  MpServices    MP_SERVICES structure.
+  @param[in]  Procedure     Pointer to the function to be run on enabled APs of the system.
+  @param[in]  SingleThread  If TRUE, then all the enabled APs execute the function specified by Procedure
+                            one by one, in ascending order of processor handle number.
+                            If FALSE, then all the enabled APs execute the function specified by Procedure
+                            simultaneously.
+  @param[in]  TimeoutInMicroSeconds Indicates the time limit in microseconds for APs to return from Procedure,
+                                    for blocking mode only. Zero means infinity.
+  @param[in]  ProcedureArgument     The parameter passed into Procedure for all APs.
+
+  @retval EFI_SUCCESS       Execute a caller provided function on all enabled APs successfully
+  @retval Others            Execute a caller provided function on all enabled APs unsuccessfully
+**/
+EFI_STATUS
+MpServicesUnitTestStartupAllAPs (
+  IN MP_SERVICES       MpServices,
+  IN EFI_AP_PROCEDURE  Procedure,
+  IN BOOLEAN           SingleThread,
+  IN UINTN             TimeoutInMicroSeconds,
+  IN VOID              *ProcedureArgument
+  )
+{
+  return MpServices.Ppi->StartupAllAPs (MpServices.Ppi, Procedure, SingleThread, TimeoutInMicroSeconds, ProcedureArgument);
+}
+
+/**
+  Get the handle number for the calling processor.
+
+  @param[in]  MpServices      MP_SERVICES structure.
+  @param[out] ProcessorNumber The handle number for the calling processor.
+
+  @retval EFI_SUCCESS       Get the handle number for the calling processor successfully.
+  @retval Others            Get the handle number for the calling processor unsuccessfully.
+**/
+EFI_STATUS
+MpServicesUnitTestWhoAmI (
+  IN MP_SERVICES  MpServices,
+  OUT UINTN       *ProcessorNumber
+  )
+{
+  return MpServices.Ppi->WhoAmI (MpServices.Ppi, ProcessorNumber);
+}
+
+/**
+  Get EDKII_PEI_MP_SERVICES2_PPI pointer.
+
+  @param[out] MpServices    Pointer to the buffer where EDKII_PEI_MP_SERVICES2_PPI is stored
+
+  @retval EFI_SUCCESS       EDKII_PEI_MP_SERVICES2_PPI interface is returned
+  @retval EFI_NOT_FOUND     EDKII_PEI_MP_SERVICES2_PPI interface is not found
+**/
+EFI_STATUS
+GetMpServices (
+  OUT MP_SERVICES  *MpServices
+  )
+{
+  return PeiServicesLocatePpi (&gEdkiiPeiMpServices2PpiGuid, 0, NULL, (VOID **)&MpServices->Ppi);
+}
+
+/**
+  Entry point of CpuExceptionHandlerPeiTest PEIM.
+
+  @param[in] FileHandle  Handle of the file being invoked.
+  @param[in] PeiServices Describes the list of possible PEI Services.
+
+  @retval EFI_SUCCESS    The PEIM executed normally.
+
+**/
+EFI_STATUS
+EFIAPI
+PeiEntryPoint (
+  IN EFI_PEI_FILE_HANDLE     FileHandle,
+  IN CONST EFI_PEI_SERVICES  **PeiServices
+  )
+{
+  EFI_STATUS                  Status;
+  UNIT_TEST_FRAMEWORK_HANDLE  Framework;
+
+  Framework = NULL;
+
+  DEBUG ((DEBUG_INFO, "%a v%a\n", UNIT_TEST_APP_NAME, UNIT_TEST_APP_VERSION));
+
+  //
+  // Start setting up the test framework for running the tests.
+  //
+  Status = InitUnitTestFramework (&Framework, UNIT_TEST_APP_NAME, gEfiCallerBaseName, UNIT_TEST_APP_VERSION);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed in InitUnitTestFramework. Status = %r\n", Status));
+    goto EXIT;
+  }
+
+  Status = AddCommonTestCase (Framework);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed in AddCommonTestCase. Status = %r\n", Status));
+    goto EXIT;
+  }
+
+  //
+  // Execute the tests.
+  //
+  Status = RunAllTestSuites (Framework);
+
+EXIT:
+  if (Framework) {
+    FreeUnitTestFramework (Framework);
+  }
+
+  return Status;
+}
-- 
2.31.1.windows.1