UefiCpuPkg: Simplify InitializeSeparateExceptionStacks

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

Cc: Eric Dong <eric.dong@intel.com>
Reviewed-by: 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>
Reviewed-by: Sami Mujawar <sami.mujawar@arm.com>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>

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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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;

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);
   }
 

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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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
   );
 
 /**

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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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;

UefiCpuPkg/CpuDxe/CpuMp.c

@@ -596,23 +596,13 @@ CollectBistDataFromHob (
   }
 }
 
-/**
+//
-  Get GDT register value.
+// Structure for InitializeSeparateExceptionStacks
-
+//
-  This function is mainly for AP purpose because AP may have different GDT
+typedef struct {
-  table than BSP.
+  VOID     *Buffer;
-
+  UINTN    *BufferSize;
-  @param[in,out] Buffer  The pointer to private data buffer.
+} EXCEPTION_STACK_SWITCH_CONTEXT;
-
-**/
-VOID
-EFIAPI
-GetGdtr (
-  IN OUT VOID  *Buffer
-  )
-{
-  AsmReadGdtr ((IA32_DESCRIPTOR *)Buffer);
-}
 
 /**
   Initializes CPU exceptions handlers for the sake of stack switch requirement.
@@ -629,27 +619,17 @@ InitializeExceptionStackSwitchHandlers (
   IN OUT VOID  *Buffer
   )
 {
-  CPU_EXCEPTION_INIT_DATA  *EssData;
+  EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
-  IA32_DESCRIPTOR          Idtr;
-  EFI_STATUS               Status;
 
-  EssData = Buffer;
+  SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
-  //
+  InitializeSeparateExceptionStacks (SwitchStackData->Buffer, SwitchStackData->BufferSize);
-  // 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);
 }
 
 /**
   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 +637,53 @@ InitializeMpExceptionStackSwitchHandlers (
   VOID
   )
 {
-  UINTN                    Index;
+  UINTN                           Index;
-  UINTN                    Bsp;
+  UINTN                           Bsp;
-  UINTN                    ExceptionNumber;
+  EXCEPTION_STACK_SWITCH_CONTEXT  SwitchStackData;
-  UINTN                    OldGdtSize;
+  UINTN                           BufferSize;
-  UINTN                    NewGdtSize;
-  UINTN                    NewStackSize;
-  IA32_DESCRIPTOR          Gdtr;
-  CPU_EXCEPTION_INIT_DATA  EssData;
-  UINT8                    *GdtBuffer;
-  UINT8                    *StackTop;
 
-  ExceptionNumber = FixedPcdGetSize (PcdCpuStackSwitchExceptionList);
+  SwitchStackData.BufferSize = &BufferSize;
-  NewStackSize    = FixedPcdGet32 (PcdCpuKnownGoodStackSize) * ExceptionNumber;
-
-  StackTop = AllocateRuntimeZeroPool (NewStackSize * mNumberOfProcessors);
-  ASSERT (StackTop != NULL);
-  StackTop += NewStackSize  * mNumberOfProcessors;
-
-  //
-  // 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;
   MpInitLibWhoAmI (&Bsp);
+
   for (Index = 0; Index < mNumberOfProcessors; ++Index) {
-    //
+    SwitchStackData.Buffer = NULL;
-    // To support stack switch, we need to re-construct GDT but not IDT.
+    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);
     }
 
-    //
+    if (BufferSize == 0) {
-    // X64 needs only one TSS of current task working for all exceptions
+      continue;
-    // 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);
+    SwitchStackData.Buffer = AllocateRuntimeZeroPool (BufferSize);
-    ASSERT (GdtBuffer != NULL);
+    ASSERT (SwitchStackData.Buffer != 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;
     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)SwitchStackData.Buffer,
+      (UINT32)BufferSize
       ));
 
     if (Index == Bsp) {
-      InitializeExceptionStackSwitchHandlers (&EssData);
+      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
     } else {
       MpInitLibStartupThisAP (
         InitializeExceptionStackSwitchHandlers,
         Index,
         NULL,
         0,
-        (VOID *)&EssData,
+        (VOID *)&SwitchStackData,
         NULL
         );
     }
-
-    StackTop -= NewStackSize;
   }
 }
 

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,23 +411,13 @@ PeiWhoAmI (
   return MpInitLibWhoAmI (ProcessorNumber);
 }
 
-/**
+//
-  Get GDT register value.
+// Structure for InitializeSeparateExceptionStacks
-
+//
-  This function is mainly for AP purpose because AP may have different GDT
+typedef struct {
-  table than BSP.
+  VOID     *Buffer;
-
+  UINTN    *BufferSize;
-  @param[in,out] Buffer  The pointer to private data buffer.
+} EXCEPTION_STACK_SWITCH_CONTEXT;
-
-**/
-VOID
-EFIAPI
-GetGdtr (
-  IN OUT VOID  *Buffer
-  )
-{
-  AsmReadGdtr ((IA32_DESCRIPTOR *)Buffer);
-}
 
 /**
   Initializes CPU exceptions handlers for the sake of stack switch requirement.
@@ -444,27 +434,17 @@ InitializeExceptionStackSwitchHandlers (
   IN OUT VOID  *Buffer
   )
 {
-  CPU_EXCEPTION_INIT_DATA  *EssData;
+  EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
-  IA32_DESCRIPTOR          Idtr;
-  EFI_STATUS               Status;
 
-  EssData = Buffer;
+  SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
-  //
+  InitializeSeparateExceptionStacks (SwitchStackData->Buffer, SwitchStackData->BufferSize);
-  // 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);
 }
 
 /**
   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,148 +452,60 @@ InitializeMpExceptionStackSwitchHandlers (
   VOID
   )
 {
-  EFI_STATUS               Status;
+  UINTN                           Index;
-  UINTN                    Index;
+  UINTN                           Bsp;
-  UINTN                    Bsp;
+  EXCEPTION_STACK_SWITCH_CONTEXT  SwitchStackData;
-  UINTN                    ExceptionNumber;
+  UINTN                           BufferSize;
-  UINTN                    OldGdtSize;
+  UINTN                           NumberOfProcessors;
-  UINTN                    NewGdtSize;
-  UINTN                    NewStackSize;
-  IA32_DESCRIPTOR          Gdtr;
-  CPU_EXCEPTION_INIT_DATA  EssData;
-  UINT8                    *GdtBuffer;
-  UINT8                    *StackTop;
-  UINTN                    NumberOfProcessors;
 
   if (!PcdGetBool (PcdCpuStackGuard)) {
     return;
   }
 
+  SwitchStackData.BufferSize = &BufferSize;
   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) {
-    //
+    SwitchStackData.Buffer = NULL;
-    // To support stack switch, we need to re-construct GDT but not IDT.
+    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);
     }
 
-    //
+    if (BufferSize == 0) {
-    // X64 needs only one TSS of current task working for all exceptions
+      continue;
-    // 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;
     }
 
-    //
+    SwitchStackData.Buffer = AllocatePages (EFI_SIZE_TO_PAGES (BufferSize));
-    // Make sure GDT table alignment
+    ASSERT (SwitchStackData.Buffer != NULL);
-    //
+    ZeroMem (SwitchStackData.Buffer, EFI_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES (BufferSize)));
-    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;
     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)SwitchStackData.Buffer,
+      (UINT32)BufferSize
       ));
 
     if (Index == Bsp) {
-      InitializeExceptionStackSwitchHandlers (&EssData);
+      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
     } else {
       MpInitLibStartupThisAP (
         InitializeExceptionStackSwitchHandlers,
         Index,
         NULL,
         0,
-        (VOID *)&EssData,
+        (VOID *)&SwitchStackData,
         NULL
         );
     }
-
-    StackTop -= NewStackSize;
   }
 }
 

UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c

@@ -104,48 +104,105 @@ 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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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) {
+  //
+  // 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. To simplify the code, we report the
+  // needed memory for IA32 case to cover both IA32 and X64 exception
+  // stack switch.
+  //
+  // Layout of memory needed 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
+  //    |                              |
+  //    --------------------------------
+  //
+  AsmReadGdtr (&Gdtr);
+  if ((Buffer == NULL) && (BufferSize == NULL)) {
     SetMem (mNewGdt, sizeof (mNewGdt), 0);
+    StackTop    = (UINTN)mNewStack + sizeof (mNewStack);
+    NewGdtTable = mNewGdt;
+  } else {
+    if (BufferSize == NULL) {
+      return EFI_INVALID_PARAMETER;
+    }
 
-    AsmReadIdtr (&Idtr);
+    //
-    AsmReadGdtr (&Gdtr);
+    // Total needed size includes stack size, new GDT table size, TSS size.
+    // Add another DESCRIPTOR size for alignment requiremet.
+    //
+    NeedBufferSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE +
+                     CPU_TSS_DESC_SIZE + Gdtr.Limit + 1 +
+                     CPU_TSS_SIZE +
+                     sizeof (IA32_TSS_DESCRIPTOR);
+    if (*BufferSize < NeedBufferSize) {
+      *BufferSize = NeedBufferSize;
+      return EFI_BUFFER_TOO_SMALL;
+    }
 
-    EssData.X64.Revision                   = CPU_EXCEPTION_INIT_DATA_REV;
+    if (Buffer == NULL) {
-    EssData.X64.KnownGoodStackTop          = (UINTN)mNewStack + sizeof (mNewStack);
+      return EFI_INVALID_PARAMETER;
-    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)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.Revision                   = CPU_EXCEPTION_INIT_DATA_REV;
+  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;
+  EssData.X64.IdtTableSize               = Idtr.Limit + 1;
+  EssData.X64.GdtTable                   = NewGdtTable;
+  EssData.X64.GdtTableSize               = CPU_TSS_DESC_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);
 }

UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c

@@ -151,25 +151,104 @@ 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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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;
+
+  //
+  // 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. To simplify the code, we report the
+  // needed memory for IA32 case to cover both IA32 and X64 exception
+  // stack switch.
+  //
+  // Layout of memory needed 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
+  //    |                              |
+  //    --------------------------------
+  //
+
+  if ((Buffer == NULL) && (BufferSize == NULL)) {
     return EFI_UNSUPPORTED;
   }
 
-  return ArchSetupExceptionStack (InitData);
+  if (BufferSize == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  AsmReadGdtr (&Gdtr);
+  //
+  // Total needed size includes stack size, new GDT table size, TSS size.
+  // Add another DESCRIPTOR size for alignment requiremet.
+  //
+  NeedBufferSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE +
+                   CPU_TSS_DESC_SIZE + Gdtr.Limit + 1 +
+                   CPU_TSS_SIZE +
+                   sizeof (IA32_TSS_DESCRIPTOR);
+  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.Revision                   = CPU_EXCEPTION_INIT_DATA_REV;
+  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;
+  EssData.X64.IdtTableSize               = Idtr.Limit + 1;
+  EssData.X64.GdtTable                   = NewGdtTable;
+  EssData.X64.GdtTableSize               = CPU_TSS_DESC_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);
 }

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

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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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;

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]       Buffer        Point to buffer used to separate exception stack.
-
+  @param[in, out]  BufferSize    On input, it indicates the byte size of Buffer.
-  @param[in]  InitData      Pointer to data optional for information about how
+                                 If the size is not enough, the return status will
-                            to assign stacks for certain exception handlers.
+                                 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;