ShellPkg/Ifconfig6: 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 a915fea68e6bea900755f85e81cdc1070b9f3dc6)

2 files changed, 55 insertions, 21 deletions

diff --git a/ShellPkg/Library/UefiShellNetwork2CommandsLib/Ifconfig6.c b/ShellPkg/Library/UefiShellNetwork2CommandsLib/Ifconfig6.c
index 42c972e1a5..32dd2845fc 100644
--- a/ShellPkg/Library/UefiShellNetwork2CommandsLib/Ifconfig6.c
+++ b/ShellPkg/Library/UefiShellNetwork2CommandsLib/Ifconfig6.c
@@ -155,6 +155,26 @@ VAR_CHECK_ITEM  mIfConfig6SetCheckList[] = {
 };

 

 /**

+  Free the ARG_LIST.

+

+  @param List Pointer to ARG_LIST to free.

+**/

+VOID

+IfConfig6FreeArgList (

+  ARG_LIST       *List

+)

+{

+  ARG_LIST       *Next;

+  while (List->Next != NULL) {

+    Next = List->Next;

+    FreePool (List);

+    List = Next;

+  }

+

+  FreePool (List);

+}

+

+/**

   Split a string with specified separator and save the substring to a list.

 

   @param[in]    String       The pointer of the input string.

@@ -181,15 +201,19 @@ IfConfig6SplitStrToList (
   //

   // Copy the CONST string to a local copy.

   //

-  Str     = AllocateCopyPool (StrSize (String), String);

-  ASSERT (Str != NULL);

+  Str = AllocateCopyPool (StrSize (String), String);

+  if (Str == NULL) {

+    return NULL;

+  }

   ArgStr  = Str;

 

   //

   // init a node for the list head.

   //

   ArgNode = (ARG_LIST *) AllocateZeroPool (sizeof (ARG_LIST));

-  ASSERT (ArgNode != NULL);

+  if (ArgNode == NULL) {

+    return NULL;

+  }

   ArgList = ArgNode;

 

   //

@@ -201,7 +225,14 @@ IfConfig6SplitStrToList (
       ArgNode->Arg  = ArgStr;

       ArgStr        = Str + 1;

       ArgNode->Next = (ARG_LIST *) AllocateZeroPool (sizeof (ARG_LIST));

-      ASSERT (ArgNode->Next != NULL);

+      if (ArgNode->Next == NULL) {

+        //

+        // Free the local copy of string stored in the first node

+        //

+        FreePool (ArgList->Arg);

+        IfConfig6FreeArgList (ArgList);

+        return NULL;

+      }

       ArgNode = ArgNode->Next;

     }

 

@@ -1668,8 +1699,6 @@ IfConfig6Cleanup (
   LIST_ENTRY                *Entry;

   LIST_ENTRY                *NextEntry;

   IFCONFIG6_INTERFACE_CB    *IfCb;

-  ARG_LIST                  *ArgNode;

-  ARG_LIST                  *ArgHead;

 

   ASSERT (Private != NULL);

 

@@ -1677,15 +1706,7 @@ IfConfig6Cleanup (
   // Clean the list which save the set config Args.

   //

   if (Private->VarArg != NULL) {

-    ArgHead = Private->VarArg;

-

-    while (ArgHead->Next != NULL) {

-      ArgNode = ArgHead->Next;

-      FreePool (ArgHead);

-      ArgHead = ArgNode;

-    }

-

-    FreePool (ArgHead);

+    IfConfig6FreeArgList (Private->VarArg);

   }

 

   if (Private->IfName != NULL)

@@ -1799,8 +1820,12 @@ ShellCommandRunIfconfig6 (
     Private->OpCode = IfConfig6OpList;

     ValueStr = ShellCommandLineGetValue (ParamPackage, L"-l");

     if (ValueStr != NULL) {

-      Str             = AllocateCopyPool (StrSize (ValueStr), ValueStr);

-      ASSERT (Str != NULL);

+      Str = AllocateCopyPool (StrSize (ValueStr), ValueStr);

+      if (Str == NULL) {

+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork2HiiHandle, L"ifconfig6");

+        ShellStatus = SHELL_OUT_OF_RESOURCES;

+        goto ON_EXIT;

+      }

       Private->IfName = Str;

     }

   }

@@ -1811,8 +1836,12 @@ ShellCommandRunIfconfig6 (
     Private->OpCode = IfConfig6OpClear;

     ValueStr = ShellCommandLineGetValue (ParamPackage, L"-r");

     if (ValueStr != NULL) {

-      Str             = AllocateCopyPool (StrSize (ValueStr), ValueStr);

-      ASSERT (Str != NULL);

+      Str = AllocateCopyPool (StrSize (ValueStr), ValueStr);

+      if (Str == NULL) {

+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork2HiiHandle, L"ifconfig6");

+        ShellStatus = SHELL_OUT_OF_RESOURCES;

+        goto ON_EXIT;

+      }

       Private->IfName = Str;

     }

   }

@@ -1830,8 +1859,12 @@ ShellCommandRunIfconfig6 (
     //

     // To split the configuration into multi-section.

     //

-    ArgList         = IfConfig6SplitStrToList (ValueStr, L' ');

-    ASSERT (ArgList != NULL);

+    ArgList = IfConfig6SplitStrToList (ValueStr, L' ');

+    if (ArgList == NULL) {

+      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork2HiiHandle, L"ifconfig6");

+      ShellStatus = SHELL_OUT_OF_RESOURCES;

+      goto ON_EXIT;

+    }

 

     Private->OpCode = IfConfig6OpSet;

     Private->IfName = ArgList->Arg;

diff --git a/ShellPkg/Library/UefiShellNetwork2CommandsLib/UefiShellNetwork2CommandsLib.uni b/ShellPkg/Library/UefiShellNetwork2CommandsLib/UefiShellNetwork2CommandsLib.uni
index ef441e1351..8f27e272d4 100644
--- a/ShellPkg/Library/UefiShellNetwork2CommandsLib/UefiShellNetwork2CommandsLib.uni
+++ b/ShellPkg/Library/UefiShellNetwork2CommandsLib/UefiShellNetwork2CommandsLib.uni
@@ -24,6 +24,7 @@
 #langdef   en-US "english"

 

 #string STR_GEN_ERR_AD                     #language en-US  "%H%s%N: Access denied.\r\n"

+#string STR_GEN_OUT_MEM                    #language en-US  "%H%s%N: Memory allocation was not successful.\r\n"

 

 #string STR_PING6_INVALID_IP               #language en-US  "%Ping6: Invalid IP6 address, %s\r\n"

 #string STR_PING6_INVALID_INPUT            #language en-US  "%Ping6: Invalid input, please type 'Ping6 -?'for help\r\n"