[PATCH 1/2] UefiCpuPkg: Simplify InitializeSeparateExceptionStacks

["Zhiguang Liu" <zhiguang.liu@intel.com>]

Hide the Exception implementation details in CpuExcetionHandlerLib and
caller only need to provide buffer

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Leif Lindholm <quic_llindhol@quicinc.com>
Cc: Dandan Bi <dandan.bi@intel.com>
Cc: Liming Gao <gaoliming@byosoft.com.cn>
Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Sami Mujawar <sami.mujawar@arm.com>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>
---
 .../Library/ArmExceptionLib/ArmExceptionLib.c |  15 +-
 MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c       |   4 +-
 .../Include/Library/CpuExceptionHandlerLib.h  |  15 +-
 .../CpuExceptionHandlerLibNull.c              |  15 +-
 UefiCpuPkg/CpuDxe/CpuMp.c                     | 157 +++-------------
 UefiCpuPkg/CpuDxe/CpuMp.h                     |  10 +-
 UefiCpuPkg/CpuMpPei/CpuMpPei.c                | 174 ++++--------------
 UefiCpuPkg/CpuMpPei/CpuMpPei.h                |  10 +-
 .../CpuExceptionHandlerLib/DxeException.c     |  77 +++++---
 .../Ia32/ArchExceptionHandler.c               |   3 +-
 .../CpuExceptionHandlerLib/PeiCpuException.c  |  62 ++++++-
 .../PeiCpuExceptionHandlerLib.inf             |   4 +-
 .../SecPeiCpuException.c                      |  15 +-
 .../CpuExceptionHandlerLib/SmmException.c     |  15 +-
 .../X64/ArchExceptionHandler.c                |   3 +-
 15 files changed, 231 insertions(+), 348 deletions(-)

