From 18bff310aae5d26d440e03742a208d7ecfa72387 Mon Sep 17 00:00:00 2001
From: Ruiyu Ni <ruiyu.ni@intel.com>
Date: Wed, 13 Jul 2016 17:39:58 +0800
Subject: 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>
---
 .../UefiShellNetwork1CommandsLib/Ifconfig.c        | 77 ++++++++++++++++------
 .../UefiShellNetwork1CommandsLib.uni               |  1 +
 2 files changed, 58 insertions(+), 20 deletions(-)

diff --git a/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c b/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c
index 92108a1..5e243d5 100644
--- a/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c
+++ b/ShellPkg/Library/UefiShellNetwork1CommandsLib/Ifconfig.c
@@ -130,6 +130,26 @@ VAR_CHECK_ITEM  mSetCheckList[] = {
 STATIC CONST CHAR16 PermanentString[10] = L"PERMANENT";
 
 /**
+  Free the ARG_LIST.
+
+  @param List Pointer to ARG_LIST to free.
+**/
+VOID
+FreeArgList (
+  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.
@@ -157,14 +177,18 @@ SplitStrToList (
   // Copy the CONST string to a local copy.
   //
   Str = AllocateCopyPool (StrSize (String), String);
-  ASSERT (Str != NULL);
+  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;
 
   //
@@ -176,7 +200,14 @@ SplitStrToList (
       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);
+        FreeArgList (ArgList);
+        return NULL;
+      }
       ArgNode = ArgNode->Next;
     }
 
@@ -1083,7 +1114,11 @@ IfConfigSetInterfaceInfo (
       }
 
       Dns   = AllocatePool (Index * sizeof (EFI_IPv4_ADDRESS));
-      ASSERT(Dns != NULL);
+      if (Dns == NULL) {
+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");
+        ShellStatus = SHELL_OUT_OF_RESOURCES;
+        goto ON_EXIT;
+      }
       Tmp   = VarArg;
       Index = 0;
       while (Tmp != NULL) {
@@ -1193,8 +1228,6 @@ IfConfigCleanup (
   LIST_ENTRY                *Entry;
   LIST_ENTRY                *NextEntry;
   IFCONFIG_INTERFACE_CB     *IfCb;
-  ARG_LIST                  *ArgNode;
-  ARG_LIST                  *ArgHead;
 
   ASSERT (Private != NULL);
 
@@ -1202,15 +1235,7 @@ IfConfigCleanup (
   // 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);
+    FreeArgList (Private->VarArg);
   }
 
   if (Private->IfName != NULL) {
@@ -1325,11 +1350,15 @@ ShellCommandRunIfconfig (
     ValueStr = ShellCommandLineGetValue (ParamPackage, L"-l");
     if (ValueStr != NULL) {
       Str = AllocateCopyPool (StrSize (ValueStr), ValueStr);
-      ASSERT (Str != NULL);
+      if (Str == NULL) {
+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");
+        ShellStatus = SHELL_OUT_OF_RESOURCES;
+        goto ON_EXIT;
+      }
       Private->IfName = Str;
     }
   }
-  
+
   //
   // To get interface name for the clear option.
   //
@@ -1338,7 +1367,11 @@ ShellCommandRunIfconfig (
     ValueStr = ShellCommandLineGetValue (ParamPackage, L"-r");
     if (ValueStr != NULL) {
       Str = AllocateCopyPool (StrSize (ValueStr), ValueStr);
-      ASSERT (Str != NULL);
+      if (Str == NULL) {
+        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");
+        ShellStatus = SHELL_OUT_OF_RESOURCES;
+        goto ON_EXIT;
+      }
       Private->IfName = Str;
     }
   }
@@ -1357,8 +1390,12 @@ ShellCommandRunIfconfig (
     //
     // To split the configuration into multi-section.
     //
-    ArgList         = SplitStrToList (ValueStr, L' ');
-    ASSERT (ArgList != NULL);
+    ArgList = SplitStrToList (ValueStr, L' ');
+    if (ArgList == NULL) {
+      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellNetwork1HiiHandle, L"ifconfig");
+      ShellStatus = SHELL_OUT_OF_RESOURCES;
+      goto ON_EXIT;
+    }
 
     Private->OpCode = IfConfigOpSet;
     Private->IfName = ArgList->Arg;
diff --git a/ShellPkg/Library/UefiShellNetwork1CommandsLib/UefiShellNetwork1CommandsLib.uni b/ShellPkg/Library/UefiShellNetwork1CommandsLib/UefiShellNetwork1CommandsLib.uni
index c9c0319..c8ab64f 100644
--- 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"
 #string STR_GEN_IS_DIR            #language en-US "%H%s%N: '%H%s%N' is a directory\r\n"
 #string STR_GEN_PROTOCOL_NF       #language en-US "%H%s%N: The protocol '%H%s%N' is required and not found (%g).\r\n"
+#string STR_GEN_OUT_MEM           #language en-US "%H%s%N: Memory allocation was not successful.\r\n"
 
 #string STR_PING_INVALID_SOURCE      #language en-US "%Ping: Require source interface option\r\n"
 #string STR_PING_CONFIG              #language en-US "Config %r\r\n"
-- 
cgit v1.1