commit 59b7dbac4ea075313af989cc20b71d44fdc7829a
author Ruiyu Ni <ruiyu.ni@intel.com> Wed Jul 13 15:45:08 2016 +0800
committer Ruiyu Ni <ruiyu.ni@intel.com> Mon Jul 18 10:55:48 2016 +0800
tree 8582a496a73b98d75ff7273d14acb0c5b27a00dd
parent 26ca45a86fa519e61c511ed27718b3f6dfe2710c

ShellPkg/For: 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/UefiShellLevel1CommandsLib/For.c

diff --git a/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c b/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
index cbf0517..6cfe8a7 100644
--- a/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
+++ b/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
@@ -2,7 +2,7 @@
   Main file for endfor and for shell level 1 functions.

 

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

-  Copyright (c) 2009 - 2014, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2009 - 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

@@ -414,7 +414,10 @@
         NewSize = StrSize(ArgSet);

         NewSize += sizeof(SHELL_FOR_INFO)+StrSize(gEfiShellParametersProtocol->Argv[1]);

         Info = AllocateZeroPool(NewSize);

-        ASSERT(Info != NULL);

+        if (Info == NULL) {

+          FreePool (ArgSet);

+          return SHELL_OUT_OF_RESOURCES;

+        }

         Info->Signature = SHELL_FOR_INFO_SIGNATURE;

         CopyMem(Info->Set, ArgSet, StrSize(ArgSet));

         NewSize = StrSize(gEfiShellParametersProtocol->Argv[1]);

@@ -458,7 +461,10 @@
         // set up for a 'run' for loop

         //

         Info = AllocateZeroPool(sizeof(SHELL_FOR_INFO)+StrSize(gEfiShellParametersProtocol->Argv[1]));

-        ASSERT(Info != NULL);

+        if (Info == NULL) {

+          FreePool (ArgSet);

+          return SHELL_OUT_OF_RESOURCES;

+        }

         Info->Signature = SHELL_FOR_INFO_SIGNATURE;

         CopyMem(Info->Set, gEfiShellParametersProtocol->Argv[1], StrSize(gEfiShellParametersProtocol->Argv[1]));

         Info->ReplacementName = Info->Set;