ShellPkg/IfConfig: Handle memory allocation failure

[mirror_edk2.git] / ShellPkg / Library / UefiShellLib / UefiShellLib.c

diff --git a/ShellPkg/Library/UefiShellLib/UefiShellLib.c b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
index dac0524fcbdd98652f4fb43368290b1bada372e4..e4e7e3d8296dab184cee91d114a6131253183f8f 100644 (file)
--- a/ShellPkg/Library/UefiShellLib/UefiShellLib.c
+++ b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
@@ -2,7 +2,7 @@
   Provides interface to shell functionality for shell commands and applications.\r
 \r
   Copyright 2016 Dell Inc.\r
-  Copyright (c) 2006 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
@@ -373,6 +373,8 @@ EFIAPI
 ShellInitialize (\r
   )\r
 {\r
+  EFI_STATUS Status;\r
+\r
   //\r
   // if auto initialize is not false then skip\r
   //\r
@@ -383,7 +385,8 @@ ShellInitialize (
   //\r
   // deinit the current stuff\r
   //\r
-  ASSERT_EFI_ERROR(ShellLibDestructor(gImageHandle, gST));\r
+  Status = ShellLibDestructor (gImageHandle, gST);\r
+  ASSERT_EFI_ERROR (Status);\r
 \r
   //\r
   // init the new stuff\r
@@ -2018,6 +2021,7 @@ InternalCommandLineParse (
   UINTN                         Count;\r
   CONST CHAR16                  *TempPointer;\r
   UINTN                         CurrentValueSize;\r
+  CHAR16                        *NewValue;\r
 \r
   CurrentItemPackage = NULL;\r
   GetItemValue = 0;\r
@@ -2116,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
@@ -3283,7 +3294,7 @@ StrnCatGrow (
   Prompt the user and return the resultant answer to the requestor.\r
 \r
   This function will display the requested question on the shell prompt and then\r
-  wait for an apropriate answer to be input from the console.\r
+  wait for an appropriate answer to be input from the console.\r
 \r
   if the SHELL_PROMPT_REQUEST_TYPE is SHELL_PROMPT_REQUEST_TYPE_YESNO, ShellPromptResponseTypeQuitContinue\r
   or SHELL_PROMPT_REQUEST_TYPE_YESNOCANCEL then *Response is of type SHELL_PROMPT_RESPONSE.\r
@@ -4071,6 +4082,9 @@ ShellFileHandleReturnLine(
     Status = ShellFileHandleReadLine(Handle, RetVal, &Size, FALSE, Ascii);\r
 \r
   }\r
+  if (Status == EFI_END_OF_FILE && RetVal != NULL && *RetVal != CHAR_NULL) {\r
+    Status = EFI_SUCCESS;\r
+  }\r
   if (EFI_ERROR(Status) && (RetVal != NULL)) {\r
     FreePool(RetVal);\r
     RetVal = NULL;\r
@@ -4176,43 +4190,27 @@ ShellFileHandleReadLine(
     //\r
     // if we have space save it...\r
     //\r
-    if ((CountSoFar + 1) * CharSize < *Size){\r
+    if ((CountSoFar+1)*sizeof(CHAR16) < *Size){\r
       ASSERT(Buffer != NULL);\r
-      if (*Ascii) {\r
-        ((CHAR8*)Buffer)[CountSoFar] = (CHAR8) CharBuffer;\r
-        ((CHAR8*)Buffer)[CountSoFar+1] = '\0';\r
-      }\r
-      else {\r
-        ((CHAR16*)Buffer)[CountSoFar] = CharBuffer;\r
-        ((CHAR16*)Buffer)[CountSoFar+1] = CHAR_NULL;\r
-      }\r
+      ((CHAR16*)Buffer)[CountSoFar] = CharBuffer;\r
+      ((CHAR16*)Buffer)[CountSoFar+1] = CHAR_NULL;\r
     }\r
   }\r
 \r
   //\r
   // if we ran out of space tell when...\r
   //\r
-  if (Status != EFI_END_OF_FILE){\r
-    if ((CountSoFar + 1) * CharSize > *Size){\r
-      *Size = (CountSoFar + 1) * CharSize;\r
-      if (!Truncate) {\r
-        gEfiShellProtocol->SetFilePosition(Handle, OriginalFilePosition);\r
-      } else {\r
-        DEBUG((DEBUG_WARN, "The line was truncated in ShellFileHandleReadLine"));\r
-      }\r
-      return (EFI_BUFFER_TOO_SMALL);\r
-    }\r
-\r
-    if (*Ascii) {\r
-      if (CountSoFar && ((CHAR8*)Buffer)[CountSoFar - 1] == '\r') {\r
-        ((CHAR8*)Buffer)[CountSoFar - 1] = '\0';\r
-      }\r
-    }\r
-    else {\r
-      if (CountSoFar && Buffer[CountSoFar - 1] == L'\r') {\r
-        Buffer[CountSoFar - 1] = CHAR_NULL;\r
-      }\r
+  if ((CountSoFar+1)*sizeof(CHAR16) > *Size){\r
+    *Size = (CountSoFar+1)*sizeof(CHAR16);\r
+    if (!Truncate) {\r
+      gEfiShellProtocol->SetFilePosition(Handle, OriginalFilePosition);\r
+    } else {\r
+      DEBUG((DEBUG_WARN, "The line was truncated in ShellFileHandleReadLine"));\r
     }\r
+    return (EFI_BUFFER_TOO_SMALL);\r
+  }\r
+  while(Buffer[StrLen(Buffer)-1] == L'\r') {\r
+    Buffer[StrLen(Buffer)-1] = CHAR_NULL;\r
   }\r
 \r
   return (Status);\r