ShellPkg/UefiShellCommandLib.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/UefiShellCommandLib/UefiShellCommandLib.c b/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c
index 48e4d4af0dca77301b278888918be9ed622482bc..35e0611a8e4dcf12f0e249e2f976fa79463a9d12 100644 (file)
--- a/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c
+++ b/ShellPkg/Library/UefiShellCommandLib/UefiShellCommandLib.c
@@ -1,7 +1,7 @@
 /** @file\r
   Provides interface to shell internal functions for shell commands.\r
 \r
-  Copyright (c) 2009 - 2014, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   (C) Copyright 2013-2015 Hewlett-Packard Development Company, L.P.<BR>
   (C) Copyright 2016 Hewlett Packard Enterprise Development LP<BR>
 
@@ -546,9 +546,14 @@ ShellCommandRegisterCommandName (
   // allocate memory for new struct\r
   //\r
   Node = AllocateZeroPool(sizeof(SHELL_COMMAND_INTERNAL_LIST_ENTRY));\r
-  ASSERT(Node != NULL);\r
+  if (Node == NULL) {\r
+    return RETURN_OUT_OF_RESOURCES;\r
+  }\r
   Node->CommandString = AllocateCopyPool(StrSize(CommandString), CommandString);\r
-  ASSERT(Node->CommandString != NULL);\r
+  if (Node->CommandString == NULL) {\r
+    FreePool (Node);\r
+    return RETURN_OUT_OF_RESOURCES;\r
+  }\r
 \r
   Node->GetManFileName  = GetManFileName;\r
   Node->CommandHandler  = CommandHandler;\r
@@ -807,11 +812,20 @@ ShellCommandRegisterAlias (
   // allocate memory for new struct\r
   //\r
   Node = AllocateZeroPool(sizeof(ALIAS_LIST));\r
-  ASSERT(Node != NULL);\r
+  if (Node == NULL) {\r
+    return RETURN_OUT_OF_RESOURCES;\r
+  }\r
   Node->CommandString = AllocateCopyPool(StrSize(Command), Command);\r
+  if (Node->CommandString == NULL) {\r
+    FreePool (Node);\r
+    return RETURN_OUT_OF_RESOURCES;\r
+  }\r
   Node->Alias = AllocateCopyPool(StrSize(Alias), Alias);\r
-  ASSERT(Node->CommandString != NULL);\r
-  ASSERT(Node->Alias != NULL);\r
+  if (Node->Alias == NULL) {\r
+    FreePool (Node->CommandString);\r
+    FreePool (Node);\r
+    return RETURN_OUT_OF_RESOURCES;\r
+  }\r
 \r
   InsertHeadList (&mAliasList.Link, &Node->Link);\r
 \r
@@ -1303,7 +1317,10 @@ ShellCommandCreateInitialMappingsAndPaths(
     // Get all Device Paths\r
     //\r
     DevicePathList = AllocateZeroPool(sizeof(EFI_DEVICE_PATH_PROTOCOL*) * Count);\r
-    ASSERT(DevicePathList != NULL);\r
+    if (DevicePathList == NULL) {\r
+      SHELL_FREE_NON_NULL (HandleList);\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
 \r
     for (Count = 0 ; HandleList[Count] != NULL ; Count++) {\r
       DevicePathList[Count] = DevicePathFromHandle(HandleList[Count]);\r
@@ -1360,7 +1377,10 @@ ShellCommandCreateInitialMappingsAndPaths(
     // Get all Device Paths\r
     //\r
     DevicePathList = AllocateZeroPool(sizeof(EFI_DEVICE_PATH_PROTOCOL*) * Count);\r
-    ASSERT(DevicePathList != NULL);\r
+    if (DevicePathList == NULL) {\r
+      SHELL_FREE_NON_NULL (HandleList);\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
 \r
     for (Count = 0 ; HandleList[Count] != NULL ; Count++) {\r
       DevicePathList[Count] = DevicePathFromHandle(HandleList[Count]);\r