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>

1 files changed, 10 insertions, 2 deletions

diff --git a/ShellPkg/Library/UefiShellLib/UefiShellLib.c b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
index 5e5e6271bf..e4e7e3d829 100644
--- a/ShellPkg/Library/UefiShellLib/UefiShellLib.c
+++ b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
@@ -2021,6 +2021,7 @@ InternalCommandLineParse (
   UINTN                         Count;

   CONST CHAR16                  *TempPointer;

   UINTN                         CurrentValueSize;

+  CHAR16                        *NewValue;

 

   CurrentItemPackage = NULL;

   GetItemValue = 0;

@@ -2119,8 +2120,15 @@ InternalCommandLineParse (
       // get the item VALUE for a previous flag

       //

       CurrentValueSize = ValueSize + StrSize(Argv[LoopCounter]) + sizeof(CHAR16);

-      CurrentItemPackage->Value = ReallocatePool(ValueSize, CurrentValueSize, CurrentItemPackage->Value);

-      ASSERT(CurrentItemPackage->Value != NULL);

+      NewValue = ReallocatePool(ValueSize, CurrentValueSize, CurrentItemPackage->Value);

+      if (NewValue == NULL) {

+        SHELL_FREE_NON_NULL (CurrentItemPackage->Value);

+        SHELL_FREE_NON_NULL (CurrentItemPackage);

+        ShellCommandLineFreeVarList (*CheckPackage);

+        *CheckPackage = NULL;

+        return EFI_OUT_OF_RESOURCES;

+      }

+      CurrentItemPackage->Value = NewValue;

       if (ValueSize == 0) {

         StrCpyS( CurrentItemPackage->Value, 

                   CurrentValueSize/sizeof(CHAR16),