ShellPkg/DMem: 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/UefiShellDebug1CommandsLib/Dmem.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/Dmem.c
index 7693835b8865739600dd99db003be26d49255a36..5491f6de578e2191626a596b864b72449543df83 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/Dmem.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/Dmem.c
@@ -1,7 +1,7 @@
 /** @file\r
   Main file for Dmem shell Debug1 function.\r
 \r
-  Copyright (c) 2010 - 2011, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>\r
   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>  \r
   (C) Copyright 2015 Hewlett Packard Enterprise Development LP<BR>  \r
   This program and the accompanying materials\r
@@ -68,7 +68,9 @@ DisplayMmioMemory(
     return (SHELL_NOT_FOUND);\r
   }\r
   Buffer = AllocateZeroPool(Size);\r
-  ASSERT(Buffer != NULL);\r
+  if (Buffer == NULL) {\r
+    return SHELL_OUT_OF_RESOURCES;\r
+  }\r
 \r
   Status = PciRbIo->Mem.Read(PciRbIo, EfiPciWidthUint8, (UINT64)(UINTN)Address, Size, Buffer);\r
   if (EFI_ERROR(Status)) {\r