MdeModulePkg/DxeCapsuleLibFmp: Add missing NULL pointer check.

[mirror_edk2.git] / MdeModulePkg / Library / DxeCapsuleLibFmp / CapsuleOnDisk.c

diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/CapsuleOnDisk.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/CapsuleOnDisk.c
index 66c9be8e1fe9f529dcf2430e042d86463ea4de6c..3193ca8f4d6ff0147de0f49327126b5695fee9e0 100644 (file)
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/CapsuleOnDisk.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/CapsuleOnDisk.c
@@ -25,6 +25,9 @@ IsCapsuleNameCapsule (
   Check the integrity of the capsule name capsule.\r
   If the capsule is vaild, return the physical address of each capsule name string.\r
 \r
+  This routine assumes the capsule has been validated by IsValidCapsuleHeader(), so\r
+  capsule memory overflow is not going to happen in this routine.\r
+\r
   @param[in]  CapsuleHeader   Pointer to the capsule header of a capsule name capsule.\r
   @param[out] CapsuleNameNum  Number of capsule name.\r
 \r
@@ -65,6 +68,9 @@ ValidateCapsuleNameCapsuleIntegrity (
   //\r
   if (((UINTN) CapsuleNameBufStart & BIT0) != 0) {\r
     CapsuleNameBufStart = AllocateCopyPool (CapsuleHeader->CapsuleImageSize - CapsuleHeader->HeaderSize, CapsuleNameBufStart);\r
+    if (CapsuleNameBufStart == NULL) {\r
+      return NULL;\r
+    }\r
   }\r
 \r
   CapsuleNameBufEnd = CapsuleNameBufStart + CapsuleHeader->CapsuleImageSize - CapsuleHeader->HeaderSize;\r