MdeModulePkg/CapsulePei: Add memory pointer check

[mirror_edk2.git] / MdeModulePkg / Universal / CapsulePei / UefiCapsule.c

diff --git a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
index 8d4ae69bb2e9f14c0257577c2518d855bf4c8aea..51afab7b05752630cf8ae267a10f26e081473a50 100644 (file)
--- a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
+++ b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
@@ -965,6 +965,10 @@ GetScatterGatherHeadEntries (
     //\r
     if ((ValidIndex + 1) >= TempListLength) {\r
       EnlargedTempList = AllocateZeroPool (TempListLength * 2);\r
+      if (EnlargedTempList == NULL) {\r
+        DEBUG ((DEBUG_ERROR, "Fail to allocate memory!\n"));\r
+        return EFI_OUT_OF_RESOURCES;\r
+      }\r
       CopyMem (EnlargedTempList, TempList, TempListLength);\r
       FreePool (TempList);\r
       TempList = EnlargedTempList;\r
@@ -1056,7 +1060,7 @@ CapsuleCoalesce (
   // Get SG list entries\r
   //\r
   Status = GetScatterGatherHeadEntries (&ListLength, &VariableArrayAddress);\r
-  if (EFI_ERROR (Status)) {\r
+  if (EFI_ERROR (Status) || VariableArrayAddress == NULL) {\r
     DEBUG ((DEBUG_ERROR, "%a failed to get Scatter Gather List Head Entries.  Status = %r\n", __FUNCTION__, Status));\r
     goto Done;\r
   }\r