ShellPkg/IsVolatileEnv: Handle memory allocation failure

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

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