ShellPkg/UefiHandleParsingLib.c: Handle memory allocation failure

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Jaben Carsey <jaben.carsey@intel.com>

diff --git a/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c b/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
index 7665fd5cdd8c57525efdc11b0c33ec77dd0339c4..c9093955317a38eee1df6a6c21f0de0ce0d7ea75 100644 (file)
--- a/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
+++ b/ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
@@ -727,8 +727,9 @@ DriverEfiVersionProtocolDumpInformation(
   ASSERT_EFI_ERROR(Status);\r
 \r
   RetVal = AllocateZeroPool(VersionStringSize);\r
-  ASSERT(RetVal != NULL);\r
-  UnicodeSPrint(RetVal, VersionStringSize, L"0x%08x", DriverEfiVersion->FirmwareVersion);\r
+  if (RetVal != NULL) {\r
+    UnicodeSPrint (RetVal, VersionStringSize, L"0x%08x", DriverEfiVersion->FirmwareVersion);\r
+  }\r
   return (RetVal);\r
 }\r
 /**\r
@@ -2217,10 +2218,11 @@ InternalShellInitHandleList(
   }\r
   for (mHandleList.NextIndex = 1 ; mHandleList.NextIndex <= HandleCount ; mHandleList.NextIndex++){\r
     ListWalker = AllocateZeroPool(sizeof(HANDLE_LIST));\r
-    ASSERT(ListWalker != NULL);\r
-    ListWalker->TheHandle = HandleBuffer[mHandleList.NextIndex-1];\r
-    ListWalker->TheIndex  = mHandleList.NextIndex;\r
-    InsertTailList(&mHandleList.List.Link,&ListWalker->Link);\r
+    if (ListWalker != NULL) {\r
+      ListWalker->TheHandle = HandleBuffer[mHandleList.NextIndex - 1];\r
+      ListWalker->TheIndex = mHandleList.NextIndex;\r
+      InsertTailList (&mHandleList.List.Link, &ListWalker->Link);\r
+    }\r
   }\r
   FreePool(HandleBuffer);\r
   return (EFI_SUCCESS);\r
@@ -2288,7 +2290,9 @@ ConvertHandleToHandleIndex(
   FreePool (ProtocolBuffer);\r
 \r
   ListWalker = AllocateZeroPool(sizeof(HANDLE_LIST));\r
-  ASSERT(ListWalker != NULL);\r
+  if (ListWalker == NULL) {\r
+    return 0;\r
+  }\r
   ListWalker->TheHandle = TheHandle;\r
   ListWalker->TheIndex  = mHandleList.NextIndex++;\r
   InsertTailList(&mHandleList.List.Link,&ListWalker->Link);\r
@@ -2415,7 +2419,10 @@ ParseHandleDatabaseByRelationshipWithType (
   }\r
 \r
   *HandleType = AllocateZeroPool (*HandleCount * sizeof (UINTN));\r
-  ASSERT(*HandleType != NULL);\r
+  if (*HandleType == NULL) {\r
+    SHELL_FREE_NON_NULL (*HandleBuffer);\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
 \r
   DriverBindingHandleIndex = -1;\r
   for (HandleIndex = 0; HandleIndex < *HandleCount; HandleIndex++) {\r
@@ -2671,26 +2678,26 @@ ParseHandleDatabaseByRelationship (
         // Allocate a handle buffer for the number of handles that matched the attributes in Mask\r
         //\r
         *MatchingHandleBuffer = AllocateZeroPool ((*MatchingHandleCount +1)* sizeof (EFI_HANDLE));\r
-        ASSERT(*MatchingHandleBuffer != NULL);\r
+        if (*MatchingHandleBuffer != NULL) {\r
+          for (HandleIndex = 0, *MatchingHandleCount = 0\r
+               ;  HandleIndex < HandleCount\r
+               ;  HandleIndex++\r
+               ) {\r
+            //\r
+            // Fill the allocated buffer with the handles that matched the attributes in Mask\r
+            //\r
+            if ((HandleType[HandleIndex] & Mask) == Mask) {\r
+              (*MatchingHandleBuffer)[(*MatchingHandleCount)++] = HandleBuffer[HandleIndex];\r
+            }\r
+          }\r
 \r
-        for (HandleIndex = 0,*MatchingHandleCount = 0\r
-          ;  HandleIndex < HandleCount\r
-          ;  HandleIndex++\r
-         ){\r
           //\r
-          // Fill the allocated buffer with the handles that matched the attributes in Mask\r
+          // Make the last one NULL\r
           //\r
-          if ((HandleType[HandleIndex] & Mask) == Mask) {\r
-            (*MatchingHandleBuffer)[(*MatchingHandleCount)++] = HandleBuffer[HandleIndex];\r
-          }\r
-        }\r
+          (*MatchingHandleBuffer)[*MatchingHandleCount] = NULL;\r
 \r
-        //\r
-        // Make the last one NULL\r
-        //\r
-        (*MatchingHandleBuffer)[*MatchingHandleCount] = NULL;\r
-\r
-        Status = EFI_SUCCESS;\r
+          Status = EFI_SUCCESS;\r
+        } // *MatchingHandleBuffer != NULL (IF)\r
       } // MacthingHandleBuffer == NULL (ELSE)\r
     } // *MatchingHandleCount  == 0 (ELSE)\r
   } // no error on ParseHandleDatabaseByRelationshipWithType\r