ShellPkg/ShellProtocol.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>

ShellPkg/Application/Shell/ShellProtocol.c

@@ -466,7 +466,10 @@ EfiShellGetFilePathFromDevicePath(
         ASSERT((PathForReturn == NULL && PathSize == 0) || (PathForReturn != NULL));
 
         AlignedNode = AllocateCopyPool (DevicePathNodeLength(FilePath), FilePath);
-        ASSERT (AlignedNode != NULL);
+        if (AlignedNode == NULL) {
+          FreePool (PathForReturn);
+          return NULL;
+        }
 
         // File Path Device Path Nodes 'can optionally add a "\" separator to
         //  the beginning and/or the end of the Path Name string.'
@@ -2142,6 +2145,14 @@ EfiShellFindFilesInDir(
       ; !EFI_ERROR(Status) && !NoFile
       ; Status = FileHandleFindNextFile(FileDirHandle, FileInfo, &NoFile)
      ){
+    if (ShellFileList == NULL) {
+      ShellFileList = (EFI_SHELL_FILE_INFO*)AllocateZeroPool(sizeof(EFI_SHELL_FILE_INFO));
+      if (ShellFileList == NULL) {
+        SHELL_FREE_NON_NULL (BasePath);
+        return EFI_OUT_OF_RESOURCES;
+      }
+      InitializeListHead(&ShellFileList->Link);
+    }
     //
     // allocate a new EFI_SHELL_FILE_INFO and populate it...
     //
@@ -2151,11 +2162,12 @@ EfiShellFindFilesInDir(
       FileInfo->FileName,
       NULL,         // no handle since not open
       FileInfo);
-
+    if (ShellFileListItem == NULL) {
-    if (ShellFileList == NULL) {
+      Status = EFI_OUT_OF_RESOURCES;
-      ShellFileList = (EFI_SHELL_FILE_INFO*)AllocateZeroPool(sizeof(EFI_SHELL_FILE_INFO));
+      //
-      ASSERT(ShellFileList != NULL);
+      // Free resources outside the loop.
-      InitializeListHead(&ShellFileList->Link);
+      //
+      break;
     }
     InsertTailList(&ShellFileList->Link, &ShellFileListItem->Link);
   }
@@ -2360,7 +2372,10 @@ ShellSearchHandle(
      ; NextFilePatternStart++);
 
   CurrentFilePattern = AllocateZeroPool((NextFilePatternStart-FilePattern+1)*sizeof(CHAR16));
-  ASSERT(CurrentFilePattern != NULL);
+  if (CurrentFilePattern == NULL) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
   StrnCpyS(CurrentFilePattern, NextFilePatternStart-FilePattern+1, FilePattern, NextFilePatternStart-FilePattern);
 
   if (CurrentFilePattern[0]   == CHAR_NULL
@@ -2469,7 +2484,6 @@ ShellSearchHandle(
             // copy the information we need into a new Node
             //
             NewShellNode = InternalDuplicateShellFileInfo(ShellInfoNode, FALSE);
-            ASSERT(NewShellNode != NULL);
             if (NewShellNode == NULL) {
               Status = EFI_OUT_OF_RESOURCES;
             }
@@ -3216,7 +3230,9 @@ EfiShellGetHelpText(
     && (Command[StrLen(Command)-4] == L'.')
     ) {
       FixCommand = AllocateZeroPool(StrSize(Command) - 4 * sizeof (CHAR16));
-      ASSERT(FixCommand != NULL);
+      if (FixCommand == NULL) {
+        return EFI_OUT_OF_RESOURCES;
+      }
 
       StrnCpyS( FixCommand, 
                 (StrSize(Command) - 4 * sizeof (CHAR16))/sizeof(CHAR16), 
@@ -3395,7 +3411,9 @@ EfiShellGetAlias(
   // Convert to lowercase to make aliases case-insensitive
   if (Alias != NULL) {
     AliasLower = AllocateCopyPool (StrSize (Alias), Alias);
-    ASSERT (AliasLower != NULL);
+    if (AliasLower == NULL) {
+      return NULL;
+    }
     ToLower (AliasLower);
 
     if (Volatile == NULL) {
@@ -3459,7 +3477,9 @@ InternalSetAlias(
   // Convert to lowercase to make aliases case-insensitive
   if (Alias != NULL) {
     AliasLower = AllocateCopyPool (StrSize (Alias), Alias);
-    ASSERT (AliasLower != NULL);
+    if (AliasLower == NULL) {
+      return EFI_OUT_OF_RESOURCES;
+    }
     ToLower (AliasLower);
   } else {
     AliasLower = NULL;
@@ -3618,6 +3638,7 @@ CreatePopulateInstallShellProtocol (
   EFI_HANDLE                  *Buffer;
   UINTN                       HandleCounter;
   SHELL_PROTOCOL_HANDLE_LIST  *OldProtocolNode;
+  EFI_SHELL_PROTOCOL          *OldShell;
 
   if (NewShell == NULL) {
     return (EFI_INVALID_PARAMETER);
@@ -3669,20 +3690,25 @@ CreatePopulateInstallShellProtocol (
     // now overwrite each of them, but save the info to restore when we end.
     //
     for (HandleCounter = 0 ; HandleCounter < (BufferSize/sizeof(EFI_HANDLE)) ; HandleCounter++) {
-      OldProtocolNode = AllocateZeroPool(sizeof(SHELL_PROTOCOL_HANDLE_LIST));
-      ASSERT(OldProtocolNode != NULL);
       Status = gBS->OpenProtocol(Buffer[HandleCounter],
                                 &gEfiShellProtocolGuid,
-                                (VOID **) &(OldProtocolNode->Interface),
+                                (VOID **) &OldShell,
                                 gImageHandle,
                                 NULL,
                                 EFI_OPEN_PROTOCOL_GET_PROTOCOL
                                );
       if (!EFI_ERROR(Status)) {
+        OldProtocolNode = AllocateZeroPool(sizeof(SHELL_PROTOCOL_HANDLE_LIST));
+        if (OldProtocolNode == NULL && !IsListEmpty (&ShellInfoObject.OldShellList.Link)) {
+          CleanUpShellProtocol (&mShellProtocol);
+          Status = EFI_OUT_OF_RESOURCES;
+          break;
+        }
         //
         // reinstall over the old one...
         //
-        OldProtocolNode->Handle = Buffer[HandleCounter];
+        OldProtocolNode->Handle    = Buffer[HandleCounter];
+        OldProtocolNode->Interface = OldShell;
         Status = gBS->ReinstallProtocolInterface(
                             OldProtocolNode->Handle,
                             &gEfiShellProtocolGuid,
@@ -3735,38 +3761,50 @@ CreatePopulateInstallShellProtocol (
   @retval EFI_SUCCESS       The operation was successful.
 **/
 EFI_STATUS
-EFIAPI
 CleanUpShellProtocol (
   IN OUT EFI_SHELL_PROTOCOL  *NewShell
   )
 {
-  EFI_STATUS                        Status;
   SHELL_PROTOCOL_HANDLE_LIST        *Node2;
-  EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *SimpleEx;
 
   //
   // if we need to restore old protocols...
   //
   if (!IsListEmpty(&ShellInfoObject.OldShellList.Link)) {
-    for (Node2 = (SHELL_PROTOCOL_HANDLE_LIST *)GetFirstNode(&ShellInfoObject.OldShellList.Link)
+    for (Node2 = (SHELL_PROTOCOL_HANDLE_LIST *) GetFirstNode (&ShellInfoObject.OldShellList.Link)
          ; !IsListEmpty (&ShellInfoObject.OldShellList.Link)
-         ; Node2 = (SHELL_PROTOCOL_HANDLE_LIST *)GetFirstNode(&ShellInfoObject.OldShellList.Link)
+         ; Node2 = (SHELL_PROTOCOL_HANDLE_LIST *) GetFirstNode (&ShellInfoObject.OldShellList.Link)
-        ){
+         ) {
-      RemoveEntryList(&Node2->Link);
+      RemoveEntryList (&Node2->Link);
-      Status = gBS->ReinstallProtocolInterface(Node2->Handle,
+      gBS->ReinstallProtocolInterface (Node2->Handle, &gEfiShellProtocolGuid, NewShell, Node2->Interface);
-                                               &gEfiShellProtocolGuid,
+      FreePool (Node2);
-                                               NewShell,
-                                               Node2->Interface);
-      FreePool(Node2);
     }
   } else {
     //
     // no need to restore
     //
-    Status = gBS->UninstallProtocolInterface(gImageHandle,
+    gBS->UninstallProtocolInterface (gImageHandle, &gEfiShellProtocolGuid, NewShell);
-                                             &gEfiShellProtocolGuid,
-                                             NewShell);
   }
+  return EFI_SUCCESS;
+}
+
+/**
+  Cleanup the shell environment.
+
+  @param[in, out] NewShell   The pointer to the new shell protocol structure.
+
+  @retval EFI_SUCCESS       The operation was successful.
+**/
+EFI_STATUS
+CleanUpShellEnvironment (
+  IN OUT EFI_SHELL_PROTOCOL  *NewShell
+  )
+{
+  EFI_STATUS                        Status;
+  EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *SimpleEx;
+  
+  CleanUpShellProtocol (NewShell);
+
   Status = gBS->CloseEvent(NewShell->ExecutionBreak);
   NewShell->ExecutionBreak = NULL;
 

ShellPkg/Application/Shell/ShellProtocol.h

@@ -3,7 +3,7 @@
   manipulation, and initialization of EFI_SHELL_PROTOCOL.
 
   (C) Copyright 2014 Hewlett-Packard Development Company, L.P.<BR>
-  Copyright (c) 2009 - 2014, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>
   This program and the accompanying materials
   are licensed and made available under the terms and conditions of the BSD License
   which accompanies this distribution.  The full text of the license may be found at
@@ -59,11 +59,22 @@ CreatePopulateInstallShellProtocol (
   @retval EFI_SUCCESS       The operation was successful.
 **/
 EFI_STATUS
-EFIAPI
 CleanUpShellProtocol (
   IN OUT EFI_SHELL_PROTOCOL  *NewShell
   );
 
+/**
+  Cleanup the shell environment.
+
+  @param[in, out] NewShell   The pointer to the new shell protocol structure.
+
+  @retval EFI_SUCCESS       The operation was successful.
+**/
+EFI_STATUS
+CleanUpShellEnvironment (
+  IN OUT EFI_SHELL_PROTOCOL  *NewShell
+  );
+
 /**
   This function creates a mapping for a device path.