Ported to GNU assembly.

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

diff --git a/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c b/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c
index a9bdf60aa1e0e1f2351e31e1684182533823c39a..ef643eed4296d2252506098a5bc5b18922425dc6 100644 (file)
--- a/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c
+++ b/MdePkg/Library/DxeMemoryAllocationLib/MemoryAllocationLib.c
@@ -168,6 +168,13 @@ InternalAllocateAlignedPages (
     //\r
     AlignmentMask  = Alignment - 1;\r
     RealPages      = Pages + EFI_SIZE_TO_PAGES (Alignment);\r
+    if (RealPages <= Pages) {\r
+      //\r
+      // This extra checking is to make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow. \r
+      //\r
+      return NULL;\r
+    }\r
+\r
     Status         = gBS->AllocatePages (AllocateAnyPages, MemoryType, RealPages, &Memory);\r
     if (EFI_ERROR (Status)) {\r
       return NULL;\r
@@ -576,8 +583,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 +600,15 @@ 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
+  if (RealAllocationSize <= AllocationSize ) {\r
+    //\r
+    // This extra checking is to make sure that AllocationSize plus OverAllocationSize does not overflow. \r
+    //\r
+    return NULL;\r
+  }\r
+  RawAddress = InternalAllocatePool (PoolType, RealAllocationSize);\r
+  if (RawAddress == NULL) {\r
     return NULL;\r
   }\r
   AlignedAddress      = ((UINTN) RawAddress + OverAllocationSize) & ~AlignmentMask;\r