Add check to the allocated pool.

[mirror_edk2.git] / MdeModulePkg / Core / Dxe / FwVolBlock / FwVolBlock.c

diff --git a/MdeModulePkg/Core/Dxe/FwVolBlock/FwVolBlock.c b/MdeModulePkg/Core/Dxe/FwVolBlock/FwVolBlock.c
index 02b8914ff82ea767295775d21a1caea43013332c..9b7b166a9949596c3c8bf649835fa772d777b2d5 100644 (file)
--- a/MdeModulePkg/Core/Dxe/FwVolBlock/FwVolBlock.c
+++ b/MdeModulePkg/Core/Dxe/FwVolBlock/FwVolBlock.c
@@ -513,6 +513,10 @@ ProduceFVBProtocolOnBuffer (
     // FV does not contains extension header, then produce MEMMAP_DEVICE_PATH\r
     //\r
     FvbDev->DevicePath = (EFI_DEVICE_PATH_PROTOCOL *) AllocateCopyPool (sizeof (FV_MEMMAP_DEVICE_PATH), &mFvMemmapDevicePathTemplate);\r
+    if (FvbDev->DevicePath == NULL) {\r
+      FreePool (FvbDev);\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
     ((FV_MEMMAP_DEVICE_PATH *) FvbDev->DevicePath)->MemMapDevPath.StartingAddress = BaseAddress;\r
     ((FV_MEMMAP_DEVICE_PATH *) FvbDev->DevicePath)->MemMapDevPath.EndingAddress   = BaseAddress + FwVolHeader->FvLength - 1;\r
   } else {\r
@@ -520,6 +524,10 @@ ProduceFVBProtocolOnBuffer (
     // FV contains extension header, then produce MEDIA_FW_VOL_DEVICE_PATH\r
     //\r
     FvbDev->DevicePath = (EFI_DEVICE_PATH_PROTOCOL *) AllocateCopyPool (sizeof (FV_PIWG_DEVICE_PATH), &mFvPIWGDevicePathTemplate);\r
+    if (FvbDev->DevicePath == NULL) {\r
+      FreePool (FvbDev);\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
     CopyGuid (\r
       &((FV_PIWG_DEVICE_PATH *)FvbDev->DevicePath)->FvDevPath.FvName, \r
       (GUID *)(UINTN)(BaseAddress + FwVolHeader->ExtHeaderOffset)\r