[PATCH v2] EmbeddedPkg/PrePiMemoryAllocationLib: Check for space on offset allocation

["Jeff Brasen" <jbrasen@nvidia.com>]

Update check for enough space to occur prior to alignment offset
modification. This prevents a case where EfiFreeMemoryTop could be
less than EfiFreeMemoryBottom

Signed-off-by: Jeff Brasen <jbrasen@nvidia.com>
---
 .../MemoryAllocationLib.c                     | 53 +++++++++++--------
 1 file changed, 30 insertions(+), 23 deletions(-)

diff --git a/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c b/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c
index 2cc2a71121..9208826565 100644
--- a/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c
+++ b/EmbeddedPkg/Library/PrePiMemoryAllocationLib/MemoryAllocationLib.c
@@ -27,37 +27,44 @@ InternalAllocatePages (
 
   Hob.Raw = GetHobList ();
 
-  // Check to see if on 4k boundary
   Offset = Hob.HandoffInformationTable->EfiFreeMemoryTop & 0xFFF;
+  //
+  // Verify that there is sufficient memory to satisfy the allocation and padding prior to updating anything
+  //
+  if ((Hob.HandoffInformationTable->EfiFreeMemoryTop - ((Pages * EFI_PAGE_SIZE) + sizeof (EFI_HOB_MEMORY_ALLOCATION)) - Offset) < Hob.HandoffInformationTable->EfiFreeMemoryBottom) {
+    if (Offset != 0) {
+      DEBUG ((DEBUG_ERROR, "Offset applied without enough space\r\n"));
+    } else {
+      DEBUG ((DEBUG_ERROR, "Out of memory\r\n"));
+    }
+
+    ASSERT (FALSE);
+    return 0;
+  }
+
+  // Check to see if on 4k boundary
   if (Offset != 0) {
     // If not aligned, make the allocation aligned.
     Hob.HandoffInformationTable->EfiFreeMemoryTop -= Offset;
   }
 
   //
-  // Verify that there is sufficient memory to satisfy the allocation
+  // Update the PHIT to reflect the memory usage
   //
-  if (Hob.HandoffInformationTable->EfiFreeMemoryTop - ((Pages * EFI_PAGE_SIZE) + sizeof (EFI_HOB_MEMORY_ALLOCATION)) < Hob.HandoffInformationTable->EfiFreeMemoryBottom) {
-    return 0;
-  } else {
-    //
-    // Update the PHIT to reflect the memory usage
-    //
-    Hob.HandoffInformationTable->EfiFreeMemoryTop -= Pages * EFI_PAGE_SIZE;
-
-    // This routine used to create a memory allocation HOB a la PEI, but that's not
-    // necessary for us.
-
-    //
-    // Create a memory allocation HOB.
-    //
-    BuildMemoryAllocationHob (
-      Hob.HandoffInformationTable->EfiFreeMemoryTop,
-      Pages * EFI_PAGE_SIZE,
-      MemoryType
-      );
-    return (VOID *)(UINTN)Hob.HandoffInformationTable->EfiFreeMemoryTop;
-  }
+  Hob.HandoffInformationTable->EfiFreeMemoryTop -= Pages * EFI_PAGE_SIZE;
+
+  // This routine used to create a memory allocation HOB a la PEI, but that's not
+  // necessary for us.
+
+  //
+  // Create a memory allocation HOB.
+  //
+  BuildMemoryAllocationHob (
+    Hob.HandoffInformationTable->EfiFreeMemoryTop,
+    Pages * EFI_PAGE_SIZE,
+    MemoryType
+    );
+  return (VOID *)(UINTN)Hob.HandoffInformationTable->EfiFreeMemoryTop;
 }
 
 /**
-- 
2.25.1