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>
(cherry picked from commit ffbc60a02757fe430448d7ad022d04cec133ba5d)

3 files changed, 43 insertions, 33 deletions

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index 9f87b9074f..77bf599eab 100644
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -442,8 +442,11 @@ ShellFindEnvVarInList (
                     including the tailing CHAR_NULL

   @param Atts       The attributes of the variable.

 

+  @retval EFI_SUCCESS  The environment variable was added to list successfully.

+  @retval others       Some errors happened.

+

 **/

-VOID

+EFI_STATUS

 ShellAddEnvVarToList (

   IN CONST CHAR16     *Key,

   IN CONST CHAR16     *Value,

@@ -452,9 +455,16 @@ ShellAddEnvVarToList (
   )

 {

   ENV_VAR_LIST      *Node;

+  CHAR16            *LocalKey;

+  CHAR16            *LocalValue;

   

   if (Key == NULL || Value == NULL || ValueSize == 0) {

-    return;

+    return EFI_INVALID_PARAMETER;

+  }

+

+  LocalValue = AllocateCopyPool (ValueSize, Value);

+  if (LocalValue == NULL) {

+    return EFI_OUT_OF_RESOURCES;

   }

 

   //

@@ -467,10 +477,8 @@ ShellAddEnvVarToList (
     if (Node->Key != NULL && StrCmp(Key, Node->Key) == 0) {

       Node->Atts = Atts;

       SHELL_FREE_NON_NULL(Node->Val);

-      Node->Val  = AllocateZeroPool (ValueSize);

-      ASSERT (Node->Val != NULL);

-      CopyMem(Node->Val, Value, ValueSize);

-      return;

+      Node->Val  = LocalValue;

+      return EFI_SUCCESS;

     }

   }

 

@@ -478,16 +486,23 @@ ShellAddEnvVarToList (
   // If the environment varialbe key doesn't exist in list just insert

   // a new node.

   //

+  LocalKey = AllocateCopyPool (StrSize(Key), Key);

+  if (LocalKey == NULL) {

+    FreePool (LocalValue);

+    return EFI_OUT_OF_RESOURCES;

+  }

   Node = (ENV_VAR_LIST*)AllocateZeroPool (sizeof(ENV_VAR_LIST));

-  ASSERT (Node != NULL);

-  Node->Key = AllocateCopyPool(StrSize(Key), Key);

-  ASSERT (Node->Key != NULL);

-  Node->Val = AllocateCopyPool(ValueSize, Value);

-  ASSERT (Node->Val != NULL);

+  if (Node == NULL) {

+    FreePool (LocalKey);

+    FreePool (LocalValue);

+    return EFI_OUT_OF_RESOURCES;

+  }

+  Node->Key = LocalKey;

+  Node->Val = LocalValue;

   Node->Atts = Atts;

   InsertTailList(&gShellEnvVarList.Link, &Node->Link);

 

-  return;

+  return EFI_SUCCESS;

 }

 

 /**

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.h b/ShellPkg/Application/Shell/ShellEnvVar.h
index 5356580759..f05ef2bd66 100644
--- 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.

   @param Value      The value of environment variable.

   @param ValueSize  The size in bytes of the environment variable

-                    including the tailing CHAR_NELL

+                    including the tailing CHAR_NULL

   @param Atts       The attributes of the variable.

 

+  @retval EFI_SUCCESS  The environment variable was added to list successfully.

+  @retval others       Some errors happened.

+

 **/

-VOID

+EFI_STATUS

 ShellAddEnvVarToList (

   IN CONST CHAR16     *Key,

   IN CONST CHAR16     *Value,

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index 39b0e78b32..a95da000f5 100644
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/ShellPkg/Application/Shell/ShellProtocol.c
@@ -2852,36 +2852,28 @@ InternalEfiShellSetEnv(
   )

 {

   EFI_STATUS      Status;

-  UINT32          Atts;

 

-  Atts = 0x0;

-  

   if (Value == NULL || StrLen(Value) == 0) {

     Status = SHELL_DELETE_ENVIRONMENT_VARIABLE(Name);

     if (!EFI_ERROR(Status)) {

       ShellRemvoeEnvVarFromList(Name);

     }

-    return Status;

   } else {

     SHELL_DELETE_ENVIRONMENT_VARIABLE(Name);

-    if (Volatile) {

-      Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value);

-      if (!EFI_ERROR(Status)) {

-        Atts   &= ~EFI_VARIABLE_NON_VOLATILE;

-        Atts   |= EFI_VARIABLE_BOOTSERVICE_ACCESS;

-        ShellAddEnvVarToList(Name, Value, StrSize(Value), Atts);

+    Status = ShellAddEnvVarToList(

+               Name, Value, StrSize(Value),

+               EFI_VARIABLE_BOOTSERVICE_ACCESS | (Volatile ? 0 : EFI_VARIABLE_NON_VOLATILE)

+               );

+    if (!EFI_ERROR (Status)) {

+      Status = Volatile

+             ? SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value)

+             : SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value);

+      if (EFI_ERROR (Status)) {

+        ShellRemvoeEnvVarFromList(Name);

       }

-      return Status;

-    } else {

-      Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value);

-      if (!EFI_ERROR(Status)) {

-        Atts   |= EFI_VARIABLE_NON_VOLATILE;

-        Atts   |= EFI_VARIABLE_BOOTSERVICE_ACCESS;

-        ShellAddEnvVarToList(Name, Value, StrSize(Value), Atts);

-      } 

-      return Status;

     }

   }

+  return Status;

 }

 

 /**