Check the buffer allocation result before using it.

[mirror_edk2.git] / ShellPkg / Library / UefiShellDriver1CommandsLib / DrvCfg.c

diff --git a/ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c b/ShellPkg/Library/UefiShellDriver1CommandsLib/DrvCfg.c
index bd5a825ad4c6916cccf948ad32b10633cc6e41c2..5203f11f9bd88ccf7d4558d62d7576ecf11979d6 100644 (file)
--- 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);\r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
     HandleBuffer = AllocateZeroPool(HandleBufferSize);\r
+    ASSERT (HandleBuffer != NULL);\r
     Status = HiiDb->ListPackageLists(HiiDb, EFI_HII_PACKAGE_DEVICE_PATH, NULL, &HandleBufferSize, HandleBuffer);\r
   }\r
   if (EFI_ERROR(Status)) {\r
@@ -69,6 +70,7 @@ FindHiiHandleViaDevPath(
     Status = HiiDb->ExportPackageLists(HiiDb, HandleBuffer[LoopVariable], &MainBufferSize, MainBuffer);\r
     if (Status == EFI_BUFFER_TOO_SMALL) {\r
       MainBuffer = AllocateZeroPool(MainBufferSize);\r
+      ASSERT (MainBuffer != NULL);\r
       Status = HiiDb->ExportPackageLists(HiiDb, HandleBuffer[LoopVariable], &MainBufferSize, MainBuffer);\r
     }\r
     //\r