ShellPkg/Shell.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/Application/Shell/Shell.c b/ShellPkg/Application/Shell/Shell.c
index 54ca76a09f85591d9836b71c4119248ec6c0f943..713c317d7411bd5baede23f6368dd4f81abc7825 100644 (file)
--- a/ShellPkg/Application/Shell/Shell.c
+++ b/ShellPkg/Application/Shell/Shell.c
@@ -1319,7 +1319,7 @@ DoShellPrompt (
 **/\r
 VOID*\r
 EFIAPI\r
-AddBufferToFreeList(\r
+AddBufferToFreeList (\r
   VOID *Buffer\r
   )\r
 {\r
@@ -1329,10 +1329,13 @@ AddBufferToFreeList(
     return (NULL);\r
   }\r
 \r
-  BufferListEntry = AllocateZeroPool(sizeof(BUFFER_LIST));\r
-  ASSERT(BufferListEntry != NULL);\r
+  BufferListEntry = AllocateZeroPool (sizeof (BUFFER_LIST));\r
+  if (BufferListEntry == NULL) {\r
+    return NULL;\r
+  }\r
+\r
   BufferListEntry->Buffer = Buffer;\r
-  InsertTailList(&ShellInfoObject.BufferToFreeList.Link, &BufferListEntry->Link);\r
+  InsertTailList (&ShellInfoObject.BufferToFreeList.Link, &BufferListEntry->Link);\r
   return (Buffer);\r
 }\r
 \r
@@ -1391,9 +1394,15 @@ AddLineToCommandHistory(
 \r
 \r
   Node = AllocateZeroPool(sizeof(BUFFER_LIST));\r
-  ASSERT(Node != NULL);\r
-  Node->Buffer = AllocateCopyPool(StrSize(Buffer), Buffer);\r
-  ASSERT(Node->Buffer != NULL);\r
+  if (Node == NULL) {\r
+    return;\r
+  }\r
+\r
+  Node->Buffer = AllocateCopyPool (StrSize (Buffer), Buffer);\r
+  if (Node->Buffer == NULL) {\r
+    FreePool (Node);\r
+    return;\r
+  }\r
 \r
   for ( Walker = (BUFFER_LIST*)GetFirstNode(&ShellInfoObject.ViewingSettings.CommandHistory.Link)\r
       ; !IsNull(&ShellInfoObject.ViewingSettings.CommandHistory.Link, &Walker->Link)\r
@@ -1721,7 +1730,9 @@ RunSplitCommand(
   // make a SPLIT_LIST item and add to list\r
   //\r
   Split = AllocateZeroPool(sizeof(SPLIT_LIST));\r
-  ASSERT(Split != NULL);\r
+  if (Split == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
   Split->SplitStdIn   = StdIn;\r
   Split->SplitStdOut  = ConvertEfiFileProtocolToShellHandle(CreateFileInterfaceMem(Unicode), NULL);\r
   ASSERT(Split->SplitStdOut != NULL);\r