ShellPkg/ShellAddEnvVarToList: 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/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index 9f87b9074fcc6ca579819862ab39bd5966822844..77bf599eab1c8f7e0b59a5b0f0551b2ba88a2ee2 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -442,8 +442,11 @@ ShellFindEnvVarInList (
                     including the tailing CHAR_NULL\r
   @param Atts       The attributes of the variable.\r
 \r
+  @retval EFI_SUCCESS  The environment variable was added to list successfully.\r
+  @retval others       Some errors happened.\r
+\r
 **/\r
-VOID\r
+EFI_STATUS\r
 ShellAddEnvVarToList (\r
   IN CONST CHAR16     *Key,\r
   IN CONST CHAR16     *Value,\r
@@ -452,9 +455,16 @@ ShellAddEnvVarToList (
   )\r
 {\r
   ENV_VAR_LIST      *Node;\r
+  CHAR16            *LocalKey;\r
+  CHAR16            *LocalValue;\r
   \r
   if (Key == NULL || Value == NULL || ValueSize == 0) {\r
-    return;\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  LocalValue = AllocateCopyPool (ValueSize, Value);\r
+  if (LocalValue == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
   //\r
@@ -467,10 +477,8 @@ ShellAddEnvVarToList (
     if (Node->Key != NULL && StrCmp(Key, Node->Key) == 0) {\r
       Node->Atts = Atts;\r
       SHELL_FREE_NON_NULL(Node->Val);\r
-      Node->Val  = AllocateZeroPool (ValueSize);\r
-      ASSERT (Node->Val != NULL);\r
-      CopyMem(Node->Val, Value, ValueSize);\r
-      return;\r
+      Node->Val  = LocalValue;\r
+      return EFI_SUCCESS;\r
     }\r
   }\r
 \r
@@ -478,16 +486,23 @@ ShellAddEnvVarToList (
   // If the environment varialbe key doesn't exist in list just insert\r
   // a new node.\r
   //\r
+  LocalKey = AllocateCopyPool (StrSize(Key), Key);\r
+  if (LocalKey == NULL) {\r
+    FreePool (LocalValue);\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
   Node = (ENV_VAR_LIST*)AllocateZeroPool (sizeof(ENV_VAR_LIST));\r
-  ASSERT (Node != NULL);\r
-  Node->Key = AllocateCopyPool(StrSize(Key), Key);\r
-  ASSERT (Node->Key != NULL);\r
-  Node->Val = AllocateCopyPool(ValueSize, Value);\r
-  ASSERT (Node->Val != NULL);\r
+  if (Node == NULL) {\r
+    FreePool (LocalKey);\r
+    FreePool (LocalValue);\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+  Node->Key = LocalKey;\r
+  Node->Val = LocalValue;\r
   Node->Atts = Atts;\r
   InsertTailList(&gShellEnvVarList.Link, &Node->Link);\r
 \r
-  return;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.h b/ShellPkg/Application/Shell/ShellEnvVar.h
index 5356580759a7adf24880a29d9f961b7300bde0e1..f05ef2bd66fff51974c648f2145ca4d48c872185 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.h
+++ b/ShellPkg/Application/Shell/ShellEnvVar.h
@@ -240,11 +240,14 @@ ShellFindEnvVarInList (
   @param Key        The name of the environment variable.\r
   @param Value      The value of environment variable.\r
   @param ValueSize  The size in bytes of the environment variable\r
-                    including the tailing CHAR_NELL\r
+                    including the tailing CHAR_NULL\r
   @param Atts       The attributes of the variable.\r
 \r
+  @retval EFI_SUCCESS  The environment variable was added to list successfully.\r
+  @retval others       Some errors happened.\r
+\r
 **/\r
-VOID\r
+EFI_STATUS\r
 ShellAddEnvVarToList (\r
   IN CONST CHAR16     *Key,\r
   IN CONST CHAR16     *Value,\r

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index 39b0e78b320aa76eb35250775fa3e2369ba9beb2..a95da000f51e8293b11946029acded08bdbde198 100644 (file)
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/ShellPkg/Application/Shell/ShellProtocol.c
@@ -2852,36 +2852,28 @@ InternalEfiShellSetEnv(
   )\r
 {\r
   EFI_STATUS      Status;\r
-  UINT32          Atts;\r
 \r
-  Atts = 0x0;\r
-  \r
   if (Value == NULL || StrLen(Value) == 0) {\r
     Status = SHELL_DELETE_ENVIRONMENT_VARIABLE(Name);\r
     if (!EFI_ERROR(Status)) {\r
       ShellRemvoeEnvVarFromList(Name);\r
     }\r
-    return Status;\r
   } else {\r
     SHELL_DELETE_ENVIRONMENT_VARIABLE(Name);\r
-    if (Volatile) {\r
-      Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value);\r
-      if (!EFI_ERROR(Status)) {\r
-        Atts   &= ~EFI_VARIABLE_NON_VOLATILE;\r
-        Atts   |= EFI_VARIABLE_BOOTSERVICE_ACCESS;\r
-        ShellAddEnvVarToList(Name, Value, StrSize(Value), Atts);\r
+    Status = ShellAddEnvVarToList(\r
+               Name, Value, StrSize(Value),\r
+               EFI_VARIABLE_BOOTSERVICE_ACCESS | (Volatile ? 0 : EFI_VARIABLE_NON_VOLATILE)\r
+               );\r
+    if (!EFI_ERROR (Status)) {\r
+      Status = Volatile\r
+             ? SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value)\r
+             : SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value);\r
+      if (EFI_ERROR (Status)) {\r
+        ShellRemvoeEnvVarFromList(Name);\r
       }\r
-      return Status;\r
-    } else {\r
-      Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value);\r
-      if (!EFI_ERROR(Status)) {\r
-        Atts   |= EFI_VARIABLE_NON_VOLATILE;\r
-        Atts   |= EFI_VARIABLE_BOOTSERVICE_ACCESS;\r
-        ShellAddEnvVarToList(Name, Value, StrSize(Value), Atts);\r
-      } \r
-      return Status;\r
     }\r
   }\r
+  return Status;\r
 }\r
 \r
 /**\r