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 <[email protected]> Signed-off-by: Chris Fernald <[email protected]>
zurcher · Nov 15, 2024 · 13fad60 · 13fad60
1 parent 843f0c1
commit 13fad60
Show file tree

Hide file tree

Showing 19 changed files with 423 additions and 73 deletions.
diff --git a/UefiCpuPkg/CpuDxe/CpuMp.c b/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,25 +765,28 @@ InitializeMpSupport (
   Status = MpInitLibInitialize ();
   ASSERT_EFI_ERROR (Status);
 
-  MpInitLibGetNumberOfProcessors (&NumberOfProcessors, &NumberOfEnabledProcessors);
-  mNumberOfProcessors = NumberOfProcessors;
-  DEBUG ((DEBUG_INFO, "Detect CPU count: %d\n", mNumberOfProcessors));
+  Status = MpInitLibGetNumberOfProcessors (&NumberOfProcessors, &NumberOfEnabledProcessors);
+  ASSERT_EFI_ERROR (Status);
+  if (!EFI_ERROR (Status)) {
+    mNumberOfProcessors = NumberOfProcessors;
+    DEBUG ((DEBUG_INFO, "Detect CPU count: %d\n", mNumberOfProcessors));
 
-  //
-  // Initialize special exception handlers for each logic processor.
-  //
-  InitializeMpExceptionHandlers ();
+    //
+    // Initialize special exception handlers for each logic processor.
+    //
+    InitializeMpExceptionHandlers ();
 
-  //
-  // Update CPU healthy information from Guided HOB
-  //
-  CollectBistDataFromHob ();
-
-  Status = gBS->InstallMultipleProtocolInterfaces (
-                  &mMpServiceHandle,
-                  &gEfiMpServiceProtocolGuid,
-                  &mMpServicesTemplate,
-                  NULL
-                  );
-  ASSERT_EFI_ERROR (Status);
+    //
+    // Update CPU healthy information from Guided HOB
+    //
+    CollectBistDataFromHob ();
+
+    Status = gBS->InstallMultipleProtocolInterfaces (
+                    &mMpServiceHandle,
+                    &gEfiMpServiceProtocolGuid,
+                    &mMpServicesTemplate,
+                    NULL
+                    );
+    ASSERT_EFI_ERROR (Status);
+  }
 }
diff --git 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) {

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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;
 }
 
 /**