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>

1 files changed, 16 insertions, 3 deletions

diff --git a/ShellPkg/Application/Shell/FileHandleWrappers.c b/ShellPkg/Application/Shell/FileHandleWrappers.c
index f64915d9ed..f6a82ee3c7 100644
--- a/ShellPkg/Application/Shell/FileHandleWrappers.c
+++ b/ShellPkg/Application/Shell/FileHandleWrappers.c
@@ -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;