[PATCH v4 01/32] gpu: nova-core: mm: retain BAR1 reservations after unmap failure

[Alok Kumar <alkumar@nvidia.com>]

From: Zhi Wang <zhiw@nvidia.com>

Vmm::unmap_pages currently returns a virtual address range to its pool
even when PTE invalidation or its TLB flush fails. Hardware can still
use that mapping, so reusing the range can alias a live mapping.

Borrow MappedRange during unmap and release its VA reservation only
after invalidation succeeds. Disarm its drop guard at that point and
reject a second unmap of a range whose guard is already disarmed.

Keep BarUserAccess's mapped state until unmap succeeds, and update its
callers to pass a borrowed range. A failed unmap leaves the VA reserved;
callers must retain the backing storage while hardware may still access
it.

Extend the MM self-test to unmap an old range after its virtual address
has been reused. Require the stale token to fail without invalidating
the new mapping.

Signed-off-by: Zhi Wang <zhiw@nvidia.com>
Signed-off-by: Alok Kumar <alkumar@nvidia.com>
---
 drivers/gpu/nova-core/mm/bar_user.rs | 35 +++++++++++++++++++---------
 drivers/gpu/nova-core/mm/vmm.rs      | 18 +++++++-------
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/drivers/gpu/nova-core/mm/bar_user.rs b/drivers/gpu/nova-core/mm/bar_user.rs
index 6e9384f06f43..622a79f1cc22 100644
--- a/drivers/gpu/nova-core/mm/bar_user.rs
+++ b/drivers/gpu/nova-core/mm/bar_user.rs
@@ -83,8 +83,7 @@ pub(crate) fn map(
 /// Borrows the [`BarUser`] managing the mapping for the lifetime of the access object.
 pub(crate) struct BarUserAccess<'a> {
     bar_user: &'a BarUser<'a>,
-    /// [`BarUserAccess::release`] [`Option::take`]s this; `Some` at
-    /// drop time means `release()` was never called.
+    /// Cleared only after PTE invalidation and its TLB flush succeed.
     mapped: Option<MappedRange>,
 }
 
@@ -92,16 +91,22 @@ pub(crate) struct BarUserAccess<'a> {
 impl BarUserAccess<'_> {
     /// Tear down the BAR1 mapping.
     pub(crate) fn release(mut self, mm: &mut GpuMm<'_>) -> Result {
-        let mapped = self.mapped.take().ok_or(EINVAL)?;
+        self.unmap(mm)
+    }
+
+    fn unmap(&mut self, mm: &mut GpuMm<'_>) -> Result {
+        let Some(mapped) = self.mapped.as_ref() else {
+            return Ok(());
+        };
         let mut vmm = self.bar_user.vmm.lock();
         vmm.unmap_pages(mm, mapped)?;
+        self.mapped = None;
         Ok(())
     }
 
     /// Returns the active mapping.
     fn mapped(&self) -> &MappedRange {
-        // `mapped` is only `None` after `take()` in `release`; hence unwrap()
-        // cannot panic here.
+        // `release` consumes the access, so no accessor can run after successful release.
         self.mapped.as_ref().unwrap()
     }
 
@@ -167,7 +172,7 @@ impl Drop for BarUserAccess<'_> {
     fn drop(&mut self) {
         if self.mapped.is_some() {
             kernel::pr_warn!(
-                "BarUserAccess dropped without calling release(). BarUser address space will leak.\n"
+                "BarUserAccess dropped with a live mapping; BAR1 VA remains reserved.\n"
             );
         }
         // The inner `MappedRange`'s own `MustUnmapGuard` will also fire,
@@ -266,19 +271,27 @@ pub(crate) fn run_self_test(
     };
 
     // Cleanup - invalidate PTE.
-    vmm.unmap_pages(mm, mapped)?;
+    vmm.unmap_pages(mm, &mapped)?;
 
     // Test 2: Two-phase prepare/execute API.
     let prepared = vmm.prepare_map(mm, 1, None)?;
     let mapped2 = vmm.execute_map(mm, prepared, &[test_pfn], true)?;
+    // An old token must not invalidate a new mapping that reused its address.
+    let repeated_unmap = vmm.unmap_pages(mm, &mapped);
     let readback = vmm.read_mapping(mm, mapped2.vfn_start)?;
-    let test2_passed = if readback == Some(test_pfn) {
+    let test2_passed = if mapped2.vfn_start == test_vfn
+        && repeated_unmap == Err(EINVAL)
+        && readback == Some(test_pfn)
+    {
         true
     } else {
-        dev_err!(dev, "MM: Test 2 FAILED - Two-phase map readback mismatch\n");
+        dev_err!(
+            dev,
+            "MM: Test 2 FAILED - Remapped page or stale-token rejection mismatch\n"
+        );
         false
     };
-    vmm.unmap_pages(mm, mapped2)?;
+    vmm.unmap_pages(mm, &mapped2)?;
 
     // Test 3: Range-constrained allocation with a hole — exercises block.size()-driven
     // BAR1 mapping. A 4K hole is punched at base+16K, then a single 32K allocation
@@ -370,7 +383,7 @@ pub(crate) fn run_self_test(
             }
         }
 
-        vmm.unmap_pages(mm, mapped)?;
+        vmm.unmap_pages(mm, &mapped)?;
     }
 
     // Verify aggregate: all returned block sizes must sum to allocation size.
diff --git a/drivers/gpu/nova-core/mm/vmm.rs b/drivers/gpu/nova-core/mm/vmm.rs
index 51b500a27233..49b1e5db1498 100644
--- a/drivers/gpu/nova-core/mm/vmm.rs
+++ b/drivers/gpu/nova-core/mm/vmm.rs
@@ -324,23 +324,25 @@ pub(crate) fn map_pages(
     }
 
     /// Unmap all pages in a [`MappedRange`] with a single TLB flush.
-    pub(crate) fn unmap_pages(&mut self, mm: &mut GpuMm<'_>, range: MappedRange) -> Result {
-        let result = self
-            .pt_map
-            .invalidate_ptes(mm, range.vfn_start, range.num_pages);
+    ///
+    /// A failed invalidation retains the VA reservation. The caller must keep
+    /// the mapped storage alive while hardware may still access it.
+    pub(crate) fn unmap_pages(&mut self, mm: &mut GpuMm<'_>, range: &MappedRange) -> Result {
+        if !range._drop_guard.armed.get() {
+            return Err(EINVAL);
+        }
+        self.pt_map
+            .invalidate_ptes(mm, range.vfn_start, range.num_pages)?;
 
         // TODO: Internal page table pages (PDE, PTE pages) are still kept around.
         // This is by design as repeated maps/unmaps will be fast. As a future TODO,
         // we can add a reclaimer here to reclaim if VRAM is short. For now, the PT
         // pages are dropped once the `Vmm` is dropped.
 
-        // Free the VA range regardless of PTE invalidation success, so that the VA
-        // range is recovered even on failure (PTEs may be stale, but that is better
-        // than leaking both PTEs and VA range).
         self.free_vfn(range.vfn_start);
 
         // Unmap complete, safe to drop `MappedRange`.
         range._drop_guard.disarm();
-        result
+        Ok(())
     }
 }
-- 
2.53.0