ShellPkg/IsVolatileEnv: 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-08 15:18:14 +08:00
parent b2c036a7f0
commit 31e5b912b9
5 changed files with 56 additions and 31 deletions

View File

@ -971,6 +971,7 @@ FileInterfaceEnvClose(
VOID* NewBuffer;
UINTN NewSize;
EFI_STATUS Status;
BOOLEAN Volatile;
//
// Most if not all UEFI commands will have an '\r\n' at the end of any output.
@ -981,6 +982,11 @@ FileInterfaceEnvClose(
NewBuffer = NULL;
NewSize = 0;
Status = IsVolatileEnv (((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &Volatile);
if (EFI_ERROR (Status)) {
return Status;
}
Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);
if (Status == EFI_BUFFER_TOO_SMALL) {
NewBuffer = AllocateZeroPool(NewSize + sizeof(CHAR16));
@ -998,8 +1004,8 @@ FileInterfaceEnvClose(
&& (((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 3] == CHAR_CARRIAGE_RETURN)) {
((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 3] = CHAR_NULL;
}
if (IsVolatileEnv(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name)) {
if (Volatile) {
Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);
} else {
Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);
@ -1157,13 +1163,20 @@ CreateFileInterfaceEnv(
IN CONST CHAR16 *EnvName
)
{
EFI_STATUS Status;
EFI_FILE_PROTOCOL_ENVIRONMENT *EnvFileInterface;
UINTN EnvNameSize;
BOOLEAN Volatile;
if (EnvName == NULL) {
return (NULL);
}
Status = IsVolatileEnv (EnvName, &Volatile);
if (EFI_ERROR (Status)) {
return NULL;
}
//
// Get some memory
//
@ -1192,7 +1205,7 @@ CreateFileInterfaceEnv(
//
// Assign the different members for Volatile and Non-Volatile variables
//
if (IsVolatileEnv(EnvName)) {
if (Volatile) {
EnvFileInterface->Write = FileInterfaceEnvVolWrite;
} else {
EnvFileInterface->Write = FileInterfaceEnvNonVolWrite;

View File

@ -26,14 +26,15 @@ ENV_VAR_LIST gShellEnvVarList;
Reports whether an environment variable is Volatile or Non-Volatile.
@param EnvVarName The name of the environment variable in question
@param Volatile Return TRUE if the environment variable is volatile
@retval TRUE This environment variable is Volatile
@retval FALSE This environment variable is NON-Volatile
@retval EFI_SUCCESS The volatile attribute is returned successfully
@retval others Some errors happened.
**/
BOOLEAN
EFIAPI
EFI_STATUS
IsVolatileEnv (
IN CONST CHAR16 *EnvVarName
IN CONST CHAR16 *EnvVarName,
OUT BOOLEAN *Volatile
)
{
EFI_STATUS Status;
@ -41,6 +42,8 @@ IsVolatileEnv (
VOID *Buffer;
UINT32 Attribs;
ASSERT (Volatile != NULL);
Size = 0;
Buffer = NULL;
@ -54,7 +57,9 @@ IsVolatileEnv (
Buffer);
if (Status == EFI_BUFFER_TOO_SMALL) {
Buffer = AllocateZeroPool(Size);
ASSERT(Buffer != NULL);
if (Buffer == NULL) {
return EFI_OUT_OF_RESOURCES;
}
Status = gRT->GetVariable((CHAR16*)EnvVarName,
&gShellVariableGuid,
&Attribs,
@ -66,21 +71,18 @@ IsVolatileEnv (
// not found means volatile
//
if (Status == EFI_NOT_FOUND) {
return (TRUE);
*Volatile = TRUE;
return EFI_SUCCESS;
}
if (EFI_ERROR (Status)) {
return Status;
}
ASSERT_EFI_ERROR(Status);
//
// check for the Non Volatile bit
//
if ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE) {
return (FALSE);
}
//
// everything else is volatile
//
return (TRUE);
*Volatile = !(BOOLEAN) ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE);
return EFI_SUCCESS;
}
/**

View File

@ -34,19 +34,18 @@ extern ENV_VAR_LIST gShellEnvVarList;
/**
Reports whether an environment variable is Volatile or Non-Volatile
This will use the Runtime Services call GetVariable to to search for the variable.
Reports whether an environment variable is Volatile or Non-Volatile.
@param EnvVarName The name of the environment variable in question
@param Volatile Return TRUE if the environment variable is volatile
@retval TRUE This environment variable is Volatile
@retval FALSE This environment variable is NON-Volatile
@retval EFI_SUCCESS The volatile attribute is returned successfully
@retval others Some errors happened.
**/
BOOLEAN
EFIAPI
EFI_STATUS
IsVolatileEnv (
IN CONST CHAR16 *EnvVarName
IN CONST CHAR16 *EnvVarName,
OUT BOOLEAN *Volatile
);
/**

View File

@ -736,6 +736,7 @@ UpdateStdInStdOutStdErr(
UINTN Size;
SPLIT_LIST *Split;
CHAR16 *FirstLocation;
BOOLEAN Volatile;
OutUnicode = TRUE;
InUnicode = TRUE;
@ -1111,8 +1112,8 @@ UpdateStdInStdOutStdErr(
//
// Check for no volatile environment variables
//
||(StdErrVarName != NULL && !IsVolatileEnv(StdErrVarName))
||(StdOutVarName != NULL && !IsVolatileEnv(StdOutVarName))
||(StdErrVarName != NULL && !EFI_ERROR (IsVolatileEnv (StdErrVarName, &Volatile)) && !Volatile)
||(StdOutVarName != NULL && !EFI_ERROR (IsVolatileEnv (StdOutVarName, &Volatile)) && !Volatile)
//
// Cant redirect during a reconnect operation.
//

View File

@ -1129,13 +1129,18 @@ EfiShellCreateFile(
{
EFI_DEVICE_PATH_PROTOCOL *DevicePath;
EFI_STATUS Status;
BOOLEAN Volatile;
//
// Is this for an environment variable
// do we start with >v
//
if (StrStr(FileName, L">v") == FileName) {
if (!IsVolatileEnv(FileName+2)) {
Status = IsVolatileEnv (FileName + 2, &Volatile);
if (EFI_ERROR (Status)) {
return Status;
}
if (!Volatile) {
return (EFI_INVALID_PARAMETER);
}
*FileHandle = CreateFileInterfaceEnv(FileName+2);
@ -1245,6 +1250,7 @@ EfiShellOpenFileByName(
{
EFI_DEVICE_PATH_PROTOCOL *DevicePath;
EFI_STATUS Status;
BOOLEAN Volatile;
*FileHandle = NULL;
@ -1304,7 +1310,11 @@ EfiShellOpenFileByName(
// do we start with >v
//
if (StrStr(FileName, L">v") == FileName) {
if (!IsVolatileEnv(FileName+2) &&
Status = IsVolatileEnv (FileName + 2, &Volatile);
if (EFI_ERROR (Status)) {
return Status;
}
if (!Volatile &&
((OpenMode & EFI_FILE_MODE_WRITE) != 0)) {
return (EFI_INVALID_PARAMETER);
}