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>

1 files changed, 68 insertions, 30 deletions

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index a95da000f5..c3331def82 100644
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/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 (ShellFileList == NULL) {

-      ShellFileList = (EFI_SHELL_FILE_INFO*)AllocateZeroPool(sizeof(EFI_SHELL_FILE_INFO));

-      ASSERT(ShellFileList != NULL);

-      InitializeListHead(&ShellFileList->Link);

+    if (ShellFileListItem == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+      //

+      // Free resources outside the loop.

+      //

+      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)

-        ){

-      RemoveEntryList(&Node2->Link);

-      Status = gBS->ReinstallProtocolInterface(Node2->Handle,

-                                               &gEfiShellProtocolGuid,

-                                               NewShell,

-                                               Node2->Interface);

-      FreePool(Node2);

+         ; Node2 = (SHELL_PROTOCOL_HANDLE_LIST *) GetFirstNode (&ShellInfoObject.OldShellList.Link)

+         ) {

+      RemoveEntryList (&Node2->Link);

+      gBS->ReinstallProtocolInterface (Node2->Handle, &gEfiShellProtocolGuid, NewShell, Node2->Interface);

+      FreePool (Node2);

     }

   } else {

     //

     // no need to restore

     //

-    Status = gBS->UninstallProtocolInterface(gImageHandle,

-                                             &gEfiShellProtocolGuid,

-                                             NewShell);

+    gBS->UninstallProtocolInterface (gImageHandle, &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;