From 9168df3dea65f707d1e9c32eba5e18ef6b84e5cd Mon Sep 17 00:00:00 2001 From: Ruiyu Ni Date: Thu, 14 Jul 2016 15:04:38 +0800 Subject: ShellPkg/ShellProtocol.c: Handle memory allocation failure Contributed-under: TianoCore Contribution Agreement 1.0 Signed-off-by: Ruiyu Ni Reviewed-by: Jaben Carsey --- ShellPkg/Application/Shell/ShellProtocol.c | 98 +++++++++++++++++++++--------- 1 file changed, 68 insertions(+), 30 deletions(-) (limited to 'ShellPkg/Application/Shell/ShellProtocol.c') 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; -- cgit v1.2.3