Merge changes from topic "aarch64-paging-0.5.0" into main

* changes:
  Use unsafe block to activate the ID map
  Map DBM data region down to pages
  Use statically typed callback function in modify_range
  Drop is_leaf() check now that modify_range() returns leaves only
  Use walk_range when traversing the page tables without modifying them
  Simplify lazy MMIO on-demand mapping
diff --git a/vmbase/src/memory/dbm.rs b/vmbase/src/memory/dbm.rs
index 401022e..108cd5d 100644
--- a/vmbase/src/memory/dbm.rs
+++ b/vmbase/src/memory/dbm.rs
@@ -14,7 +14,7 @@
 
 //! Hardware management of the access flag and dirty state.
 
-use super::page_table::{is_leaf_pte, PageTable};
+use super::page_table::PageTable;
 use super::util::flush_region;
 use crate::{dsb, isb, read_sysreg, tlbi, write_sysreg};
 use aarch64_paging::paging::{Attributes, Descriptor, MemoryRegion};
@@ -52,14 +52,10 @@
 /// Flushes a memory range the descriptor refers to, if the descriptor is in writable-dirty state.
 pub(super) fn flush_dirty_range(
     va_range: &MemoryRegion,
-    desc: &mut Descriptor,
-    level: usize,
+    desc: &Descriptor,
+    _level: usize,
 ) -> Result<(), ()> {
-    // Only flush ranges corresponding to dirty leaf PTEs.
     let flags = desc.flags().ok_or(())?;
-    if !is_leaf_pte(&flags, level) {
-        return Ok(());
-    }
     if !flags.contains(Attributes::READ_ONLY) {
         flush_region(va_range.start().0, va_range.len());
     }
@@ -71,12 +67,9 @@
 pub(super) fn mark_dirty_block(
     va_range: &MemoryRegion,
     desc: &mut Descriptor,
-    level: usize,
+    _level: usize,
 ) -> Result<(), ()> {
     let flags = desc.flags().ok_or(())?;
-    if !is_leaf_pte(&flags, level) {
-        return Ok(());
-    }
     if flags.contains(Attributes::DBM) {
         assert!(flags.contains(Attributes::READ_ONLY), "unexpected PTE writable state");
         desc.modify_flags(Attributes::empty(), Attributes::READ_ONLY);
diff --git a/vmbase/src/memory/page_table.rs b/vmbase/src/memory/page_table.rs
index e067e96..dc346e7 100644
--- a/vmbase/src/memory/page_table.rs
+++ b/vmbase/src/memory/page_table.rs
@@ -16,7 +16,7 @@
 
 use crate::read_sysreg;
 use aarch64_paging::idmap::IdMap;
-use aarch64_paging::paging::{Attributes, MemoryRegion, PteUpdater};
+use aarch64_paging::paging::{Attributes, Constraints, Descriptor, MemoryRegion};
 use aarch64_paging::MapError;
 use core::result;
 
@@ -83,7 +83,9 @@
     /// code being currently executed. Otherwise, the Rust execution model (on which the borrow
     /// checker relies) would be violated.
     pub unsafe fn activate(&mut self) {
-        self.idmap.activate()
+        // SAFETY: the caller of this unsafe function asserts that switching to a different
+        // translation is safe
+        unsafe { self.idmap.activate() }
     }
 
     /// Maps the given range of virtual addresses to the physical addresses as lazily mapped
@@ -107,7 +109,15 @@
     /// Maps the given range of virtual addresses to the physical addresses as non-executable,
     /// read-only and writable-clean normal memory.
     pub fn map_data_dbm(&mut self, range: &MemoryRegion) -> Result<()> {
-        self.idmap.map_range(range, DATA_DBM)
+        // Map the region down to pages to minimize the size of the regions that will be marked
+        // dirty once a store hits them, but also to ensure that we can clear the read-only
+        // attribute while the mapping is live without causing break-before-make (BBM) violations.
+        // The latter implies that we must avoid the use of the contiguous hint as well.
+        self.idmap.map_range_with_constraints(
+            range,
+            DATA_DBM,
+            Constraints::NO_BLOCK_MAPPINGS | Constraints::NO_CONTIGUOUS_HINT,
+        )
     }
 
     /// Maps the given range of virtual addresses to the physical addresses as read-only
@@ -124,18 +134,20 @@
 
     /// Applies the provided updater function to a number of PTEs corresponding to a given memory
     /// range.
-    pub fn modify_range(&mut self, range: &MemoryRegion, f: &PteUpdater) -> Result<()> {
+    pub fn modify_range<F>(&mut self, range: &MemoryRegion, f: &F) -> Result<()>
+    where
+        F: Fn(&MemoryRegion, &mut Descriptor, usize) -> result::Result<(), ()>,
+    {
         self.idmap.modify_range(range, f)
     }
-}
 
-/// Checks whether a PTE at given level is a page or block descriptor.
-#[inline]
-pub(super) fn is_leaf_pte(flags: &Attributes, level: usize) -> bool {
-    const LEAF_PTE_LEVEL: usize = 3;
-    if flags.contains(Attributes::TABLE_OR_PAGE) {
-        level == LEAF_PTE_LEVEL
-    } else {
-        level < LEAF_PTE_LEVEL
+    /// Applies the provided callback function to a number of PTEs corresponding to a given memory
+    /// range.
+    pub fn walk_range<F>(&self, range: &MemoryRegion, f: &F) -> Result<()>
+    where
+        F: Fn(&MemoryRegion, &Descriptor, usize) -> result::Result<(), ()>,
+    {
+        let mut callback = |mr: &MemoryRegion, d: &Descriptor, l: usize| f(mr, d, l);
+        self.idmap.walk_range(range, &mut callback)
     }
 }
diff --git a/vmbase/src/memory/shared.rs b/vmbase/src/memory/shared.rs
index 6c8a844..dd433d4 100644
--- a/vmbase/src/memory/shared.rs
+++ b/vmbase/src/memory/shared.rs
@@ -16,12 +16,14 @@
 
 use super::dbm::{flush_dirty_range, mark_dirty_block, set_dbm_enabled};
 use super::error::MemoryTrackerError;
-use super::page_table::{is_leaf_pte, PageTable, MMIO_LAZY_MAP_FLAG};
+use super::page_table::{PageTable, MMIO_LAZY_MAP_FLAG};
 use super::util::{page_4kb_of, virt_to_phys};
 use crate::dsb;
 use crate::exceptions::HandleExceptionError;
 use crate::util::RangeExt as _;
-use aarch64_paging::paging::{Attributes, Descriptor, MemoryRegion as VaRange, VirtualAddress};
+use aarch64_paging::paging::{
+    Attributes, Descriptor, MemoryRegion as VaRange, VirtualAddress, BITS_PER_LEVEL, PAGE_SIZE,
+};
 use alloc::alloc::{alloc_zeroed, dealloc, handle_alloc_error};
 use alloc::boxed::Box;
 use alloc::vec::Vec;
@@ -253,7 +255,7 @@
         if get_mmio_guard().is_some() {
             for range in &self.mmio_regions {
                 self.page_table
-                    .modify_range(&get_va_range(range), &mmio_guard_unmap_page)
+                    .walk_range(&get_va_range(range), &mmio_guard_unmap_page)
                     .map_err(|_| MemoryTrackerError::FailedToUnmap)?;
             }
         }
@@ -319,14 +321,24 @@
     /// table entry and MMIO guard mapping the block. Breaks apart a block entry if required.
     fn handle_mmio_fault(&mut self, addr: VirtualAddress) -> Result<()> {
         let page_start = VirtualAddress(page_4kb_of(addr.0));
+        assert_eq!(page_start.0 % MMIO_GUARD_GRANULE_SIZE, 0);
         let page_range: VaRange = (page_start..page_start + MMIO_GUARD_GRANULE_SIZE).into();
         let mmio_guard = get_mmio_guard().unwrap();
+        // This must be safe and free from break-before-make (BBM) violations, given that the
+        // initial lazy mapping has the valid bit cleared, and each newly created valid descriptor
+        // created inside the mapping has the same size and alignment.
         self.page_table
-            .modify_range(&page_range, &verify_lazy_mapped_block)
+            .modify_range(&page_range, &|_: &VaRange, desc: &mut Descriptor, _: usize| {
+                let flags = desc.flags().expect("Unsupported PTE flags set");
+                if flags.contains(MMIO_LAZY_MAP_FLAG) && !flags.contains(Attributes::VALID) {
+                    desc.modify_flags(Attributes::VALID, Attributes::empty());
+                    Ok(())
+                } else {
+                    Err(())
+                }
+            })
             .map_err(|_| MemoryTrackerError::InvalidPte)?;
-        mmio_guard.map(page_start.0)?;
-        // Maps a single device page, breaking up block mappings if necessary.
-        self.page_table.map_device(&page_range).map_err(|_| MemoryTrackerError::FailedToMap)
+        Ok(mmio_guard.map(page_start.0)?)
     }
 
     /// Flush all memory regions marked as writable-dirty.
@@ -340,7 +352,7 @@
         // Now flush writable-dirty pages in those regions.
         for range in writable_regions.chain(self.payload_range.as_ref().into_iter()) {
             self.page_table
-                .modify_range(&get_va_range(range), &flush_dirty_range)
+                .walk_range(&get_va_range(range), &flush_dirty_range)
                 .map_err(|_| MemoryTrackerError::FlushRegionFailed)?;
         }
         Ok(())
@@ -467,33 +479,13 @@
     }
 }
 
-/// Checks whether block flags indicate it should be MMIO guard mapped.
-fn verify_lazy_mapped_block(
-    _range: &VaRange,
-    desc: &mut Descriptor,
-    level: usize,
-) -> result::Result<(), ()> {
-    let flags = desc.flags().expect("Unsupported PTE flags set");
-    if !is_leaf_pte(&flags, level) {
-        return Ok(()); // Skip table PTEs as they aren't tagged with MMIO_LAZY_MAP_FLAG.
-    }
-    if flags.contains(MMIO_LAZY_MAP_FLAG) && !flags.contains(Attributes::VALID) {
-        Ok(())
-    } else {
-        Err(())
-    }
-}
-
 /// MMIO guard unmaps page
 fn mmio_guard_unmap_page(
     va_range: &VaRange,
-    desc: &mut Descriptor,
+    desc: &Descriptor,
     level: usize,
 ) -> result::Result<(), ()> {
     let flags = desc.flags().expect("Unsupported PTE flags set");
-    if !is_leaf_pte(&flags, level) {
-        return Ok(());
-    }
     // This function will be called on an address range that corresponds to a device. Only if a
     // page has been accessed (written to or read from), will it contain the VALID flag and be MMIO
     // guard mapped. Therefore, we can skip unmapping invalid pages, they were never MMIO guard
@@ -503,9 +495,11 @@
             flags.contains(MMIO_LAZY_MAP_FLAG),
             "Attempting MMIO guard unmap for non-device pages"
         );
+        const MMIO_GUARD_GRANULE_SHIFT: u32 = MMIO_GUARD_GRANULE_SIZE.ilog2() - PAGE_SIZE.ilog2();
+        const MMIO_GUARD_GRANULE_LEVEL: usize =
+            3 - (MMIO_GUARD_GRANULE_SHIFT as usize / BITS_PER_LEVEL);
         assert_eq!(
-            va_range.len(),
-            MMIO_GUARD_GRANULE_SIZE,
+            level, MMIO_GUARD_GRANULE_LEVEL,
             "Failed to break down block mapping before MMIO guard mapping"
         );
         let page_base = va_range.start().0;