[Patch V3 1/4] UefiCpuPkg: Add Unit tests for DxeCpuExceptionHandlerLib

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

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.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             | 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