[PATCH 2/4] UefiCpuPkg/CpuExceptionHandlerLib: support stack switch for PEI exceptions

[Jian J Wang <jian.j.wang@intel.com>]

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>
---
 .../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 6f271983f2..5687c98c0d 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;
 
@@ -208,5 +209,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..0ea44b3052 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuExceptionHandlerLib.inf
@@ -60,3 +60,7 @@
   HobLib
   MemoryAllocationLib
   SynchronizationLib
+
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard
+
-- 
2.16.2.windows.1