diff --git a/ArmPkg/Library/ArmExceptionLib/ArmExceptionLib.c b/ArmPkg/Library/ArmExceptionLib/ArmExceptionLib.c
index 2c7bc66aa7..a521c33f32 100644
--- a/ArmPkg/Library/ArmExceptionLib/ArmExceptionLib.c
+++ b/ArmPkg/Library/ArmExceptionLib/ArmExceptionLib.c
@@ -288,20 +288,23 @@ CommonCExceptionHandler (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   )
 {
   return EFI_SUCCESS;
diff --git a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
index 0a1f3d79e2..5733f0c8ec 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
+++ b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
@@ -1,7 +1,7 @@
 /** @file
   DXE Core Main Entry Point
 
-Copyright (c) 2006 - 2018, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2022, Intel Corporation. All rights reserved.<BR>
 SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -260,7 +260,7 @@ DxeMain (
   // Setup Stack Guard
   //
   if (PcdGetBool (PcdCpuStackGuard)) {
-    Status = InitializeSeparateExceptionStacks (NULL);
+    Status = InitializeSeparateExceptionStacks (NULL, NULL);
     ASSERT_EFI_ERROR (Status);
   }
 
diff --git a/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h b/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h
index 9a495081f7..8d44ed916a 100644
--- a/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h
+++ b/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h
@@ -104,20 +104,23 @@ InitializeCpuExceptionHandlers (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   );
 
 /**
diff --git a/MdeModulePkg/Library/CpuExceptionHandlerLibNull/CpuExceptionHandlerLibNull.c b/MdeModulePkg/Library/CpuExceptionHandlerLibNull/CpuExceptionHandlerLibNull.c
index 8aeedcb4d1..74908a379b 100644
--- a/MdeModulePkg/Library/CpuExceptionHandlerLibNull/CpuExceptionHandlerLibNull.c
+++ b/MdeModulePkg/Library/CpuExceptionHandlerLibNull/CpuExceptionHandlerLibNull.c
@@ -83,20 +83,23 @@ DumpCpuContext (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   )
 {
   return EFI_UNSUPPORTED;
diff --git a/UefiCpuPkg/CpuDxe/CpuMp.c b/UefiCpuPkg/CpuDxe/CpuMp.c
index e385f585c7..286ef2d3fd 100644
--- a/UefiCpuPkg/CpuDxe/CpuMp.c
+++ b/UefiCpuPkg/CpuDxe/CpuMp.c
@@ -596,24 +596,6 @@ CollectBistDataFromHob (
   }
 }
 
-/**
-  Get GDT register value.
-
-  This function is mainly for AP purpose because AP may have different GDT
-  table than BSP.
-
-  @param[in,out] Buffer  The pointer to private data buffer.
-
-**/
-VOID
-EFIAPI
-GetGdtr (
-  IN OUT VOID  *Buffer
-  )
-{
-  AsmReadGdtr ((IA32_DESCRIPTOR *)Buffer);
-}
-
 /**
   Initializes CPU exceptions handlers for the sake of stack switch requirement.
 
@@ -629,27 +611,17 @@ InitializeExceptionStackSwitchHandlers (
   IN OUT VOID  *Buffer
   )
 {
-  CPU_EXCEPTION_INIT_DATA  *EssData;
-  IA32_DESCRIPTOR          Idtr;
-  EFI_STATUS               Status;
+  SWITCH_STACK_DATA  *SwitchStackData;
 
-  EssData = Buffer;
-  //
-  // We don't plan to replace IDT table with a new one, but we should not assume
-  // the AP's IDT is the same as BSP's IDT either.
-  //
-  AsmReadIdtr (&Idtr);
-  EssData->Ia32.IdtTable     = (VOID *)Idtr.Base;
-  EssData->Ia32.IdtTableSize = Idtr.Limit + 1;
-  Status                     = InitializeSeparateExceptionStacks (EssData);
-  ASSERT_EFI_ERROR (Status);
+  SwitchStackData = (SWITCH_STACK_DATA *)Buffer;
+  InitializeSeparateExceptionStacks (SwitchStackData->Buffer, SwitchStackData->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 CPU_EXCEPTION_INIT_DATA to each logic processor.
+  and pass them through SwitchStackData to each logic processor.
 
 **/
 VOID
@@ -657,129 +629,52 @@ InitializeMpExceptionStackSwitchHandlers (
   VOID
   )
 {
-  UINTN                    Index;
-  UINTN                    Bsp;
-  UINTN                    ExceptionNumber;
-  UINTN                    OldGdtSize;
-  UINTN                    NewGdtSize;
-  UINTN                    NewStackSize;
-  IA32_DESCRIPTOR          Gdtr;
-  CPU_EXCEPTION_INIT_DATA  EssData;
-  UINT8                    *GdtBuffer;
-  UINT8                    *StackTop;
-
-  ExceptionNumber = FixedPcdGetSize (PcdCpuStackSwitchExceptionList);
-  NewStackSize    = FixedPcdGet32 (PcdCpuKnownGoodStackSize) * ExceptionNumber;
-
-  StackTop = AllocateRuntimeZeroPool (NewStackSize * mNumberOfProcessors);
-  ASSERT (StackTop != NULL);
-  StackTop += NewStackSize  * mNumberOfProcessors;
+  UINTN              Index;
+  UINTN              Bsp;
+  UINT8              *Buffer;
+  SWITCH_STACK_DATA  SwitchStackData;
+  UINTN              BufferSize;
 
-  //
-  // The default exception handlers must have been initialized. Let's just skip
-  // it in this method.
-  //
-  EssData.Ia32.Revision            = CPU_EXCEPTION_INIT_DATA_REV;
-  EssData.Ia32.InitDefaultHandlers = FALSE;
-
-  EssData.Ia32.StackSwitchExceptions      = FixedPcdGetPtr (PcdCpuStackSwitchExceptionList);
-  EssData.Ia32.StackSwitchExceptionNumber = ExceptionNumber;
-  EssData.Ia32.KnownGoodStackSize         = FixedPcdGet32 (PcdCpuKnownGoodStackSize);
-
-  //
-  // Initialize Gdtr to suppress incorrect compiler/analyzer warnings.
-  //
-  Gdtr.Base  = 0;
-  Gdtr.Limit = 0;
+  SwitchStackData.BufferSize = &BufferSize;
   MpInitLibWhoAmI (&Bsp);
+
   for (Index = 0; Index < mNumberOfProcessors; ++Index) {
-    //
-    // To support stack switch, we need to re-construct GDT but not IDT.
-    //
+    SwitchStackData.Buffer = NULL;
+    BufferSize             = 0;
+
     if (Index == Bsp) {
-      GetGdtr (&Gdtr);
+      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
     } else {
       //
-      // AP might have different size of GDT from BSP.
+      // AP might need different buffer size from BSP.
       //
-      MpInitLibStartupThisAP (GetGdtr, Index, NULL, 0, (VOID *)&Gdtr, NULL);
+      MpInitLibStartupThisAP (InitializeExceptionStackSwitchHandlers, Index, NULL, 0, (VOID *)&SwitchStackData, NULL);
     }
 
-    //
-    // X64 needs only one TSS of current task working for all exceptions
-    // because of its IST feature. IA32 needs one TSS for each exception
-    // in addition to current task. Since AP is not supposed to allocate
-    // memory, we have to do it in BSP. To simplify the code, we allocate
-    // memory for IA32 case to cover both IA32 and X64 exception stack
-    // switch.
-    //
-    // Layout of memory to allocate for each processor:
-    //    --------------------------------
-    //    |            Alignment         |  (just in case)
-    //    --------------------------------
-    //    |                              |
-    //    |        Original GDT          |
-    //    |                              |
-    //    --------------------------------
-    //    |    Current task descriptor   |
-    //    --------------------------------
-    //    |                              |
-    //    |  Exception task descriptors  |  X ExceptionNumber
-    //    |                              |
-    //    --------------------------------
-    //    |  Current task-state segment  |
-    //    --------------------------------
-    //    |                              |
-    //    | Exception task-state segment |  X ExceptionNumber
-    //    |                              |
-    //    --------------------------------
-    //
-    OldGdtSize                        = Gdtr.Limit + 1;
-    EssData.Ia32.ExceptionTssDescSize = sizeof (IA32_TSS_DESCRIPTOR) *
-                                        (ExceptionNumber + 1);
-    EssData.Ia32.ExceptionTssSize = sizeof (IA32_TASK_STATE_SEGMENT) *
-                                    (ExceptionNumber + 1);
-    NewGdtSize = sizeof (IA32_TSS_DESCRIPTOR) +
-                 OldGdtSize +
-                 EssData.Ia32.ExceptionTssDescSize +
-                 EssData.Ia32.ExceptionTssSize;
-
-    GdtBuffer = AllocateRuntimeZeroPool (NewGdtSize);
-    ASSERT (GdtBuffer != NULL);
-
-    //
-    // Make sure GDT table alignment
-    //
-    EssData.Ia32.GdtTable     = ALIGN_POINTER (GdtBuffer, sizeof (IA32_TSS_DESCRIPTOR));
-    NewGdtSize               -= ((UINT8 *)EssData.Ia32.GdtTable - GdtBuffer);
-    EssData.Ia32.GdtTableSize = NewGdtSize;
-
-    EssData.Ia32.ExceptionTssDesc = ((UINT8 *)EssData.Ia32.GdtTable + OldGdtSize);
-    EssData.Ia32.ExceptionTss     = ((UINT8 *)EssData.Ia32.GdtTable + OldGdtSize +
-                                     EssData.Ia32.ExceptionTssDescSize);
-
-    EssData.Ia32.KnownGoodStackTop = (UINTN)StackTop;
+    ASSERT (BufferSize != 0);
+    Buffer = AllocateRuntimeZeroPool (BufferSize);
+    ASSERT (Buffer != NULL);
+    SwitchStackData.Buffer = Buffer;
     DEBUG ((
       DEBUG_INFO,
-      "Exception stack top[cpu%lu]: 0x%lX\n",
+      "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlers: 0x%lX with size 0x%x\n",
       (UINT64)(UINTN)Index,
-      (UINT64)(UINTN)StackTop
+      (UINT64)(UINTN)Buffer,
+      (UINT32)BufferSize
       ));
 
     if (Index == Bsp) {
-      InitializeExceptionStackSwitchHandlers (&EssData);
+      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
     } else {
       MpInitLibStartupThisAP (
         InitializeExceptionStackSwitchHandlers,
         Index,
         NULL,
         0,
-        (VOID *)&EssData,
+        (VOID *)&SwitchStackData,
         NULL
         );
     }
-
-    StackTop -= NewStackSize;
   }
 }
 
diff --git a/UefiCpuPkg/CpuDxe/CpuMp.h b/UefiCpuPkg/CpuDxe/CpuMp.h
index b461753510..c545a711b8 100644
--- a/UefiCpuPkg/CpuDxe/CpuMp.h
+++ b/UefiCpuPkg/CpuDxe/CpuMp.h
@@ -1,7 +1,7 @@
 /** @file
   CPU DXE MP support
 
-  Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2006 - 2022, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -9,6 +9,14 @@
 #ifndef _CPU_MP_H_
 #define _CPU_MP_H_
 
+//
+// Structure for InitializeSeparateExceptionStacks
+//
+typedef struct {
+  VOID     *Buffer;
+  UINTN    *BufferSize;
+} SWITCH_STACK_DATA;
+
 /**
   Initialize Multi-processor support.
 
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.c b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
index d4786979fa..dfc0128361 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.c
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
@@ -1,7 +1,7 @@
 /** @file
   CPU PEI Module installs CPU Multiple Processor PPI.
 
-  Copyright (c) 2015 - 2021, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2015 - 2022, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -411,24 +411,6 @@ PeiWhoAmI (
   return MpInitLibWhoAmI (ProcessorNumber);
 }
 
-/**
-  Get GDT register value.
-
-  This function is mainly for AP purpose because AP may have different GDT
-  table than BSP.
-
-  @param[in,out] Buffer  The pointer to private data buffer.
-
-**/
-VOID
-EFIAPI
-GetGdtr (
-  IN OUT VOID  *Buffer
-  )
-{
-  AsmReadGdtr ((IA32_DESCRIPTOR *)Buffer);
-}
-
 /**
   Initializes CPU exceptions handlers for the sake of stack switch requirement.
 
@@ -444,27 +426,17 @@ InitializeExceptionStackSwitchHandlers (
   IN OUT VOID  *Buffer
   )
 {
-  CPU_EXCEPTION_INIT_DATA  *EssData;
-  IA32_DESCRIPTOR          Idtr;
-  EFI_STATUS               Status;
+  SWITCH_STACK_DATA  *SwitchStackData;
 
-  EssData = Buffer;
-  //
-  // We don't plan to replace IDT table with a new one, but we should not assume
-  // the AP's IDT is the same as BSP's IDT either.
-  //
-  AsmReadIdtr (&Idtr);
-  EssData->Ia32.IdtTable     = (VOID *)Idtr.Base;
-  EssData->Ia32.IdtTableSize = Idtr.Limit + 1;
-  Status                     = InitializeSeparateExceptionStacks (EssData);
-  ASSERT_EFI_ERROR (Status);
+  SwitchStackData = (SWITCH_STACK_DATA *)Buffer;
+  InitializeSeparateExceptionStacks (SwitchStackData->Buffer, SwitchStackData->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 CPU_EXCEPTION_INIT_DATA to each logic processor.
+  and pass them through SwitchStackData to each logic processor.
 
 **/
 VOID
@@ -472,18 +444,14 @@ InitializeMpExceptionStackSwitchHandlers (
   VOID
   )
 {
-  EFI_STATUS               Status;
-  UINTN                    Index;
-  UINTN                    Bsp;
-  UINTN                    ExceptionNumber;
-  UINTN                    OldGdtSize;
-  UINTN                    NewGdtSize;
-  UINTN                    NewStackSize;
-  IA32_DESCRIPTOR          Gdtr;
-  CPU_EXCEPTION_INIT_DATA  EssData;
-  UINT8                    *GdtBuffer;
-  UINT8                    *StackTop;
-  UINTN                    NumberOfProcessors;
+  UINTN              Index;
+  UINTN              Bsp;
+  UINT8              *Buffer;
+  SWITCH_STACK_DATA  SwitchStackData;
+  UINTN              BufferSize;
+
+  SwitchStackData.BufferSize = &BufferSize;
+  UINTN  NumberOfProcessors;
 
   if (!PcdGetBool (PcdCpuStackGuard)) {
     return;
@@ -492,128 +460,48 @@ InitializeMpExceptionStackSwitchHandlers (
   MpInitLibGetNumberOfProcessors (&NumberOfProcessors, NULL);
   MpInitLibWhoAmI (&Bsp);
 
-  ExceptionNumber = FixedPcdGetSize (PcdCpuStackSwitchExceptionList);
-  NewStackSize    = FixedPcdGet32 (PcdCpuKnownGoodStackSize) * ExceptionNumber;
-
-  StackTop = AllocatePages (EFI_SIZE_TO_PAGES (NewStackSize * NumberOfProcessors));
-  ASSERT (StackTop != NULL);
-  if (StackTop == NULL) {
-    return;
-  }
-
-  StackTop += NewStackSize  * NumberOfProcessors;
-
-  //
-  // The default exception handlers must have been initialized. Let's just skip
-  // it in this method.
-  //
-  EssData.Ia32.Revision            = CPU_EXCEPTION_INIT_DATA_REV;
-  EssData.Ia32.InitDefaultHandlers = FALSE;
-
-  EssData.Ia32.StackSwitchExceptions      = FixedPcdGetPtr (PcdCpuStackSwitchExceptionList);
-  EssData.Ia32.StackSwitchExceptionNumber = ExceptionNumber;
-  EssData.Ia32.KnownGoodStackSize         = FixedPcdGet32 (PcdCpuKnownGoodStackSize);
-
-  //
-  // Initialize Gdtr to suppress incorrect compiler/analyzer warnings.
-  //
-  Gdtr.Base  = 0;
-  Gdtr.Limit = 0;
   for (Index = 0; Index < NumberOfProcessors; ++Index) {
-    //
-    // To support stack switch, we need to re-construct GDT but not IDT.
-    //
+    SwitchStackData.Buffer = NULL;
+    BufferSize             = 0;
+
     if (Index == Bsp) {
-      GetGdtr (&Gdtr);
+      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
     } else {
       //
-      // AP might have different size of GDT from BSP.
+      // AP might need different buffer size from BSP.
       //
-      MpInitLibStartupThisAP (GetGdtr, Index, NULL, 0, (VOID *)&Gdtr, NULL);
+      MpInitLibStartupThisAP (InitializeExceptionStackSwitchHandlers, Index, NULL, 0, (VOID *)&SwitchStackData, NULL);
     }
 
-    //
-    // X64 needs only one TSS of current task working for all exceptions
-    // because of its IST feature. IA32 needs one TSS for each exception
-    // in addition to current task. Since AP is not supposed to allocate
-    // memory, we have to do it in BSP. To simplify the code, we allocate
-    // memory for IA32 case to cover both IA32 and X64 exception stack
-    // switch.
-    //
-    // Layout of memory to allocate for each processor:
-    //    --------------------------------
-    //    |            Alignment         |  (just in case)
-    //    --------------------------------
-    //    |                              |
-    //    |        Original GDT          |
-    //    |                              |
-    //    --------------------------------
-    //    |    Current task descriptor   |
-    //    --------------------------------
-    //    |                              |
-    //    |  Exception task descriptors  |  X ExceptionNumber
-    //    |                              |
-    //    --------------------------------
-    //    |  Current task-state segment  |
-    //    --------------------------------
-    //    |                              |
-    //    | Exception task-state segment |  X ExceptionNumber
-    //    |                              |
-    //    --------------------------------
-    //
-    OldGdtSize                        = Gdtr.Limit + 1;
-    EssData.Ia32.ExceptionTssDescSize = sizeof (IA32_TSS_DESCRIPTOR) *
-                                        (ExceptionNumber + 1);
-    EssData.Ia32.ExceptionTssSize = sizeof (IA32_TASK_STATE_SEGMENT) *
-                                    (ExceptionNumber + 1);
-    NewGdtSize = sizeof (IA32_TSS_DESCRIPTOR) +
-                 OldGdtSize +
-                 EssData.Ia32.ExceptionTssDescSize +
-                 EssData.Ia32.ExceptionTssSize;
-
-    Status = PeiServicesAllocatePool (
-               NewGdtSize,
-               (VOID **)&GdtBuffer
-               );
-    ASSERT (GdtBuffer != NULL);
-    if (EFI_ERROR (Status)) {
-      ASSERT_EFI_ERROR (Status);
-      return;
+    if (BufferSize == 0 ) {
+      continue;
     }
 
-    //
-    // Make sure GDT table alignment
-    //
-    EssData.Ia32.GdtTable     = ALIGN_POINTER (GdtBuffer, sizeof (IA32_TSS_DESCRIPTOR));
-    NewGdtSize               -= ((UINT8 *)EssData.Ia32.GdtTable - GdtBuffer);
-    EssData.Ia32.GdtTableSize = NewGdtSize;
-
-    EssData.Ia32.ExceptionTssDesc = ((UINT8 *)EssData.Ia32.GdtTable + OldGdtSize);
-    EssData.Ia32.ExceptionTss     = ((UINT8 *)EssData.Ia32.GdtTable + OldGdtSize +
-                                     EssData.Ia32.ExceptionTssDescSize);
-
-    EssData.Ia32.KnownGoodStackTop = (UINTN)StackTop;
+    ASSERT (BufferSize != 0);
+    Buffer = AllocatePages (EFI_SIZE_TO_PAGES (BufferSize));
+    ASSERT (Buffer != NULL);
+    ZeroMem (Buffer, EFI_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES (BufferSize)));
+    SwitchStackData.Buffer = Buffer;
     DEBUG ((
       DEBUG_INFO,
-      "Exception stack top[cpu%lu]: 0x%lX\n",
+      "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlers: 0x%lX with size 0x%x\n",
       (UINT64)(UINTN)Index,
-      (UINT64)(UINTN)StackTop
+      (UINT64)(UINTN)Buffer,
+      (UINT32)BufferSize
       ));
 
     if (Index == Bsp) {
-      InitializeExceptionStackSwitchHandlers (&EssData);
+      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
     } else {
       MpInitLibStartupThisAP (
         InitializeExceptionStackSwitchHandlers,
         Index,
         NULL,
         0,
-        (VOID *)&EssData,
+        (VOID *)&SwitchStackData,
         NULL
         );
     }
-
-    StackTop -= NewStackSize;
   }
 }
 
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.h b/UefiCpuPkg/CpuMpPei/CpuMpPei.h
index 0649c48d14..f4fe7a3d39 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.h
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.h
@@ -1,7 +1,7 @@
 /** @file
   Definitions to install Multiple Processor PPI.
 
-  Copyright (c) 2015 - 2021, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2015 - 2022, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -31,6 +31,14 @@
 
 extern EFI_PEI_PPI_DESCRIPTOR  mPeiCpuMpPpiDesc;
 
+//
+// Structure for InitializeSeparateExceptionStacks
+//
+typedef struct {
+  VOID     *Buffer;
+  UINTN    *BufferSize;
+} SWITCH_STACK_DATA;
+
 /**
   This service retrieves the number of logical processor in the platform
   and the number of those logical processors that are enabled on this boot.
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
index e62bb5e6c0..4b75c42f1e 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
@@ -104,48 +104,71 @@ RegisterCpuInterruptHandler (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   )
 {
   CPU_EXCEPTION_INIT_DATA  EssData;
   IA32_DESCRIPTOR          Idtr;
   IA32_DESCRIPTOR          Gdtr;
+  UINTN                    NeedBufferSize;
+  UINTN                    StackTop;
+  UINT8                    *NewGdtTable;
 
-  if (InitData == NULL) {
+  AsmReadGdtr (&Gdtr);
+  if ((Buffer == NULL) && (BufferSize == NULL)) {
     SetMem (mNewGdt, sizeof (mNewGdt), 0);
-
-    AsmReadIdtr (&Idtr);
-    AsmReadGdtr (&Gdtr);
-
-    EssData.X64.Revision                   = CPU_EXCEPTION_INIT_DATA_REV;
-    EssData.X64.KnownGoodStackTop          = (UINTN)mNewStack + sizeof (mNewStack);
-    EssData.X64.KnownGoodStackSize         = CPU_KNOWN_GOOD_STACK_SIZE;
-    EssData.X64.StackSwitchExceptions      = CPU_STACK_SWITCH_EXCEPTION_LIST;
-    EssData.X64.StackSwitchExceptionNumber = CPU_STACK_SWITCH_EXCEPTION_NUMBER;
-    EssData.X64.IdtTable                   = (VOID *)Idtr.Base;
-    EssData.X64.IdtTableSize               = Idtr.Limit + 1;
-    EssData.X64.GdtTable                   = mNewGdt;
-    EssData.X64.GdtTableSize               = sizeof (mNewGdt);
-    EssData.X64.ExceptionTssDesc           = mNewGdt + Gdtr.Limit + 1;
-    EssData.X64.ExceptionTssDescSize       = CPU_TSS_DESC_SIZE;
-    EssData.X64.ExceptionTss               = mNewGdt + Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE;
-    EssData.X64.ExceptionTssSize           = CPU_TSS_SIZE;
-
-    InitData = &EssData;
+    StackTop    = (UINTN)mNewStack + sizeof (mNewStack);
+    NewGdtTable = mNewGdt;
+  } else {
+    if (BufferSize == NULL) {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    NeedBufferSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE + // Stack size
+                     CPU_TSS_DESC_SIZE + CPU_TSS_SIZE + Gdtr.Limit + 1 +             // GDT size
+                     sizeof (IA32_TSS_DESCRIPTOR);                                   // Make sure GDT table alignment
+    if (*BufferSize < NeedBufferSize) {
+      *BufferSize = NeedBufferSize;
+      return EFI_BUFFER_TOO_SMALL;
+    }
+
+    if (Buffer == NULL) {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    StackTop    = (UINTN)Buffer + CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE;
+    NewGdtTable = ALIGN_POINTER (StackTop, sizeof (IA32_TSS_DESCRIPTOR));
   }
 
-  return ArchSetupExceptionStack (InitData);
+  AsmReadIdtr (&Idtr);
+  EssData.X64.KnownGoodStackTop          = StackTop;
+  EssData.X64.KnownGoodStackSize         = CPU_KNOWN_GOOD_STACK_SIZE;
+  EssData.X64.StackSwitchExceptions      = CPU_STACK_SWITCH_EXCEPTION_LIST;
+  EssData.X64.StackSwitchExceptionNumber = CPU_STACK_SWITCH_EXCEPTION_NUMBER;
+  EssData.X64.IdtTable                   = (VOID *)Idtr.Base; // Won't change IDT table in this function
+  EssData.X64.IdtTableSize               = Idtr.Limit + 1;
+  EssData.X64.GdtTable                   = NewGdtTable;
+  EssData.X64.GdtTableSize               = CPU_TSS_DESC_SIZE + CPU_TSS_SIZE + Gdtr.Limit + 1;
+  EssData.X64.ExceptionTssDesc           = NewGdtTable + Gdtr.Limit + 1;
+  EssData.X64.ExceptionTssDescSize       = CPU_TSS_DESC_SIZE;
+  EssData.X64.ExceptionTss               = NewGdtTable + Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE;
+  EssData.X64.ExceptionTssSize           = CPU_TSS_SIZE;
+
+  return ArchSetupExceptionStack (&EssData);
 }
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
index f13e8e7020..fa62074023 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
@@ -1,7 +1,7 @@
 /** @file
   IA32 CPU Exception Handler functons.
 
-  Copyright (c) 2012 - 2019, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2012 - 2022, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -132,7 +132,6 @@ ArchSetupExceptionStack (
   EXCEPTION_HANDLER_TEMPLATE_MAP  TemplateMap;
 
   if ((StackSwitchData == NULL) ||
-      (StackSwitchData->Ia32.Revision != CPU_EXCEPTION_INIT_DATA_REV) ||
       (StackSwitchData->Ia32.KnownGoodStackTop == 0) ||
       (StackSwitchData->Ia32.KnownGoodStackSize == 0) ||
       (StackSwitchData->Ia32.StackSwitchExceptions == NULL) ||
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
index 494c2ab433..2868560855 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
@@ -151,25 +151,71 @@ InitializeCpuExceptionHandlers (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   )
 {
-  if (InitData == NULL) {
+  CPU_EXCEPTION_INIT_DATA  EssData;
+  IA32_DESCRIPTOR          Idtr;
+  IA32_DESCRIPTOR          Gdtr;
+  UINTN                    NeedBufferSize;
+  UINTN                    StackTop;
+  UINT8                    *NewGdtTable;
+
+  if ((Buffer == NULL) && (BufferSize == NULL)) {
     return EFI_UNSUPPORTED;
   }
 
-  return ArchSetupExceptionStack (InitData);
+  if (BufferSize == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  AsmReadGdtr (&Gdtr);
+
+  NeedBufferSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE +  // Stack size
+                   CPU_TSS_DESC_SIZE + CPU_TSS_SIZE + Gdtr.Limit + 1 +              // GDT size
+                   sizeof (IA32_TSS_DESCRIPTOR);                                    // Make sure GDT table alignment
+
+  if (*BufferSize < NeedBufferSize) {
+    *BufferSize = NeedBufferSize;
+    return EFI_BUFFER_TOO_SMALL;
+  }
+
+  if (Buffer == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  StackTop    = (UINTN)Buffer + CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE;
+  NewGdtTable = ALIGN_POINTER (StackTop, sizeof (IA32_TSS_DESCRIPTOR));
+
+  AsmReadIdtr (&Idtr);
+  EssData.X64.KnownGoodStackTop          = StackTop;
+  EssData.X64.KnownGoodStackSize         = CPU_KNOWN_GOOD_STACK_SIZE;
+  EssData.X64.StackSwitchExceptions      = CPU_STACK_SWITCH_EXCEPTION_LIST;
+  EssData.X64.StackSwitchExceptionNumber = CPU_STACK_SWITCH_EXCEPTION_NUMBER;
+  EssData.X64.IdtTable                   = (VOID *)Idtr.Base; // Won't change IDT table in this function
+  EssData.X64.IdtTableSize               = Idtr.Limit + 1;
+  EssData.X64.GdtTable                   = NewGdtTable;
+  EssData.X64.GdtTableSize               = CPU_TSS_DESC_SIZE + CPU_TSS_SIZE + Gdtr.Limit + 1;
+  EssData.X64.ExceptionTssDesc           = NewGdtTable + Gdtr.Limit + 1;
+  EssData.X64.ExceptionTssDescSize       = CPU_TSS_DESC_SIZE;
+  EssData.X64.ExceptionTss               = NewGdtTable + Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE;
+  EssData.X64.ExceptionTssSize           = CPU_TSS_SIZE;
+
+  return ArchSetupExceptionStack (&EssData);
 }
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
index cf5bfe4083..7c2ec3b2db 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
@@ -1,7 +1,7 @@
 ## @file
 #  CPU Exception Handler library instance for PEI module.
 #
-#  Copyright (c) 2016 - 2018, Intel Corporation. All rights reserved.<BR>
+#  Copyright (c) 2016 - 2022, Intel Corporation. All rights reserved.<BR>
 #  SPDX-License-Identifier: BSD-2-Clause-Patent
 #
 ##
@@ -56,6 +56,8 @@
 
 [Pcd]
   gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard    # CONSUMES
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuKnownGoodStackSize
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuStackSwitchExceptionList
 
 [FeaturePcd]
   gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmStackGuard                    ## CONSUMES
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c
index 4313cc5582..ad5e0e9ed4 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c
@@ -201,20 +201,23 @@ RegisterCpuInterruptHandler (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   )
 {
   return EFI_UNSUPPORTED;
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c
index 1c97dab926..46a86ad2c6 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c
@@ -97,20 +97,23 @@ RegisterCpuInterruptHandler (
 
 /**
   Setup separate stacks for certain exception handlers.
+  If the input Buffer and BufferSize are both NULL, use global variable if possible.
 
-  InitData is optional and processor arch dependent.
-
-  @param[in]  InitData      Pointer to data optional for information about how
-                            to assign stacks for certain exception handlers.
+  @param[in]       Buffer        Point to buffer used to separate exception stack.
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
+                                 If the size is not enough, the return status will
+                                 be EFI_BUFFER_TOO_SMALL, and output BufferSize
+                                 will be the size it needs.
 
   @retval EFI_SUCCESS             The stacks are assigned successfully.
   @retval EFI_UNSUPPORTED         This function is not supported.
-
+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.
 **/
 EFI_STATUS
 EFIAPI
 InitializeSeparateExceptionStacks (
-  IN CPU_EXCEPTION_INIT_DATA  *InitData OPTIONAL
+  IN     VOID   *Buffer,
+  IN OUT UINTN  *BufferSize
   )
 {
   return EFI_UNSUPPORTED;
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
index cd7dccd481..ff0dde4f12 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
@@ -1,7 +1,7 @@
 /** @file
   x64 CPU Exception Handler.
 
-  Copyright (c) 2012 - 2019, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2012 - 2022, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -136,7 +136,6 @@ ArchSetupExceptionStack (
   UINTN                     GdtSize;
 
   if ((StackSwitchData == NULL) ||
-      (StackSwitchData->Ia32.Revision != CPU_EXCEPTION_INIT_DATA_REV) ||
       (StackSwitchData->X64.KnownGoodStackTop == 0) ||
       (StackSwitchData->X64.KnownGoodStackSize == 0) ||
       (StackSwitchData->X64.StackSwitchExceptions == NULL) ||
-- 
2.26.2.windows.1