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:
@ -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;
|
||||
|
Reference in New Issue
Block a user