[Patch V2 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 of test cases are created in this module:
a.Test if exception handler can be registered/unregistered
for no error code exception.In the test case, only no error
code exception is triggered and tested by INTn instruction.

b.Test if exception handler can be registered/unregistered
for GP and PF. In the test case, GP exception is triggered
and tested by setting CR4_RESERVED_BIT to 1. PF exception
is triggered by writting to not-present or RO address.

c.Test if CpuContext is consistent before and after exception.
In this test case:
1.Set Cpu register to mExpectedContextInHandler before
exception. 2.Trigger exception specified by ExceptionType.
3.Store SystemContext in mActualContextInHandler and set
SystemContext to mExpectedContextAfterException in handler.
4.After return from exception, store Cpu registers in
mActualContextAfterException.
The expectation is:
1.Register values in mActualContextInHandler are the same
with register values in mExpectedContextInHandler.
2.Register values in mActualContextAfterException are the
same with register values mActualContextAfterException.

d.Test if stack overflow can be captured by CpuStackGuard
in both Bsp and AP. In this test case, stack overflow is
triggered by a funtion which calls itself continuously.
This test case triggers stack overflow in both BSP and AP.
All AP use same Idt with Bsp. The expectation is:
1. PF exception is triggered (leading to a DF if sepereated
stack is not prepared for PF) when Esp<=StackBase+SIZE_4KB
since [StackBase, StackBase + SIZE_4KB] is guarded in page
table when PcdCpuStackGuard is TRUE.
2. Stack for PF/DF exception handler in both Bsp and AP is
succussfully switched by InitializeSeparateExceptionStacks.

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             |  21 ++++++++++++++++-----
 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, 624 insertions(+), 5 deletions(-)

diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
index 3034257d9c..7d9a07a491 100644
--- a/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
+++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
@@ -122,11 +122,22 @@ typedef struct {
   UINT64    R15Register;
 } GENERAL_REGISTER;
 
-extern UINTN               mFaultInstructionLength;
-extern EFI_EXCEPTION_TYPE  mExceptionType;
-extern UINTN               mRspAddress[];
-extern GENERAL_REGISTER    mExpectedContextInHandler;
-extern GENERAL_REGISTER    mActualContextAfterException;
+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[];
+extern GENERAL_REGISTER       mExpectedContextInHandler;
+extern GENERAL_REGISTER       mActualContextAfterException;
+extern GENERAL_REGISTER_IA32  mIa32ExpectedContextInHandler;
+extern GENERAL_REGISTER_IA32  mIa32ActualContextAfterException;
 
 /**
   Initialize Bsp Idt with a new Idt table and return the IA32_DESCRIPTOR buffer.
diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c b/UefiCpuPkg/CpuExceptionHandlerUnitTest/Ia32/ArchExceptionHandlerTest.c
new file mode 100644
index 0000000000..cb83dad215
--- /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  mIa32ActualContextInHandler;
+GENERAL_REGISTER_IA32  mIa32ActualContextAfterException;
+
+//
+// In TestCpuContextConsistency, Cpu registers will be set to mIa32ExpectedContextInHandler/mIa32ExpectedContextAfterException.
+// Ecx in mIa32ExpectedContextInHandler 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  mIa32ExpectedContextInHandler      = { 1, 2, 3, 4, 5, 0 };
+GENERAL_REGISTER_IA32  mIa32ExpectedContextAfterException = { 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.
+  //
+  mIa32ActualContextInHandler.Edi = SystemContext.SystemContextIa32->Edi;
+  mIa32ActualContextInHandler.Esi = SystemContext.SystemContextIa32->Esi;
+  mIa32ActualContextInHandler.Ebx = SystemContext.SystemContextIa32->Ebx;
+  mIa32ActualContextInHandler.Edx = SystemContext.SystemContextIa32->Edx;
+  mIa32ActualContextInHandler.Ecx = SystemContext.SystemContextIa32->Ecx;
+  mIa32ActualContextInHandler.Eax = SystemContext.SystemContextIa32->Eax;
+
+  //
+  // Modify cpu context. These registers will be stored in mIa32ActualContextAfterException.
+  // Do not handle Esp and Ebp in SystemContext. CpuExceptionHandlerLib doesn't set Esp and
+  // Esp register to the value in SystemContext.
+  //
+  SystemContext.SystemContextIa32->Edi = mIa32ExpectedContextAfterException.Edi;
+  SystemContext.SystemContextIa32->Esi = mIa32ExpectedContextAfterException.Esi;
+  SystemContext.SystemContextIa32->Ebx = mIa32ExpectedContextAfterException.Ebx;
+  SystemContext.SystemContextIa32->Edx = mIa32ExpectedContextAfterException.Edx;
+  SystemContext.SystemContextIa32->Ecx = mIa32ExpectedContextAfterException.Ecx;
+  SystemContext.SystemContextIa32->Eax = mIa32ExpectedContextAfterException.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 mIa32ActualContextInHandler with mIa32ExpectedContextInHandler.
+  2.Compare mIa32ActualContextAfterException with mIa32ExpectedContextAfterException.
+
+  @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 (mIa32ActualContextInHandler.Edi, mIa32ExpectedContextInHandler.Edi);
+  UT_ASSERT_EQUAL (mIa32ActualContextInHandler.Esi, mIa32ExpectedContextInHandler.Esi);
+  UT_ASSERT_EQUAL (mIa32ActualContextInHandler.Ebx, mIa32ExpectedContextInHandler.Ebx);
+  UT_ASSERT_EQUAL (mIa32ActualContextInHandler.Edx, mIa32ExpectedContextInHandler.Edx);
+  UT_ASSERT_EQUAL (mIa32ActualContextInHandler.Ecx, mIa32ExpectedContextInHandler.Ecx);
+  UT_ASSERT_EQUAL (mIa32ActualContextInHandler.Eax, mIa32ExpectedContextInHandler.Eax);
+
+  UT_ASSERT_EQUAL (mIa32ActualContextAfterException.Edi, mIa32ExpectedContextAfterException.Edi);
+  UT_ASSERT_EQUAL (mIa32ActualContextAfterException.Esi, mIa32ExpectedContextAfterException.Esi);
+  UT_ASSERT_EQUAL (mIa32ActualContextAfterException.Ebx, mIa32ExpectedContextAfterException.Ebx);
+  UT_ASSERT_EQUAL (mIa32ActualContextAfterException.Edx, mIa32ExpectedContextAfterException.Edx);
+  UT_ASSERT_EQUAL (mIa32ActualContextAfterException.Ecx, mIa32ExpectedContextAfterException.Ecx);
+  UT_ASSERT_EQUAL (mIa32ActualContextAfterException.Eax, mIa32ExpectedContextAfterException.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..45186bd5c6
--- /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(mIa32ExpectedContextInHandler)
+extern ASM_PFX(mIa32ActualContextAfterException)
+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 mIa32ExpectedContextInHandler for different exception.
+;------------------------------------------------------------------------------
+global ASM_PFX(ModifyEcxInGlobalBeforeException)
+ASM_PFX(ModifyEcxInGlobalBeforeException):
+    push eax
+    lea  eax, [ASM_PFX(mIa32ExpectedContextInHandler)]
+    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 mIa32ExpectedContextInHandler. Do not handle Esp and Ebp.
+    ; CpuExceptionHandlerLib doesn't set Esp and Esp register to the value in SystemContext.
+    ;
+    lea eax, [ASM_PFX(mIa32ExpectedContextInHandler)]
+    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 mIa32ExpectedContextInHandler.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 mIa32ExpectedContextInHandler.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 mIa32ExpectedContextInHandler.Ecx to ExceptionType.
+    push ecx                                        ; Push ExceptionType into stack.
+    call ASM_PFX(TriggerINTnException)
+
+AfterException:
+    ;
+    ; Save register in mIa32ActualContextAfterException.
+    ;
+    push eax
+    lea  eax, [ASM_PFX(mIa32ActualContextAfterException)]
+    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