UefiCpuPkg: Simplify the implementation when separate exception stacks

The API of InitializeSeparateExceptionStacks is just changed before, and
makes the struct CPU_EXCEPTION_INIT_DATA an internal definition.
Furthermore, we can even remove the struct to make core simpler.

Cc: Eric Dong <eric.dong@intel.com>
Reviewed-by: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>

8 files changed, 173 insertions, 383 deletions

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h b/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h
index 11a5624f51..4593c204a6 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h
@@ -49,61 +49,6 @@
 

 #define CPU_TSS_GDT_SIZE  (SIZE_2KB + CPU_TSS_DESC_SIZE + CPU_TSS_SIZE)

 

-typedef struct {

-  //

-  // The address of top of known good stack reserved for *ALL* exceptions

-  // listed in field StackSwitchExceptions.

-  //

-  UINTN    KnownGoodStackTop;

-  //

-  // The size of known good stack for *ONE* exception only.

-  //

-  UINTN    KnownGoodStackSize;

-  //

-  // Buffer of exception vector list for stack switch.

-  //

-  UINT8    *StackSwitchExceptions;

-  //

-  // Number of exception vectors in StackSwitchExceptions.

-  //

-  UINTN    StackSwitchExceptionNumber;

-  //

-  // Buffer of IDT table. It must be type of IA32_IDT_GATE_DESCRIPTOR.

-  // Normally there's no need to change IDT table size.

-  //

-  VOID     *IdtTable;

-  //

-  // Size of buffer for IdtTable.

-  //

-  UINTN    IdtTableSize;

-  //

-  // Buffer of GDT table. It must be type of IA32_SEGMENT_DESCRIPTOR.

-  //

-  VOID     *GdtTable;

-  //

-  // Size of buffer for GdtTable.

-  //

-  UINTN    GdtTableSize;

-  //

-  // Pointer to start address of descriptor of exception task gate in the

-  // GDT table. It must be type of IA32_TSS_DESCRIPTOR.

-  //

-  VOID     *ExceptionTssDesc;

-  //

-  // Size of buffer for ExceptionTssDesc.

-  //

-  UINTN    ExceptionTssDescSize;

-  //

-  // Buffer of task-state segment for exceptions. It must be type of

-  // IA32_TASK_STATE_SEGMENT.

-  //

-  VOID     *ExceptionTss;

-  //

-  // Size of buffer for ExceptionTss.

-  //

-  UINTN    ExceptionTssSize;

-} CPU_EXCEPTION_INIT_DATA;

-

 //

 // Record exception handler information

 //

