Fixed GCC ld issue

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

diff --git a/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c b/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c
index a9bdf60aa1e0e1f2351e31e1684182533823c39a..15a419bc566af3dc8e6f178e5c4afb1418ee51a9 100644 (file)
--- a/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c
+++ b/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c
@@ -168,6 +168,11 @@ InternalAllocateAlignedPages (
     //\r
     AlignmentMask  = Alignment - 1;\r
     RealPages      = Pages + EFI_SIZE_TO_PAGES (Alignment);\r
+    //\r
+    // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow.\r
+    //\r
+    ASSERT (RealPages > Pages);\r
+ \r
     Status         = gBS->AllocatePages (AllocateAnyPages, MemoryType, RealPages, &Memory);\r
     if (EFI_ERROR (Status)) {\r
       return NULL;\r
@@ -469,6 +474,9 @@ InternalAllocateCopyPool (
 {\r
   VOID  *Memory;\r
 \r
+  ASSERT (Buffer != NULL);\r
+  ASSERT (AllocationSize <= (MAX_ADDRESS - (UINTN) Buffer + 1));\r
+\r
   Memory = InternalAllocatePool (PoolType, AllocationSize);\r
   if (Memory != NULL) {\r
      Memory = CopyMem (Memory, Buffer, AllocationSize);\r
@@ -576,8 +584,8 @@ InternalAllocateAlignedPool (
   UINTN       AlignedAddress;\r
   UINTN       AlignmentMask;\r
   UINTN       OverAllocationSize;\r
+  UINTN       RealAllocationSize;\r
   VOID        **FreePointer;\r
-  EFI_STATUS  Status;\r
 \r
   //\r
   // Alignment must be a power of two or zero.\r
@@ -593,8 +601,14 @@ InternalAllocateAlignedPool (
   // Calculate the extra memory size, over-allocate memory pool and get the aligned memory address. \r
   //\r
   OverAllocationSize  = sizeof (RawAddress) + AlignmentMask;\r
-  Status = gBS->AllocatePool (PoolType, AllocationSize + OverAllocationSize, &RawAddress);\r
-  if (EFI_ERROR (Status)) {\r
+  RealAllocationSize  = AllocationSize + OverAllocationSize;\r
+  //\r
+  // Make sure that AllocationSize plus OverAllocationSize does not overflow. \r
+  //\r
+  ASSERT (RealAllocationSize > AllocationSize); \r
+\r
+  RawAddress = InternalAllocatePool (PoolType, RealAllocationSize);\r
+  if (RawAddress == NULL) {\r
     return NULL;\r
   }\r
   AlignedAddress      = ((UINTN) RawAddress + OverAllocationSize) & ~AlignmentMask;\r
@@ -782,6 +796,9 @@ InternalAllocateAlignedCopyPool (
 {\r
   VOID  *Memory;\r
   \r
+  ASSERT (Buffer != NULL);\r
+  ASSERT (AllocationSize <= (MAX_ADDRESS - (UINTN) Buffer + 1));\r
+\r
   Memory = InternalAllocateAlignedPool (PoolType, AllocationSize, Alignment);\r
   if (Memory != NULL) {\r
     Memory = CopyMem (Memory, Buffer, AllocationSize);\r