pci: align VFIO MMIO mmap to a huge page boundary

mmap(NULL, ...) only guarantees page-sized alignment, so the device
MMIO regions mapped for VFIO passthrough were not necessarily aligned
on a huge page boundary and thus could not take fast paths in the
kernel.

Align the address to the largest possible size. This allows VFIO to take
fast paths when walking the page table.

Signed-off-by: Wei Liu <liuwe@microsoft.com>
Assisted-by: Copilot:Opus-4.8
Assisted-by: Copilot:GPT-5.5

pci/src/mmap.rs

@@ -4,13 +4,165 @@
 
 //! Helpers for `mmap()`
 
-use core::ffi::c_int;
+use core::ffi::{c_int, c_void};
 use core::ptr::null_mut;
 use std::io::{self, Error, ErrorKind};
 use std::os::fd::{AsRawFd as _, BorrowedFd};
 
 use libc::size_t;
 use log::warn;
+use vm_allocator::page_size::get_page_size;
+
+const ONE_GIB: usize = 1024 * 1024 * 1024;
+
+/// # SAFETY
+///
+/// Callers must guarantee that the range can be passed to munmap().
+unsafe fn munmap(addr: *mut c_void, len: size_t) {
+    // SAFETY: see function comment
+    let ret = unsafe { libc::munmap(addr, len) };
+
+    if ret != 0 {
+        warn!(
+            "Failed to munmap region address {:p} length 0x{:x}, {}, leaking...",
+            addr,
+            len,
+            Error::last_os_error()
+        );
+    }
+}
+
+fn find_alignment(len: size_t) -> usize {
+    let page_size = get_page_size() as usize;
+    if len >= ONE_GIB {
+        ONE_GIB
+    } else {
+        len.next_power_of_two().max(page_size)
+    }
+}
+
+fn reserve_mapping(reserve: size_t, alignment: usize) -> io::Result<(*mut c_void, *mut c_void)> {
+    // SAFETY: FFI call. Reserving address space with a NULL hint and no fd.
+    let base = unsafe {
+        libc::mmap(
+            null_mut(),
+            reserve,
+            libc::PROT_NONE,
+            libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_NORESERVE,
+            -1,
+            0,
+        )
+    };
+    if base == libc::MAP_FAILED {
+        return Err(Error::last_os_error());
+    }
+
+    let aligned = (base as usize).next_multiple_of(alignment) as *mut c_void;
+
+    Ok((base, aligned))
+}
+
+fn tail_trim_range(
+    base_addr: usize,
+    reserve: size_t,
+    aligned_addr: usize,
+    len: size_t,
+) -> (usize, size_t) {
+    let page_size = get_page_size() as usize;
+    let mapped_end = aligned_addr + len;
+    // munmap() requires a page-aligned address, so trim from the next page.
+    let tail_addr = mapped_end.next_multiple_of(page_size);
+    let tail_len = (base_addr + reserve) - tail_addr;
+    assert!(tail_len > 0);
+
+    (tail_addr, tail_len)
+}
+
+/// # SAFETY
+///
+/// Callers must make sure base and aligned come from the same valid mapping
+/// from mmap.
+unsafe fn trim_mapping(base: *mut c_void, reserve: size_t, aligned: *mut c_void, len: size_t) {
+    let base_addr = base as usize;
+    let aligned_addr = aligned as usize;
+
+    let head_len = aligned_addr - base_addr;
+    if head_len > 0 {
+        // SAFETY: [base, base+head_len) is a valid region
+        unsafe {
+            munmap(base, head_len);
+        }
+    }
+
+    let (tail_addr, tail_len) = tail_trim_range(base_addr, reserve, aligned_addr, len);
+    // SAFETY: [tail_addr, tail_addr+tail_len) is a valid region
+    unsafe {
+        munmap(tail_addr as *mut c_void, tail_len);
+    }
+}
+
+fn allocate_aligned_mapping(
+    len: size_t,
+    prot: c_int,
+    fd: BorrowedFd,
+    offset: libc::off_t,
+) -> io::Result<*mut c_void> {
+    let alignment = find_alignment(len);
+
+    if alignment == get_page_size() as usize {
+        // SAFETY: FFI call with correct parameters.
+        let addr = unsafe {
+            libc::mmap(
+                null_mut(),
+                len,
+                prot,
+                libc::MAP_SHARED,
+                fd.as_raw_fd(),
+                offset,
+            )
+        };
+        if addr == libc::MAP_FAILED {
+            return Err(Error::last_os_error());
+        }
+        return Ok(addr);
+    }
+
+    let Some(reserve) = len.checked_add(alignment) else {
+        return Err(Error::new(
+            ErrorKind::InvalidInput,
+            "Length+Alignment overflows",
+        ));
+    };
+
+    let (base, aligned) = reserve_mapping(reserve, alignment)?;
+    // SAFETY: base and aligned are from the same mmap region
+    unsafe {
+        trim_mapping(base, reserve, aligned, len);
+    }
+
+    // SAFETY: FFI call. MAP_FIXED is safe here because it only replaces the
+    // remaining anonymous reservation we just created.
+    let addr = unsafe {
+        libc::mmap(
+            aligned,
+            len,
+            prot,
+            libc::MAP_SHARED | libc::MAP_FIXED,
+            fd.as_raw_fd(),
+            offset,
+        )
+    };
+    if addr == libc::MAP_FAILED {
+        let err = Error::last_os_error();
+        // SAFETY: [aligned, aligned+len) is a valid region
+        unsafe {
+            munmap(aligned, len);
+        };
+        return Err(err);
+    }
+
+    Ok(addr)
+}
 
 /// A region of `mmap()`-allocated memory that calls `munmap()` when dropped.
 /// This guarantees that the buffer is valid and that its address space
