From 92e488dba2b01cf1ace162c14d1ef0931b6ebafe Mon Sep 17 00:00:00 2001
From: Ken Lautner <klautner@microsoft.com>
Date: Thu, 13 Jul 2023 12:03:40 -0700
Subject: [PATCH] Added additional CodeQL fixes for 202302

---
 .../UnitTest/Library/BaseCryptLib/X509Tests.c |  9 +++++-
 UefiCpuPkg/CpuMpPei/CpuMpPei.c                | 29 +++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/CryptoPkg/Test/UnitTest/Library/BaseCryptLib/X509Tests.c b/CryptoPkg/Test/UnitTest/Library/BaseCryptLib/X509Tests.c
index b3d11ea330..376298a093 100644
--- a/CryptoPkg/Test/UnitTest/Library/BaseCryptLib/X509Tests.c
+++ b/CryptoPkg/Test/UnitTest/Library/BaseCryptLib/X509Tests.c
@@ -573,7 +573,14 @@ TestVerifyX509 (
   Status      = X509GetIssuerName (mTestCert, sizeof (mTestCert), NULL, &SubjectSize);
   UT_ASSERT_TRUE (!Status);
   Subject = AllocatePool (SubjectSize);
-  Status  = X509GetIssuerName (mTestCert, sizeof (mTestCert), Subject, &SubjectSize);
+  // MU_CHANGE [BEGIN] - CodeQL change
+  if (Subject == NULL) {
+    ASSERT (Subject != NULL);
+    return UNIT_TEST_ERROR_TEST_FAILED;
+  }
+
+  // MU_CHANGE [END] - CodeQL change
+  Status = X509GetIssuerName (mTestCert, sizeof (mTestCert), Subject, &SubjectSize);
   UT_ASSERT_TRUE (Status);
   FreePool (Subject);
 
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.c b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
index 1a5e31c0f5..d80c1b9eb1 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.c
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
@@ -541,9 +541,25 @@ InitializeMpExceptionStackSwitchHandlers (
   // }
   // MU_CHANGE END
 
-  MpInitLibGetNumberOfProcessors (&NumberOfProcessors, NULL);
+  // MU_CHANGE [BEGIN] - CodeQL change
+  Status = MpInitLibGetNumberOfProcessors (&NumberOfProcessors, NULL);
+  if (EFI_ERROR (Status)) {
+    ASSERT (NumberOfProcessors != NULL);
+    DEBUG ((DEBUG_ERROR, "%a - Failed to get number of processors.  Status = %r\n", __FUNCTION__, Status));
+    return;
+  }
+
+  // MU_CHANGE [END] - CodeQL change
+
   SwitchStackData = AllocatePages (EFI_SIZE_TO_PAGES (NumberOfProcessors * sizeof (EXCEPTION_STACK_SWITCH_CONTEXT)));
-  ASSERT (SwitchStackData != NULL);
+  // MU_CHANGE [BEGIN] - CodeQL change
+  if (SwitchStackData == NULL) {
+    ASSERT (SwitchStackData != NULL);
+    DEBUG ((DEBUG_ERROR, "%a - Failed to allocate Switch Stack pages.\n", __FUNCTION__));
+    return;
+  }
+
+  // MU_CHANGE [END] - CodeQL change
   ZeroMem (SwitchStackData, NumberOfProcessors * sizeof (EXCEPTION_STACK_SWITCH_CONTEXT));
   for (Index = 0; Index < NumberOfProcessors; ++Index) {
     //
@@ -573,7 +589,14 @@ InitializeMpExceptionStackSwitchHandlers (
 
   if (BufferSize != 0) {
     Buffer = AllocatePages (EFI_SIZE_TO_PAGES (BufferSize));
-    ASSERT (Buffer != NULL);
+    // MU_CHANGE [BEGIN] - CodeQL change
+    if (Buffer == NULL) {
+      ASSERT (Buffer != NULL);
+      DEBUG ((DEBUG_ERROR, "%a - Failed to allocate Buffer pages.\n", __FUNCTION__));
+      return;
+    }
+
+    // MU_CHANGE [END] - CodeQL change
     BufferSize = 0;
     for (Index = 0; Index < NumberOfProcessors; ++Index) {
       if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {