From 31e5b912b99e0fb39e81f70bc24a4be589191abb Mon Sep 17 00:00:00 2001
From: Ruiyu Ni <ruiyu.ni@intel.com>
Date: Fri, 8 Jul 2016 15:18:14 +0800
Subject: 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 | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

(limited to 'ShellPkg/Application/Shell/FileHandleWrappers.c')

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;
-- 
cgit v1.2.3