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