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>

diff --git a/ShellPkg/Application/Shell/FileHandleWrappers.c b/ShellPkg/Application/Shell/FileHandleWrappers.c
index f64915d9ed33f2501cd75ee268c166dbb8779872..f6a82ee3c760de910ef46ba554fa62564af835da 100644 (file)
--- a/ShellPkg/Application/Shell/FileHandleWrappers.c
+++ b/ShellPkg/Application/Shell/FileHandleWrappers.c
@@ -971,6 +971,7 @@ FileInterfaceEnvClose(
   VOID*       NewBuffer;\r
   UINTN       NewSize;\r
   EFI_STATUS  Status;\r
+  BOOLEAN     Volatile;\r
 \r
   //\r
   // Most if not all UEFI commands will have an '\r\n' at the end of any output. \r
@@ -981,6 +982,11 @@ FileInterfaceEnvClose(
   NewBuffer   = NULL;\r
   NewSize     = 0;\r
 \r
+  Status = IsVolatileEnv (((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &Volatile);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
   Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);\r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
     NewBuffer = AllocateZeroPool(NewSize + sizeof(CHAR16));\r
@@ -998,8 +1004,8 @@ FileInterfaceEnvClose(
            && (((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 3] == CHAR_CARRIAGE_RETURN)) {\r
         ((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 3] = CHAR_NULL;   \r
       }\r
-    \r
-      if (IsVolatileEnv(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name)) {\r
+\r
+      if (Volatile) {\r
         Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);\r
       } else {\r
         Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);\r
@@ -1157,13 +1163,20 @@ CreateFileInterfaceEnv(
   IN CONST CHAR16 *EnvName\r
   )\r
 {\r
+  EFI_STATUS                     Status;\r
   EFI_FILE_PROTOCOL_ENVIRONMENT  *EnvFileInterface;\r
   UINTN                          EnvNameSize;\r
+  BOOLEAN                        Volatile;\r
 \r
   if (EnvName == NULL) {\r
     return (NULL);\r
   }\r
 \r
+  Status = IsVolatileEnv (EnvName, &Volatile);\r
+  if (EFI_ERROR (Status)) {\r
+    return NULL;\r
+  }\r
+\r
   //\r
   // Get some memory\r
   //\r
@@ -1192,7 +1205,7 @@ CreateFileInterfaceEnv(
   //\r
   // Assign the different members for Volatile and Non-Volatile variables\r
   //\r
-  if (IsVolatileEnv(EnvName)) {\r
+  if (Volatile) {\r
     EnvFileInterface->Write       = FileInterfaceEnvVolWrite;\r
   } else {\r
     EnvFileInterface->Write       = FileInterfaceEnvNonVolWrite;\r

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index 5eb382a5866f551e511848c5f892d43dd682547f..9f87b9074fcc6ca579819862ab39bd5966822844 100644 (file)
--- 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.\r
 \r
   @param EnvVarName             The name of the environment variable in question\r
+  @param Volatile               Return TRUE if the environment variable is volatile\r
 \r
-  @retval TRUE                  This environment variable is Volatile\r
-  @retval FALSE                 This environment variable is NON-Volatile\r
+  @retval EFI_SUCCESS           The volatile attribute is returned successfully\r
+  @retval others                Some errors happened.\r
 **/\r
-BOOLEAN\r
-EFIAPI\r
+EFI_STATUS\r
 IsVolatileEnv (\r
-  IN CONST CHAR16 *EnvVarName\r
+  IN CONST CHAR16 *EnvVarName,\r
+  OUT BOOLEAN     *Volatile\r
   )\r
 {\r
   EFI_STATUS  Status;\r
@@ -41,6 +42,8 @@ IsVolatileEnv (
   VOID        *Buffer;\r
   UINT32      Attribs;\r
 \r
+  ASSERT (Volatile != NULL);\r
+\r
   Size = 0;\r
   Buffer = NULL;\r
 \r
@@ -54,7 +57,9 @@ IsVolatileEnv (
                             Buffer);\r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
     Buffer = AllocateZeroPool(Size);\r
-    ASSERT(Buffer != NULL);\r
+    if (Buffer == NULL) {\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
     Status = gRT->GetVariable((CHAR16*)EnvVarName,\r
                               &gShellVariableGuid,\r
                               &Attribs,\r
@@ -66,21 +71,18 @@ IsVolatileEnv (
   // not found means volatile\r
   //\r
   if (Status == EFI_NOT_FOUND) {\r
-    return (TRUE);\r
+    *Volatile = TRUE;\r
+    return EFI_SUCCESS;\r
   }\r
-  ASSERT_EFI_ERROR(Status);\r
-\r
-  //\r
-  // check for the Non Volatile bit\r
-  //\r
-  if ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE) {\r
-    return (FALSE);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
   }\r
 \r
   //\r
-  // everything else is volatile\r
+  // check for the Non Volatile bit\r
   //\r
-  return (TRUE);\r
+  *Volatile = !(BOOLEAN) ((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE);\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.h b/ShellPkg/Application/Shell/ShellEnvVar.h
index dd88b29fa040e23083219435bed8d566a66f28b1..5356580759a7adf24880a29d9f961b7300bde0e1 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.h
+++ b/ShellPkg/Application/Shell/ShellEnvVar.h
@@ -34,19 +34,18 @@ extern ENV_VAR_LIST    gShellEnvVarList;
 \r
 \r
 /**\r
-  Reports whether an environment variable is Volatile or Non-Volatile\r
-\r
-  This will use the Runtime Services call GetVariable to to search for the variable.\r
+  Reports whether an environment variable is Volatile or Non-Volatile.\r
 \r
   @param EnvVarName             The name of the environment variable in question\r
+  @param Volatile               Return TRUE if the environment variable is volatile\r
 \r
-  @retval TRUE                  This environment variable is Volatile\r
-  @retval FALSE                 This environment variable is NON-Volatile\r
+  @retval EFI_SUCCESS           The volatile attribute is returned successfully\r
+  @retval others                Some errors happened.\r
 **/\r
-BOOLEAN\r
-EFIAPI\r
+EFI_STATUS\r
 IsVolatileEnv (\r
-  IN CONST CHAR16 *EnvVarName\r
+  IN CONST CHAR16 *EnvVarName,\r
+  OUT BOOLEAN     *Volatile\r
   );\r
 \r
 /**\r

diff --git a/ShellPkg/Application/Shell/ShellParametersProtocol.c b/ShellPkg/Application/Shell/ShellParametersProtocol.c
index b3767bbbf2b9b7d407a896b2d188f08cbda27dd8..3684f9cd827b175d2abcde6b07e042c15b6d351e 100644 (file)
--- a/ShellPkg/Application/Shell/ShellParametersProtocol.c
+++ b/ShellPkg/Application/Shell/ShellParametersProtocol.c
@@ -736,6 +736,7 @@ UpdateStdInStdOutStdErr(
   UINTN             Size;\r
   SPLIT_LIST        *Split;\r
   CHAR16            *FirstLocation;\r
+  BOOLEAN           Volatile;\r
 \r
   OutUnicode      = TRUE;\r
   InUnicode       = TRUE;\r
@@ -1111,8 +1112,8 @@ UpdateStdInStdOutStdErr(
       //\r
       // Check for no volatile environment variables\r
       //\r
-      ||(StdErrVarName  != NULL && !IsVolatileEnv(StdErrVarName))\r
-      ||(StdOutVarName  != NULL && !IsVolatileEnv(StdOutVarName))\r
+      ||(StdErrVarName  != NULL && !EFI_ERROR (IsVolatileEnv (StdErrVarName, &Volatile)) && !Volatile)\r
+      ||(StdOutVarName  != NULL && !EFI_ERROR (IsVolatileEnv (StdOutVarName, &Volatile)) && !Volatile)\r
       //\r
       // Cant redirect during a reconnect operation.\r
       //\r

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index dc65b7d09aa26dc2eeb85d38dcbe0f642333b385..39b0e78b320aa76eb35250775fa3e2369ba9beb2 100644 (file)
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/ShellPkg/Application/Shell/ShellProtocol.c
@@ -1129,13 +1129,18 @@ EfiShellCreateFile(
 {\r
   EFI_DEVICE_PATH_PROTOCOL  *DevicePath;\r
   EFI_STATUS                Status;\r
+  BOOLEAN                   Volatile;\r
 \r
   //\r
   // Is this for an environment variable\r
   // do we start with >v\r
   //\r
   if (StrStr(FileName, L">v") == FileName) {\r
-    if (!IsVolatileEnv(FileName+2)) {\r
+    Status = IsVolatileEnv (FileName + 2, &Volatile);\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+    if (!Volatile) {\r
       return (EFI_INVALID_PARAMETER);\r
     }\r
     *FileHandle = CreateFileInterfaceEnv(FileName+2);\r
@@ -1245,6 +1250,7 @@ EfiShellOpenFileByName(
 {\r
   EFI_DEVICE_PATH_PROTOCOL        *DevicePath;\r
   EFI_STATUS                      Status;\r
+  BOOLEAN                         Volatile;\r
 \r
   *FileHandle = NULL;\r
 \r
@@ -1304,7 +1310,11 @@ EfiShellOpenFileByName(
   // do we start with >v\r
   //\r
   if (StrStr(FileName, L">v") == FileName) {\r
-    if (!IsVolatileEnv(FileName+2) &&\r
+    Status = IsVolatileEnv (FileName + 2, &Volatile);\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+    if (!Volatile &&\r
         ((OpenMode & EFI_FILE_MODE_WRITE) != 0)) {\r
       return (EFI_INVALID_PARAMETER);\r
     }\r