@@ -27,15 +179,8 @@ pub struct MmapRegion {
 impl Drop for MmapRegion {
     fn drop(&mut self) {
         // SAFETY: guaranteed by type validity invariant
-        let ret = unsafe { libc::munmap(self.addr.cast(), self.len) };
-
-        if ret != 0 {
-            warn!(
-                "Failed to munmap region address {:p} length 0x{:x}, {}, leaking...",
-                self.addr,
-                self.len,
-                Error::last_os_error()
-            );
+        unsafe {
+            munmap(self.addr.cast(), self.len);
         }
     }
 }
@@ -86,14 +231,90 @@ in both isize and libc::size_t";
             (prot & !(libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC)) == 0,
             "bad protection"
         );
-        let flags = libc::MAP_SHARED;
-        // SAFETY: FFI call with correct parameters.
-        let addr = unsafe { libc::mmap(null_mut(), len, prot, flags, fd.as_raw_fd(), offset) };
-        if addr == libc::MAP_FAILED {
-            Err(Error::last_os_error())
-        } else {
-            let addr = addr.cast();
-            Ok(Self { addr, len })
-        }
+
+        let addr = allocate_aligned_mapping(len, prot, fd, offset)?;
+
+        Ok(Self {
+            addr: addr.cast(),
+            len,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn tail_trim_range_from_len(base_addr: usize, len: size_t) -> (usize, size_t) {
+        let alignment = find_alignment(len);
+        let reserve = len.checked_add(alignment).unwrap();
+        let aligned_addr = base_addr.next_multiple_of(alignment);
+
+        tail_trim_range(base_addr, reserve, aligned_addr, len)
+    }
+
+    #[test]
+    fn find_alignment_scales_with_length() {
+        let page_size = get_page_size() as usize;
+
+        // Sub-page and page-sized mappings only need page alignment.
+        assert_eq!(find_alignment(1), page_size);
+        assert_eq!(find_alignment(page_size), page_size);
+
+        // Larger mappings round up to the next power of two.
+        assert_eq!(
+            find_alignment(page_size + 1),
+            (page_size + 1).next_power_of_two()
+        );
+        assert!(find_alignment(page_size + 1) > page_size);
+
+        // The alignment is capped at 1 GiB.
+        assert_eq!(find_alignment(ONE_GIB), ONE_GIB);
+        assert_eq!(find_alignment(ONE_GIB + 1), ONE_GIB);
+        assert_eq!(find_alignment(2 * ONE_GIB), ONE_GIB);
+    }
+
+    #[test]
+    fn tail_trim_range_rounds_non_page_sized_len_up() {
+        let page_size = get_page_size() as usize;
+        let base_addr = 0x1000_0000;
+        let len = page_size * 3 + 1;
+        let alignment = find_alignment(len);
+
+        let (tail_addr, tail_len) = tail_trim_range_from_len(base_addr, len);
+
+        // The tail starts at the first page boundary past the mapping ...
+        assert_eq!(tail_addr, base_addr + page_size * 4);
+        assert_eq!(tail_addr % page_size, 0);
+        // ... and reaches the end of the reserved region.
+        assert_eq!(tail_addr + tail_len, base_addr + len + alignment);
+    }
+
+    #[test]
+    fn tail_trim_range_uses_exact_end_for_page_sized_len() {
+        let page_size = get_page_size() as usize;
+        let base_addr = 0x1000_0000;
+        let len = page_size * 3;
+        let alignment = find_alignment(len);
+
+        let (tail_addr, tail_len) = tail_trim_range_from_len(base_addr, len);
+
+        // A page-aligned end needs no rounding.
+        assert_eq!(tail_addr, base_addr + len);
+        assert_eq!(tail_addr + tail_len, base_addr + len + alignment);
+    }
+
+    #[test]
+    fn tail_trim_range_trims_after_partial_final_page() {
+        let page_size = get_page_size() as usize;
+        let base_addr = 0x1000_0000;
+        let len = page_size * 512 + 1;
+        let alignment = find_alignment(len);
+
+        let (tail_addr, tail_len) = tail_trim_range_from_len(base_addr, len);
+
+        assert_eq!(tail_addr, base_addr + page_size * 513);
+        assert_eq!(tail_addr % page_size, 0);
+        assert_eq!(tail_addr + tail_len, base_addr + len + alignment);
     }
 }