Reviewed-by: Ray Ni <ray...@intel.com>

> -----Original Message-----
> From: Tan, Dun <dun....@intel.com>
> Sent: Friday, October 14, 2022 5:19 PM
> To: devel@edk2.groups.io
> Cc: Dong, Eric <eric.d...@intel.com>; Ni, Ray <ray...@intel.com>; Kumar, 
> Rahul R <rahul.r.ku...@intel.com>
> Subject: [Patch V3 1/4] UefiCpuPkg: Add Unit tests for 
> DxeCpuExceptionHandlerLib
> 
> Add target based unit tests for the DxeCpuExceptionHandlerLib.
> A DXE driver is created to test DxeCpuExceptionHandlerLib.
> 
> Four test cases are created in this Unit Test 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 Rsp<=StackBase+SIZE_4KB
> since [StackBase, StackBase + SIZE_4KB] is marked as not
> present 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....@intel.com>
> Cc: Eric Dong <eric.d...@intel.com>
> Cc: Ray Ni <ray...@intel.com>
> Cc: Rahul Kumar <rahul1.ku...@intel.com>
> ---
>  UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h             
> | 336
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++
>  UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTestCommon.c       
> | 852
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++
>  UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerLibUnitTest.inf 
> |  58
> ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
>  UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerUnitTest.c      
> | 196
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> ++++++++++
>  UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTest.c        
> | 166
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
>  UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTestAsm.nasm  
> | 256
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
> ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
>  6 files changed, 1864 insertions(+)
> 
> diff --git a/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
> new file mode 100644
> index 0000000000..936098fde8
> --- /dev/null
> +++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTest.h
> @@ -0,0 +1,336 @@
> +/** @file
> +
> +  Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
> +  SPDX-License-Identifier: BSD-2-Clause-Patent
> +
> +  Four test cases are created in this Unit Test module.
> +  a.Test if exception handler can be registered/unregistered for no error 
> code exception
> +    In this test case, only no error code exception is triggered and tested 
> by INTn instruction.
> +    The special hanlder for these exception will modify a global variable 
> for check.
> +
> +  b.Test if exception handler can be registered/unregistered for GP and PF.
> +    In this test case, GP exception is triggered and tested by setting 
> CR4_RESERVED_BIT to 1.
> +    PF exception is triggered and tested by writting to not-present or RO 
> addres.
> +    The special hanlder for these exceptions will set a global vartiable for 
> check and adjust Rip to return from fault exception.
> +
> +  c.Test if Cpu Context 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 Rsp <= StackBase +
> SIZE_4KB
> +         since [StackBase, StackBase + SIZE_4KB] is marked as not present in 
> page table when PcdCpuStackGuard is TRUE.
> +      2. Stack for PF/DF exception handler in both Bsp and AP is 
> succussfully switched by InitializeSeparateExceptionStacks.
> +
> +**/
> +
> +#ifndef CPU_EXCEPTION_HANDLER_TEST_H_
> +#define CPU_EXCEPTION_HANDLER_TEST_H_
> +
> +#include <Uefi.h>
> +#include <Library/BaseLib.h>
> +#include <Library/BaseMemoryLib.h>
> +#include <Library/DebugLib.h>
> +#include <Library/UnitTestLib.h>
> +#include <Library/MemoryAllocationLib.h>
> +#include <Library/UnitTestHostBaseLib.h>
> +#include <Library/CpuExceptionHandlerLib.h>
> +#include <Library/UefiLib.h>
> +#include <Library/SerialPortLib.h>
> +#include <Library/HobLib.h>
> +#include <Library/CpuPageTableLib.h>
> +#include <Guid/MemoryAllocationHob.h>
> +#include <Protocol/MpService.h>
> +#include <PiPei.h>
> +#include <Ppi/MpServices2.h>
> +
> +#define UNIT_TEST_APP_NAME     "Cpu Exception Handler Lib Unit Tests"
> +#define UNIT_TEST_APP_VERSION  "1.0"
> +
> +#define  CPU_INTERRUPT_NUM       256
> +#define  SPEC_MAX_EXCEPTION_NUM  22
> +#define  CR4_RESERVED_BIT        BIT15
> +
> +typedef struct {
> +  IA32_DESCRIPTOR    OriginalGdtr;
> +  IA32_DESCRIPTOR    OriginalIdtr;
> +  UINT16             Tr;
> +} CPU_REGISTER_BUFFER;
> +
> +typedef union {
> +  EDKII_PEI_MP_SERVICES2_PPI    *Ppi;
> +  EFI_MP_SERVICES_PROTOCOL      *Protocol;
> +} MP_SERVICES;
> +
> +typedef struct {
> +  VOID          *Buffer;
> +  UINTN         BufferSize;
> +  EFI_STATUS    Status;
> +} EXCEPTION_STACK_SWITCH_CONTEXT;
> +
> +typedef struct {
> +  UINT64    Rdi;
> +  UINT64    Rsi;
> +  UINT64    Rbx;
> +  UINT64    Rdx;
> +  UINT64    Rcx;
> +  UINT64    Rax;
> +  UINT64    R8;
> +  UINT64    R9;
> +  UINT64    R10;
> +  UINT64    R11;
> +  UINT64    R12;
> +  UINT64    R13;
> +  UINT64    R14;
> +  UINT64    R15;
> +} GENERAL_REGISTER;
> +
> +extern UINTN               mFaultInstructionLength;
> +extern EFI_EXCEPTION_TYPE  mExceptionType;
> +extern UINTN               mRspAddress[];
> +
> +/**
> +  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
> +  );
> +
> +/**
> +  Trigger no error code exception by INT n instruction.
> +
> +  @param[in]  ExceptionType  No error code exception type.
> +**/
> +VOID
> +EFIAPI
> +TriggerINTnException (
> +  IN  EFI_EXCEPTION_TYPE  ExceptionType
> +  );
> +
> +/**
> +  Trigger GP exception by setting CR4_RESERVED_BIT to 1.
> +
> +  @param[in]  Cr4ReservedBit  Cr4 reserved bit.
> +**/
> +VOID
> +EFIAPI
> +TriggerGPException (
> +  UINTN  Cr4ReservedBit
> +  );
> +
> +/**
> +  Trigger PF exception by write to not present or ReadOnly address.
> +
> +  @param[in]  PFAddress  Not present or ReadOnly address in page table.
> +**/
> +VOID
> +EFIAPI
> +TriggerPFException (
> +  UINTN  PFAddress
> +  );
> +
> +/**
> +  Special handler for fault exception.
> +  This handler sets Rip/Eip in SystemContext to the instruction address 
> 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
> +  );
> +
> +/**
> +  Test consistency of Cpu context. Four steps:
> +  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.
> +
> +  Rcx/Ecx in mExpectedContextInHandler is decided by different exception 
> type runtime since Rcx/Ecx is needed in
> assembly code.
> +  For GP and PF, Rcx/Ecx is set to FaultParameter. For other exception 
> triggered by INTn, Rcx/Ecx is set to ExceptionType.
> +
> +  @param[in] ExceptionType   Exception type.
> +  @param[in] FaultParameter  Parameter for GP and PF. OPTIONAL
> +**/
> +VOID
> +EFIAPI
> +AsmTestConsistencyOfCpuContext (
> +  IN  EFI_EXCEPTION_TYPE  ExceptionType,
> +  IN  UINTN               FaultParameter   OPTIONAL
> +  );
> +
> +/**
> +  Special handler for ConsistencyOfCpuContext test case. General register in 
> SystemContext
> +  is modified to mExpectedContextInHandler in this handler.
> +
> +  @param ExceptionType  Exception type.
> +  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
> +**/
> +VOID
> +EFIAPI
> +AdjustCpuContextHandler (
> +  IN EFI_EXCEPTION_TYPE  ExceptionType,
> +  IN EFI_SYSTEM_CONTEXT  SystemContext
> +  );
> +
> +/**
> +  Compare cpu context in ConsistencyOfCpuContext test case.
> +  1.Compare mActualContextInHandler with mExpectedContextInHandler.
> +  2.Compare mActualContextAfterException with mActualContextAfterException.
> +
> +  @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
> +  );
> +
> +/**
> +  Get EFI_MP_SERVICES_PROTOCOL/EDKII_PEI_MP_SERVICES2_PPI pointer.
> +
> +  @param[out] MpServices    Pointer to the MP_SERVICES buffer
> +
> +  @retval EFI_SUCCESS       EFI_MP_SERVICES_PROTOCOL/PPI interface is 
> returned
> +  @retval EFI_NOT_FOUND     EFI_MP_SERVICES_PROTOCOL/PPI interface is not 
> found
> +**/
> +EFI_STATUS
> +GetMpServices (
> +  OUT MP_SERVICES  *MpServices
> +  );
> +
> +/**
> +  Create CpuExceptionLibUnitTestSuite and add test case.
> +
> +  @param[in]  FrameworkHandle    Unit test framework.
> +
> +  @return  EFI_SUCCESS           The unit test suite was created.
> +  @retval  EFI_OUT_OF_RESOURCES  There are not enough resources available to
> +                                 initialize the unit test suite.
> +**/
> +EFI_STATUS
> +AddCommonTestCase (
> +  IN  UNIT_TEST_FRAMEWORK_HANDLE  Framework
> +  );
> +
> +/**
> +  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
> +  );
> +
> +/**
> +  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
> +  );
> +
> +/**
> +  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
> +  );
> +
> +/**
> +  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
> +  );
> +
> +/**
> +  Trigger stack overflow by calling itself continuously.
> +**/
> +VOID
> +EFIAPI
> +TriggerStackOverflow (
> +  VOID
> +  );
> +
> +/**
> +  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
> +  );
> +
> +#endif
> diff --git 
> a/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTestCommon.c
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTestCommon.c
> new file mode 100644
> index 0000000000..17afb592d3
> --- /dev/null
> +++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/CpuExceptionHandlerTestCommon.c
> @@ -0,0 +1,852 @@
> +/** @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"
> +
> +//
> +// Length of the assembly falut instruction.
> +//
> +UINTN               mFaultInstructionLength = 0;
> +EFI_EXCEPTION_TYPE  mExceptionType          = 256;
> +UINTN               mNumberOfProcessors     = 1;
> +UINTN               mRspAddress[2]          = { 0 };
> +
> +//
> +// Error code flag indicating whether or not an error code will be
> +// pushed on the stack if an exception occurs.
> +//
> +// 1 means an error code will be pushed, otherwise 0
> +//
> +CONST UINT32  mErrorCodeExceptionFlag = 0x20227d00;
> +
> +/**
> +  Special handler for exception triggered by INTn instruction.
> +  This hanlder only modifies a global variable for check.
> +
> +  @param ExceptionType  Exception type.
> +  @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
> +**/
> +VOID
> +EFIAPI
> +INTnExceptionHandler (
> +  IN EFI_EXCEPTION_TYPE  ExceptionType,
> +  IN EFI_SYSTEM_CONTEXT  SystemContext
> +  )
> +{
> +  mExceptionType = ExceptionType;
> +}
> +
> +/**
> +  Restore cpu original registers before exit test case.
> +
> +  @param[in] Buffer  Argument of the procedure.
> +**/
> +VOID
> +EFIAPI
> +RestoreRegistersPerCpu (
> +  IN VOID  *Buffer
> +  )
> +{
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer;
> +  UINT16               Tr;
> +  IA32_TSS_DESCRIPTOR  *Tss;
> +
> +  CpuOriginalRegisterBuffer = (CPU_REGISTER_BUFFER *)Buffer;
> +
> +  AsmWriteGdtr (&(CpuOriginalRegisterBuffer->OriginalGdtr));
> +  AsmWriteIdtr (&(CpuOriginalRegisterBuffer->OriginalIdtr));
> +  Tr = CpuOriginalRegisterBuffer->Tr;
> +  if ((Tr != 0) && (Tr < CpuOriginalRegisterBuffer->OriginalGdtr.Limit)) {
> +    Tss = (IA32_TSS_DESCRIPTOR 
> *)(CpuOriginalRegisterBuffer->OriginalGdtr.Base + Tr);
> +    if (Tss->Bits.P == 1) {
> +      //
> +      // Clear busy bit of TSS before write Tr
> +      //
> +      Tss->Bits.Type &= 0xD;
> +      AsmWriteTr (Tr);
> +    }
> +  }
> +}
> +
> +/**
> +  Restore cpu original registers before exit test case.
> +
> +  @param[in] MpServices                 MpServices.
> +  @param[in] CpuOriginalRegisterBuffer  Address of CpuOriginalRegisterBuffer.
> +  @param[in] BspProcessorNum            Bsp processor number.
> +**/
> +VOID
> +RestoreAllCpuRegisters (
> +  MP_SERVICES *MpServices, OPTIONAL
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer,
> +  UINTN                         BspProcessorNum
> +  )
> +{
> +  UINTN       Index;
> +  EFI_STATUS  Status;
> +
> +  for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +    if (Index == BspProcessorNum) {
> +      RestoreRegistersPerCpu ((VOID *)&CpuOriginalRegisterBuffer[Index]);
> +      continue;
> +    }
> +
> +    ASSERT (MpServices != NULL);
> +    Status = MpServicesUnitTestStartupThisAP (
> +               *MpServices,
> +               (EFI_AP_PROCEDURE)RestoreRegistersPerCpu,
> +               Index,
> +               0,
> +               (VOID *)&CpuOriginalRegisterBuffer[Index]
> +               );
> +    ASSERT_EFI_ERROR (Status);
> +  }
> +}
> +
> +/**
> +  Store cpu registers before the test case starts.
> +
> +  @param[in] Buffer  Argument of the procedure.
> +**/
> +VOID
> +EFIAPI
> +SaveRegisterPerCpu (
> +  IN VOID  *Buffer
> +  )
> +{
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer;
> +  IA32_DESCRIPTOR      Gdtr;
> +  IA32_DESCRIPTOR      Idtr;
> +
> +  CpuOriginalRegisterBuffer = (CPU_REGISTER_BUFFER *)Buffer;
> +
> +  AsmReadGdtr (&Gdtr);
> +  AsmReadIdtr (&Idtr);
> +  CpuOriginalRegisterBuffer->OriginalGdtr.Base  = Gdtr.Base;
> +  CpuOriginalRegisterBuffer->OriginalGdtr.Limit = Gdtr.Limit;
> +  CpuOriginalRegisterBuffer->OriginalIdtr.Base  = Idtr.Base;
> +  CpuOriginalRegisterBuffer->OriginalIdtr.Limit = Idtr.Limit;
> +  CpuOriginalRegisterBuffer->Tr                 = AsmReadTr ();
> +}
> +
> +/**
> +  Store cpu registers before the test case starts.
> +
> +  @param[in] MpServices       MpServices.
> +  @param[in] BspProcessorNum  Bsp processor number.
> +
> +  @return Pointer to the allocated CPU_REGISTER_BUFFER.
> +**/
> +CPU_REGISTER_BUFFER *
> +SaveAllCpuRegisters (
> +  MP_SERVICES *MpServices, OPTIONAL
> +  UINTN       BspProcessorNum
> +  )
> +{
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer;
> +  EFI_STATUS           Status;
> +  UINTN                Index;
> +
> +  CpuOriginalRegisterBuffer = AllocateZeroPool (mNumberOfProcessors * sizeof 
> (CPU_REGISTER_BUFFER));
> +  ASSERT (CpuOriginalRegisterBuffer != NULL);
> +
> +  for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +    if (Index == BspProcessorNum) {
> +      SaveRegisterPerCpu ((VOID *)&CpuOriginalRegisterBuffer[Index]);
> +      continue;
> +    }
> +
> +    ASSERT (MpServices != NULL);
> +    Status = MpServicesUnitTestStartupThisAP (
> +               *MpServices,
> +               (EFI_AP_PROCEDURE)SaveRegisterPerCpu,
> +               Index,
> +               0,
> +               (VOID *)&CpuOriginalRegisterBuffer[Index]
> +               );
> +    ASSERT_EFI_ERROR (Status);
> +  }
> +
> +  return CpuOriginalRegisterBuffer;
> +}
> +
> +/**
> +  Initialize Ap Idt Procedure.
> +
> +  @param[in] Buffer  Argument of the procedure.
> +**/
> +VOID
> +EFIAPI
> +InitializeIdtPerAp (
> +  IN VOID  *Buffer
> +  )
> +{
> +  AsmWriteIdtr (Buffer);
> +}
> +
> +/**
> +  Initialize all Ap Idt.
> +
> +  @param[in] MpServices MpServices.
> +  @param[in] BspIdtr    Pointer to IA32_DESCRIPTOR allocated by Bsp.
> +**/
> +VOID
> +InitializeApIdt (
> +  MP_SERVICES  MpServices,
> +  VOID         *BspIdtr
> +  )
> +{
> +  EFI_STATUS  Status;
> +
> +  Status = MpServicesUnitTestStartupAllAPs (
> +             MpServices,
> +             (EFI_AP_PROCEDURE)InitializeIdtPerAp,
> +             FALSE,
> +             0,
> +             BspIdtr
> +             );
> +  ASSERT_EFI_ERROR (Status);
> +}
> +
> +/**
> +  Check if exception handler can registered/unregistered for no error code 
> exception.
> +
> +  @param[in]  Context    [Optional] An optional parameter that enables:
> +                         1) test-case reuse with varied parameters and
> +                         2) test-case re-entry for Target tests that need a
> +                         reboot.  This parameter is a VOID* and it is the
> +                         responsibility of the test author to ensure that the
> +                         contents are well understood by all test cases that 
> may
> +                         consume it.
> +
> +  @retval  UNIT_TEST_PASSED             The Unit test has completed and the 
> test
> +                                        case was successful.
> +  @retval  UNIT_TEST_ERROR_TEST_FAILED  A test case assertion has failed.
> +**/
> +UNIT_TEST_STATUS
> +EFIAPI
> +TestRegisterHandlerForNoErrorCodeException (
> +  IN UNIT_TEST_CONTEXT  Context
> +  )
> +{
> +  EFI_STATUS           Status;
> +  UINTN                Index;
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer;
> +  VOID                 *NewIdtr;
> +
> +  CpuOriginalRegisterBuffer = SaveAllCpuRegisters (NULL, 0);
> +  NewIdtr                   = InitializeBspIdt ();
> +  Status                    = InitializeCpuExceptionHandlers (NULL);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +  for (Index = 0; Index < SPEC_MAX_EXCEPTION_NUM; Index++) {
> +    //
> +    // Only test no error code exception by INT n instruction.
> +    //
> +    if ((mErrorCodeExceptionFlag & (1 << Index)) != 0) {
> +      continue;
> +    }
> +
> +    DEBUG ((DEBUG_INFO, "TestCase1: ExceptionType is %d\n", Index));
> +    Status = RegisterCpuInterruptHandler (Index, INTnExceptionHandler);
> +    UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +    TriggerINTnException (Index);
> +    UT_ASSERT_EQUAL (mExceptionType, Index);
> +    Status = RegisterCpuInterruptHandler (Index, NULL);
> +    UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  }
> +
> +  RestoreAllCpuRegisters (NULL, CpuOriginalRegisterBuffer, 0);
> +  FreePool (CpuOriginalRegisterBuffer);
> +  FreePool (NewIdtr);
> +  return UNIT_TEST_PASSED;
> +}
> +
> +/**
> +  Get Bsp stack base.
> +
> +  @param[out] StackBase  Pointer to stack base of BSP.
> +**/
> +VOID
> +GetBspStackBase (
> +  OUT UINTN  *StackBase
> +  )
> +{
> +  EFI_PEI_HOB_POINTERS       Hob;
> +  EFI_HOB_MEMORY_ALLOCATION  *MemoryHob;
> +
> +  //
> +  // Get the base of stack from Hob.
> +  //
> +  ASSERT (StackBase != NULL);
> +  Hob.Raw = GetHobList ();
> +  while ((Hob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, Hob.Raw)) != 
> NULL) {
> +    MemoryHob = Hob.MemoryAllocation;
> +    if (CompareGuid (&gEfiHobMemoryAllocStackGuid, 
> &MemoryHob->AllocDescriptor.Name)) {
> +      DEBUG ((
> +        DEBUG_INFO,
> +        "%a: Bsp StackBase = 0x%016lx  StackSize = 0x%016lx\n",
> +        __FUNCTION__,
> +        MemoryHob->AllocDescriptor.MemoryBaseAddress,
> +        MemoryHob->AllocDescriptor.MemoryLength
> +        ));
> +
> +      *StackBase = (UINTN)MemoryHob->AllocDescriptor.MemoryBaseAddress;
> +      //
> +      // Ensure the base of the stack is page-size aligned.
> +      //
> +      ASSERT ((*StackBase & EFI_PAGE_MASK) == 0);
> +      break;
> +    }
> +
> +    Hob.Raw = GET_NEXT_HOB (Hob);
> +  }
> +
> +  ASSERT (*StackBase != 0);
> +}
> +
> +/**
> +  Get Ap stack base procedure.
> +
> +  @param[out] ApStackBase  Pointer to Ap stack base.
> +**/
> +VOID
> +EFIAPI
> +GetStackBasePerAp (
> +  OUT VOID  *ApStackBase
> +  )
> +{
> +  UINTN  ApTopOfStack;
> +
> +  ApTopOfStack          = ALIGN_VALUE ((UINTN)&ApTopOfStack, (UINTN)PcdGet32 
> (PcdCpuApStackSize));
> +  *(UINTN *)ApStackBase = ApTopOfStack - (UINTN)PcdGet32 (PcdCpuApStackSize);
> +}
> +
> +/**
> +  Get all Cpu stack base.
> +
> +  @param[in] MpServices       MpServices.
> +  @param[in] BspProcessorNum  Bsp processor number.
> +
> +  @return Pointer to the allocated CpuStackBaseBuffer.
> +**/
> +UINTN *
> +GetAllCpuStackBase (
> +  MP_SERVICES  *MpServices,
> +  UINTN        BspProcessorNum
> +  )
> +{
> +  UINTN       *CpuStackBaseBuffer;
> +  EFI_STATUS  Status;
> +  UINTN       Index;
> +
> +  CpuStackBaseBuffer = AllocateZeroPool (mNumberOfProcessors * sizeof 
> (UINTN));
> +  ASSERT (CpuStackBaseBuffer != NULL);
> +
> +  for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +    if (Index == BspProcessorNum) {
> +      GetBspStackBase (&CpuStackBaseBuffer[Index]);
> +      continue;
> +    }
> +
> +    ASSERT (MpServices != NULL);
> +    Status = MpServicesUnitTestStartupThisAP (
> +               *MpServices,
> +               (EFI_AP_PROCEDURE)GetStackBasePerAp,
> +               Index,
> +               0,
> +               (VOID *)&CpuStackBaseBuffer[Index]
> +               );
> +    ASSERT_EFI_ERROR (Status);
> +    DEBUG ((DEBUG_INFO, "AP[%d] StackBase = 0x%x\n", Index, 
> CpuStackBaseBuffer[Index]));
> +  }
> +
> +  return CpuStackBaseBuffer;
> +}
> +
> +/**
> +  Find not present or ReadOnly address in page table.
> +
> +  @param[out] PFAddress  Access to the address which is not permitted will 
> trigger PF exceptions.
> +
> +  @retval TRUE   Found not present or ReadOnly address in page table.
> +  @retval FALSE  Failed to found PFAddress in page table.
> +**/
> +BOOLEAN
> +FindPFAddressInPageTable (
> +  OUT UINTN  *PFAddress
> +  )
> +{
> +  IA32_CR0        Cr0;
> +  IA32_CR4        Cr4;
> +  UINTN           PageTable;
> +  PAGING_MODE     PagingMode;
> +  BOOLEAN         Enable5LevelPaging;
> +  RETURN_STATUS   Status;
> +  IA32_MAP_ENTRY  *Map;
> +  UINTN           MapCount;
> +  UINTN           Index;
> +  UINTN           PreviousAddress;
> +
> +  ASSERT (PFAddress != NULL);
> +
> +  Cr0.UintN = AsmReadCr0 ();
> +  if (Cr0.Bits.PG == 0) {
> +    return FALSE;
> +  }
> +
> +  PageTable = AsmReadCr3 ();
> +  Cr4.UintN = AsmReadCr4 ();
> +  if (sizeof (UINTN) == sizeof (UINT32)) {
> +    ASSERT (Cr4.Bits.PAE == 1);
> +    PagingMode = PagingPae;
> +  } else {
> +    Enable5LevelPaging = (BOOLEAN)(Cr4.Bits.LA57 == 1);
> +    PagingMode         = Enable5LevelPaging ? Paging5Level : Paging4Level;
> +  }
> +
> +  MapCount = 0;
> +  Status   = PageTableParse (PageTable, PagingMode, NULL, &MapCount);
> +  ASSERT (Status == RETURN_BUFFER_TOO_SMALL);
> +  Map    = AllocatePages (EFI_SIZE_TO_PAGES (MapCount * sizeof 
> (IA32_MAP_ENTRY)));
> +  Status = PageTableParse (PageTable, PagingMode, Map, &MapCount);
> +  ASSERT (Status == RETURN_SUCCESS);
> +
> +  PreviousAddress = 0;
> +  for (Index = 0; Index < MapCount; Index++) {
> +    DEBUG ((
> +      DEBUG_ERROR,
> +      "%02d: %016lx - %016lx, %016lx\n",
> +      Index,
> +      Map[Index].LinearAddress,
> +      Map[Index].LinearAddress + Map[Index].Length,
> +      Map[Index].Attribute.Uint64
> +      ));
> +
> +    //
> +    // Not present address in page table.
> +    //
> +    if (Map[Index].LinearAddress > PreviousAddress) {
> +      *PFAddress = PreviousAddress;
> +      return TRUE;
> +    }
> +
> +    PreviousAddress = (UINTN)(Map[Index].LinearAddress + Map[Index].Length);
> +
> +    //
> +    // ReadOnly address in page table.
> +    //
> +    if ((Cr0.Bits.WP != 0) && (Map[Index].Attribute.Bits.ReadWrite == 0)) {
> +      *PFAddress = (UINTN)Map[Index].LinearAddress;
> +      return TRUE;
> +    }
> +  }
> +
> +  return FALSE;
> +}
> +
> +/**
> +  Test if exception handler can registered/unregistered for GP and PF.
> +
> +  @param[in]  Context    [Optional] An optional parameter that enables:
> +                         1) test-case reuse with varied parameters and
> +                         2) test-case re-entry for Target tests that need a
> +                         reboot.  This parameter is a VOID* and it is the
> +                         responsibility of the test author to ensure that the
> +                         contents are well understood by all test cases that 
> may
> +                         consume it.
> +
> +  @retval  UNIT_TEST_PASSED             The Unit test has completed and the 
> test
> +                                        case was successful.
> +  @retval  UNIT_TEST_ERROR_TEST_FAILED  A test case assertion has failed.
> +**/
> +UNIT_TEST_STATUS
> +EFIAPI
> +TestRegisterHandlerForGPAndPF (
> +  IN UNIT_TEST_CONTEXT  Context
> +  )
> +{
> +  EFI_STATUS           Status;
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer;
> +  UINTN                PFAddress;
> +  VOID                 *NewIdtr;
> +
> +  PFAddress                 = 0;
> +  CpuOriginalRegisterBuffer = SaveAllCpuRegisters (NULL, 0);
> +  NewIdtr                   = InitializeBspIdt ();
> +  Status                    = InitializeCpuExceptionHandlers (NULL);
> +
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +  //
> +  // GP exception.
> +  //
> +  DEBUG ((DEBUG_INFO, "TestCase2: ExceptionType is %d\n", 
> EXCEPT_IA32_GP_FAULT));
> +  Status = RegisterCpuInterruptHandler (EXCEPT_IA32_GP_FAULT, 
> AdjustRipForFaultHandler);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +  TriggerGPException (CR4_RESERVED_BIT);
> +  UT_ASSERT_EQUAL (mExceptionType, EXCEPT_IA32_GP_FAULT);
> +  Status = RegisterCpuInterruptHandler (EXCEPT_IA32_GP_FAULT, NULL);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +  //
> +  // PF exception.
> +  //
> +  if (FindPFAddressInPageTable (&PFAddress)) {
> +    DEBUG ((DEBUG_INFO, "TestCase2: ExceptionType is %d\n", 
> EXCEPT_IA32_PAGE_FAULT));
> +    Status = RegisterCpuInterruptHandler (EXCEPT_IA32_PAGE_FAULT, 
> AdjustRipForFaultHandler);
> +    UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +    TriggerPFException (PFAddress);
> +
> +    UT_ASSERT_EQUAL (mExceptionType, EXCEPT_IA32_PAGE_FAULT);
> +    Status = RegisterCpuInterruptHandler (EXCEPT_IA32_PAGE_FAULT, NULL);
> +    UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  }
> +
> +  RestoreAllCpuRegisters (NULL, CpuOriginalRegisterBuffer, 0);
> +  FreePool (CpuOriginalRegisterBuffer);
> +  FreePool (NewIdtr);
> +  return UNIT_TEST_PASSED;
> +}
> +
> +/**
> +  Test if Cpu Context is consistent before and after exception.
> +
> +  @param[in]  Context    [Optional] An optional parameter that enables:
> +                         1) test-case reuse with varied parameters and
> +                         2) test-case re-entry for Target tests that need a
> +                         reboot.  This parameter is a VOID* and it is the
> +                         responsibility of the test author to ensure that the
> +                         contents are well understood by all test cases that 
> may
> +                         consume it.
> +
> +  @retval  UNIT_TEST_PASSED             The Unit test has completed and the 
> test
> +                                        case was successful.
> +  @retval  UNIT_TEST_ERROR_TEST_FAILED  A test case assertion has failed.
> +**/
> +UNIT_TEST_STATUS
> +EFIAPI
> +TestCpuContextConsistency (
> +  IN UNIT_TEST_CONTEXT  Context
> +  )
> +{
> +  EFI_STATUS           Status;
> +  UINTN                Index;
> +  CPU_REGISTER_BUFFER  *CpuOriginalRegisterBuffer;
> +  UINTN                FaultParameter;
> +  VOID                 *NewIdtr;
> +
> +  FaultParameter            = 0;
> +  CpuOriginalRegisterBuffer = SaveAllCpuRegisters (NULL, 0);
> +  NewIdtr                   = InitializeBspIdt ();
> +  Status                    = InitializeCpuExceptionHandlers (NULL);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +  for (Index = 0; Index < 22; Index++) {
> +    if (Index == EXCEPT_IA32_PAGE_FAULT) {
> +      if (!FindPFAddressInPageTable (&FaultParameter)) {
> +        continue;
> +      }
> +    } else if (Index == EXCEPT_IA32_GP_FAULT) {
> +      FaultParameter = CR4_RESERVED_BIT;
> +    } else {
> +      if ((mErrorCodeExceptionFlag & (1 << Index)) != 0) {
> +        continue;
> +      }
> +    }
> +
> +    DEBUG ((DEBUG_INFO, "TestCase3: ExceptionType is %d\n", Index));
> +    Status = RegisterCpuInterruptHandler (Index, AdjustCpuContextHandler);
> +    UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +    //
> +    // Trigger different type exception and compare different stage cpu 
> context.
> +    //
> +    AsmTestConsistencyOfCpuContext (Index, FaultParameter);
> +    CompareCpuContext ();
> +    Status = RegisterCpuInterruptHandler (Index, NULL);
> +    UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  }
> +
> +  RestoreAllCpuRegisters (NULL, CpuOriginalRegisterBuffer, 0);
> +  FreePool (CpuOriginalRegisterBuffer);
> +  FreePool (NewIdtr);
> +  return UNIT_TEST_PASSED;
> +}
> +
> +/**
> +  Initializes CPU exceptions handlers for the sake of stack switch 
> requirement.
> +
> +  This function is a wrapper of InitializeSeparateExceptionStacks. It's 
> mainly
> +  for the sake of AP's init because of EFI_AP_PROCEDURE API requirement.
> +
> +  @param[in,out] Buffer  The pointer to private data buffer.
> +
> +**/
> +VOID
> +EFIAPI
> +InitializeExceptionStackSwitchHandlersPerAp (
> +  IN OUT VOID  *Buffer
> +  )
> +{
> +  EXCEPTION_STACK_SWITCH_CONTEXT  *CpuSwitchStackData;
> +
> +  CpuSwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
> +
> +  //
> +  // This may be called twice for each Cpu. Only run 
> InitializeSeparateExceptionStacks
> +  // if this is the first call or the first call failed because of size too 
> small.
> +  //
> +  if ((CpuSwitchStackData->Status == EFI_NOT_STARTED) || 
> (CpuSwitchStackData->Status == EFI_BUFFER_TOO_SMALL)) {
> +    CpuSwitchStackData->Status = InitializeSeparateExceptionStacks 
> (CpuSwitchStackData->Buffer, &CpuSwitchStackData-
> >BufferSize);
> +  }
> +}
> +
> +/**
> +  Initializes MP exceptions handlers for the sake of stack switch 
> requirement.
> +
> +  This function will allocate required resources required to setup stack 
> switch
> +  and pass them through SwitchStackData to each logic processor.
> +
> +  @param[in, out] MpServices       MpServices.
> +  @param[in, out] BspProcessorNum  Bsp processor number.
> +
> +  @return Pointer to the allocated SwitchStackData.
> +**/
> +EXCEPTION_STACK_SWITCH_CONTEXT  *
> +InitializeMpExceptionStackSwitchHandlers (
> +  MP_SERVICES  MpServices,
> +  UINTN        BspProcessorNum
> +  )
> +{
> +  UINTN                           Index;
> +  EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
> +  UINTN                           BufferSize;
> +  EFI_STATUS                      Status;
> +  UINT8                           *Buffer;
> +
> +  SwitchStackData = AllocateZeroPool (mNumberOfProcessors * sizeof 
> (EXCEPTION_STACK_SWITCH_CONTEXT));
> +  ASSERT (SwitchStackData != NULL);
> +  for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +    //
> +    // Because the procedure may runs multiple times, use the status 
> EFI_NOT_STARTED
> +    // to indicate the procedure haven't been run yet.
> +    //
> +    SwitchStackData[Index].Status = EFI_NOT_STARTED;
> +    if (Index == BspProcessorNum) {
> +      InitializeExceptionStackSwitchHandlersPerAp ((VOID 
> *)&SwitchStackData[Index]);
> +      continue;
> +    }
> +
> +    Status = MpServicesUnitTestStartupThisAP (
> +               MpServices,
> +               InitializeExceptionStackSwitchHandlersPerAp,
> +               Index,
> +               0,
> +               (VOID *)&SwitchStackData[Index]
> +               );
> +    ASSERT_EFI_ERROR (Status);
> +  }
> +
> +  BufferSize = 0;
> +  for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +    if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {
> +      ASSERT (SwitchStackData[Index].BufferSize != 0);
> +      BufferSize += SwitchStackData[Index].BufferSize;
> +    } else {
> +      ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);
> +      ASSERT (SwitchStackData[Index].BufferSize == 0);
> +    }
> +  }
> +
> +  if (BufferSize != 0) {
> +    Buffer = AllocateZeroPool (BufferSize);
> +    ASSERT (Buffer != NULL);
> +    BufferSize = 0;
> +    for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +      if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {
> +        SwitchStackData[Index].Buffer = (VOID *)(&Buffer[BufferSize]);
> +        BufferSize                   += SwitchStackData[Index].BufferSize;
> +        DEBUG ((
> +          DEBUG_INFO,
> +          "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlersPerAp: 
> 0x%lX with size 0x%lX\n",
> +          (UINT64)(UINTN)Index,
> +          (UINT64)(UINTN)SwitchStackData[Index].Buffer,
> +          (UINT64)(UINTN)SwitchStackData[Index].BufferSize
> +          ));
> +      }
> +    }
> +
> +    for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +      if (Index == BspProcessorNum) {
> +        InitializeExceptionStackSwitchHandlersPerAp ((VOID 
> *)&SwitchStackData[Index]);
> +        continue;
> +      }
> +
> +      Status = MpServicesUnitTestStartupThisAP (
> +                 MpServices,
> +                 InitializeExceptionStackSwitchHandlersPerAp,
> +                 Index,
> +                 0,
> +                 (VOID *)&SwitchStackData[Index]
> +                 );
> +      ASSERT_EFI_ERROR (Status);
> +    }
> +
> +    for (Index = 0; Index < mNumberOfProcessors; ++Index) {
> +      ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);
> +    }
> +  }
> +
> +  return SwitchStackData;
> +}
> +
> +/**
> +  Test if stack overflow is captured by CpuStackGuard in both Bsp and AP.
> +
> +  @param[in]  Context    [Optional] An optional parameter that enables:
> +                         1) test-case reuse with varied parameters and
> +                         2) test-case re-entry for Target tests that need a
> +                         reboot.  This parameter is a VOID* and it is the
> +                         responsibility of the test author to ensure that the
> +                         contents are well understood by all test cases that 
> may
> +                         consume it.
> +
> +  @retval  UNIT_TEST_PASSED             The Unit test has completed and the 
> test
> +                                        case was successful.
> +  @retval  UNIT_TEST_ERROR_TEST_FAILED  A test case assertion has failed.
> +**/
> +UNIT_TEST_STATUS
> +EFIAPI
> +TestCpuStackGuardInBspAndAp (
> +  IN UNIT_TEST_CONTEXT  Context
> +  )
> +{
> +  EFI_STATUS                      Status;
> +  UINTN                           OriginalStackBase;
> +  UINTN                           NewStackTop;
> +  UINTN                           NewStackBase;
> +  EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
> +  MP_SERVICES                     MpServices;
> +  UINTN                           ProcessorNumber;
> +  UINTN                           EnabledProcessorNum;
> +  CPU_REGISTER_BUFFER             *CpuOriginalRegisterBuffer;
> +  UINTN                           Index;
> +  UINTN                           BspProcessorNum;
> +  VOID                            *NewIdtr;
> +  UINTN                           *CpuStackBaseBuffer;
> +
> +  if (!PcdGetBool (PcdCpuStackGuard)) {
> +    return UNIT_TEST_PASSED;
> +  }
> +
> +  //
> +  // Get MP Service Protocol
> +  //
> +  Status = GetMpServices (&MpServices);
> +  Status = MpServicesUnitTestGetNumberOfProcessors (MpServices, 
> &ProcessorNumber, &EnabledProcessorNum);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  Status = MpServicesUnitTestWhoAmI (MpServices, &BspProcessorNum);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  mNumberOfProcessors = ProcessorNumber;
> +
> +  CpuOriginalRegisterBuffer = SaveAllCpuRegisters (&MpServices, 
> BspProcessorNum);
> +
> +  //
> +  // Initialize Bsp and AP Idt.
> +  // Idt buffer should not be empty or it will hang in MP API.
> +  //
> +  NewIdtr = InitializeBspIdt ();
> +  Status  = InitializeCpuExceptionHandlers (NULL);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  InitializeApIdt (MpServices, NewIdtr);
> +
> +  //
> +  // Get BSP and AP original stack base.
> +  //
> +  CpuStackBaseBuffer = GetAllCpuStackBase (&MpServices, BspProcessorNum);
> +
> +  //
> +  // InitializeMpExceptionStackSwitchHandlers and register exception handler.
> +  //
> +  SwitchStackData = InitializeMpExceptionStackSwitchHandlers (MpServices, 
> BspProcessorNum);
> +  Status          = RegisterCpuInterruptHandler (EXCEPT_IA32_PAGE_FAULT, 
> CpuStackGuardExceptionHandler);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  Status = RegisterCpuInterruptHandler (EXCEPT_IA32_DOUBLE_FAULT, 
> AdjustRipForFaultHandler);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +
> +  for (Index = 0; Index < mNumberOfProcessors; Index++) {
> +    OriginalStackBase = CpuStackBaseBuffer[Index];
> +    NewStackTop       = (UINTN)(SwitchStackData[Index].Buffer) + 
> SwitchStackData[Index].BufferSize;
> +    NewStackBase      = (UINTN)(SwitchStackData[Index].Buffer);
> +    if (Index == BspProcessorNum) {
> +      TriggerStackOverflow ();
> +    } else {
> +      MpServicesUnitTestStartupThisAP (
> +        MpServices,
> +        (EFI_AP_PROCEDURE)TriggerStackOverflow,
> +        Index,
> +        0,
> +        NULL
> +        );
> +    }
> +
> +    DEBUG ((DEBUG_INFO, "TestCase4: mRspAddress[0] is 0x%x, mRspAddress[1] 
> is 0x%x\n", mRspAddress[0],
> mRspAddress[1]));
> +    UT_ASSERT_TRUE ((mRspAddress[0] >= OriginalStackBase) && (mRspAddress[0] 
> <= (OriginalStackBase + SIZE_4KB)));
> +    UT_ASSERT_TRUE ((mRspAddress[1] >= NewStackBase) && (mRspAddress[1] < 
> NewStackTop));
> +  }
> +
> +  Status = RegisterCpuInterruptHandler (EXCEPT_IA32_PAGE_FAULT, NULL);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  Status = RegisterCpuInterruptHandler (EXCEPT_IA32_DOUBLE_FAULT, NULL);
> +  UT_ASSERT_EQUAL (Status, EFI_SUCCESS);
> +  RestoreAllCpuRegisters (&MpServices, CpuOriginalRegisterBuffer, 
> BspProcessorNum);
> +  FreePool (SwitchStackData);
> +  FreePool (CpuOriginalRegisterBuffer);
> +  FreePool (NewIdtr);
> +
> +  return UNIT_TEST_PASSED;
> +}
> +
> +/**
> +  Create CpuExceptionLibUnitTestSuite and add test case.
> +
> +  @param[in]  FrameworkHandle    Unit test framework.
> +
> +  @return  EFI_SUCCESS           The unit test suite was created.
> +  @retval  EFI_OUT_OF_RESOURCES  There are not enough resources available to
> +                                 initialize the unit test suite.
> +**/
> +EFI_STATUS
> +AddCommonTestCase (
> +  IN  UNIT_TEST_FRAMEWORK_HANDLE  Framework
> +  )
> +{
> +  EFI_STATUS              Status;
> +  UNIT_TEST_SUITE_HANDLE  CpuExceptionLibUnitTestSuite;
> +
> +  //
> +  // Populate the Manual Test Cases.
> +  //
> +  Status = CreateUnitTestSuite (&CpuExceptionLibUnitTestSuite, Framework, 
> "Test CpuExceptionHandlerLib",
> "CpuExceptionHandlerLib.Manual", NULL, NULL);
> +  if (EFI_ERROR (Status)) {
> +    DEBUG ((DEBUG_ERROR, "Failed in CreateUnitTestSuite for 
> CpuExceptionHandlerLib Test Cases\n"));
> +    Status = EFI_OUT_OF_RESOURCES;
> +    return Status;
> +  }
> +
> +  AddTestCase (CpuExceptionLibUnitTestSuite, "Check if exception handler can 
> be registered/unregistered for no error
> code exception", "TestRegisterHandlerForNoErrorCodeException", 
> TestRegisterHandlerForNoErrorCodeException, NULL,
> NULL, NULL);
> +  AddTestCase (CpuExceptionLibUnitTestSuite, "Check if exception handler can 
> be registered/unregistered for GP and PF",
> "TestRegisterHandlerForGPAndPF", TestRegisterHandlerForGPAndPF, NULL, NULL, 
> NULL);
> +
> +  AddTestCase (CpuExceptionLibUnitTestSuite, "Check if Cpu Context is 
> consistent before and after exception.",
> "TestCpuContextConsistency", TestCpuContextConsistency, NULL, NULL, NULL);
> +  AddTestCase (CpuExceptionLibUnitTestSuite, "Check if stack overflow is 
> captured by CpuStackGuard in Bsp and AP",
> "TestCpuStackGuardInBspAndAp", TestCpuStackGuardInBspAndAp, NULL, NULL, NULL);
> +
> +  return EFI_SUCCESS;
> +}
> diff --git 
> a/UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerLibUnitTest.inf
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerLibUnitTest.inf
> new file mode 100644
> index 0000000000..e3dbe7b9ab
> --- /dev/null
> +++ 
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerLibUnitTest.inf
> @@ -0,0 +1,58 @@
> +## @file
> +# Unit tests of the DxeCpuExceptionHandlerLib instance.
> +#
> +# Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
> +# SPDX-License-Identifier: BSD-2-Clause-Patent
> +##
> +
> +[Defines]
> +  INF_VERSION                    = 0x00010005
> +  BASE_NAME                      = CpuExceptionHandlerDxeTest
> +  FILE_GUID                      = D76BFD9C-0B6D-46BD-AD66-2BBB6FA7031A
> +  MODULE_TYPE                    = DXE_DRIVER
> +  VERSION_STRING                 = 1.0
> +  ENTRY_POINT                    = CpuExceptionHandlerTestEntry
> +
> +#
> +# The following information is for reference only and not required by the 
> build tools.
> +#
> +#  VALID_ARCHITECTURES           = X64
> +#
> +[Sources.X64]
> +  X64/ArchExceptionHandlerTestAsm.nasm
> +  X64/ArchExceptionHandlerTest.c
> +
> +[Sources.common]
> +  CpuExceptionHandlerTest.h
> +  CpuExceptionHandlerTestCommon.c
> +  DxeCpuExceptionHandlerUnitTest.c
> +
> +[Packages]
> +  MdePkg/MdePkg.dec
> +  MdeModulePkg/MdeModulePkg.dec
> +  UefiCpuPkg/UefiCpuPkg.dec
> +
> +[LibraryClasses]
> +  BaseLib
> +  BaseMemoryLib
> +  DebugLib
> +  UnitTestLib
> +  MemoryAllocationLib
> +  CpuExceptionHandlerLib
> +  UefiDriverEntryPoint
> +  HobLib
> +  UefiBootServicesTableLib
> +  CpuPageTableLib
> +
> +[Guids]
> +  gEfiHobMemoryAllocStackGuid
> +
> +[Pcd]
> +  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard       ## CONSUMES
> +  gUefiCpuPkgTokenSpaceGuid.PcdCpuApStackSize           ## CONSUMES
> +
> +[Protocols]
> +  gEfiMpServiceProtocolGuid
> +
> +[Depex]
> +  gEfiMpServiceProtocolGuid
> diff --git 
> a/UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerUnitTest.c
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerUnitTest.c
> new file mode 100644
> index 0000000000..917fc549bf
> --- /dev/null
> +++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/DxeCpuExceptionHandlerUnitTest.c
> @@ -0,0 +1,196 @@
> +/** @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/UefiBootServicesTableLib.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);
> +  ASSERT (NewIdtTable != NULL);
> +  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.Protocol->GetNumberOfProcessors (MpServices.Protocol, 
> NumberOfProcessors,
> NumberOfEnabledProcessors);
> +}
> +
> +/**
> +  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.Protocol->WhoAmI (MpServices.Protocol, ProcessorNumber);
> +}
> +
> +/**
> +  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.Protocol->StartupThisAP (MpServices.Protocol, Procedure, 
> ProcessorNumber, NULL,
> TimeoutInMicroSeconds, ProcedureArgument, NULL);
> +}
> +
> +/**
> +  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.Protocol->StartupAllAPs (MpServices.Protocol, Procedure, 
> SingleThread, NULL,
> TimeoutInMicroSeconds, ProcedureArgument, NULL);
> +}
> +
> +/**
> +  Get EFI_MP_SERVICES_PROTOCOL pointer.
> +
> +  @param[out] MpServices    Pointer to the buffer where 
> EFI_MP_SERVICES_PROTOCOL is stored
> +
> +  @retval EFI_SUCCESS       EFI_MP_SERVICES_PROTOCOL interface is returned
> +  @retval EFI_NOT_FOUND     EFI_MP_SERVICES_PROTOCOL interface is not found
> +**/
> +EFI_STATUS
> +GetMpServices (
> +  OUT MP_SERVICES  *MpServices
> +  )
> +{
> +  return gBS->LocateProtocol (&gEfiMpServiceProtocolGuid, NULL, (VOID 
> **)&MpServices->Protocol);
> +}
> +
> +/**
> +  Entry for CpuExceptionHandlerDxeTest driver.
> +
> +  @param ImageHandle     Image handle this driver.
> +  @param SystemTable     Pointer to the System Table.
> +
> +  @retval EFI_SUCCESS    The driver executed normally.
> +
> +**/
> +EFI_STATUS
> +EFIAPI
> +CpuExceptionHandlerTestEntry (
> +  IN EFI_HANDLE        ImageHandle,
> +  IN EFI_SYSTEM_TABLE  *SystemTable
> +  )
> +{
> +  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;
> +}
> diff --git 
> a/UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTest.c
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTest.c
> new file mode 100644
> index 0000000000..c0d962f26d
> --- /dev/null
> +++ b/UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTest.c
> @@ -0,0 +1,166 @@
> +/** @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  mActualContextInHandler;
> +GENERAL_REGISTER  mActualContextAfterException;
> +
> +//
> +// In TestCpuContextConsistency, Cpu registers will be set to
> mExpectedContextInHandler/mExpectedContextAfterException.
> +// Rcx in mExpectedContextInHandler is set runtime since Rcx is needed in 
> assembly code.
> +// For GP and PF, Rcx is set to FaultParameter. For other exception 
> triggered by INTn, Rcx is set to ExceptionType.
> +//
> +GENERAL_REGISTER  mExpectedContextInHandler      = { 1, 2, 3, 4, 5, 0, 7, 8, 
> 9, 0xa, 0xb, 0xc, 0xd, 0xe };
> +GENERAL_REGISTER  mExpectedContextAfterException = { 0x11, 0x12, 0x13, 0x14, 
> 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
> 0x1c, 0x1d, 0x1e };
> +
> +/**
> +  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.SystemContextX64->Rip += 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 mActualContextInHandler.
> +  //
> +  mActualContextInHandler.Rdi = SystemContext.SystemContextX64->Rdi;
> +  mActualContextInHandler.Rsi = SystemContext.SystemContextX64->Rsi;
> +  mActualContextInHandler.Rbx = SystemContext.SystemContextX64->Rbx;
> +  mActualContextInHandler.Rdx = SystemContext.SystemContextX64->Rdx;
> +  mActualContextInHandler.Rcx = SystemContext.SystemContextX64->Rcx;
> +  mActualContextInHandler.Rax = SystemContext.SystemContextX64->Rax;
> +  mActualContextInHandler.R8  = SystemContext.SystemContextX64->R8;
> +  mActualContextInHandler.R9  = SystemContext.SystemContextX64->R9;
> +  mActualContextInHandler.R10 = SystemContext.SystemContextX64->R10;
> +  mActualContextInHandler.R11 = SystemContext.SystemContextX64->R11;
> +  mActualContextInHandler.R12 = SystemContext.SystemContextX64->R12;
> +  mActualContextInHandler.R13 = SystemContext.SystemContextX64->R13;
> +  mActualContextInHandler.R14 = SystemContext.SystemContextX64->R14;
> +  mActualContextInHandler.R15 = SystemContext.SystemContextX64->R15;
> +
> +  //
> +  // Modify cpu context. These registers will be stored in 
> mActualContextAfterException.
> +  // Do not handle Rsp and Rbp. CpuExceptionHandlerLib doesn't set Rsp and 
> Rbp register
> +  // to the value in SystemContext.
> +  //
> +  SystemContext.SystemContextX64->Rdi = mExpectedContextAfterException.Rdi;
> +  SystemContext.SystemContextX64->Rsi = mExpectedContextAfterException.Rsi;
> +  SystemContext.SystemContextX64->Rbx = mExpectedContextAfterException.Rbx;
> +  SystemContext.SystemContextX64->Rdx = mExpectedContextAfterException.Rdx;
> +  SystemContext.SystemContextX64->Rcx = mExpectedContextAfterException.Rcx;
> +  SystemContext.SystemContextX64->Rax = mExpectedContextAfterException.Rax;
> +  SystemContext.SystemContextX64->R8  = mExpectedContextAfterException.R8;
> +  SystemContext.SystemContextX64->R9  = mExpectedContextAfterException.R9;
> +  SystemContext.SystemContextX64->R10 = mExpectedContextAfterException.R10;
> +  SystemContext.SystemContextX64->R11 = mExpectedContextAfterException.R11;
> +  SystemContext.SystemContextX64->R12 = mExpectedContextAfterException.R12;
> +  SystemContext.SystemContextX64->R13 = mExpectedContextAfterException.R13;
> +  SystemContext.SystemContextX64->R14 = mExpectedContextAfterException.R14;
> +  SystemContext.SystemContextX64->R15 = mExpectedContextAfterException.R15;
> +
> +  //
> +  // 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 mActualContextAfterException.
> +
> +  @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.Rdi, 
> mExpectedContextInHandler.Rdi);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.Rsi, 
> mExpectedContextInHandler.Rsi);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.Rbx, 
> mExpectedContextInHandler.Rbx);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.Rdx, 
> mExpectedContextInHandler.Rdx);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.Rcx, 
> mExpectedContextInHandler.Rcx);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.Rax, 
> mExpectedContextInHandler.Rax);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R8, mExpectedContextInHandler.R8);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R9, mExpectedContextInHandler.R9);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R10, 
> mExpectedContextInHandler.R10);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R11, 
> mExpectedContextInHandler.R11);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R12, 
> mExpectedContextInHandler.R12);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R13, 
> mExpectedContextInHandler.R13);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R14, 
> mExpectedContextInHandler.R14);
> +  UT_ASSERT_EQUAL (mActualContextInHandler.R15, 
> mExpectedContextInHandler.R15);
> +
> +  UT_ASSERT_EQUAL (mActualContextAfterException.Rdi, 
> mExpectedContextAfterException.Rdi);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.Rsi, 
> mExpectedContextAfterException.Rsi);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.Rbx, 
> mExpectedContextAfterException.Rbx);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.Rdx, 
> mExpectedContextAfterException.Rdx);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.Rcx, 
> mExpectedContextAfterException.Rcx);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.Rax, 
> mExpectedContextAfterException.Rax);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R8, 
> mExpectedContextAfterException.R8);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R9, 
> mExpectedContextAfterException.R9);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R10, 
> mExpectedContextAfterException.R10);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R11, 
> mExpectedContextAfterException.R11);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R12, 
> mExpectedContextAfterException.R12);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R13, 
> mExpectedContextAfterException.R13);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R14, 
> mExpectedContextAfterException.R14);
> +  UT_ASSERT_EQUAL (mActualContextAfterException.R15, 
> mExpectedContextAfterException.R15);
> +  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.SystemContextX64->Rsp;
> +  mRspAddress[1] = (UINTN)(&LocalVariable);
> +
> +  return;
> +}
> diff --git 
> a/UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTestAsm.nasm
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTestAsm.nasm
> new file mode 100644
> index 0000000000..e229dbed00
> --- /dev/null
> +++ 
> b/UefiCpuPkg/CpuExceptionHandlerUnitTest/X64/ArchExceptionHandlerTestAsm.nasm
> @@ -0,0 +1,256 @@
> +;------------------------------------------------------------------------------
> +;
> +; Copyright (c) 2022, Intel Corporation. All rights reserved.<BR>
> +; SPDX-License-Identifier: BSD-2-Clause-Patent
> +;
> +; Module Name:
> +;
> +;   ArchExceptionHandlerTestAsm.nasm
> +;
> +; Abstract:
> +;
> +;   x64 CPU Exception Handler Lib Unit test
> +;
> +;------------------------------------------------------------------------------
> +
> +    DEFAULT REL
> +    SECTION .text
> +
> +struc GENERAL_REGISTER
> +  .Rdi:    resq    1
> +  .Rsi:    resq    1
> +  .Rbx:    resq    1
> +  .Rdx:    resq    1
> +  .Rcx:    resq    1
> +  .Rax:    resq    1
> +  .R8:     resq    1
> +  .R9:     resq    1
> +  .R10:    resq    1
> +  .R11:    resq    1
> +  .R12:    resq    1
> +  .R13:    resq    1
> +  .R14:    resq    1
> +  .R15:    resq    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
> +    ;
> +    push rcx
> +    lea  rcx, [ASM_PFX(mFaultInstructionLength)]
> +    mov  qword[rcx], TriggerGPExceptionAfter - TriggerGPExceptionBefore
> +    pop  rcx
> +TriggerGPExceptionBefore:
> +    mov  cr4, rcx
> +TriggerGPExceptionAfter:
> +    ret
> +
> +;------------------------------------------------------------------------------
> +; VOID
> +; EFIAPI
> +; TriggerPFException (
> +;  UINTN  PFAddress
> +;  );
> +;------------------------------------------------------------------------------
> +global ASM_PFX(TriggerPFException)
> +ASM_PFX(TriggerPFException):
> +    push rcx
> +    lea  rcx, [ASM_PFX(mFaultInstructionLength)]
> +    mov  qword[rcx], TriggerPFExceptionAfter - TriggerPFExceptionBefore
> +    pop  rcx
> +TriggerPFExceptionBefore:
> +    mov  qword[rcx], 0x1
> +TriggerPFExceptionAfter:
> +    ret
> +
> +;------------------------------------------------------------------------------
> +; ModifyRcxInGlobalBeforeException;
> +; This function is writed by assebly code because it's only called in this 
> file.
> +; It's used to set Rcx in mExpectedContextInHandler for different exception.
> +;------------------------------------------------------------------------------
> +global ASM_PFX(ModifyRcxInGlobalBeforeException)
> +ASM_PFX(ModifyRcxInGlobalBeforeException):
> +    push rax
> +    lea  rax, [ASM_PFX(mExpectedContextInHandler)]
> +    mov  [rax + GENERAL_REGISTER.Rcx], rcx
> +    pop  rax
> +    ret
> +
> +;------------------------------------------------------------------------------
> +;VOID
> +;EFIAPI
> +;AsmTestConsistencyOfCpuContext (
> +;  IN  EFI_EXCEPTION_TYPE ExceptionType
> +;  IN  UINTN              FaultParameter   OPTIONAL
> +;  );
> +;------------------------------------------------------------------------------
> +global ASM_PFX(AsmTestConsistencyOfCpuContext)
> +ASM_PFX(AsmTestConsistencyOfCpuContext):
> +    ;
> +    ; Push original register
> +    ;
> +    push r15
> +    push r14
> +    push r13
> +    push r12
> +    push r11
> +    push r10
> +    push r9
> +    push r8
> +    push rax
> +    push rcx
> +    push rbx
> +    push rsi
> +    push rdi
> +    push rdx
> +    push rdx
> +
> +    ;
> +    ; Modify registers to mExpectedContextInHandler. Do not handle Rsp and 
> Rbp.
> +    ; CpuExceptionHandlerLib doesn't set Rsp and Rsp register to the value 
> in SystemContext.
> +    ;
> +    lea r15, [ASM_PFX(mExpectedContextInHandler)]
> +    mov rdi, [r15 + GENERAL_REGISTER.Rdi]
> +    mov rsi, [r15 + GENERAL_REGISTER.Rsi]
> +    mov rbx, [r15 + GENERAL_REGISTER.Rbx]
> +    mov rdx, [r15 + GENERAL_REGISTER.Rdx]
> +    mov rax, [r15 + GENERAL_REGISTER.Rax]
> +    mov r8,  [r15 + GENERAL_REGISTER.R8]
> +    mov r9,  [r15 + GENERAL_REGISTER.R9]
> +    mov r10, [r15 + GENERAL_REGISTER.R10]
> +    mov r11, [r15 + GENERAL_REGISTER.R11]
> +    mov r12, [r15 + GENERAL_REGISTER.R12]
> +    mov r13, [r15 + GENERAL_REGISTER.R13]
> +    mov r14, [r15 + GENERAL_REGISTER.R14]
> +    mov r15, [r15 + GENERAL_REGISTER.R15]
> +
> +    cmp  rcx, 0xd
> +    jz   GPException
> +    cmp  rcx, 0xe
> +    jz   PFException
> +    jmp  INTnException
> +
> +PFException:
> +    pop  rcx                                       ; Pop rdx(PFAddress) to 
> rcx.
> +    call ASM_PFX(ModifyRcxInGlobalBeforeException) ; Set 
> mExpectedContextInHandler.Rcx to PFAddress.
> +    call ASM_PFX(TriggerPFException)
> +    jmp  AfterException
> +
> +GPException:
> +    pop  rcx                                       ; Pop rdx(Cr4ReservedBit) 
> to rcx.
> +    call ASM_PFX(ModifyRcxInGlobalBeforeException) ; Set 
> mExpectedContextInHandler.Rcx to Cr4ReservedBit.
> +    call ASM_PFX(TriggerGPException)
> +    jmp  AfterException
> +
> +INTnException:
> +    ;
> +    ; Modify Rcx in mExpectedContextInHandler.
> +    ;
> +    add  Rsp, 8                                    ; Discard the extra Rdx 
> in stack. Rcx is ExceptionType now.
> +    call ASM_PFX(ModifyRcxInGlobalBeforeException) ; Set 
> mExpectedContextInHandler.Rcx to ExceptionType.
> +    call ASM_PFX(TriggerINTnException)
> +
> +AfterException:
> +    ;
> +    ; Save registers in mActualContextAfterException
> +    ;
> +    push rax
> +    lea  rax, [ASM_PFX(mActualContextAfterException)]
> +    mov  [rax + GENERAL_REGISTER.Rdi], rdi
> +    mov  [rax + GENERAL_REGISTER.Rsi], rsi
> +    mov  [rax + GENERAL_REGISTER.Rbx], rbx
> +    mov  [rax + GENERAL_REGISTER.Rdx], rdx
> +    mov  [rax + GENERAL_REGISTER.Rcx], rcx
> +    pop  rcx
> +    mov  [rax + GENERAL_REGISTER.Rax], rcx
> +    mov  [rax + GENERAL_REGISTER.R8],  r8
> +    mov  [rax + GENERAL_REGISTER.R9],  r9
> +    mov  [rax + GENERAL_REGISTER.R10], r10
> +    mov  [rax + GENERAL_REGISTER.R11], r11
> +    mov  [rax + GENERAL_REGISTER.R12], r12
> +    mov  [rax + GENERAL_REGISTER.R13], r13
> +    mov  [rax + GENERAL_REGISTER.R14], r14
> +    mov  [rax + GENERAL_REGISTER.R15], r15
> +
> +    ;
> +    ; restore original register
> +    ;
> +    pop rdx
> +    pop rdi
> +    pop rsi
> +    pop rbx
> +    pop rcx
> +    pop rax
> +    pop r8
> +    pop r9
> +    pop r10
> +    pop r11
> +    pop r12
> +    pop r13
> +    pop r14
> +    pop r15
> +
> +    ret
> +
> +;------------------------------------------------------------------------------
> +; VOID
> +; EFIAPI
> +; TriggerStackOverflow (
> +;  VOID
> +;  );
> +;------------------------------------------------------------------------------
> +global ASM_PFX(TriggerStackOverflow)
> +ASM_PFX(TriggerStackOverflow):
> +    push rcx
> +    lea  rcx, [ASM_PFX(mFaultInstructionLength)]
> +    mov  qword[rcx], TriggerCpuStackGuardAfter - TriggerCpuStackGuardBefore
> +    pop  rcx
> +TriggerCpuStackGuardBefore:
> +    call TriggerCpuStackGuardBefore
> +TriggerCpuStackGuardAfter:
> +    ret
> +
> +;------------------------------------------------------------------------------
> +; VOID
> +; EFIAPI
> +; TriggerINTnException (
> +;  IN  EFI_EXCEPTION_TYPE ExceptionType
> +;  );
> +;------------------------------------------------------------------------------
> +global ASM_PFX(TriggerINTnException)
> +ASM_PFX(TriggerINTnException):
> +    push rax
> +    push rdx
> +    push rcx
> +    lea  rax, [AsmTriggerException1 - AsmTriggerException0]
> +    mul  rcx
> +    mov  rcx, AsmTriggerException0
> +    add  rax, rcx
> +    pop  rcx
> +    pop  rdx
> +    jmp  rax
> +    ;
> +    ; rax = AsmTriggerException0 + (AsmTriggerException1 - 
> AsmTriggerException0) * rcx
> +    ;
> +%assign Vector 0
> +%rep  22
> +AsmTriggerException %+ Vector:
> +    pop rax
> +    INT Vector
> +    ret
> +%assign Vector Vector+1
> +%endrep
> --
> 2.31.1.windows.1



-=-=-=-=-=-=-=-=-=-=-=-
Groups.io Links: You receive all messages sent to this group.
View/Reply Online (#95206): https://edk2.groups.io/g/devel/message/95206
Mute This Topic: https://groups.io/mt/94323058/21656
Group Owner: devel+ow...@edk2.groups.io
Unsubscribe: https://edk2.groups.io/g/devel/unsub [arch...@mail-archive.com]
-=-=-=-=-=-=-=-=-=-=-=-



Reply via email to