ShellPkg/UefiShellLib: 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/UefiShellLib/UefiShellLib.c b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
index 5e5e6271bf24d418aaa7f04ff29925aee0be1754..e4e7e3d8296dab184cee91d114a6131253183f8f 100644 (file)
--- a/ShellPkg/Library/UefiShellLib/UefiShellLib.c
+++ b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
@@ -2021,6 +2021,7 @@ InternalCommandLineParse (
   UINTN                         Count;\r
   CONST CHAR16                  *TempPointer;\r
   UINTN                         CurrentValueSize;\r
+  CHAR16                        *NewValue;\r
 \r
   CurrentItemPackage = NULL;\r
   GetItemValue = 0;\r
@@ -2119,8 +2120,15 @@ InternalCommandLineParse (
       // get the item VALUE for a previous flag\r
       //\r
       CurrentValueSize = ValueSize + StrSize(Argv[LoopCounter]) + sizeof(CHAR16);\r
-      CurrentItemPackage->Value = ReallocatePool(ValueSize, CurrentValueSize, CurrentItemPackage->Value);\r
-      ASSERT(CurrentItemPackage->Value != NULL);\r
+      NewValue = ReallocatePool(ValueSize, CurrentValueSize, CurrentItemPackage->Value);\r
+      if (NewValue == NULL) {\r
+        SHELL_FREE_NON_NULL (CurrentItemPackage->Value);\r
+        SHELL_FREE_NON_NULL (CurrentItemPackage);\r
+        ShellCommandLineFreeVarList (*CheckPackage);\r
+        *CheckPackage = NULL;\r
+        return EFI_OUT_OF_RESOURCES;\r
+      }\r
+      CurrentItemPackage->Value = NewValue;\r
       if (ValueSize == 0) {\r
         StrCpyS( CurrentItemPackage->Value, \r
                   CurrentValueSize/sizeof(CHAR16), \r