From c80c222198a88847f3e1ce3e482af8e4c0bd36fa Mon Sep 17 00:00:00 2001
From: Oliver Smith-Denny <osde@microsoft.com>
Date: Thu, 3 Oct 2024 10:31:12 -0700
Subject: ShellPkg: UefiShellCommandLib: CodeQL Fixes

Includes changes across the module for the following CodeQL rules:
 - cpp/comparison-with-wider-type
 - cpp/overflow-buffer
 - cpp/redundant-null-check-param
 - cpp/uselesstest

Co-authored-by: Taylor Beebe <taylor.d.beebe@gmail.com>

Signed-off-by: Oliver Smith-Denny <osde@linux.microsoft.com>
---
 .../Library/UefiShellCommandLib/ConsistMapping.c   |  4 ++
 .../UefiShellCommandLib/UefiShellCommandLib.c      | 67 +++++++++++++++-------
 ShellPkg/Library/UefiShellLevel1CommandsLib/For.c  |  2 +
 3 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/ShellPkg/Library/UefiShellCommandLib/ConsistMapping.c b/ShellPkg/Library/UefiShellCommandLib/ConsistMapping.c
index 1767a7dc4a..9ef6a0192f 100755
--- a/ShellPkg/Library/UefiShellCommandLib/ConsistMapping.c
+++ b/ShellPkg/Library/UefiShellCommandLib/ConsistMapping.c
@@ -1404,6 +1404,10 @@ GetHIDevicePath (
   NonHIDevicePathNodeCount = 0;
 
   HIDevicePath = AllocateZeroPool (sizeof (EFI_DEVICE_PATH_PROTOCOL));
+  if (HIDevicePath == NULL) {
+    return NULL;
+  }
+
   SetDevicePathEndNode (HIDevicePath);
 
   Node.DevPath.Type      = END_DEVICE_PATH_TYPE;
diff --git a/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c b/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c
index 75fbd50cbe..b0c77e4e80 100644
--- a/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c
+++ b/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c
@@ -1142,7 +1142,7 @@ DeleteScriptFileStruct (
   IN SCRIPT_FILE  *Script
   )
 {
-  UINT8  LoopVar;
+  UINTN  LoopVar;
 
   if (Script == NULL) {
     return;
@@ -1263,6 +1263,10 @@ ShellCommandCreateNewMappingName (
   String = NULL;
 
   String = AllocateZeroPool (PcdGet8 (PcdShellMapNameLength) * sizeof (String[0]));
+  if (String == NULL) {
+    return (NULL);
+  }
+
   UnicodeSPrint (
     String,
     PcdGet8 (PcdShellMapNameLength) * sizeof (String[0]),
@@ -1459,29 +1463,39 @@ ShellCommandCreateInitialMappingsAndPaths (
     //
     PerformQuickSort (DevicePathList, Count, sizeof (EFI_DEVICE_PATH_PROTOCOL *), DevicePathCompare);
 
-    if (!EFI_ERROR (ShellCommandConsistMappingInitialize (&ConsistMappingTable))) {
+    Status = ShellCommandConsistMappingInitialize (&ConsistMappingTable);
+    if (EFI_ERROR (Status)) {
+      SHELL_FREE_NON_NULL (HandleList);
+      SHELL_FREE_NON_NULL (DevicePathList);
+      return Status;
+    }
+
+    //
+    // Assign new Mappings to all...
+    //
+    for (Count = 0; HandleList[Count] != NULL; Count++) {
       //
-      // Assign new Mappings to all...
+      // Get default name first
       //
-      for (Count = 0; HandleList[Count] != NULL; Count++) {
-        //
-        // Get default name first
-        //
-        NewDefaultName = ShellCommandCreateNewMappingName (MappingTypeFileSystem);
+      NewDefaultName = ShellCommandCreateNewMappingName (MappingTypeFileSystem);
+      if (NewDefaultName == NULL) {
         ASSERT (NewDefaultName != NULL);
-        Status = ShellCommandAddMapItemAndUpdatePath (NewDefaultName, DevicePathList[Count], 0, TRUE);
-        ASSERT_EFI_ERROR (Status);
-        FreePool (NewDefaultName);
+        Status = EFI_OUT_OF_RESOURCES;
+        break;
+      }
 
-        //
-        // Now do consistent name
-        //
-        NewConsistName = ShellCommandConsistMappingGenMappingName (DevicePathList[Count], ConsistMappingTable);
-        if (NewConsistName != NULL) {
-          Status = ShellCommandAddMapItemAndUpdatePath (NewConsistName, DevicePathList[Count], 0, FALSE);
-          ASSERT_EFI_ERROR (Status);
-          FreePool (NewConsistName);
-        }
+      Status = ShellCommandAddMapItemAndUpdatePath (NewDefaultName, DevicePathList[Count], 0, TRUE);
+      ASSERT_EFI_ERROR (Status);
+      FreePool (NewDefaultName);
+
+      //
+      // Now do consistent name
+      //
+      NewConsistName = ShellCommandConsistMappingGenMappingName (DevicePathList[Count], ConsistMappingTable);
+      if (NewConsistName != NULL) {
+        Status = ShellCommandAddMapItemAndUpdatePath (NewConsistName, DevicePathList[Count], 0, FALSE);
+        ASSERT_EFI_ERROR (Status);
+        FreePool (NewConsistName);
       }
     }
 
@@ -1561,7 +1575,13 @@ ShellCommandCreateInitialMappingsAndPaths (
       // Get default name first
       //
       NewDefaultName = ShellCommandCreateNewMappingName (MappingTypeBlockIo);
-      ASSERT (NewDefaultName != NULL);
+      if (NewDefaultName == NULL) {
+        ASSERT (NewDefaultName != NULL);
+        SHELL_FREE_NON_NULL (HandleList);
+        SHELL_FREE_NON_NULL (DevicePathList);
+        return EFI_OUT_OF_RESOURCES;
+      }
+
       Status = ShellCommandAddMapItemAndUpdatePath (NewDefaultName, DevicePathList[Count], 0, FALSE);
       ASSERT_EFI_ERROR (Status);
       FreePool (NewDefaultName);
@@ -1631,6 +1651,11 @@ ShellCommandUpdateMapping (
     PerformQuickSort (DevicePathList, Count, sizeof (EFI_DEVICE_PATH_PROTOCOL *), DevicePathCompare);
 
     Status = ShellCommandConsistMappingInitialize (&ConsistMappingTable);
+    if (EFI_ERROR (Status)) {
+      SHELL_FREE_NON_NULL (HandleList);
+      SHELL_FREE_NON_NULL (DevicePathList);
+      return Status;
+    }
 
     //
     // Assign new Mappings to remainders
diff --git a/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c b/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
index 40bb3d59b5..9cd45c2f4c 100644
--- a/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
+++ b/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
@@ -691,6 +691,8 @@ ShellCommandRunFor (
 
       TempString = AllocateZeroPool (50*sizeof (CHAR16));
       if (TempString == NULL) {
+        SHELL_FREE_NON_NULL (ArgSet);
+        SHELL_FREE_NON_NULL (Info);
         return (SHELL_OUT_OF_RESOURCES);
       }
 
-- 
cgit