From fc0e7fd5e8aad193fcc3479634d5c69996f6460a Mon Sep 17 00:00:00 2001
From: Jian J Wang <jian.j.wang@intel.com>
Date: Mon, 3 Sep 2018 10:36:21 +0800
Subject: [PATCH] UefiCpuPkg/CpuExceptionHandlerLib: support stack switch for
 PEI exceptions

Stack Guard needs to setup stack switch capability to allow exception
handler to be called with good stack if stack overflow is detected.
This patch update InitializeCpuExceptionHandlersEx() to allow pass
extra initialization data used to setup exception stack switch for
specified exceptions.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: "Ware, Ryan R" <ryan.r.ware@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
Regression-tested-by: Laszlo Ersek <lersek@redhat.com>
Reviewed-by: Eric Dong <eric.dong@intel.com>
---
 .../CpuExceptionHandlerLib/PeiCpuException.c  | 27 ++++++++++++++++++-
 .../PeiCpuExceptionHandlerLib.inf             |  4 +++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
index 5dd8423d2f..658f1087a2 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
@@ -17,6 +17,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/DebugLib.h>
 #include <Library/HobLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/PcdLib.h>
 
 CONST UINTN    mDoFarReturnFlag  = 0;
 
@@ -239,5 +240,29 @@ InitializeCpuExceptionHandlersEx (
   IN CPU_EXCEPTION_INIT_DATA            *InitData OPTIONAL
   )
 {
-  return InitializeCpuExceptionHandlers (VectorInfo);
+  EFI_STATUS                        Status;
+
+  //
+  // To avoid repeat initialization of default handlers, the caller should pass
+  // an extended init data with InitDefaultHandlers set to FALSE. There's no
+  // need to call this method to just initialize default handlers. Call non-ex
+  // version instead; or this method must be implemented as a simple wrapper of
+  // non-ex version of it, if this version has to be called.
+  //
+  if (InitData == NULL || InitData->Ia32.InitDefaultHandlers) {
+    Status = InitializeCpuExceptionHandlers (VectorInfo);
+  } else {
+    Status = EFI_SUCCESS;
+  }
+
+  if (!EFI_ERROR (Status)) {
+    //
+    // Initializing stack switch is only necessary for Stack Guard functionality.
+    //
+    if (PcdGetBool (PcdCpuStackGuard) && InitData != NULL) {
+      Status = ArchSetupExcpetionStack (InitData);
+    }
+  }
+
+  return  Status;
 }
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
index 783260e39a..e192641db2 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
@@ -60,3 +60,7 @@
   HobLib
   MemoryAllocationLib
   SynchronizationLib
+
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard    # CONSUMES
+