UefiCpuPkg: Fix unchecked returns and potential integer overflows

Resolves several issues in UefiCpuPkg related to:

1. Unchecked returns leading to potential NULL or uninitialized access.
2. Potential unchecked integer overflows.
3. Incorrect comparison between integers of different sizes.

Co-authored-by: kenlautner <85201046+kenlautner@users.noreply.github.com>
Signed-off-by: Chris Fernald <chfernal@microsoft.com>

UefiCpuPkg/CpuDxe/CpuMp.c

@@ -622,8 +622,15 @@ InitializeExceptionStackSwitchHandlers (
 {
   EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
   UINTN                           Index;
+  EFI_STATUS                      Status;
+
+  Status = MpInitLibWhoAmI (&Index);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  The exception stack was not initialized.\n", __func__));
+    return;
+  }
 
-  MpInitLibWhoAmI (&Index);
   SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
 
   //
@@ -758,7 +765,9 @@ InitializeMpSupport (
   Status = MpInitLibInitialize ();
   ASSERT_EFI_ERROR (Status);
 
-  MpInitLibGetNumberOfProcessors (&NumberOfProcessors, &NumberOfEnabledProcessors);
+  Status = MpInitLibGetNumberOfProcessors (&NumberOfProcessors, &NumberOfEnabledProcessors);
+  ASSERT_EFI_ERROR (Status);
+  if (!EFI_ERROR (Status)) {
     mNumberOfProcessors = NumberOfProcessors;
     DEBUG ((DEBUG_INFO, "Detect CPU count: %d\n", mNumberOfProcessors));
 
@@ -779,4 +788,5 @@ InitializeMpSupport (
                     NULL
                     );
     ASSERT_EFI_ERROR (Status);
+  }
 }

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)));
+  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));
+    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) {

UefiCpuPkg/Library/CpuCommonFeaturesLib/MachineCheck.c

@@ -174,7 +174,7 @@ McaInitialize (
     }
 
     if (PcdGetBool (PcdIsPowerOnReset)) {
-      for (BankIndex = 0; BankIndex < (UINTN)McgCap.Bits.Count; BankIndex++) {
+      for (BankIndex = 0; BankIndex < (UINT32)McgCap.Bits.Count; BankIndex++) {
         CPU_REGISTER_TABLE_WRITE64 (
           ProcessorNumber,
           Msr,

UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c

@@ -71,7 +71,11 @@ SetExceptionHandlerData (
   IdtTable = (IA32_IDT_GATE_DESCRIPTOR *)IdtDescriptor.Base;
 
   Exception0StubHeader = AllocatePool (sizeof (*Exception0StubHeader));
+  if (Exception0StubHeader == NULL) {
     ASSERT (Exception0StubHeader != NULL);
+    return;
+  }
+
   CopyMem (
     Exception0StubHeader->ExceptionStubHeader,
     (VOID *)ArchGetIdtHandler (&IdtTable[0]),
@@ -165,10 +169,18 @@ InitializeCpuExceptionHandlers (
   RESERVED_VECTORS_DATA   *ReservedVectors;
 
   ReservedVectors = AllocatePool (sizeof (RESERVED_VECTORS_DATA) * CPU_EXCEPTION_NUM);
+  if (ReservedVectors == NULL) {
     ASSERT (ReservedVectors != NULL);
+    return EFI_OUT_OF_RESOURCES;
+  }
 
   ExceptionHandlerData = AllocatePool (sizeof (EXCEPTION_HANDLER_DATA));
+  if (ExceptionHandlerData == NULL) {
     ASSERT (ExceptionHandlerData != NULL);
+    FreePool (ReservedVectors);
+    return EFI_OUT_OF_RESOURCES;
+  }
+
   ExceptionHandlerData->IdtEntryCount            = CPU_EXCEPTION_NUM;
   ExceptionHandlerData->ReservedVectors          = ReservedVectors;
   ExceptionHandlerData->ExternalInterruptHandler = AllocateZeroPool (sizeof (EFI_CPU_INTERRUPT_HANDLER) * ExceptionHandlerData->IdtEntryCount);

UefiCpuPkg/Library/MpInitLib/AmdSev.c

@@ -25,7 +25,9 @@ GetProtectedMode16CS (
   IA32_DESCRIPTOR          GdtrDesc;
   IA32_SEGMENT_DESCRIPTOR  *GdtEntry;
   UINTN                    GdtEntryCount;
-  UINT16                   Index;
+  UINTN                    Index;
+  UINT16                   CodeSegmentValue;
+  EFI_STATUS               Status;
 
   Index = (UINT16)-1;
   AsmReadGdtr (&GdtrDesc);
@@ -42,8 +44,19 @@ GetProtectedMode16CS (
     GdtEntry++;
   }
 
-  ASSERT (Index != GdtEntryCount);
+  Status = SafeUintnToUint16 (Index, &CodeSegmentValue);
-  return Index * 8;
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  Status = SafeUint16Mult (CodeSegmentValue, 8, &CodeSegmentValue);
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  return CodeSegmentValue;
 }
 
 /**
@@ -61,7 +74,9 @@ GetProtectedMode32CS (
   IA32_DESCRIPTOR          GdtrDesc;
   IA32_SEGMENT_DESCRIPTOR  *GdtEntry;
   UINTN                    GdtEntryCount;
-  UINT16                   Index;
+  UINTN                    Index;
+  UINT16                   CodeSegmentValue;
+  EFI_STATUS               Status;
 
   Index = (UINT16)-1;
   AsmReadGdtr (&GdtrDesc);
@@ -79,7 +94,19 @@ GetProtectedMode32CS (
   }
 
   ASSERT (Index != GdtEntryCount);
-  return Index * 8;
+  Status = SafeUintnToUint16 (Index, &CodeSegmentValue);
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  Status = SafeUint16Mult (CodeSegmentValue, 8, &CodeSegmentValue);
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  return CodeSegmentValue;
 }
 
 /**

UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf

@@ -63,6 +63,7 @@
 
 [LibraryClasses.IA32, LibraryClasses.X64]
   AmdSvsmLib
+  SafeIntLib
   CcExitLib
   LocalApicLib
   MicrocodeLib

UefiCpuPkg/Library/MpInitLib/DxeMpLib.c

@@ -313,7 +313,9 @@ GetProtectedMode16CS (
   IA32_DESCRIPTOR          GdtrDesc;
   IA32_SEGMENT_DESCRIPTOR  *GdtEntry;
   UINTN                    GdtEntryCount;
-  UINT16                   Index;
+  UINTN                    Index;
+  UINT16                   CodeSegmentValue;
+  EFI_STATUS               Status;
 
   Index = (UINT16)-1;
   AsmReadGdtr (&GdtrDesc);
@@ -329,8 +331,19 @@ GetProtectedMode16CS (
     GdtEntry++;
   }
 
-  ASSERT (Index != GdtEntryCount);
+  Status = SafeUintnToUint16 (Index, &CodeSegmentValue);
-  return Index * 8;
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  Status = SafeUint16Mult (CodeSegmentValue, 8, &CodeSegmentValue);
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  return CodeSegmentValue;
 }
 
 /**
@@ -346,7 +359,9 @@ GetProtectedModeCS (
   IA32_DESCRIPTOR          GdtrDesc;
   IA32_SEGMENT_DESCRIPTOR  *GdtEntry;
   UINTN                    GdtEntryCount;
-  UINT16                   Index;
+  UINTN                    Index;
+  UINT16                   CodeSegmentValue;
+  EFI_STATUS               Status;
 
   AsmReadGdtr (&GdtrDesc);
   GdtEntryCount = (GdtrDesc.Limit + 1) / sizeof (IA32_SEGMENT_DESCRIPTOR);
@@ -361,8 +376,19 @@ GetProtectedModeCS (
     GdtEntry++;
   }
 
-  ASSERT (Index != GdtEntryCount);
+  Status = SafeUintnToUint16 (Index, &CodeSegmentValue);
-  return Index * 8;
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  Status = SafeUint16Mult (CodeSegmentValue, 8, &CodeSegmentValue);
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  return CodeSegmentValue;
 }
 
 /**

UefiCpuPkg/Library/MpInitLib/MpLib.c

@@ -1657,6 +1657,11 @@ ResetProcessorToIdleState (
   CpuMpData = GetCpuMpData ();
 
   CpuMpData->WakeUpByInitSipiSipi = TRUE;
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.  Aborting the AP reset to idle.\n", __func__));
+    return;
+  }
+
   WakeUpAP (CpuMpData, FALSE, ProcessorNumber, NULL, NULL, TRUE);
   while (CpuMpData->FinishedCount < 1) {
     CpuPause ();
@@ -1686,6 +1691,11 @@ GetNextWaitingProcessorNumber (
 
   CpuMpData = GetCpuMpData ();
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   for (ProcessorNumber = 0; ProcessorNumber < CpuMpData->CpuCount; ProcessorNumber++) {
     if (CpuMpData->CpuData[ProcessorNumber].Waiting) {
       *NextProcessorNumber = ProcessorNumber;
@@ -1716,6 +1726,12 @@ CheckThisAP (
   CPU_AP_DATA  *CpuData;
 
   CpuMpData = GetCpuMpData ();
+
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   CpuData = &CpuMpData->CpuData[ProcessorNumber];
 
   //
@@ -1778,6 +1794,11 @@ CheckAllAPs (
 
   CpuMpData = GetCpuMpData ();
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   NextProcessorNumber = 0;
 
   //
@@ -2097,6 +2118,10 @@ MpInitLibInitialize (
   BufferSize += (sizeof (CPU_AP_DATA) + sizeof (CPU_INFO_IN_HOB))* MaxLogicalProcessorNumber;
   MpBuffer    = AllocatePages (EFI_SIZE_TO_PAGES (BufferSize));
   ASSERT (MpBuffer != NULL);
+  if (MpBuffer == NULL) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
   ZeroMem (MpBuffer, BufferSize);
   Buffer = ALIGN_VALUE ((UINTN)MpBuffer, ApStackSize);
 
@@ -2457,8 +2482,15 @@ MpInitLibGetProcessorInfo (
   UINTN            CallerNumber;
   CPU_INFO_IN_HOB  *CpuInfoInHob;
   UINTN            OriginalProcessorNumber;
+  EFI_STATUS       Status;
 
   CpuMpData = GetCpuMpData ();
+
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;
 
   //
@@ -2470,7 +2502,13 @@ MpInitLibGetProcessorInfo (
   //
   // Check whether caller processor is BSP
   //
-  MpInitLibWhoAmI (&CallerNumber);
+  Status = MpInitLibWhoAmI (&CallerNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Failed to get MpInit Processor info.\n", __func__));
+    return Status;
+  }
+
   if (CallerNumber != CpuMpData->BspNumber) {
     return EFI_DEVICE_ERROR;
   }
@@ -2551,6 +2589,7 @@ SwitchBSPWorker (
   MSR_IA32_APIC_BASE_REGISTER  ApicBaseMsr;
   BOOLEAN                      OldInterruptState;
   BOOLEAN                      OldTimerInterruptState;
+  EFI_STATUS                   Status;
 
   //
   // Save and Disable Local APIC timer interrupt
@@ -2573,10 +2612,21 @@ SwitchBSPWorker (
 
   CpuMpData = GetCpuMpData ();
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   //
   // Check whether caller processor is BSP
   //
-  MpInitLibWhoAmI (&CallerNumber);
+  Status = MpInitLibWhoAmI (&CallerNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Failed to get MpInit Processor info.\n", __func__));
+    return Status;
+  }
+
   if (CallerNumber != CpuMpData->BspNumber) {
     return EFI_DEVICE_ERROR;
   }
@@ -2700,13 +2750,25 @@ EnableDisableApWorker (
 {
   CPU_MP_DATA  *CpuMpData;
   UINTN        CallerNumber;
+  EFI_STATUS   Status;
 
   CpuMpData = GetCpuMpData ();
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   //
   // Check whether caller processor is BSP
   //
-  MpInitLibWhoAmI (&CallerNumber);
+  Status = MpInitLibWhoAmI (&CallerNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Failed to get MpInit Processor info.\n", __func__));
+    return Status;
+  }
+
   if (CallerNumber != CpuMpData->BspNumber) {
     return EFI_DEVICE_ERROR;
   }
@@ -2763,6 +2825,11 @@ MpInitLibWhoAmI (
 
   CpuMpData = GetCpuMpData ();
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   return GetProcessorNumber (CpuMpData, ProcessorNumber);
 }
 
@@ -2798,9 +2865,15 @@ MpInitLibGetNumberOfProcessors (
   UINTN        ProcessorNumber;
   UINTN        EnabledProcessorNumber;
   UINTN        Index;
+  EFI_STATUS   Status;
 
   CpuMpData = GetCpuMpData ();
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   if ((NumberOfProcessors == NULL) && (NumberOfEnabledProcessors == NULL)) {
     return EFI_INVALID_PARAMETER;
   }
@@ -2808,7 +2881,13 @@ MpInitLibGetNumberOfProcessors (
   //
   // Check whether caller processor is BSP
   //
-  MpInitLibWhoAmI (&CallerNumber);
+  Status = MpInitLibWhoAmI (&CallerNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Failed to get MpInit Processor info.\n", __func__));
+    return Status;
+  }
+
   if (CallerNumber != CpuMpData->BspNumber) {
     return EFI_DEVICE_ERROR;
   }
@@ -2890,6 +2969,11 @@ StartupAllCPUsWorker (
     *FailedCpuList = NULL;
   }
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   if ((CpuMpData->CpuCount == 1) && ExcludeBsp) {
     return EFI_NOT_STARTED;
   }
@@ -2901,7 +2985,13 @@ StartupAllCPUsWorker (
   //
   // Check whether caller processor is BSP
   //
-  MpInitLibWhoAmI (&CallerNumber);
+  Status = MpInitLibWhoAmI (&CallerNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Failed to get MpInit Processor info.\n", __func__));
+    return Status;
+  }
+
   if (CallerNumber != CpuMpData->BspNumber) {
     return EFI_DEVICE_ERROR;
   }
@@ -3043,10 +3133,21 @@ StartupThisAPWorker (
     *Finished = FALSE;
   }
 
+  if (CpuMpData == NULL) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));
+    return EFI_LOAD_ERROR;
+  }
+
   //
   // Check whether caller processor is BSP
   //
-  MpInitLibWhoAmI (&CallerNumber);
+  Status = MpInitLibWhoAmI (&CallerNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Failed to get MpInit Processor info.\n", __func__));
+    return Status;
+  }
+
   if (CallerNumber != CpuMpData->BspNumber) {
     return EFI_DEVICE_ERROR;
   }
@@ -3268,8 +3369,15 @@ RelocateApLoop (
   BOOLEAN      MwaitSupport;
   UINTN        ProcessorNumber;
   UINTN        StackStart;
+  EFI_STATUS   Status;
+
+  Status = MpInitLibWhoAmI (&ProcessorNumber);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get processor number.  Aborting AP sync.\n", __func__));
+    return;
+  }
 
-  MpInitLibWhoAmI (&ProcessorNumber);
   CpuMpData    = GetCpuMpData ();
   MwaitSupport = IsMwaitSupport ();
   if (CpuMpData->UseSevEsAPMethod) {

UefiCpuPkg/Library/MpInitLib/MpLib.h

@@ -34,6 +34,7 @@
 #include <Library/PcdLib.h>
 #include <Library/MicrocodeLib.h>
 #include <Library/CpuPageTableLib.h>
+#include <Library/SafeIntLib.h>
 #include <ConfidentialComputingGuestAttr.h>
 
 #include <Register/Amd/SevSnpMsr.h>

UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf

@@ -62,6 +62,7 @@
 
 [LibraryClasses.IA32, LibraryClasses.X64]
   AmdSvsmLib
+  SafeIntLib
   CcExitLib
   LocalApicLib
   MicrocodeLib

UefiCpuPkg/Library/MpInitLib/PeiMpLib.c

@@ -230,7 +230,7 @@ GetWakeupBuffer (
           WakeupBufferEnd = BASE_1MB;
         }
 
-        while (WakeupBufferEnd > WakeupBufferSize) {
+        while (WakeupBufferEnd > (UINT64)WakeupBufferSize) {
           //
           // Wakeup buffer should be aligned on 4KB
           //

UefiCpuPkg/Library/MtrrLib/MtrrLib.c

@@ -1703,8 +1703,8 @@ MtrrLibCalculateMtrrs (
   }
 
   for (TypeCount = 2; TypeCount <= 3; TypeCount++) {
-    for (Start = 0; Start < VertexCount; Start++) {
+    for (Start = 0; (UINT32)Start < VertexCount; Start++) {
-      for (Stop = Start + 2; Stop < VertexCount; Stop++) {
+      for (Stop = Start + 2; (UINT32)Stop < VertexCount; Stop++) {
         ASSERT (Vertices[Stop].Address > Vertices[Start].Address);
         Length = Vertices[Stop].Address - Vertices[Start].Address;
         if (Length > Vertices[Start].Alignment) {
@@ -2139,13 +2139,13 @@ MtrrLibSetMemoryRanges (
   //
   BiggestScratchSize = 0;
 
-  for (Index = 0; Index < RangeCount;) {
+  for (Index = 0; (UINTN)Index < RangeCount;) {
     Base0 = Ranges[Index].BaseAddress;
 
     //
     // Full step is optimal
     //
-    while (Index < RangeCount) {
+    while ((UINTN)Index < RangeCount) {
       ASSERT (Ranges[Index].BaseAddress == Base0);
       Alignment = MtrrLibBiggestAlignment (Base0, A0);
       while (Base0 + Alignment <= Ranges[Index].BaseAddress + Ranges[Index].Length) {
@@ -2193,7 +2193,7 @@ MtrrLibSetMemoryRanges (
     CompatibleTypes = MtrrLibGetCompatibleTypes (&Ranges[Index], RangeCount - Index);
 
     End = Index; // End points to last one that matches the CompatibleTypes.
-    while (End + 1 < RangeCount) {
+    while ((UINTN)(End + 1) < RangeCount) {
       if (((1 << Ranges[End + 1].Type) & CompatibleTypes) == 0) {
         break;
       }
@@ -2209,7 +2209,7 @@ MtrrLibSetMemoryRanges (
     // Base1 may not in Ranges[End]. Update End to the range Base1 belongs to.
     //
     End = Index;
-    while (End + 1 < RangeCount) {
+    while ((UINTN)(End + 1) < RangeCount) {
       if (Base1 <= Ranges[End + 1].BaseAddress) {
         break;
       }

UefiCpuPkg/Library/RegisterCpuFeaturesLib/CpuFeaturesInitialize.c

@@ -95,6 +95,7 @@ CpuInitDataInitialize (
   EFI_STATUS                 Status;
   UINTN                      ProcessorNumber;
   EFI_PROCESSOR_INFORMATION  ProcessorInfoBuffer;
+  CPU_STATUS_INFORMATION     CpuStatusBackupBuffer;
   CPU_FEATURES_ENTRY         *CpuFeature;
   CPU_FEATURES_INIT_ORDER    *InitOrder;
   CPU_FEATURES_DATA          *CpuFeaturesData;
@@ -120,7 +121,19 @@ CpuInitDataInitialize (
   Package = 0;
   Thread  = 0;
 
+  CpuFeaturesData       = NULL;
+  CpuStatus             = NULL;
+  FirstCore             = NULL;
+  InitOrder             = NULL;
+  Location              = NULL;
+  ThreadCountPerCore    = NULL;
+  ThreadCountPerPackage = NULL;
+
   CpuFeaturesData = GetCpuFeaturesData ();
+  if (CpuFeaturesData == NULL) {
+    ASSERT (CpuFeaturesData != NULL);
+    return;
+  }
 
   //
   // Initialize CpuFeaturesData->MpService as early as possile, so later function can use it.
@@ -130,7 +143,11 @@ CpuInitDataInitialize (
   GetNumberOfProcessor (&NumberOfCpus, &NumberOfEnabledProcessors);
 
   CpuFeaturesData->InitOrder = AllocatePages (EFI_SIZE_TO_PAGES (sizeof (CPU_FEATURES_INIT_ORDER) * NumberOfCpus));
+  if (CpuFeaturesData->InitOrder == NULL) {
     ASSERT (CpuFeaturesData->InitOrder != NULL);
+    return;
+  }
+
   ZeroMem (CpuFeaturesData->InitOrder, sizeof (CPU_FEATURES_INIT_ORDER) * NumberOfCpus);
 
   //
@@ -150,19 +167,32 @@ CpuInitDataInitialize (
   CpuFeaturesData->NumberOfCpus = (UINT32)NumberOfCpus;
 
   AcpiCpuData = GetAcpiCpuData ();
+  if (AcpiCpuData == NULL) {
     ASSERT (AcpiCpuData != NULL);
+    goto ExitOnError;
+  }
+
   CpuFeaturesData->AcpiCpuData = AcpiCpuData;
 
   CpuStatus = &AcpiCpuData->CpuFeatureInitData.CpuStatus;
+  CopyMem (&CpuStatusBackupBuffer, CpuStatus, sizeof (CpuStatusBackupBuffer));
   Location = AllocatePages (EFI_SIZE_TO_PAGES (sizeof (EFI_CPU_PHYSICAL_LOCATION) * NumberOfCpus));
+  if (Location == NULL) {
     ASSERT (Location != NULL);
+    goto ExitOnError;
+  }
+
   ZeroMem (Location, sizeof (EFI_CPU_PHYSICAL_LOCATION) * NumberOfCpus);
   AcpiCpuData->CpuFeatureInitData.ApLocation = (EFI_PHYSICAL_ADDRESS)(UINTN)Location;
 
   for (ProcessorNumber = 0; ProcessorNumber < NumberOfCpus; ProcessorNumber++) {
     InitOrder                        = &CpuFeaturesData->InitOrder[ProcessorNumber];
     InitOrder->FeaturesSupportedMask = AllocateZeroPool (CpuFeaturesData->BitMaskSize);
+    if (InitOrder->FeaturesSupportedMask == NULL) {
       ASSERT (InitOrder->FeaturesSupportedMask != NULL);
+      goto ExitOnError;
+    }
+
     InitializeListHead (&InitOrder->OrderList);
     Status = GetProcessorInformation (ProcessorNumber, &ProcessorInfoBuffer);
     ASSERT_EFI_ERROR (Status);
@@ -214,12 +244,20 @@ CpuInitDataInitialize (
   // Collect valid core count in each package because not all cores are valid.
   //
   ThreadCountPerPackage = AllocatePages (EFI_SIZE_TO_PAGES (sizeof (UINT32) * CpuStatus->PackageCount));
+  if (ThreadCountPerPackage == NULL) {
     ASSERT (ThreadCountPerPackage != NULL);
+    goto ExitOnError;
+  }
+
   ZeroMem (ThreadCountPerPackage, sizeof (UINT32) * CpuStatus->PackageCount);
   CpuStatus->ThreadCountPerPackage = (EFI_PHYSICAL_ADDRESS)(UINTN)ThreadCountPerPackage;
 
   ThreadCountPerCore = AllocatePages (EFI_SIZE_TO_PAGES (sizeof (UINT8) * CpuStatus->PackageCount * CpuStatus->MaxCoreCount));
+  if (ThreadCountPerCore == NULL) {
     ASSERT (ThreadCountPerCore != NULL);
+    goto ExitOnError;
+  }
+
   ZeroMem (ThreadCountPerCore, sizeof (UINT8) * CpuStatus->PackageCount * CpuStatus->MaxCoreCount);
   CpuStatus->ThreadCountPerCore = (EFI_PHYSICAL_ADDRESS)(UINTN)ThreadCountPerCore;
 
@@ -247,9 +285,16 @@ CpuInitDataInitialize (
   }
 
   CpuFeaturesData->CpuFlags.CoreSemaphoreCount = AllocateZeroPool (sizeof (UINT32) * CpuStatus->PackageCount * CpuStatus->MaxCoreCount * CpuStatus->MaxThreadCount);
+  if (CpuFeaturesData->CpuFlags.CoreSemaphoreCount == NULL) {
     ASSERT (CpuFeaturesData->CpuFlags.CoreSemaphoreCount != NULL);
+    goto ExitOnError;
+  }
+
   CpuFeaturesData->CpuFlags.PackageSemaphoreCount = AllocateZeroPool (sizeof (UINT32) * CpuStatus->PackageCount * CpuStatus->MaxCoreCount * CpuStatus->MaxThreadCount);
+  if (CpuFeaturesData->CpuFlags.PackageSemaphoreCount == NULL) {
     ASSERT (CpuFeaturesData->CpuFlags.PackageSemaphoreCount != NULL);
+    goto ExitOnError;
+  }
 
   //
   // Initialize CpuFeaturesData->InitOrder[].CpuInfo.First
@@ -257,7 +302,11 @@ CpuInitDataInitialize (
   //
   Pages     = EFI_SIZE_TO_PAGES (CpuStatus->PackageCount * sizeof (UINT32) + CpuStatus->PackageCount * CpuStatus->MaxCoreCount * sizeof (UINT32));
   FirstCore = AllocatePages (Pages);
+  if (FirstCore == NULL) {
     ASSERT (FirstCore != NULL);
+    goto ExitOnError;
+  }
+
   FirstThread = FirstCore + CpuStatus->PackageCount;
 
   //
@@ -317,6 +366,60 @@ CpuInitDataInitialize (
   }
 
   FreePages (FirstCore, Pages);
+
+  return;
+
+ExitOnError:
+  if (FirstCore != NULL) {
+    FreePages (FirstCore, Pages);
+  }
+
+  if ((CpuFeaturesData != NULL) && (CpuFeaturesData->CpuFlags.PackageSemaphoreCount != NULL)) {
+    FreePool ((VOID *)CpuFeaturesData->CpuFlags.PackageSemaphoreCount);
+    CpuFeaturesData->CpuFlags.PackageSemaphoreCount = NULL;
+  }
+
+  if ((CpuFeaturesData != NULL) && (CpuFeaturesData->CpuFlags.CoreSemaphoreCount != NULL)) {
+    FreePool ((VOID *)CpuFeaturesData->CpuFlags.CoreSemaphoreCount);
+    CpuFeaturesData->CpuFlags.CoreSemaphoreCount = NULL;
+  }
+
+  if ((ThreadCountPerCore != NULL) && (CpuStatus != NULL)) {
+    FreePages (
+      ThreadCountPerCore,
+      EFI_SIZE_TO_PAGES (sizeof (UINT8) * CpuStatus->PackageCount * CpuStatus->MaxCoreCount)
+      );
+  }
+
+  if ((ThreadCountPerPackage != NULL) && (CpuStatus != NULL)) {
+    FreePages (
+      ThreadCountPerPackage,
+      EFI_SIZE_TO_PAGES (sizeof (UINT32) * CpuStatus->PackageCount)
+      );
+  }
+
+  if (InitOrder != NULL) {
+    for (ProcessorNumber = 0; ProcessorNumber < NumberOfCpus; ProcessorNumber++) {
+      InitOrder = &CpuFeaturesData->InitOrder[ProcessorNumber];
+      if (InitOrder->FeaturesSupportedMask != NULL) {
+        FreePool (InitOrder->FeaturesSupportedMask);
+        InitOrder->FeaturesSupportedMask = NULL;
+      }
+    }
+  }
+
+  if (Location != NULL) {
+    FreePages (Location, EFI_SIZE_TO_PAGES (sizeof (EFI_CPU_PHYSICAL_LOCATION) * NumberOfCpus));
+  }
+
+  if (CpuFeaturesData->InitOrder != NULL) {
+    FreePages (CpuFeaturesData->InitOrder, EFI_SIZE_TO_PAGES (sizeof (CPU_FEATURES_INIT_ORDER) * NumberOfCpus));
+    CpuFeaturesData->InitOrder = NULL;
+  }
+
+  if (CpuStatus != NULL) {
+    CopyMem (CpuStatus, &CpuStatusBackupBuffer, sizeof (*CpuStatus));
+  }
 }
 
 /**

UefiCpuPkg/Library/RegisterCpuFeaturesLib/RegisterCpuFeaturesLib.c

@@ -966,6 +966,8 @@ RegisterCpuFeature (
   Return ACPI_CPU_DATA data.
 
   @return  Pointer to ACPI_CPU_DATA data.
+           NULL if the ACPI CPU data structure cannot be allocated.
+
 **/
 ACPI_CPU_DATA *
 GetAcpiCpuData (
@@ -984,7 +986,11 @@ GetAcpiCpuData (
   AcpiCpuData = (ACPI_CPU_DATA *)(UINTN)PcdGet64 (PcdCpuS3DataAddress);
   if (AcpiCpuData == NULL) {
     AcpiCpuData = AllocatePages (EFI_SIZE_TO_PAGES (sizeof (ACPI_CPU_DATA)));
+    if (AcpiCpuData == NULL) {
       ASSERT (AcpiCpuData != NULL);
+      return NULL;
+    }
+
     ZeroMem (AcpiCpuData, sizeof (ACPI_CPU_DATA));
 
     //
@@ -1006,7 +1012,12 @@ GetAcpiCpuData (
     NumberOfCpus  = AcpiCpuData->NumberOfCpus;
     TableSize     = 2 * NumberOfCpus * sizeof (CPU_REGISTER_TABLE);
     RegisterTable = AllocatePages (EFI_SIZE_TO_PAGES (TableSize));
+    if (RegisterTable == NULL) {
+      // Leave the AcpiCpuData data buffer allocated since it was assigned to a dynamic PCD
+      // which could have invoked PCD set callbacks that may have cached the buffer.
       ASSERT (RegisterTable != NULL);
+      return NULL;
+    }
 
     for (Index = 0; Index < NumberOfCpus; Index++) {
       Status = GetProcessorInformation (Index, &ProcessorInfoBuffer);
@@ -1111,7 +1122,12 @@ CpuRegisterTableWriteWorker (
   CpuFeaturesData = GetCpuFeaturesData ();
   if (CpuFeaturesData->RegisterTable == NULL) {
     AcpiCpuData = GetAcpiCpuData ();
-    ASSERT ((AcpiCpuData != NULL) && (AcpiCpuData->CpuFeatureInitData.RegisterTable != 0));
+    if (AcpiCpuData == NULL) {
+      ASSERT (AcpiCpuData != NULL);
+      return;
+    }
+
+    ASSERT (AcpiCpuData->CpuFeatureInitData.RegisterTable != 0);
     CpuFeaturesData->RegisterTable       = (CPU_REGISTER_TABLE *)(UINTN)AcpiCpuData->CpuFeatureInitData.RegisterTable;
     CpuFeaturesData->PreSmmRegisterTable = (CPU_REGISTER_TABLE *)(UINTN)AcpiCpuData->CpuFeatureInitData.PreSmmInitRegisterTable;
   }

UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c

@@ -1006,10 +1006,17 @@ AllocateTokenBuffer (
   // Separate the Spin_lock and Proc_token because the alignment requires by Spin_Lock.
   //
   SpinLockBuffer = AllocatePool (SpinLockSize * TokenCountPerChunk);
+  if (SpinLockBuffer == NULL) {
     ASSERT (SpinLockBuffer != NULL);
+    return NULL;
+  }
 
   ProcTokens = AllocatePool (sizeof (PROCEDURE_TOKEN) * TokenCountPerChunk);
+  if (ProcTokens == NULL) {
     ASSERT (ProcTokens != NULL);
+    FreePool (SpinLockBuffer);
+    return NULL;
+  }
 
   for (Index = 0; Index < TokenCountPerChunk; Index++) {
     SpinLock = (SPIN_LOCK *)(SpinLockBuffer + SpinLockSize * Index);
@@ -1820,7 +1827,11 @@ InitializeSmmCpuSemaphores (
   DEBUG ((DEBUG_INFO, "Total Semaphores Size = 0x%x\n", TotalSize));
   Pages          = EFI_SIZE_TO_PAGES (TotalSize);
   SemaphoreBlock = AllocatePages (Pages);
+  if (SemaphoreBlock == NULL) {
     ASSERT (SemaphoreBlock != NULL);
+    return;
+  }
+
   ZeroMem (SemaphoreBlock, TotalSize);
 
   SemaphoreAddr                                   = (UINTN)SemaphoreBlock;

UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c

@@ -136,7 +136,11 @@ GetSmiCommandPort (
   Fadt = (EFI_ACPI_2_0_FIXED_ACPI_DESCRIPTION_TABLE *)EfiLocateFirstAcpiTable (
                                                         EFI_ACPI_2_0_FIXED_ACPI_DESCRIPTION_TABLE_SIGNATURE
                                                         );
+
+  if (Fadt == NULL) {
     ASSERT (Fadt != NULL);
+    return;
+  }
 
   mSmiCommandPort = Fadt->SmiCmd;
   DEBUG ((DEBUG_INFO, "mSmiCommandPort = %x\n", mSmiCommandPort));

UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c

@@ -184,6 +184,7 @@ CalculateMaximumSupportAddress (
   Create PageTable for SMM use.
 
   @return The address of PML4 (to set CR3).
+          Zero if any error occurs.
 
 **/
 UINT32
@@ -201,6 +202,9 @@ SmmInitPageTable (
   UINT64                    *Pml4Entry;
   UINT64                    *Pml5Entry;
 
+  Pml4Entry = NULL;
+  Pml5Entry = NULL;
+
   //
   // Initialize spin lock
   //
@@ -254,7 +258,16 @@ SmmInitPageTable (
     // Add pages to page pool
     //
     FreePage = (LIST_ENTRY *)AllocatePageTableMemory (PAGE_TABLE_PAGES);
+    if (FreePage == NULL) {
+      FreePages (Pml4Entry, 1);
+      if (Pml5Entry != NULL) {
+        FreePages (Pml5Entry, 1);
+      }
+
       ASSERT (FreePage != NULL);
+      return 0;
+    }
+
     for (Index = 0; Index < PAGE_TABLE_PAGES; Index++) {
       InsertTailList (&mPagePool, FreePage);
       FreePage += EFI_PAGE_SIZE / sizeof (*FreePage);

UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmFuncsArch.c

@@ -118,7 +118,7 @@ GetProtectedModeCS (
   AsmReadGdtr (&GdtrDesc);
   GdtEntryCount = (GdtrDesc.Limit + 1) / sizeof (IA32_SEGMENT_DESCRIPTOR);
   GdtEntry      = (IA32_SEGMENT_DESCRIPTOR *)GdtrDesc.Base;
-  for (Index = 0; Index < GdtEntryCount; Index++) {
+  for (Index = 0; (UINTN)Index < GdtEntryCount; Index++) {
     if (GdtEntry->Bits.L == 0) {
       if ((GdtEntry->Bits.Type > 8) && (GdtEntry->Bits.DB == 1)) {
         break;

UefiCpuPkg/UefiCpuPkg.dsc

@@ -27,6 +27,7 @@
 [LibraryClasses]
   BaseLib|MdePkg/Library/BaseLib/BaseLib.inf
   BaseMemoryLib|MdePkg/Library/BaseMemoryLib/BaseMemoryLib.inf
+  SafeIntLib|MdePkg/Library/BaseSafeIntLib/BaseSafeIntLib.inf
   CpuLib|MdePkg/Library/BaseCpuLib/BaseCpuLib.inf
   DebugLib|MdePkg/Library/BaseDebugLibNull/BaseDebugLibNull.inf
   SerialPortLib|MdePkg/Library/BaseSerialPortLibNull/BaseSerialPortLibNull.inf