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>

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;

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 EFI_SUCCESS           The volatile attribute is returned successfully
-  @retval FALSE                 This environment variable is NON-Volatile
+  @retval others                Some errors happened.
 **/
-BOOLEAN
+EFI_STATUS
-EFIAPI
 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) {
+  *Volatile = !(BOOLEAN) ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE);
-    return (FALSE);
+  return EFI_SUCCESS;
-  }
-
-  //
-  // everything else is volatile
-  //
-  return (TRUE);
 }
 
 /**

ShellPkg/Application/Shell/ShellEnvVar.h

@@ -34,19 +34,18 @@ extern ENV_VAR_LIST    gShellEnvVarList;
 
 
 /**
-  Reports whether an environment variable is Volatile or Non-Volatile
+  Reports whether an environment variable is Volatile or Non-Volatile.
-
-  This will use the Runtime Services call GetVariable to to search for the variable.
 
   @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 EFI_SUCCESS           The volatile attribute is returned successfully
-  @retval FALSE                 This environment variable is NON-Volatile
+  @retval others                Some errors happened.
 **/
-BOOLEAN
+EFI_STATUS
-EFIAPI
 IsVolatileEnv (
-  IN CONST CHAR16 *EnvVarName
+  IN CONST CHAR16 *EnvVarName,
+  OUT BOOLEAN     *Volatile
   );
 
 /**

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))
+      ||(StdErrVarName  != NULL && !EFI_ERROR (IsVolatileEnv (StdErrVarName, &Volatile)) && !Volatile)
-      ||(StdOutVarName  != NULL && !IsVolatileEnv(StdOutVarName))
+      ||(StdOutVarName  != NULL && !EFI_ERROR (IsVolatileEnv (StdOutVarName, &Volatile)) && !Volatile)
       //
       // Cant redirect during a reconnect operation.
       //

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);
     }