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>

5 files changed, 56 insertions, 31 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;

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index 5eb382a586..9f87b9074f 100644
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -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;

   }

-  ASSERT_EFI_ERROR(Status);

-

-  //

-  // check for the Non Volatile bit

-  //

-  if ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE) {

-    return (FALSE);

+  if (EFI_ERROR (Status)) {

+    return Status;

   }

 

   //

-  // everything else is volatile

+  // check for the Non Volatile bit

   //

-  return (TRUE);

+  *Volatile = !(BOOLEAN) ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE);

+  return EFI_SUCCESS;

 }

 

 /**

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.h b/ShellPkg/Application/Shell/ShellEnvVar.h
index dd88b29fa0..5356580759 100644
--- a/ShellPkg/Application/Shell/ShellEnvVar.h
+++ b/ShellPkg/Application/Shell/ShellEnvVar.h
@@ -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

   );

 

 /**

diff --git a/ShellPkg/Application/Shell/ShellParametersProtocol.c b/ShellPkg/Application/Shell/ShellParametersProtocol.c
index b3767bbbf2..3684f9cd82 100644
--- a/ShellPkg/Application/Shell/ShellParametersProtocol.c
+++ b/ShellPkg/Application/Shell/ShellParametersProtocol.c
@@ -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.

       //

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index dc65b7d09a..39b0e78b32 100644
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/ShellPkg/Application/Shell/ShellProtocol.c
@@ -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);

     }