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>
This commit is contained in:
Ruiyu Ni
2016-07-14 15:04:38 +08:00
parent ffbc60a027
commit 9168df3dea
2 changed files with 82 additions and 33 deletions

View File

@ -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;