From cd8462ea9d8aa0edb8116c43fd26918f8e155c31 Mon Sep 17 00:00:00 2001
From: ydong10 <ydong10@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Mon, 13 Jun 2011 01:44:03 +0000
Subject: Check the buffer allocation result before using it.

Signed-off-by:ydong10
Reviewed-by:jljusten


git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@11811 6f19259b-4bc3-4df7-8a09-765794883524
---
 ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c | 2 ++
 1 file changed, 2 insertions(+)

(limited to 'ShellPkg/Library/UefiShellDriver1CommandsLib')

diff --git a/ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c b/ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c
index bd5a825ad4..5203f11f9b 100644
--- a/ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c
+++ b/ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c
@@ -56,6 +56,7 @@ FindHiiHandleViaDevPath(
   Status = HiiDb->ListPackageLists(HiiDb, EFI_HII_PACKAGE_DEVICE_PATH, NULL, &HandleBufferSize, HandleBuffer);
   if (Status == EFI_BUFFER_TOO_SMALL) {
     HandleBuffer = AllocateZeroPool(HandleBufferSize);
+    ASSERT (HandleBuffer != NULL);
     Status = HiiDb->ListPackageLists(HiiDb, EFI_HII_PACKAGE_DEVICE_PATH, NULL, &HandleBufferSize, HandleBuffer);
   }
   if (EFI_ERROR(Status)) {
@@ -69,6 +70,7 @@ FindHiiHandleViaDevPath(
     Status = HiiDb->ExportPackageLists(HiiDb, HandleBuffer[LoopVariable], &MainBufferSize, MainBuffer);
     if (Status == EFI_BUFFER_TOO_SMALL) {
       MainBuffer = AllocateZeroPool(MainBufferSize);
+      ASSERT (MainBuffer != NULL);
       Status = HiiDb->ExportPackageLists(HiiDb, HandleBuffer[LoopVariable], &MainBufferSize, MainBuffer);
     }
     //
-- 
cgit v1.2.3