From c84ccd98880b8561b357e8524630cfb00bdfce4b Mon Sep 17 00:00:00 2001
From: Ruiyu Ni <ruiyu.ni@intel.com>
Date: Tue, 12 Jul 2016 13:51:54 +0800
Subject: 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>
(cherry picked from commit aa3276c171e3bd2f29745c4b48589f0faded4322)
---
 .../UefiHandleParsingLib/UefiHandleParsingLib.c    | 55 ++++++++++++----------
 1 file changed, 31 insertions(+), 24 deletions(-)

(limited to 'ShellPkg/Library')

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