fix comments.

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

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index cd55fa907be69a6d3f0ac3df4fc066f3ffc3c10f..519181b56ddbb4e5cb998d4e4b65b40891df8554 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -13,6 +13,7 @@
 **/\r
 \r
 #include <Uefi.h>\r
+#include <ShellBase.h>\r
 \r
 #include <Guid/ShellVariableGuid.h>\r
 \r
@@ -24,7 +25,6 @@
 \r
 #include "ShellEnvVar.h"\r
 \r
-\r
 /**\r
   Reports whether an environment variable is Volatile or Non-Volatile.\r
 \r
@@ -169,18 +169,31 @@ GetEnvironmentVariableList(
     }\r
     if (!EFI_ERROR(Status) && CompareGuid(&Guid, &gShellVariableGuid)){\r
       VarList = AllocateZeroPool(sizeof(ENV_VAR_LIST));\r
-      ValSize = 0;\r
-      Status = SHELL_GET_ENVIRONMENT_VARIABLE_AND_ATTRIBUTES(VariableName, &VarList->Atts, &ValSize, VarList->Val);\r
-      if (Status == EFI_BUFFER_TOO_SMALL){\r
-        VarList->Val = AllocatePool(ValSize);\r
-        ASSERT(VarList->Val != NULL);\r
+      if (VarList == NULL) {\r
+        Status = EFI_OUT_OF_RESOURCES;\r
+      } else {\r
+        ValSize = 0;\r
         Status = SHELL_GET_ENVIRONMENT_VARIABLE_AND_ATTRIBUTES(VariableName, &VarList->Atts, &ValSize, VarList->Val);\r
-      }\r
-      if (!EFI_ERROR(Status)) {\r
-        VarList->Key = AllocatePool(StrSize(VariableName));\r
-        ASSERT(VarList->Key != NULL);\r
-        StrCpy(VarList->Key, VariableName);\r
-        InsertTailList(ListHead, &VarList->Link);\r
+        if (Status == EFI_BUFFER_TOO_SMALL){\r
+          VarList->Val = AllocatePool(ValSize);\r
+          if (VarList->Val == NULL) {\r
+            SHELL_FREE_NON_NULL(VarList);\r
+            Status = EFI_OUT_OF_RESOURCES;\r
+          } else {\r
+            Status = SHELL_GET_ENVIRONMENT_VARIABLE_AND_ATTRIBUTES(VariableName, &VarList->Atts, &ValSize, VarList->Val);\r
+          }\r
+        }\r
+        if (!EFI_ERROR(Status) && VarList != NULL) {\r
+          VarList->Key = AllocatePool(StrSize(VariableName));\r
+          if (VarList->Key == NULL) {\r
+            SHELL_FREE_NON_NULL(VarList->Val);\r
+            SHELL_FREE_NON_NULL(VarList);\r
+            Status = EFI_OUT_OF_RESOURCES;\r
+          } else {\r
+            StrCpy(VarList->Key, VariableName);\r
+            InsertTailList(ListHead, &VarList->Link);\r
+          }\r
+        }\r
       }\r
     } // compare guid\r
   } // while\r