EmbeddedPkg/PrePi: Check for enough space before aligning heap pointer

[mirror_edk2.git] / EmbeddedPkg / Library / PrePiMemoryAllocationLib / MemoryAllocationLib.c

diff --git a/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c b/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c
index 2cc2a711219790279cad73bc879e158ca9cce548..08a0add340cfb50eeeefdf93022a379127d6f169 100644 (file)
--- a/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c
+++ b/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c
@@ -23,41 +23,35 @@ InternalAllocatePages (
   )\r
 {\r
   EFI_PEI_HOB_POINTERS  Hob;\r
-  EFI_PHYSICAL_ADDRESS  Offset;\r
+  EFI_PHYSICAL_ADDRESS  NewTop;\r
 \r
   Hob.Raw = GetHobList ();\r
 \r
-  // Check to see if on 4k boundary\r
-  Offset = Hob.HandoffInformationTable->EfiFreeMemoryTop & 0xFFF;\r
-  if (Offset != 0) {\r
-    // If not aligned, make the allocation aligned.\r
-    Hob.HandoffInformationTable->EfiFreeMemoryTop -= Offset;\r
-  }\r
+  NewTop  = Hob.HandoffInformationTable->EfiFreeMemoryTop & ~(EFI_PHYSICAL_ADDRESS)EFI_PAGE_MASK;\r
+  NewTop -= Pages * EFI_PAGE_SIZE;\r
 \r
   //\r
   // Verify that there is sufficient memory to satisfy the allocation\r
   //\r
-  if (Hob.HandoffInformationTable->EfiFreeMemoryTop - ((Pages * EFI_PAGE_SIZE) + sizeof (EFI_HOB_MEMORY_ALLOCATION)) < Hob.HandoffInformationTable->EfiFreeMemoryBottom) {\r
-    return 0;\r
-  } else {\r
-    //\r
-    // Update the PHIT to reflect the memory usage\r
-    //\r
-    Hob.HandoffInformationTable->EfiFreeMemoryTop -= Pages * EFI_PAGE_SIZE;\r
-\r
-    // This routine used to create a memory allocation HOB a la PEI, but that's not\r
-    // necessary for us.\r
-\r
-    //\r
-    // Create a memory allocation HOB.\r
-    //\r
-    BuildMemoryAllocationHob (\r
-      Hob.HandoffInformationTable->EfiFreeMemoryTop,\r
-      Pages * EFI_PAGE_SIZE,\r
-      MemoryType\r
-      );\r
-    return (VOID *)(UINTN)Hob.HandoffInformationTable->EfiFreeMemoryTop;\r
+  if (NewTop < (Hob.HandoffInformationTable->EfiFreeMemoryBottom + sizeof (EFI_HOB_MEMORY_ALLOCATION))) {\r
+    return NULL;\r
   }\r
+\r
+  //\r
+  // Update the PHIT to reflect the memory usage\r
+  //\r
+  Hob.HandoffInformationTable->EfiFreeMemoryTop = NewTop;\r
+\r
+  //\r
+  // Create a memory allocation HOB.\r
+  //\r
+  BuildMemoryAllocationHob (\r
+    Hob.HandoffInformationTable->EfiFreeMemoryTop,\r
+    Pages * EFI_PAGE_SIZE,\r
+    MemoryType\r
+    );\r
+\r
+  return (VOID *)(UINTN)Hob.HandoffInformationTable->EfiFreeMemoryTop;\r
 }\r
 \r
 /**\r