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;