@@ -345,18 +290,22 @@ CommonExceptionHandlerWorker (
   );

 

 /**

-  Setup separate stack for specific exceptions.

+  Setup separate stacks for certain exception handlers.

 

-  @param[in] StackSwitchData      Pointer to data required for setuping up

-                                  stack switch.

+  @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 exceptions have been successfully

-                                  initialized with new stack.

-  @retval EFI_INVALID_PARAMETER   StackSwitchData contains invalid content.

+  @retval EFI_SUCCESS             The stacks are assigned successfully.

+  @retval EFI_BUFFER_TOO_SMALL    This BufferSize is too small.

+  @retval EFI_UNSUPPORTED         This function is not supported.

 **/

 EFI_STATUS

 ArchSetupExceptionStack (

-  IN CPU_EXCEPTION_INIT_DATA  *StackSwitchData

+  IN     VOID   *Buffer,

+  IN OUT UINTN  *BufferSize

   );

 

 /**

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
index d90c607bd7..ee989bf079 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
@@ -23,9 +23,8 @@ EXCEPTION_HANDLER_DATA     mExceptionHandlerData = {
   mExternalInterruptHandlerTable

 };

 

-UINT8  mNewStack[CPU_STACK_SWITCH_EXCEPTION_NUMBER *

-                 CPU_KNOWN_GOOD_STACK_SIZE];

-UINT8  mNewGdt[CPU_TSS_GDT_SIZE];

+UINT8  mBuffer[CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE

+               + CPU_TSS_GDT_SIZE];

 

 /**

   Common exception handler.

@@ -123,85 +122,16 @@ InitializeSeparateExceptionStacks (
   IN OUT UINTN  *BufferSize

   )

 {

-  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

-  //    |                              |

-  //    --------------------------------

-  //

-  AsmReadGdtr (&Gdtr);

+  UINTN       LocalBufferSize;

+  EFI_STATUS  Status;

+

   if ((Buffer == NULL) && (BufferSize == NULL)) {

-    SetMem (mNewGdt, sizeof (mNewGdt), 0);

-    StackTop    = (UINTN)mNewStack + sizeof (mNewStack);

-    NewGdtTable = mNewGdt;

+    SetMem (mBuffer, sizeof (mBuffer), 0);

+    LocalBufferSize = sizeof (mBuffer);

+    Status          = ArchSetupExceptionStack (mBuffer, &LocalBufferSize);

+    ASSERT_EFI_ERROR (Status);

+    return Status;

   } else {

-    if (BufferSize == NULL) {

-      return EFI_INVALID_PARAMETER;

-    }

-

-    //

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

+    return ArchSetupExceptionStack (Buffer, BufferSize);

   }

-

-  AsmReadIdtr (&Idtr);

-  EssData.KnownGoodStackTop          = StackTop;

-  EssData.KnownGoodStackSize         = CPU_KNOWN_GOOD_STACK_SIZE;

-  EssData.StackSwitchExceptions      = CPU_STACK_SWITCH_EXCEPTION_LIST;

-  EssData.StackSwitchExceptionNumber = CPU_STACK_SWITCH_EXCEPTION_NUMBER;

-  EssData.IdtTable                   = (VOID *)Idtr.Base;

-  EssData.IdtTableSize               = Idtr.Limit + 1;

-  EssData.GdtTable                   = NewGdtTable;

-  EssData.GdtTableSize               = CPU_TSS_DESC_SIZE + Gdtr.Limit + 1;

-  EssData.ExceptionTssDesc           = NewGdtTable + Gdtr.Limit + 1;

-  EssData.ExceptionTssDescSize       = CPU_TSS_DESC_SIZE;

-  EssData.ExceptionTss               = NewGdtTable + Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE;

-  EssData.ExceptionTssSize           = CPU_TSS_SIZE;

-

-  return ArchSetupExceptionStack (&EssData);

 }

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
index 194d3a499b..8c398ebc5b 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
@@ -104,108 +104,97 @@ ArchRestoreExceptionContext (
 }

 

 /**

-  Setup separate stack for given exceptions.

+  Setup separate stacks for certain exception handlers.

 

-  @param[in] StackSwitchData      Pointer to data required for setuping up

-                                  stack switch.

-

-  @retval EFI_SUCCESS             The exceptions have been successfully

-                                  initialized with new stack.

-  @retval EFI_INVALID_PARAMETER   StackSwitchData contains invalid content.

+  @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_BUFFER_TOO_SMALL    This BufferSize is too small.

 **/

 EFI_STATUS

 ArchSetupExceptionStack (

-  IN CPU_EXCEPTION_INIT_DATA  *StackSwitchData

+  IN     VOID   *Buffer,

+  IN OUT UINTN  *BufferSize

   )

 {

   IA32_DESCRIPTOR                 Gdtr;

   IA32_DESCRIPTOR                 Idtr;

   IA32_IDT_GATE_DESCRIPTOR        *IdtTable;

   IA32_TSS_DESCRIPTOR             *TssDesc;

+  IA32_TSS_DESCRIPTOR             *TssDescBase;

   IA32_TASK_STATE_SEGMENT         *Tss;

+  VOID                            *NewGdtTable;

   UINTN                           StackTop;

   UINTN                           Index;

   UINTN                           Vector;

   UINTN                           TssBase;

-  UINTN                           GdtSize;

+  UINT8                           *StackSwitchExceptions;

+  UINTN                           NeedBufferSize;

   EXCEPTION_HANDLER_TEMPLATE_MAP  TemplateMap;

 

-  if ((StackSwitchData == NULL) ||

-      (StackSwitchData->KnownGoodStackTop == 0) ||

-      (StackSwitchData->KnownGoodStackSize == 0) ||

-      (StackSwitchData->StackSwitchExceptions == NULL) ||

-      (StackSwitchData->StackSwitchExceptionNumber == 0) ||

-      (StackSwitchData->StackSwitchExceptionNumber > CPU_EXCEPTION_NUM) ||

-      (StackSwitchData->GdtTable == NULL) ||

-      (StackSwitchData->IdtTable == NULL) ||

-      (StackSwitchData->ExceptionTssDesc == NULL) ||

-      (StackSwitchData->ExceptionTss == NULL))

-  {

+  if (BufferSize == NULL) {

     return EFI_INVALID_PARAMETER;

   }

 

   //

-  // The caller is responsible for that the GDT table, no matter the existing

-  // one or newly allocated, has enough space to hold descriptors for exception

-  // task-state segments.

+  // Total needed size includes stack size, new GDT table size, TSS size.

+  // Add another DESCRIPTOR size for alignment requiremet.

   //

-  if (((UINTN)StackSwitchData->GdtTable & (IA32_GDT_ALIGNMENT - 1)) != 0) {

-    return EFI_INVALID_PARAMETER;

-  }

-

-  if ((UINTN)StackSwitchData->ExceptionTssDesc < (UINTN)(StackSwitchData->GdtTable)) {

-    return EFI_INVALID_PARAMETER;

-  }

-

-  if ((UINTN)StackSwitchData->ExceptionTssDesc + StackSwitchData->ExceptionTssDescSize >

-      ((UINTN)(StackSwitchData->GdtTable) + StackSwitchData->GdtTableSize))

-  {

-    return EFI_INVALID_PARAMETER;

-  }

-

+  // Layout of memory needed for each processor:

+  //    --------------------------------

+  //    |                              |

+  //    |          Stack Size          |  X ExceptionNumber

+  //    |                              |

+  //    --------------------------------

+  //    |          Alignment           |  (just in case)

+  //    --------------------------------

+  //    |                              |

+  //    |         Original GDT         |

+  //    |                              |

+  //    --------------------------------

+  //    |    Current task descriptor   |

+  //    --------------------------------

+  //    |                              |

+  //    |  Exception task descriptors  |  X ExceptionNumber

+  //    |                              |

+  //    --------------------------------

+  //    |  Current task-state segment  |

+  //    --------------------------------

+  //    |                              |

+  //    | Exception task-state segment |  X ExceptionNumber

+  //    |                              |

+  //    --------------------------------

   //

-  // We need one descriptor and one TSS for current task and every exception

-  // specified.

-  //

-  if (StackSwitchData->ExceptionTssDescSize <

-      sizeof (IA32_TSS_DESCRIPTOR) * (StackSwitchData->StackSwitchExceptionNumber + 1))

-  {

-    return EFI_INVALID_PARAMETER;

+  AsmReadGdtr (&Gdtr);

+  NeedBufferSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE +

+                   sizeof (IA32_TSS_DESCRIPTOR) +

+                   Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE +

+                   CPU_TSS_SIZE;

+

+  if (*BufferSize < NeedBufferSize) {

+    *BufferSize = NeedBufferSize;

+    return EFI_BUFFER_TOO_SMALL;

   }

 

-  if (StackSwitchData->ExceptionTssSize <

-      sizeof (IA32_TASK_STATE_SEGMENT) * (StackSwitchData->StackSwitchExceptionNumber + 1))

-  {

+  if (Buffer == NULL) {

     return EFI_INVALID_PARAMETER;

   }

 

-  TssDesc = StackSwitchData->ExceptionTssDesc;

-  Tss     = StackSwitchData->ExceptionTss;

-

-  //

-  // Initialize new GDT table and/or IDT table, if any

-  //

   AsmReadIdtr (&Idtr);

-  AsmReadGdtr (&Gdtr);

-

-  GdtSize = (UINTN)TssDesc +

-            sizeof (IA32_TSS_DESCRIPTOR) *

-            (StackSwitchData->StackSwitchExceptionNumber + 1) -

-            (UINTN)(StackSwitchData->GdtTable);

-  if ((UINTN)StackSwitchData->GdtTable != Gdtr.Base) {

-    CopyMem (StackSwitchData->GdtTable, (VOID *)Gdtr.Base, Gdtr.Limit + 1);

-    Gdtr.Base  = (UINTN)StackSwitchData->GdtTable;

-    Gdtr.Limit = (UINT16)GdtSize - 1;

-  }

-

-  if ((UINTN)StackSwitchData->IdtTable != Idtr.Base) {

-    Idtr.Base = (UINTN)StackSwitchData->IdtTable;

-  }

+  StackSwitchExceptions = CPU_STACK_SWITCH_EXCEPTION_LIST;

+  StackTop              = (UINTN)Buffer + CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE;

+  NewGdtTable           = ALIGN_POINTER (StackTop, sizeof (IA32_TSS_DESCRIPTOR));

+  TssDesc               = (IA32_TSS_DESCRIPTOR *)((UINTN)NewGdtTable + Gdtr.Limit + 1);

+  Tss                   = (IA32_TASK_STATE_SEGMENT *)((UINTN)TssDesc + CPU_TSS_DESC_SIZE);

+  TssDescBase           = TssDesc;

 

-  if (StackSwitchData->IdtTableSize > 0) {

-    Idtr.Limit = (UINT16)(StackSwitchData->IdtTableSize - 1);

-  }

+  CopyMem (NewGdtTable, (VOID *)Gdtr.Base, Gdtr.Limit + 1);

+  Gdtr.Base  = (UINTN)NewGdtTable;

+  Gdtr.Limit = (UINT16)(Gdtr.Limit + CPU_TSS_DESC_SIZE);

 

   //

   // Fixup current task descriptor. Task-state segment for current task will

@@ -226,10 +215,10 @@ ArchSetupExceptionStack (
   // Fixup exception task descriptor and task-state segment

   //

   AsmGetTssTemplateMap (&TemplateMap);

-  StackTop = StackSwitchData->KnownGoodStackTop - CPU_STACK_ALIGNMENT;

+  StackTop = StackTop - CPU_STACK_ALIGNMENT;

   StackTop = (UINTN)ALIGN_POINTER (StackTop, CPU_STACK_ALIGNMENT);

-  IdtTable = StackSwitchData->IdtTable;

-  for (Index = 0; Index < StackSwitchData->StackSwitchExceptionNumber; ++Index) {

+  IdtTable = (IA32_IDT_GATE_DESCRIPTOR  *)Idtr.Base;

+  for (Index = 0; Index < CPU_STACK_SWITCH_EXCEPTION_NUMBER; ++Index) {

     TssDesc += 1;

     Tss     += 1;

 

@@ -250,7 +239,7 @@ ArchSetupExceptionStack (
     //

     // Fixup TSS

     //

-    Vector = StackSwitchData->StackSwitchExceptions[Index];

+    Vector = StackSwitchExceptions[Index];

     if ((Vector >= CPU_EXCEPTION_NUM) ||

         (Vector >= (Idtr.Limit + 1) / sizeof (IA32_IDT_GATE_DESCRIPTOR)))

     {

@@ -270,7 +259,7 @@ ArchSetupExceptionStack (
     Tss->FS     = AsmReadFs ();

     Tss->GS     = AsmReadGs ();

 

-    StackTop -= StackSwitchData->KnownGoodStackSize;

+    StackTop -= CPU_KNOWN_GOOD_STACK_SIZE;

 

     //

     // Update IDT to use Task Gate for given exception

@@ -290,12 +279,7 @@ ArchSetupExceptionStack (
   //

   // Load current task

   //

-  AsmWriteTr ((UINT16)((UINTN)StackSwitchData->ExceptionTssDesc - Gdtr.Base));

-

-  //

-  // Publish IDT

-  //

-  AsmWriteIdtr (&Idtr);

+  AsmWriteTr ((UINT16)((UINTN)TssDescBase - Gdtr.Base));

 

   return EFI_SUCCESS;

 }

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
index 5952295126..940d83a92f 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
@@ -170,84 +170,9 @@ InitializeSeparateExceptionStacks (
   IN OUT UINTN  *BufferSize

   )

 {

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

   }

 

-  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.KnownGoodStackTop          = StackTop;

-  EssData.KnownGoodStackSize         = CPU_KNOWN_GOOD_STACK_SIZE;

-  EssData.StackSwitchExceptions      = CPU_STACK_SWITCH_EXCEPTION_LIST;

-  EssData.StackSwitchExceptionNumber = CPU_STACK_SWITCH_EXCEPTION_NUMBER;

-  EssData.IdtTable                   = (VOID *)Idtr.Base;

-  EssData.IdtTableSize               = Idtr.Limit + 1;

-  EssData.GdtTable                   = NewGdtTable;

-  EssData.GdtTableSize               = CPU_TSS_DESC_SIZE + Gdtr.Limit + 1;

-  EssData.ExceptionTssDesc           = NewGdtTable + Gdtr.Limit + 1;

-  EssData.ExceptionTssDescSize       = CPU_TSS_DESC_SIZE;

-  EssData.ExceptionTss               = NewGdtTable + Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE;

-  EssData.ExceptionTssSize           = CPU_TSS_SIZE;

-

-  return ArchSetupExceptionStack (&EssData);

+  return ArchSetupExceptionStack (Buffer, BufferSize);

 }

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuExceptionHandlerLib.inf b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuExceptionHandlerLib.inf
index 8ae4feae62..6a170286c8 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuExceptionHandlerLib.inf
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuExceptionHandlerLib.inf
@@ -1,7 +1,7 @@
 ## @file

 #  CPU Exception Handler library instance for SEC/PEI modules.

 #

-#  Copyright (c) 2012 - 2018, Intel Corporation. All rights reserved.<BR>

+#  Copyright (c) 2012 - 2022, Intel Corporation. All rights reserved.<BR>

 #  SPDX-License-Identifier: BSD-2-Clause-Patent

 #

 ##

@@ -50,6 +50,11 @@
   PeCoffGetEntryPointLib

   VmgExitLib

 

+[Pcd]

+  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard

+  gUefiCpuPkgTokenSpaceGuid.PcdCpuStackSwitchExceptionList

+  gUefiCpuPkgTokenSpaceGuid.PcdCpuKnownGoodStackSize

+

 [FeaturePcd]

   gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmStackGuard                    ## CONSUMES

 

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmCpuExceptionHandlerLib.inf b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmCpuExceptionHandlerLib.inf
index c9f20da058..9dde07612a 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmCpuExceptionHandlerLib.inf
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmCpuExceptionHandlerLib.inf
@@ -1,7 +1,7 @@
 ## @file

 #  CPU Exception Handler library instance for SMM modules.

 #

-#  Copyright (c) 2013 - 2018, Intel Corporation. All rights reserved.<BR>

+#  Copyright (c) 2013 - 2022, Intel Corporation. All rights reserved.<BR>

 #  SPDX-License-Identifier: BSD-2-Clause-Patent

 #

 ##

@@ -53,6 +53,11 @@
   DebugLib

   VmgExitLib

 

+[Pcd]

+  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard

+  gUefiCpuPkgTokenSpaceGuid.PcdCpuStackSwitchExceptionList

+  gUefiCpuPkgTokenSpaceGuid.PcdCpuKnownGoodStackSize

+

 [FeaturePcd]

   gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmStackGuard                    ## CONSUMES

 

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
index c14ac66c43..80e9f08e5b 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
@@ -109,19 +109,22 @@ ArchRestoreExceptionContext (
 }

 

 /**

-  Setup separate stack for given exceptions.

+  Setup separate stacks for certain exception handlers.

 

-  @param[in] StackSwitchData      Pointer to data required for setuping up

-                                  stack switch.

-

-  @retval EFI_SUCCESS             The exceptions have been successfully

-                                  initialized with new stack.

-  @retval EFI_INVALID_PARAMETER   StackSwitchData contains invalid content.

+  @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_BUFFER_TOO_SMALL    This BufferSize is too small.

+  @retval EFI_UNSUPPORTED         This function is not supported.

 **/

 EFI_STATUS

 ArchSetupExceptionStack (

-  IN CPU_EXCEPTION_INIT_DATA  *StackSwitchData

+  IN     VOID   *Buffer,

+  IN OUT UINTN  *BufferSize

   )

 {

   IA32_DESCRIPTOR           Gdtr;

@@ -129,86 +132,75 @@ ArchSetupExceptionStack (
   IA32_IDT_GATE_DESCRIPTOR  *IdtTable;

   IA32_TSS_DESCRIPTOR       *TssDesc;

   IA32_TASK_STATE_SEGMENT   *Tss;

+  VOID                      *NewGdtTable;

   UINTN                     StackTop;

   UINTN                     Index;

   UINTN                     Vector;

   UINTN                     TssBase;

-  UINTN                     GdtSize;

-

-  if ((StackSwitchData == NULL) ||

-      (StackSwitchData->KnownGoodStackTop == 0) ||

-      (StackSwitchData->KnownGoodStackSize == 0) ||

-      (StackSwitchData->StackSwitchExceptions == NULL) ||

-      (StackSwitchData->StackSwitchExceptionNumber == 0) ||

-      (StackSwitchData->StackSwitchExceptionNumber > CPU_EXCEPTION_NUM) ||

-      (StackSwitchData->GdtTable == NULL) ||

-      (StackSwitchData->IdtTable == NULL) ||

-      (StackSwitchData->ExceptionTssDesc == NULL) ||

-      (StackSwitchData->ExceptionTss == NULL))

-  {

-    return EFI_INVALID_PARAMETER;

-  }

-

-  //

-  // The caller is responsible for that the GDT table, no matter the existing

-  // one or newly allocated, has enough space to hold descriptors for exception

-  // task-state segments.

-  //

-  if (((UINTN)StackSwitchData->GdtTable & (IA32_GDT_ALIGNMENT - 1)) != 0) {

-    return EFI_INVALID_PARAMETER;

-  }

-

-  if ((UINTN)StackSwitchData->ExceptionTssDesc < (UINTN)(StackSwitchData->GdtTable)) {

-    return EFI_INVALID_PARAMETER;

-  }

-

-  if (((UINTN)StackSwitchData->ExceptionTssDesc + StackSwitchData->ExceptionTssDescSize) >

-      ((UINTN)(StackSwitchData->GdtTable) + StackSwitchData->GdtTableSize))

-  {

-    return EFI_INVALID_PARAMETER;

-  }

+  UINT8                     *StackSwitchExceptions;

+  UINTN                     NeedBufferSize;

 

-  //

-  // One task gate descriptor and one task-state segment are needed.

-  //

-  if (StackSwitchData->ExceptionTssDescSize < sizeof (IA32_TSS_DESCRIPTOR)) {

-    return EFI_INVALID_PARAMETER;

-  }

-

-  if (StackSwitchData->ExceptionTssSize < sizeof (IA32_TASK_STATE_SEGMENT)) {

+  if (BufferSize == NULL) {

     return EFI_INVALID_PARAMETER;

   }

 

   //

   // Interrupt stack table supports only 7 vectors.

   //

-  TssDesc = StackSwitchData->ExceptionTssDesc;

-  Tss     = StackSwitchData->ExceptionTss;

-  if (StackSwitchData->StackSwitchExceptionNumber > ARRAY_SIZE (Tss->IST)) {

-    return EFI_INVALID_PARAMETER;

+  if (CPU_STACK_SWITCH_EXCEPTION_NUMBER > ARRAY_SIZE (Tss->IST)) {

+    return EFI_UNSUPPORTED;

   }

 

   //

-  // Initialize new GDT table and/or IDT table, if any

+  // Total needed size includes stack size, new GDT table size, TSS size.

+  // Add another DESCRIPTOR size for alignment requiremet.

+  //

+  // Layout of memory needed for each processor:

+  //    --------------------------------

+  //    |                              |

+  //    |          Stack Size          |  X ExceptionNumber

+  //    |                              |

+  //    --------------------------------

+  //    |          Alignment           |  (just in case)

+  //    --------------------------------

+  //    |                              |

+  //    |         Original GDT         |

+  //    |                              |

+  //    --------------------------------

+  //    |                              |

+  //    |  Exception task descriptors  |  X 1

+  //    |                              |

+  //    --------------------------------

+  //    |                              |

+  //    | Exception task-state segment |  X 1

+  //    |                              |

+  //    --------------------------------

   //

-  AsmReadIdtr (&Idtr);

   AsmReadGdtr (&Gdtr);

-

-  GdtSize = (UINTN)TssDesc + sizeof (IA32_TSS_DESCRIPTOR) -

-            (UINTN)(StackSwitchData->GdtTable);

-  if ((UINTN)StackSwitchData->GdtTable != Gdtr.Base) {

-    CopyMem (StackSwitchData->GdtTable, (VOID *)Gdtr.Base, Gdtr.Limit + 1);

-    Gdtr.Base  = (UINTN)StackSwitchData->GdtTable;

-    Gdtr.Limit = (UINT16)GdtSize - 1;

+  NeedBufferSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE +

+                   sizeof (IA32_TSS_DESCRIPTOR) +

+                   Gdtr.Limit + 1 + CPU_TSS_DESC_SIZE +

+                   CPU_TSS_SIZE;

+

+  if (*BufferSize < NeedBufferSize) {

+    *BufferSize = NeedBufferSize;

+    return EFI_BUFFER_TOO_SMALL;

   }

 

-  if ((UINTN)StackSwitchData->IdtTable != Idtr.Base) {

-    Idtr.Base = (UINTN)StackSwitchData->IdtTable;

+  if (Buffer == NULL) {

+    return EFI_INVALID_PARAMETER;

   }

 

-  if (StackSwitchData->IdtTableSize > 0) {

-    Idtr.Limit = (UINT16)(StackSwitchData->IdtTableSize - 1);

-  }

+  AsmReadIdtr (&Idtr);

+  StackSwitchExceptions = CPU_STACK_SWITCH_EXCEPTION_LIST;

+  StackTop              = (UINTN)Buffer + CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE;

+  NewGdtTable           = ALIGN_POINTER (StackTop, sizeof (IA32_TSS_DESCRIPTOR));

+  TssDesc               = (IA32_TSS_DESCRIPTOR *)((UINTN)NewGdtTable + Gdtr.Limit + 1);

+  Tss                   = (IA32_TASK_STATE_SEGMENT *)((UINTN)TssDesc + CPU_TSS_DESC_SIZE);

+

+  CopyMem (NewGdtTable, (VOID *)Gdtr.Base, Gdtr.Limit + 1);

+  Gdtr.Base  = (UINTN)NewGdtTable;

+  Gdtr.Limit = (UINT16)(Gdtr.Limit + CPU_TSS_DESC_SIZE);

 

   //

   // Fixup current task descriptor. Task-state segment for current task will

@@ -231,20 +223,20 @@ ArchSetupExceptionStack (
   // Fixup exception task descriptor and task-state segment

   //

   ZeroMem (Tss, sizeof (*Tss));

-  StackTop = StackSwitchData->KnownGoodStackTop - CPU_STACK_ALIGNMENT;

+  StackTop = StackTop - CPU_STACK_ALIGNMENT;

   StackTop = (UINTN)ALIGN_POINTER (StackTop, CPU_STACK_ALIGNMENT);

-  IdtTable = StackSwitchData->IdtTable;

-  for (Index = 0; Index < StackSwitchData->StackSwitchExceptionNumber; ++Index) {

+  IdtTable = (IA32_IDT_GATE_DESCRIPTOR  *)Idtr.Base;

+  for (Index = 0; Index < CPU_STACK_SWITCH_EXCEPTION_NUMBER; ++Index) {

     //

     // Fixup IST

     //

     Tss->IST[Index] = StackTop;

-    StackTop       -= StackSwitchData->KnownGoodStackSize;

+    StackTop       -= CPU_KNOWN_GOOD_STACK_SIZE;

 

     //

     // Set the IST field to enable corresponding IST

     //

-    Vector = StackSwitchData->StackSwitchExceptions[Index];

+    Vector = StackSwitchExceptions[Index];

     if ((Vector >= CPU_EXCEPTION_NUM) ||

         (Vector >= (Idtr.Limit + 1) / sizeof (IA32_IDT_GATE_DESCRIPTOR)))

     {

@@ -262,12 +254,7 @@ ArchSetupExceptionStack (
   //

   // Load current task

   //

-  AsmWriteTr ((UINT16)((UINTN)StackSwitchData->ExceptionTssDesc - Gdtr.Base));

-

-  //

-  // Publish IDT

-  //

-  AsmWriteIdtr (&Idtr);

+  AsmWriteTr ((UINT16)((UINTN)TssDesc - Gdtr.Base));

 

   return EFI_SUCCESS;

 }

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Xcode5SecPeiCpuExceptionHandlerLib.inf b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Xcode5SecPeiCpuExceptionHandlerLib.inf
index a15f125d5b..6d2f66504a 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Xcode5SecPeiCpuExceptionHandlerLib.inf
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Xcode5SecPeiCpuExceptionHandlerLib.inf
@@ -2,7 +2,7 @@
 #  CPU Exception Handler library instance for SEC/PEI modules.

 #

 #  Copyright (C) 2020, Advanced Micro Devices, Inc. All rights reserved.<BR>

-#  Copyright (c) 2012 - 2018, Intel Corporation. All rights reserved.<BR>

+#  Copyright (c) 2012 - 2022, Intel Corporation. All rights reserved.<BR>

 #  SPDX-License-Identifier: BSD-2-Clause-Patent

 #

 #  This is the XCODE5 variant of the SEC/PEI CpuExceptionHandlerLib. This

@@ -55,6 +55,11 @@
   PeCoffGetEntryPointLib

   VmgExitLib

 

+[Pcd]

+  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard

+  gUefiCpuPkgTokenSpaceGuid.PcdCpuStackSwitchExceptionList

+  gUefiCpuPkgTokenSpaceGuid.PcdCpuKnownGoodStackSize

+

 [FeaturePcd]

   gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmStackGuard                    ## CONSUMES