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>

12 files changed, 339 insertions, 44 deletions

diff --git a/UefiCpuPkg/Library/CpuCommonFeaturesLib/MachineCheck.c b/UefiCpuPkg/Library/CpuCommonFeaturesLib/MachineCheck.c
index cb569769a1..6bfb2388e7 100644
--- a/UefiCpuPkg/Library/CpuCommonFeaturesLib/MachineCheck.c
+++ b/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,

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
index 3e38676b23..994e3917fb 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
@@ -71,7 +71,11 @@ SetExceptionHandlerData (
   IdtTable = (IA32_IDT_GATE_DESCRIPTOR *)IdtDescriptor.Base;

 

   Exception0StubHeader = AllocatePool (sizeof (*Exception0StubHeader));

-  ASSERT (Exception0StubHeader != NULL);

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

-  ASSERT (ReservedVectors != NULL);

+  if (ReservedVectors == NULL) {

+    ASSERT (ReservedVectors != NULL);

+    return EFI_OUT_OF_RESOURCES;

+  }

 

   ExceptionHandlerData = AllocatePool (sizeof (EXCEPTION_HANDLER_DATA));

-  ASSERT (ExceptionHandlerData != NULL);

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

diff --git a/UefiCpuPkg/Library/MpInitLib/AmdSev.c b/UefiCpuPkg/Library/MpInitLib/AmdSev.c
index d34f9513e0..75429e3dae 100644
--- a/UefiCpuPkg/Library/MpInitLib/AmdSev.c
+++ b/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);

-  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;

 }

 

 /**

@@ -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;

 }

 

 /**

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
index 41d5a80bc9..922c7b12d3 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
@@ -63,6 +63,7 @@
 

 [LibraryClasses.IA32, LibraryClasses.X64]

   AmdSvsmLib

+  SafeIntLib

   CcExitLib

   LocalApicLib

   MicrocodeLib

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index a8d884f607..8c1428c2fc 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/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);

-  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;

 }

 

 /**

@@ -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);

-  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;

 }

 

 /**

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 67e8556a4a..fdcc21d794 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/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,7 +1726,13 @@ CheckThisAP (
   CPU_AP_DATA  *CpuData;

 

   CpuMpData = GetCpuMpData ();

-  CpuData   = &CpuMpData->CpuData[ProcessorNumber];

+

+  if (CpuMpData == NULL) {

+    DEBUG ((DEBUG_ERROR, "[%a] - Failed to get CpuMpData.\n", __func__));

+    return EFI_LOAD_ERROR;

+  }

+

+  CpuData = &CpuMpData->CpuData[ProcessorNumber];

 

   //

   //  Check the CPU state of AP. If it is CpuStateIdle, then the AP has finished its task.

@@ -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;

+  }

 

-  CpuMpData    = GetCpuMpData ();

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

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 690b7b0e1b..145538b6ee 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/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>

diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
index d2b6a43cdc..d1e8312c02 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
@@ -62,6 +62,7 @@
 

 [LibraryClasses.IA32, LibraryClasses.X64]

   AmdSvsmLib

+  SafeIntLib

   CcExitLib

   LocalApicLib

   MicrocodeLib

diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
index 495dc108b2..bb287940e1 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
+++ b/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

           //

diff --git a/UefiCpuPkg/Library/MtrrLib/MtrrLib.c b/UefiCpuPkg/Library/MtrrLib/MtrrLib.c
index 61af77d9de..bfe02f3cb9 100644
--- a/UefiCpuPkg/Library/MtrrLib/MtrrLib.c
+++ b/UefiCpuPkg/Library/MtrrLib/MtrrLib.c
@@ -1703,8 +1703,8 @@ MtrrLibCalculateMtrrs (
   }

 

   for (TypeCount = 2; TypeCount <= 3; TypeCount++) {

-    for (Start = 0; Start < VertexCount; Start++) {

-      for (Stop = Start + 2; Stop < VertexCount; Stop++) {

+    for (Start = 0; (UINT32)Start < VertexCount; Start++) {

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

       }

diff --git a/UefiCpuPkg/Library/RegisterCpuFeaturesLib/CpuFeaturesInitialize.c b/UefiCpuPkg/Library/RegisterCpuFeaturesLib/CpuFeaturesInitialize.c
index 552fdab417..8fe91696ba 100644
--- a/UefiCpuPkg/Library/RegisterCpuFeaturesLib/CpuFeaturesInitialize.c
+++ b/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));

-  ASSERT (CpuFeaturesData->InitOrder != NULL);

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

-  ASSERT (AcpiCpuData != NULL);

+  if (AcpiCpuData == NULL) {

+    ASSERT (AcpiCpuData != NULL);

+    goto ExitOnError;

+  }

+

   CpuFeaturesData->AcpiCpuData = AcpiCpuData;

 

   CpuStatus = &AcpiCpuData->CpuFeatureInitData.CpuStatus;

-  Location  = AllocatePages (EFI_SIZE_TO_PAGES (sizeof (EFI_CPU_PHYSICAL_LOCATION) * NumberOfCpus));

-  ASSERT (Location != NULL);

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

-    ASSERT (InitOrder->FeaturesSupportedMask != NULL);

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

-  ASSERT (ThreadCountPerPackage != NULL);

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

-  ASSERT (ThreadCountPerCore != NULL);

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

-  ASSERT (CpuFeaturesData->CpuFlags.CoreSemaphoreCount != NULL);

+  if (CpuFeaturesData->CpuFlags.CoreSemaphoreCount == NULL) {

+    ASSERT (CpuFeaturesData->CpuFlags.CoreSemaphoreCount != NULL);

+    goto ExitOnError;

+  }

+

   CpuFeaturesData->CpuFlags.PackageSemaphoreCount = AllocateZeroPool (sizeof (UINT32) * CpuStatus->PackageCount * CpuStatus->MaxCoreCount * CpuStatus->MaxThreadCount);

-  ASSERT (CpuFeaturesData->CpuFlags.PackageSemaphoreCount != NULL);

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

-  ASSERT (FirstCore != NULL);

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

+  }

 }

 

 /**

diff --git a/UefiCpuPkg/Library/RegisterCpuFeaturesLib/RegisterCpuFeaturesLib.c b/UefiCpuPkg/Library/RegisterCpuFeaturesLib/RegisterCpuFeaturesLib.c
index 0285aaf5c9..b5dceba552 100644
--- a/UefiCpuPkg/Library/RegisterCpuFeaturesLib/RegisterCpuFeaturesLib.c
+++ b/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)));

-    ASSERT (AcpiCpuData != NULL);

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

-    ASSERT (RegisterTable != NULL);

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

   }