commit 55b1d63de7677f89178f200404d8d05a23150a4a
author Ruiyu Ni <ruiyu.ni@intel.com> Wed Jul 13 15:43:55 2016 +0800
committer Ruiyu Ni <ruiyu.ni@intel.com> Mon Jul 18 10:55:38 2016 +0800
tree a2ecbffff331b776ae1c1d69b21f436fd72d83fe
parent 3e9442a564a767f7bd60113117d7b993a4e29a9b

ShellPkg/DevTree: 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>

ShellPkg/Library/UefiShellDriver1CommandsLib/DevTree.c

diff --git a/ShellPkg/Library/UefiShellDriver1CommandsLib/DevTree.c b/ShellPkg/Library/UefiShellDriver1CommandsLib/DevTree.c
index fb9de10..237cb85 100644
--- a/ShellPkg/Library/UefiShellDriver1CommandsLib/DevTree.c
+++ b/ShellPkg/Library/UefiShellDriver1CommandsLib/DevTree.c
@@ -2,7 +2,7 @@
   Main file for DevTree shell Driver1 function.

 

   (C) Copyright 2014-2015 Hewlett-Packard Development Company, L.P.<BR>

-  Copyright (c) 2010 - 2012, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

   which accompanies this distribution.  The full text of the license may be found at

@@ -57,7 +57,8 @@
   ChildHandleBuffer   = NULL;

   ChildCount          = 0;

 

-  ASSERT(TheHandle    != NULL);

+  ASSERT (TheHandle != NULL);

+  ASSERT (HiiString != NULL);

   

   if (ShellGetExecutionBreakFlag()) {

     ShellStatus = SHELL_ABORTED;

@@ -91,10 +92,10 @@
     return SHELL_SUCCESS;

   }

 

-  FormatString        = AllocateZeroPool(StrSize(HiiString) + (10)*sizeof(FormatString[0]));

-

-  ASSERT(HiiString    != NULL);

-  ASSERT(FormatString != NULL);

+  FormatString = AllocateZeroPool(StrSize(HiiString) + (10)*sizeof(FormatString[0]));

+  if (FormatString == NULL) {

+    return SHELL_OUT_OF_RESOURCES;

+  }

 

   //

   // we generate the format string on the fly so that we can control the