ShellPkg/IsVolatileEnv: Handle memory allocation failure

[mirror_edk2.git] / ShellPkg / Application / Shell / ShellEnvVar.c

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