[PATCH 2/3] UefiCpuPkg: Add Unit tests for PeiCpuExceptionHandlerLib

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

Add target based unit tests for the PeiCpuExceptionHandlerLib.
A PEIM is created to test PeiCpuExceptionHandlerLib. Four kinds
of test cases are created in this module:
1.Test if exception handler can be registered/unregistered for
no error code exception. 2.Test if exception handler can be
registered/unregistered for GP and PF. 3.Test if Cpu Context is
consistent before and after exception. 4.Test if stack overflow
can be captured by CpuStackGuard in both Bsp and AP.

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/Ia32/ArchExceptionHandlerTest.c       | 133 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTestAsm.nasm | 217 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerLibUnitTest.inf |  61 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/CpuExceptionHandlerUnitTest/PeiCpuExceptionHandlerUnitTest.c      | 204 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 615 insertions(+)

diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c
new file mode 100644
index 0000000000..32188b9a81
--- /dev/null
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c
@@ -0,0 +1,133 @@
+/** @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  mIa32RegisterForCheck[2];
+
+//
+// In TestCpuContextConsistency, Cpu registers will be set to mAdjustRegisterBeforeException/mAdjustRegisterInsideException.
+// Ecx in mAdjustRegisterInsideException 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  mAdjustIa32RegisterBeforeException = { 1, 2, 3, 4, 5, 0, 7 };
+GENERAL_REGISTER_IA32  mAdjustIa32RegisterInsideException = { 11, 12, 13, 14, 15, 16, 17 };
+
+/**
+  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
+  )
+{
+  //
+  // Do not handle Esp and ebp. They will not be restored.
+  // Store SystemContext modified before exception.
+  //
+  mIa32RegisterForCheck[0].Edi = SystemContext.SystemContextIa32->Edi;
+  mIa32RegisterForCheck[0].Esi = SystemContext.SystemContextIa32->Esi;
+  mIa32RegisterForCheck[0].Ebx = SystemContext.SystemContextIa32->Ebx;
+  mIa32RegisterForCheck[0].Edx = SystemContext.SystemContextIa32->Edx;
+  mIa32RegisterForCheck[0].Ecx = SystemContext.SystemContextIa32->Ecx;
+  mIa32RegisterForCheck[0].Eax = SystemContext.SystemContextIa32->Eax;
+
+  //
+  // Modify cpu context which will be stored in mIa32RegisterForCheck[1].
+  //
+  SystemContext.SystemContextIa32->Edi = mAdjustIa32RegisterInsideException.Edi;
+  SystemContext.SystemContextIa32->Esi = mAdjustIa32RegisterInsideException.Esi;
+  SystemContext.SystemContextIa32->Ebx = mAdjustIa32RegisterInsideException.Ebx;
+  SystemContext.SystemContextIa32->Edx = mAdjustIa32RegisterInsideException.Edx;
+  SystemContext.SystemContextIa32->Ecx = mAdjustIa32RegisterInsideException.Ecx;
+  SystemContext.SystemContextIa32->Eax = mAdjustIa32RegisterInsideException.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 mRegisterForCheck[0] with mAdjustRegisterBeforeException.
+  2.Compare mRegisterForCheck[1] with mAdjustRegisterAfterException.
+
+  @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 (mIa32RegisterForCheck[0].Edi, mAdjustIa32RegisterBeforeException.Edi);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[0].Esi, mAdjustIa32RegisterBeforeException.Esi);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[0].Ebx, mAdjustIa32RegisterBeforeException.Ebx);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[0].Edx, mAdjustIa32RegisterBeforeException.Edx);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[0].Ecx, mAdjustIa32RegisterBeforeException.Ecx);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[0].Eax, mAdjustIa32RegisterBeforeException.Eax);
+
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[1].Edi, mAdjustIa32RegisterInsideException.Edi);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[1].Esi, mAdjustIa32RegisterInsideException.Esi);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[1].Ebx, mAdjustIa32RegisterInsideException.Ebx);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[1].Edx, mAdjustIa32RegisterInsideException.Edx);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[1].Ecx, mAdjustIa32RegisterInsideException.Ecx);
+  UT_ASSERT_EQUAL (mIa32RegisterForCheck[1].Eax, mAdjustIa32RegisterInsideException.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..ccadd8765f
--- /dev/null
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTestAsm.nasm
@@ -0,0 +1,217 @@
+;------------------------------------------------------------------------------
+;
+; 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
+  .Ebp:    resd    1
+  .Ebx:    resd    1
+  .Edx:    resd    1
+  .Ecx:    resd    1
+  .Eax:    resd    1
+
+endstruc
+
+extern ASM_PFX(mIa32RegisterForCheck)
+extern ASM_PFX(mAdjustIa32RegisterBeforeException)
+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
+
+global ASM_PFX(ModifyEcxInGlobalBeforeException)
+ASM_PFX(ModifyEcxInGlobalBeforeException):
+    push eax
+    lea  eax, [ASM_PFX(mAdjustIa32RegisterBeforeException)]
+    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 Modify register to mAdjustRegisterBeforeException
+    ;
+    lea eax, [ASM_PFX(mAdjustIa32RegisterBeforeException)]
+    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:
+    ;
+    ; Set ecx to PFAddress and Modify Ecx in mAdjustRegisterBeforeException.
+    ; Push PfAddress into stack.
+    ;
+    mov  ecx, dword [esp + 0x28]
+    call ASM_PFX(ModifyEcxInGlobalBeforeException)
+    push ecx
+    call ASM_PFX(TriggerPFException)
+    jmp  AfterException
+
+GPException:
+    ;
+    ; Set ecx to CR4_RESERVED_BIT and Modify Ecx in mAdjustRegisterBeforeException.
+    ; Push CR4_RESERVED_BIT into stack
+    ;
+    mov  ecx, dword [esp + 0x28]
+    call ASM_PFX(ModifyEcxInGlobalBeforeException)
+    push ecx
+    call ASM_PFX(TriggerGPException)
+    jmp  AfterException
+
+INTnException:
+    ;
+    ; Modify Ecx in mAdjustRegisterBeforeException.
+    ; Push exception index into stack
+    ;
+    call ASM_PFX(ModifyEcxInGlobalBeforeException)
+    push ecx
+    call ASM_PFX(TriggerINTnException)
+
+AfterException:
+    ;
+    ; Save register in mRegisterForCheck[1].
+    ; Esp and ebp will not be synced to the value in SystemContext. Do not test esp and ebp.
+    ;
+    push eax
+    lea  eax, [ASM_PFX(mIa32RegisterForCheck)]
+    add  eax, GENERAL_REGISTER_IA32_size
+    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
+; TriggerStackOverflowbyCpuStackGuard (
+;  VOID
+;  );
+;------------------------------------------------------------------------------
+global ASM_PFX(TriggerStackOverflowbyCpuStackGuard)
+ASM_PFX(TriggerStackOverflowbyCpuStackGuard):
+    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