[PATCH v3] 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 | 20 +++++++++++++++++++-
 2 files changed, 20 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..f95b5e5 100644
--- a/MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c
+++ b/MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c
@@ -567,6 +567,7 @@ UsbHcAllocateAlignedPages (
   EFI_PHYSICAL_ADDRESS  DeviceMemory;
   UINTN                 AlignedDeviceMemory;
   UINTN                 RealPages;
+  UINTN                 UnalignedPages;
 
   //
   // Alignment must be a power of two or zero.
@@ -593,7 +594,7 @@ UsbHcAllocateAlignedPages (
     ASSERT (RealPages > Pages);
 
     Status = IoMmuAllocateBuffer (
-               Pages,
+               RealPages,
                &Memory,
                &DeviceMemory,
                Mapping
@@ -603,6 +604,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