ShellPkg/IfConfig: 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/Library/UefiShellNetwork1CommandsLib/Ifconfig.c b/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c
index 92108a14ade22c4895055ca84fda1314c98ddc70..5e243d58843017810d57dd7c688ca4748cf75e30 100644 (file)
--- a/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c
+++ b/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c
@@ -129,6 +129,26 @@ VAR_CHECK_ITEM  mSetCheckList[] = {
 \r
 STATIC CONST CHAR16 PermanentString[10] = L"PERMANENT";\r
 \r
+/**\r
+  Free the ARG_LIST.\r
+\r
+  @param List Pointer to ARG_LIST to free.\r
+**/\r
+VOID\r
+FreeArgList (\r
+  ARG_LIST       *List\r
+)\r
+{\r
+  ARG_LIST       *Next;\r
+  while (List->Next != NULL) {\r
+    Next = List->Next;\r
+    FreePool (List);\r
+    List = Next;\r
+  }\r
+\r
+  FreePool (List);\r
+}\r
+\r
 /**\r
   Split a string with specified separator and save the substring to a list.\r
 \r
@@ -157,14 +177,18 @@ SplitStrToList (
   // Copy the CONST string to a local copy.\r
   //\r
   Str = AllocateCopyPool (StrSize (String), String);\r
-  ASSERT (Str != NULL);\r
+  if (Str == NULL) {\r
+    return NULL;\r
+  }\r
   ArgStr  = Str;\r
 \r
   //\r
   // init a node for the list head.\r
   //\r
   ArgNode = (ARG_LIST *) AllocateZeroPool (sizeof (ARG_LIST));\r
-  ASSERT (ArgNode != NULL);\r
+  if (ArgNode == NULL) {\r
+    return NULL;\r
+  }\r
   ArgList = ArgNode;\r
 \r
   //\r
@@ -176,7 +200,14 @@ SplitStrToList (
       ArgNode->Arg  = ArgStr;\r
       ArgStr        = Str + 1;\r
       ArgNode->Next = (ARG_LIST *) AllocateZeroPool (sizeof (ARG_LIST));\r
-      ASSERT (ArgNode->Next != NULL);\r
+      if (ArgNode->Next == NULL) {\r
+        //\r
+        // Free the local copy of string stored in the first node\r
+        //\r
+        FreePool (ArgList->Arg);\r
+        FreeArgList (ArgList);\r
+        return NULL;\r
+      }\r
       ArgNode = ArgNode->Next;\r
     }\r
 \r
@@ -1083,7 +1114,11 @@ IfConfigSetInterfaceInfo (
       }\r
 \r
       Dns   = AllocatePool (Index * sizeof (EFI_IPv4_ADDRESS));\r
-      ASSERT(Dns != NULL);\r
+      if (Dns == NULL) {\r
+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");\r
+        ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+        goto ON_EXIT;\r
+      }\r
       Tmp   = VarArg;\r
       Index = 0;\r
       while (Tmp != NULL) {\r
@@ -1193,8 +1228,6 @@ IfConfigCleanup (
   LIST_ENTRY                *Entry;\r
   LIST_ENTRY                *NextEntry;\r
   IFCONFIG_INTERFACE_CB     *IfCb;\r
-  ARG_LIST                  *ArgNode;\r
-  ARG_LIST                  *ArgHead;\r
 \r
   ASSERT (Private != NULL);\r
 \r
@@ -1202,15 +1235,7 @@ IfConfigCleanup (
   // Clean the list which save the set config Args.\r
   //\r
   if (Private->VarArg != NULL) {\r
-    ArgHead = Private->VarArg;\r
-\r
-    while (ArgHead->Next != NULL) {\r
-      ArgNode = ArgHead->Next;\r
-      FreePool (ArgHead);\r
-      ArgHead = ArgNode;\r
-    }\r
-\r
-    FreePool (ArgHead);\r
+    FreeArgList (Private->VarArg);\r
   }\r
 \r
   if (Private->IfName != NULL) {\r
@@ -1325,11 +1350,15 @@ ShellCommandRunIfconfig (
     ValueStr = ShellCommandLineGetValue (ParamPackage, L"-l");\r
     if (ValueStr != NULL) {\r
       Str = AllocateCopyPool (StrSize (ValueStr), ValueStr);\r
-      ASSERT (Str != NULL);\r
+      if (Str == NULL) {\r
+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");\r
+        ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+        goto ON_EXIT;\r
+      }\r
       Private->IfName = Str;\r
     }\r
   }\r
-  \r
+\r
   //\r
   // To get interface name for the clear option.\r
   //\r
@@ -1338,7 +1367,11 @@ ShellCommandRunIfconfig (
     ValueStr = ShellCommandLineGetValue (ParamPackage, L"-r");\r
     if (ValueStr != NULL) {\r
       Str = AllocateCopyPool (StrSize (ValueStr), ValueStr);\r
-      ASSERT (Str != NULL);\r
+      if (Str == NULL) {\r
+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");\r
+        ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+        goto ON_EXIT;\r
+      }\r
       Private->IfName = Str;\r
     }\r
   }\r
@@ -1357,8 +1390,12 @@ ShellCommandRunIfconfig (
     //\r
     // To split the configuration into multi-section.\r
     //\r
-    ArgList         = SplitStrToList (ValueStr, L' ');\r
-    ASSERT (ArgList != NULL);\r
+    ArgList = SplitStrToList (ValueStr, L' ');\r
+    if (ArgList == NULL) {\r
+      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");\r
+      ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+      goto ON_EXIT;\r
+    }\r
 \r
     Private->OpCode = IfConfigOpSet;\r
     Private->IfName = ArgList->Arg;\r

diff --git a/ShellPkg/Library/UefiShellNetwork1CommandsLib/UefiShellNetwork1CommandsLib.uni b/ShellPkg/Library/UefiShellNetwork1CommandsLib/UefiShellNetwork1CommandsLib.uni
index c9c031951e85dcd8ccbfae59dbf0475fac08c086..c8ab64fc55f44cad1d2077336495758b6e73eae6 100644 (file)
--- a/ShellPkg/Library/UefiShellNetwork1CommandsLib/UefiShellNetwork1CommandsLib.uni
+++ b/ShellPkg/Library/UefiShellNetwork1CommandsLib/UefiShellNetwork1CommandsLib.uni
@@ -43,6 +43,7 @@
 #string STR_GEN_FILE_NF           #language en-US "%H%s%N: File not found - '%H%s%N'\r\n"\r
 #string STR_GEN_IS_DIR            #language en-US "%H%s%N: '%H%s%N' is a directory\r\n"\r
 #string STR_GEN_PROTOCOL_NF       #language en-US "%H%s%N: The protocol '%H%s%N' is required and not found (%g).\r\n"\r
+#string STR_GEN_OUT_MEM           #language en-US "%H%s%N: Memory allocation was not successful.\r\n"\r
 \r
 #string STR_PING_INVALID_SOURCE      #language en-US "%Ping: Require source interface option\r\n"\r
 #string STR_PING_CONFIG              #language en-US "Config %r\r\n"\r