[PATCH v2] MdeModulePkg/XhciDxe: Fix Aligned Page Allocation

["Ashish Singhal" <ashishsingha@nvidia.com>]

While allocating pages aligned at an alignment higher than
4K, allocate memory taking into consideration the padding
required for that alignment. The calls to free pages takes
care of this already.

Signed-off-by: Ashish Singhal <ashishsingha@nvidia.com>
---
 MdeModulePkg/Bus/Pci/XhciDxe/UsbHcMem.c |  2 +-
 MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c | 19 ++++++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/MdeModulePkg/Bus/Pci/XhciDxe/UsbHcMem.c b/MdeModulePkg/Bus/Pci/XhciDxe/UsbHcMem.c
index fd79988..aa69c47 100644
--- a/MdeModulePkg/Bus/Pci/XhciDxe/UsbHcMem.c
+++ b/MdeModulePkg/Bus/Pci/XhciDxe/UsbHcMem.c
@@ -656,7 +656,7 @@ UsbHcAllocateAlignedPages (
                       PciIo,
                       AllocateAnyPages,
                       EfiBootServicesData,
-                      Pages,
+                      RealPages,
                       &Memory,
                       0
                       );
diff --git a/MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c b/MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c
index 56c0b90..074beee 100644
--- a/MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c
+++ b/MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c
@@ -593,7 +593,7 @@ UsbHcAllocateAlignedPages (
     ASSERT (RealPages > Pages);
 
     Status = IoMmuAllocateBuffer (
-               Pages,
+               RealPages,
                &Memory,
                &DeviceMemory,
                Mapping
@@ -603,6 +603,23 @@ UsbHcAllocateAlignedPages (
     }
     AlignedMemory = ((UINTN) Memory + AlignmentMask) & ~AlignmentMask;
     AlignedDeviceMemory = ((UINTN) DeviceMemory + AlignmentMask) & ~AlignmentMask;
+    UnalignedPages = EFI_SIZE_TO_PAGES (AlignedMemory - (UINTN) Memory);
+    if (UnalignedPages > 0) {
+      //
+      // Free first unaligned page(s).
+      //
+      Status = IoMmuFreeBuffer (UnalignedPages, Memory, Mapping);
+      ASSERT_EFI_ERROR (Status);
+    }
+    Memory = (VOID *)(UINTN)(AlignedMemory + EFI_PAGES_TO_SIZE (Pages));
+    UnalignedPages = RealPages - Pages - UnalignedPages;
+    if (UnalignedPages > 0) {
+      //
+      // Free last unaligned page(s).
+      //
+      Status = IoMmuFreeBuffer (UnalignedPages, Memory, Mapping);
+      ASSERT_EFI_ERROR (Status);
+    }
   } else {
     //
     // Do not over-allocate pages in this case.
-- 
2.7.4