1 files changed, 19 insertions, 3 deletions
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.c b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
index 2ce4d6ab50..d01fbe047e 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.c
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
@@ -437,8 +437,14 @@ InitializeExceptionStackSwitchHandlers (
 {
   EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
   UINTN                           Index;
+  EFI_STATUS                      Status;
+
+  Status = MpInitLibWhoAmI (&Index);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a - Failed to allocate Switch Stack pages.\n", __func__));
+    return;
+  }
 
-  MpInitLibWhoAmI (&Index);
   SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
 
   //
@@ -481,7 +487,12 @@ InitializeMpExceptionStackSwitchHandlers (
   }
 
   SwitchStackData = AllocatePages (EFI_SIZE_TO_PAGES (NumberOfProcessors * sizeof (EXCEPTION_STACK_SWITCH_CONTEXT)));
-  ASSERT (SwitchStackData != NULL);
+  if (SwitchStackData == NULL) {
+    ASSERT (SwitchStackData != NULL);
+    DEBUG ((DEBUG_ERROR, "%a - Failed to allocate Switch Stack pages.\n", __func__));
+    return;
+  }
+
   ZeroMem (SwitchStackData, NumberOfProcessors * sizeof (EXCEPTION_STACK_SWITCH_CONTEXT));
   for (Index = 0; Index < NumberOfProcessors; ++Index) {
     //
@@ -511,7 +522,12 @@ InitializeMpExceptionStackSwitchHandlers (
 
   if (BufferSize != 0) {
     Buffer = AllocatePages (EFI_SIZE_TO_PAGES (BufferSize));
-    ASSERT (Buffer != NULL);
+    if (Buffer == NULL) {
+      ASSERT (Buffer != NULL);
+      DEBUG ((DEBUG_ERROR, "%a - Failed to allocate Buffer pages.\n", __func__));
+      return;
+    }
+
     BufferSize = 0;
     for (Index = 0; Index < NumberOfProcessors; ++Index) {
       if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {