MemoryAllocationLib: For boundary case: “AllocationSize + OverAllocationSize >= 4G”

[mirror_edk2.git] / MdePkg / Library / PeiMemoryAllocationLib / MemoryAllocationLib.c

diff --git a/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c b/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c
index 63641569ffe65d2450d59ee6c6e25ae8994faee9..274287cf035e216cc537f3c6aff59750efcac32a 100644 (file)
--- a/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c
+++ b/MdePkg/Library/PeiMemoryAllocationLib/MemoryAllocationLib.c
@@ -158,6 +158,10 @@ InternalAllocateAlignedPages (
     return NULL;\r
   }\r
   //\r
+  // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow.\r
+  //\r
+  ASSERT (Pages <= (MAX_ADDRESS - EFI_SIZE_TO_PAGES (Alignment)));\r
+  //\r
   // We would rather waste some memory to save PEI code size.\r
   //\r
   Memory = InternalAllocatePages (MemoryType, Pages + EFI_SIZE_TO_PAGES (Alignment));\r
@@ -569,6 +573,10 @@ InternalAllocateAlignedPool (
   } else {\r
     AlignmentMask = Alignment - 1;\r
   }\r
+  //\r
+  // Make sure that AllocationSize plus AlignmentMask does not overflow.\r
+  //\r
+  ASSERT (AllocationSize <= (MAX_ADDRESS - AlignmentMask));\r
 \r
   RawAddress      = InternalAllocatePool (PoolType, AllocationSize + AlignmentMask);\r
   \r
@@ -609,6 +617,11 @@ AllocateAlignedPool (
     AlignmentMask = Alignment - 1;\r
   }\r
 \r
+  //\r
+  // Make sure that AllocationSize plus AlignmentMask does not overflow.\r
+  //\r
+  ASSERT (AllocationSize <= (MAX_ADDRESS - AlignmentMask));\r
+\r
   RawAddress      = AllocatePool (AllocationSize + AlignmentMask);\r
   \r
   AlignedAddress  = ((UINTN) RawAddress + AlignmentMask) & ~AlignmentMask;\r