supermachine 0.7.46

// DAX session — orchestrates SETUPMAPPING/REMOVEMAPPING by combining a
// FUSE filesystem backend (which knows how to mmap a file region into
// the host process) with an HVF mapper (which knows how to plumb that
// host VA into the guest's DAX window via `hv_vm_map`).
//
// This module is deliberately VM-substrate-agnostic: it takes an
// `Arc<dyn HvfMapper>` so tests can drive the full session machinery
// without a real VM. The real impl lives in `crate::hvf::dax_mapper`
// (added in a subsequent slice).

use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;

use super::backend::{Errno, FsBackend};
use super::protocol::{
    RemoveMappingIn, RemoveMappingOne, SetupMappingIn, FUSE_SETUPMAPPING_FLAG_READ,
    FUSE_SETUPMAPPING_FLAG_WRITE,
};

/// Permission bits we pass through to the underlying HVF mapping.
/// We deliberately keep the bit shape stable across HVF / KVM so the
/// trait is portable.
pub const DAX_PROT_READ: u32 = 1 << 0;
pub const DAX_PROT_WRITE: u32 = 1 << 1;

/// Abstracts `hv_vm_map` / `hv_vm_unmap` so the DAX session can be
/// tested without a real HVF VM. The real impl wraps applevisor-sys
/// in `crate::hvf` and is wired into the worker once we plumb the
/// device into the VMM device list.
pub trait HvfMapper: Send + Sync {
    /// Map `len` bytes from `host_va` into guest physical at
    /// `gpa` with `prot` (DAX_PROT_*).
    fn map(&self, host_va: *mut u8, gpa: u64, len: u64, prot: u32) -> Result<(), Errno>;

    /// Unmap a previously-established region. `gpa` + `len` must be
    /// 16 KiB-aligned and match an existing mapping; HVF tolerates
    /// unmapping the exact range we mapped.
    fn unmap(&self, gpa: u64, len: u64) -> Result<(), Errno>;
}

/// A no-op mapper used by tests + the not-yet-wired-up production
/// path. Records every map/unmap call so tests can inspect them.
pub struct MockHvfMapper {
    log: Mutex<Vec<MockCall>>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum MockCall {
    Map { host_va: usize, gpa: u64, len: u64, prot: u32 },
    Unmap { gpa: u64, len: u64 },
}

impl MockHvfMapper {
    pub fn new() -> Self {
        Self {
            log: Mutex::new(Vec::new()),
        }
    }
    pub fn calls(&self) -> Vec<MockCall> {
        self.log.lock().unwrap().clone()
    }
}

impl Default for MockHvfMapper {
    fn default() -> Self {
        Self::new()
    }
}

impl HvfMapper for MockHvfMapper {
    fn map(&self, host_va: *mut u8, gpa: u64, len: u64, prot: u32) -> Result<(), Errno> {
        self.log.lock().unwrap().push(MockCall::Map {
            host_va: host_va as usize,
            gpa,
            len,
            prot,
        });
        Ok(())
    }
    fn unmap(&self, gpa: u64, len: u64) -> Result<(), Errno> {
        self.log.lock().unwrap().push(MockCall::Unmap { gpa, len });
        Ok(())
    }
}

/// One active DAX mapping. Recorded by `DaxSession` so REMOVEMAPPING
/// can find the right host VA + len to tear down — and so a restored
/// session can re-establish the host mmap on first stage-2 fault.
///
/// `host_va = None` means the slot was hydrated from snapshot metadata
/// but has NOT yet been re-bound. The backing mmap (and its host VA)
/// is a process-local thing that can't survive daemon restart, so we
/// only persist the FUSE-level metadata (`nodeid` / `foffset` / `prot`)
/// — enough to reproduce the SETUPMAPPING that originally created
/// the slot. The first guest fault into the slot's GPA range triggers
/// `handle_stage2_fault`, which calls `backend.dax_map` + `mapper.map`
/// and flips `host_va` to `Some(...)`.
#[derive(Clone, Debug)]
struct Slot {
    /// None = slot is metadata-only (post-restore, not yet bound).
    /// Some = actively mapped host VA.
    host_va: Option<*mut u8>,
    len: u64,
    /// Fields needed to lazily re-mmap on first stage-2 fault.
    nodeid: u64,
    foffset: u64,
    prot: u32,
}

// SAFETY: host_va is a process-local pointer used only by the
// owning DaxSession (which itself is Send+Sync). We don't dereference
// it across threads — only HvfMapper does the actual map/unmap.
// `Option<*mut u8>` is `!Send` by default (because raw pointers are),
// so the unsafe impl is still required after the `Option` wrap.
unsafe impl Send for Slot {}
unsafe impl Sync for Slot {}

/// Per-mount DAX state. Owned by the FuseServer when a DAX mount is
/// configured.
pub struct DaxSession {
    /// Guest physical base of the DAX window. All FUSE moffset values
    /// are relative to this.
    dax_base_gpa: u64,
    /// Window length. SETUPMAPPING with moffset+len > window_len is rejected.
    dax_window_len: u64,
    /// Aggregate cap on the sum of currently-mapped byte counts.
    /// SETUPMAPPING with `currently_mapped + req.len > cap` returns
    /// ENOSPC. Defends against a hostile guest churning
    /// map/remap/leak cycles to consume host stage-2 page-table
    /// memory. Default = `dax_window_len` (the natural upper bound).
    cap_bytes: u64,
    /// Running sum of `Slot.len` across all active entries. Atomic so
    /// concurrent SETUPMAPPING calls don't race.
    current_mapped_bytes: AtomicU64,
    /// Backend that knows how to mmap file regions into host VA space.
    backend: std::sync::Arc<dyn FsBackend>,
    /// HVF mapper.
    mapper: std::sync::Arc<dyn HvfMapper>,
    /// Active mappings, keyed by moffset (window-relative).
    active: Mutex<BTreeMap<u64, Slot>>,
}

impl DaxSession {
    pub fn new(
        dax_base_gpa: u64,
        dax_window_len: u64,
        backend: std::sync::Arc<dyn FsBackend>,
        mapper: std::sync::Arc<dyn HvfMapper>,
    ) -> Self {
        Self::with_cap(dax_base_gpa, dax_window_len, dax_window_len, backend, mapper)
    }

    /// Construct with an explicit aggregate cap. The cap bounds the
    /// SUM of currently-active mapping lengths. Pass a smaller value
    /// than the window to tighten the limit for hostile-tenant
    /// deployments — e.g. cap at 2 GiB while the window is 8 GiB.
    pub fn with_cap(
        dax_base_gpa: u64,
        dax_window_len: u64,
        cap_bytes: u64,
        backend: std::sync::Arc<dyn FsBackend>,
        mapper: std::sync::Arc<dyn HvfMapper>,
    ) -> Self {
        Self {
            dax_base_gpa,
            dax_window_len,
            cap_bytes,
            current_mapped_bytes: AtomicU64::new(0),
            backend,
            mapper,
            active: Mutex::new(BTreeMap::new()),
        }
    }

    /// Handle FUSE_SETUPMAPPING. Returns Ok(()) on success, negated
    /// errno on failure.
    pub fn setup(&self, nodeid: u64, req: &SetupMappingIn) -> Result<(), Errno> {
        log_dax(|| format!(
            "SETUPMAPPING: nodeid={nodeid} fh={} foffset={:#x} len={:#x} flags={:#x} moffset={:#x}",
            req.fh, req.foffset, req.len, req.flags, req.moffset
        ));
        // Bounds + alignment checks. Apple Silicon page granule is
        // 16 KiB; FUSE_INIT advertised map_alignment=14 to the guest.
        if (req.moffset & 0x3FFF) != 0 || (req.len & 0x3FFF) != 0 {
            return Err(super::backend::EINVAL);
        }
        if req.moffset.checked_add(req.len).map_or(true, |end| end > self.dax_window_len) {
            return Err(super::backend::EINVAL);
        }
        // Aggregate-bytes cap. Pre-reserve the increment with
        // compare-exchange so concurrent SETUPMAPPINGs can't both
        // observe room then race past the cap. On reject the cap
        // returns ENOSPC; the guest's virtio-fs driver retries
        // after dropping references (REMOVEMAPPING).
        let mut current = self.current_mapped_bytes.load(Ordering::Acquire);
        loop {
            let proposed = match current.checked_add(req.len) {
                Some(v) if v <= self.cap_bytes => v,
                _ => return Err(super::backend::ENOSPC),
            };
            match self.current_mapped_bytes.compare_exchange_weak(
                current,
                proposed,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => break,
                Err(observed) => current = observed,
            }
        }
        // Permissions: virtio-fs sets at least one of READ / WRITE.
        let mut prot = 0u32;
        if req.flags & FUSE_SETUPMAPPING_FLAG_READ != 0 {
            prot |= DAX_PROT_READ;
        }
        if req.flags & FUSE_SETUPMAPPING_FLAG_WRITE != 0 {
            prot |= DAX_PROT_WRITE;
        }
        if prot == 0 {
            // Release the cap reservation we just took.
            self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
            return Err(super::backend::EINVAL);
        }

        // Ask the backend for a host VA covering the file region.
        let host_va = match self
            .backend
            .dax_map(nodeid, req.fh, req.foffset, req.len, prot)
        {
            Ok(va) => va,
            Err(e) => {
                log_dax(|| format!("backend.dax_map FAILED: errno={e}"));
                self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
                return Err(e);
            }
        };

        // Plumb into the guest.
        let gpa = self.dax_base_gpa + req.moffset;
        log_dax(|| format!(
            "hv_vm_map: host_va={:#x} gpa={gpa:#x} len={:#x} prot={prot:#x}",
            host_va as usize, req.len
        ));
        self.mapper.map(host_va, gpa, req.len, prot).map_err(|e| {
            log_dax(|| format!("mapper.map FAILED: errno={e}"));
            // Failed to map into the guest — release the backend's
            // mapping so we don't leak.
            let _ = self
                .backend
                .dax_unmap(nodeid, host_va, req.len);
            self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
            e
        })?;

        // Record.
        let mut active = self.active.lock().unwrap();
        if active.contains_key(&req.moffset) {
            // Overlap with an existing slot. Refuse — the guest is
            // expected to REMOVEMAPPING before re-mapping the same
            // moffset.
            let _ = self.mapper.unmap(gpa, req.len);
            let _ = self.backend.dax_unmap(nodeid, host_va, req.len);
            self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
            return Err(super::backend::EEXIST);
        }
        active.insert(
            req.moffset,
            Slot {
                host_va: Some(host_va),
                len: req.len,
                nodeid,
                foffset: req.foffset,
                prot,
            },
        );
        Ok(())
    }

    /// Handle FUSE_REMOVEMAPPING. The payload contains a count followed
    /// by `count` RemoveMappingOne entries; the FuseServer parses the
    /// count and feeds them to us one by one.
    pub fn remove(&self, nodeid: u64, entry: &RemoveMappingOne) -> Result<(), Errno> {
        let mut active = self.active.lock().unwrap();
        let slot = active.remove(&entry.moffset).ok_or(super::backend::EINVAL)?;
        if slot.len != entry.len {
            // The kernel asked us to unmap a different len than we
            // mapped. Refuse — re-insert the slot.
            active.insert(entry.moffset, slot);
            return Err(super::backend::EINVAL);
        }
        let gpa = self.dax_base_gpa + entry.moffset;
        let released_len = slot.len;
        drop(active);
        // Best-effort: even if either step fails, we've removed the
        // slot from `active` so the next SETUPMAPPING at the same
        // moffset can proceed. Release the cap reservation
        // unconditionally — the slot is gone regardless.
        //
        // If the slot was metadata-only (host_va = None, restored from
        // snapshot but never faulted into), there's no host mmap or
        // stage-2 mapping to tear down. The cap reservation still
        // needs to be released (we accounted for `slot.len` at
        // setup-time AND at restore-time).
        let mapper_res = if slot.host_va.is_some() {
            self.mapper.unmap(gpa, entry.len)
        } else {
            Ok(())
        };
        let backend_res = if let Some(host_va) = slot.host_va {
            self.backend.dax_unmap(nodeid, host_va, slot.len)
        } else {
            Ok(())
        };
        self.current_mapped_bytes
            .fetch_sub(released_len, Ordering::AcqRel);
        mapper_res.and(backend_res)
    }

    /// Bytes currently consumed by active mappings. Used by tests +
    /// observability. Cheap atomic load.
    pub fn current_mapped_bytes(&self) -> u64 {
        self.current_mapped_bytes.load(Ordering::Acquire)
    }

    /// Number of currently-active DAX slots. Used by tests + density
    /// metrics.
    pub fn active_slot_count(&self) -> usize {
        self.active.lock().unwrap().len()
    }

    /// Capture the session's metadata for snapshot. Returns a binary
    /// blob suitable for `restore_state`.
    ///
    /// We persist FUSE-level metadata ONLY — `nodeid`, `foffset`,
    /// `prot`, `len`, `moffset`. The actual host VA (from `mmap`) is
    /// process-local and would be meaningless across daemon restart;
    /// `host_va` is therefore deliberately NOT serialised. On restore,
    /// `host_va = None` on every entry, and the first guest stage-2
    /// fault into the slot's GPA range triggers a re-`dax_map` via
    /// `handle_stage2_fault`.
    ///
    /// This makes snapshot capture for DAX state ~O(slot_count) with
    /// no syscalls — capturing 10k slots takes microseconds, not
    /// milliseconds. The first-fault cost (~25µs/page on M-series)
    /// amortises over the natural guest fault stream post-restore,
    /// matching fresh-boot semantics.
    pub fn snapshot_state(&self) -> Vec<u8> {
        // Layout (little-endian throughout):
        //   "DAXS" (magic, 4 bytes)
        //   u32 version              (currently 1)
        //   u64 dax_base_gpa         (consistency check on restore)
        //   u64 dax_window_len       (consistency check on restore)
        //   u64 cap_bytes
        //   u64 current_mapped_bytes (rebuilt as sum of slot.len; stored for audit)
        //   u64 entry_count
        //   per entry:
        //     u64 moffset
        //     u64 len
        //     u64 nodeid
        //     u64 foffset
        //     u32 prot
        //     u32 reserved (= 0)
        let active = self.active.lock().unwrap();
        let entry_count = active.len() as u64;
        let mut out = Vec::with_capacity(8 + 4 + 8 * 5 + (active.len() * 48));
        out.extend_from_slice(b"DAXS");
        out.extend_from_slice(&1u32.to_le_bytes());
        out.extend_from_slice(&self.dax_base_gpa.to_le_bytes());
        out.extend_from_slice(&self.dax_window_len.to_le_bytes());
        out.extend_from_slice(&self.cap_bytes.to_le_bytes());
        out.extend_from_slice(
            &self.current_mapped_bytes.load(Ordering::Acquire).to_le_bytes(),
        );
        out.extend_from_slice(&entry_count.to_le_bytes());
        for (moffset, slot) in active.iter() {
            out.extend_from_slice(&moffset.to_le_bytes());
            out.extend_from_slice(&slot.len.to_le_bytes());
            out.extend_from_slice(&slot.nodeid.to_le_bytes());
            out.extend_from_slice(&slot.foffset.to_le_bytes());
            out.extend_from_slice(&slot.prot.to_le_bytes());
            out.extend_from_slice(&0u32.to_le_bytes());
        }
        out
    }

    /// Hydrate the session's active-slot table from a blob produced by
    /// `snapshot_state`. All restored entries have `host_va = None`
    /// (metadata-only) and will be lazily re-bound on first guest
    /// stage-2 fault via `handle_stage2_fault`.
    ///
    /// Refuses (`InvalidData`) if the snapshot's window layout
    /// (`dax_base_gpa`, `dax_window_len`) doesn't match `self` — that
    /// would mean the snapshot was baked for a different DAX slice
    /// layout (different mount count, for example), and the slot
    /// GPAs would point into the wrong device's window.
    ///
    /// Idempotent: clearing happens before re-population, so calling
    /// twice with the same blob produces the same state.
    pub fn restore_state(&self, blob: &[u8]) -> Result<(), std::io::Error> {
        fn take<'a>(b: &'a [u8], p: &mut usize, n: usize) -> Result<&'a [u8], std::io::Error> {
            if *p + n > b.len() {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    "dax snapshot truncated",
                ));
            }
            let s = &b[*p..*p + n];
            *p += n;
            Ok(s)
        }
        fn read_u32(b: &[u8], p: &mut usize) -> Result<u32, std::io::Error> {
            let s = take(b, p, 4)?;
            Ok(u32::from_le_bytes([s[0], s[1], s[2], s[3]]))
        }
        fn read_u64(b: &[u8], p: &mut usize) -> Result<u64, std::io::Error> {
            let s = take(b, p, 8)?;
            Ok(u64::from_le_bytes([
                s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7],
            ]))
        }
        let mut p = 0usize;
        let magic = take(blob, &mut p, 4)?;
        if magic != b"DAXS" {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "dax snapshot bad magic",
            ));
        }
        let version = read_u32(blob, &mut p)?;
        if version != 1 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("dax snapshot version {version} unsupported"),
            ));
        }
        let snap_base = read_u64(blob, &mut p)?;
        let snap_window = read_u64(blob, &mut p)?;
        if snap_base != self.dax_base_gpa || snap_window != self.dax_window_len {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!(
                    "dax snapshot window mismatch: snap=({:#x}, {:#x}) self=({:#x}, {:#x})",
                    snap_base, snap_window, self.dax_base_gpa, self.dax_window_len
                ),
            ));
        }
        // cap_bytes + current_mapped_bytes are read for forward
        // compat / audit but the authoritative cap_bytes is whatever
        // the runtime DaxSession was constructed with; current_mapped
        // is recomputed as the sum of restored slot.len below.
        let _cap_bytes = read_u64(blob, &mut p)?;
        let _stored_current = read_u64(blob, &mut p)?;
        let entry_count = read_u64(blob, &mut p)?;
        let mut new_entries: BTreeMap<u64, Slot> = BTreeMap::new();
        let mut sum_len: u64 = 0;
        for _ in 0..entry_count {
            let moffset = read_u64(blob, &mut p)?;
            let len = read_u64(blob, &mut p)?;
            let nodeid = read_u64(blob, &mut p)?;
            let foffset = read_u64(blob, &mut p)?;
            let prot = read_u32(blob, &mut p)?;
            let _reserved = read_u32(blob, &mut p)?;
            // Re-validate window bounds — defensive against a
            // tampered or cross-layout sidecar.
            if moffset.checked_add(len).map_or(true, |end| end > self.dax_window_len) {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    format!(
                        "dax snapshot entry out of window: moffset={:#x} len={:#x} window={:#x}",
                        moffset, len, self.dax_window_len
                    ),
                ));
            }
            sum_len = sum_len.saturating_add(len);
            new_entries.insert(
                moffset,
                Slot {
                    host_va: None,
                    len,
                    nodeid,
                    foffset,
                    prot,
                },
            );
        }
        // Swap in atomically under the active lock so the restore is
        // observable to subsequent fault handlers in one step.
        let mut active = self.active.lock().unwrap();
        active.clear();
        for (k, v) in new_entries {
            active.insert(k, v);
        }
        drop(active);
        self.current_mapped_bytes.store(sum_len, Ordering::Release);
        Ok(())
    }

    /// Re-establish the host-side mmap + HVF stage-2 mapping for a
    /// restored slot on first guest fault, if `gpa` belongs to one.
    ///
    /// Returns:
    ///   * `Ok(false)` — `gpa` is not in this session's window OR no
    ///     slot covers it. Caller should try the next session (or
    ///     fall through to the existing MMIO dispatch).
    ///   * `Ok(true)` — fault was claimed by this session. Either we
    ///     just re-established the mapping, OR another thread did so
    ///     concurrently (`host_va` was already `Some`). In either case
    ///     the vCPU should resume at the same PC and re-execute the
    ///     load/store — HVF will hit the now-present stage-2 entry
    ///     and complete the access naturally.
    ///   * `Err(_)` — backend or HVF mapper rejected the re-map. The
    ///     slot is left with `host_va = None` so a subsequent fault
    ///     can retry; the vCPU caller is expected to surface this as
    ///     a fatal data abort.
    ///
    /// Design note: we re-execute rather than emulate because the
    /// trapped instruction is a normal load/store against guest VA
    /// that just happens to be DAX-mapped. HVF doesn't decode it for
    /// us in the fault path, and reproducing the addressing-mode
    /// decode for every aarch64 load/store variant is a pile of work
    /// the architecture itself does for free on retry. The cost is
    /// one extra HVF entry/exit cycle (~µs) — negligible compared to
    /// the ~25µs cost of mmap + hv_vm_map we just paid.
    pub fn handle_stage2_fault(&self, gpa: u64) -> Result<bool, Errno> {
        // Fast-path window check before touching the lock — most
        // stage-2 faults are MMIO or genuine paging events outside
        // any DAX window, and we don't want to serialise on
        // `self.active` for every one of them.
        if gpa < self.dax_base_gpa {
            return Ok(false);
        }
        let moffset_fault = gpa - self.dax_base_gpa;
        if moffset_fault >= self.dax_window_len {
            return Ok(false);
        }
        // Find the slot whose [moffset, moffset+len) contains the
        // fault. `range(..=moffset_fault).next_back()` gives us the
        // largest key <= moffset_fault (since BTreeMap is ordered
        // by key).
        let active = self.active.lock().unwrap();
        let (&slot_moffset, slot) = match active.range(..=moffset_fault).next_back() {
            Some(kv) => kv,
            None => return Ok(false),
        };
        if moffset_fault >= slot_moffset + slot.len {
            // Largest slot <= fault doesn't cover the fault — there
            // is no DAX mapping at this GPA. Not our problem; let
            // the caller fall through to MMIO dispatch (which will
            // surface as an unhandled access).
            return Ok(false);
        }
        if slot.host_va.is_some() {
            // Already bound — another vCPU thread won the race or
            // the slot was never restore-pending. Tell the caller
            // to resume; HVF will retry the load and succeed.
            return Ok(true);
        }
        // Clone the metadata we need to release the lock for the
        // (potentially slow) backend.dax_map / mapper.map calls.
        // We don't want to hold `self.active` across a syscall —
        // it would serialise concurrent stage-2 faults on disjoint
        // slots unnecessarily.
        let nodeid = slot.nodeid;
        let foffset = slot.foffset;
        let len = slot.len;
        let prot = slot.prot;
        let slot_gpa_start = self.dax_base_gpa + slot_moffset;
        drop(active);

        log_dax(|| format!(
            "restore-fault: re-map nodeid={nodeid} foffset={foffset:#x} \
             len={len:#x} prot={prot:#x} gpa={slot_gpa_start:#x}",
        ));

        // Re-mmap via the backend. `u64::MAX` for `fh` triggers the
        // on-demand `host_path_of(nodeid)` open path in PosixFs
        // (see posix.rs:2181-2197) — the original SETUPMAPPING's
        // `fh` is not part of post-restore state (handle tables
        // don't survive snapshot/restore).
        let host_va = self.backend.dax_map(nodeid, u64::MAX, foffset, len, prot)?;

        // Plumb into the guest.
        if let Err(e) = self.mapper.map(host_va, slot_gpa_start, len, prot) {
            // Release the backend's mmap on failure so we don't leak
            // host VA. Leave `host_va = None` so a future fault can
            // retry. The cap was already accounted for at restore-
            // time so don't touch `current_mapped_bytes`.
            let _ = self.backend.dax_unmap(nodeid, host_va, len);
            return Err(e);
        }

        // Re-acquire the lock to publish the host_va. Re-check that
        // the slot still exists at the same moffset — between drop
        // and reacquire a REMOVEMAPPING could have removed it. In
        // that rare race we leak the mapping we just made (the
        // guest is unmapping the very page it just faulted on,
        // which is a guest bug, but we don't want to crash on it).
        let mut active = self.active.lock().unwrap();
        if let Some(s) = active.get_mut(&slot_moffset) {
            // Another concurrent fault may have raced and bound the
            // slot first — if so, undo our work to avoid leaking.
            if s.host_va.is_some() {
                drop(active);
                let _ = self.mapper.unmap(slot_gpa_start, len);
                let _ = self.backend.dax_unmap(nodeid, host_va, len);
                return Ok(true);
            }
            s.host_va = Some(host_va);
        } else {
            // Slot vanished while we were in the syscall — guest
            // raced with REMOVEMAPPING. Tear down what we made.
            drop(active);
            let _ = self.mapper.unmap(slot_gpa_start, len);
            let _ = self.backend.dax_unmap(nodeid, host_va, len);
            return Ok(true);
        }
        Ok(true)
    }
}

fn log_dax<F: FnOnce() -> String>(make_msg: F) {
    use std::io::Write;
    let Some(target) = crate::trace::fuse_target() else { return };
    let s = make_msg();
    let target = target.to_string_lossy().into_owned();
    if target == "1" || target == "stderr" {
        eprintln!("[dax] {s}");
        return;
    }
    if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(&target) {
        let _ = writeln!(f, "[dax] {s}");
    }
}

/// Parse a FUSE_REMOVEMAPPING payload into a list of RemoveMappingOne
/// entries. Header carries `count`; entries follow contiguously.
pub fn parse_remove_payload(payload: &[u8]) -> Result<Vec<RemoveMappingOne>, Errno> {
    if payload.len() < core::mem::size_of::<RemoveMappingIn>() {
        return Err(super::backend::EINVAL);
    }
    let hdr: RemoveMappingIn = unsafe { core::ptr::read_unaligned(payload.as_ptr() as *const RemoveMappingIn) };
    let count = hdr.count as usize;
    let one_size = core::mem::size_of::<RemoveMappingOne>();
    let needed = core::mem::size_of::<RemoveMappingIn>() + count * one_size;
    if payload.len() < needed {
        return Err(super::backend::EINVAL);
    }
    let mut out = Vec::with_capacity(count);
    let base = payload.as_ptr();
    for i in 0..count {
        let off = core::mem::size_of::<RemoveMappingIn>() + i * one_size;
        let one: RemoveMappingOne =
            unsafe { core::ptr::read_unaligned(base.add(off) as *const RemoveMappingOne) };
        out.push(one);
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fuse::backend::MemoryFs;
    use std::sync::Arc;

    /// A backend with a dax_map impl that returns a synthetic pointer
    /// so SETUPMAPPING orchestration can run end-to-end.
    struct StubDaxBackend {
        inner: MemoryFs,
        next_va: Mutex<usize>,
    }
    impl StubDaxBackend {
        fn new() -> Self {
            Self {
                inner: MemoryFs::new(),
                next_va: Mutex::new(0x40_0000_0000usize),
            }
        }
    }

    impl FsBackend for StubDaxBackend {
        fn lookup(&self, p: u64, n: &std::ffi::OsStr) -> Result<crate::fuse::backend::Entry, Errno> { self.inner.lookup(p, n) }
        fn forget(&self, n: u64, nl: u64) { self.inner.forget(n, nl) }
        fn getattr(&self, n: u64, f: Option<u64>) -> Result<crate::fuse::protocol::Attr, Errno> { self.inner.getattr(n, f) }
        fn open(&self, n: u64, f: u32) -> Result<u64, Errno> { self.inner.open(n, f) }
        fn read(&self, n: u64, h: u64, o: u64, s: u32) -> Result<Vec<u8>, Errno> { self.inner.read(n, h, o, s) }
        fn release(&self, n: u64, h: u64) -> Result<(), Errno> { self.inner.release(n, h) }
        fn opendir(&self, n: u64, f: u32) -> Result<u64, Errno> { self.inner.opendir(n, f) }
        fn readdir(&self, n: u64, h: u64, o: u64, s: u32) -> Result<Vec<crate::fuse::backend::DirEntry>, Errno> { self.inner.readdir(n, h, o, s) }
        fn releasedir(&self, n: u64, h: u64) -> Result<(), Errno> { self.inner.releasedir(n, h) }
        fn statfs(&self, n: u64) -> Result<crate::fuse::backend::StatFs, Errno> { self.inner.statfs(n) }
        fn dax_map(&self, _nodeid: u64, _fh: u64, _foffset: u64, len: u64, _prot: u32) -> Result<*mut u8, Errno> {
            let mut g = self.next_va.lock().unwrap();
            let va = *g;
            *g += len as usize;
            Ok(va as *mut u8)
        }
        fn dax_unmap(&self, _nodeid: u64, _host_va: *mut u8, _len: u64) -> Result<(), Errno> {
            Ok(())
        }
    }

    fn make_session() -> (DaxSession, Arc<MockHvfMapper>) {
        let backend = Arc::new(StubDaxBackend::new());
        let mapper = Arc::new(MockHvfMapper::new());
        let session = DaxSession::new(0x80_0000_0000, 0x4_0000_0000, backend, mapper.clone());
        (session, mapper)
    }

    #[test]
    fn setup_records_mapping_and_calls_mapper() {
        let (sess, m) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0,
        };
        sess.setup(7, &req).unwrap();
        assert_eq!(sess.active_slot_count(), 1);
        let calls = m.calls();
        assert_eq!(calls.len(), 1);
        match &calls[0] {
            MockCall::Map { gpa, len, prot, .. } => {
                assert_eq!(*gpa, 0x80_0000_0000);
                assert_eq!(*len, 0x4000);
                assert_eq!(*prot, DAX_PROT_READ);
            }
            _ => panic!("expected Map"),
        }
    }

    #[test]
    fn remove_undoes_a_prior_setup() {
        let (sess, m) = make_session();
        let setup = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0x1_0000,
        };
        sess.setup(7, &setup).unwrap();
        let rem = RemoveMappingOne {
            moffset: 0x1_0000,
            len: 0x4000,
        };
        sess.remove(7, &rem).unwrap();
        assert_eq!(sess.active_slot_count(), 0);
        let calls = m.calls();
        assert_eq!(calls.len(), 2);
        match &calls[1] {
            MockCall::Unmap { gpa, len } => {
                assert_eq!(*gpa, 0x80_0000_0000 + 0x1_0000);
                assert_eq!(*len, 0x4000);
            }
            _ => panic!("expected Unmap"),
        }
    }

    #[test]
    fn setup_rejects_unaligned() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0x100, // not 16 KiB aligned
        };
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EINVAL);
    }

    #[test]
    fn setup_rejects_out_of_window() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            // 4 GiB window; ask for the last 16 KiB + 1 KiB → overflow
            moffset: 0x4_0000_0000 - 0x4000 + 0x4000,
        };
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EINVAL);
    }

    #[test]
    fn setup_rejects_overlapping_existing_moffset() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0,
        };
        sess.setup(7, &req).unwrap();
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EEXIST);
    }

    #[test]
    fn setup_rejects_zero_perms() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: 0, // no R, no W
            moffset: 0,
        };
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EINVAL);
    }

    #[test]
    fn setup_rejects_when_cap_would_be_exceeded() {
        // Cap of 8 KiB (2 pages). Two 4 KiB mappings fit; third rejects.
        let backend = Arc::new(StubDaxBackend::new());
        let mapper = Arc::new(MockHvfMapper::new());
        let session = DaxSession::with_cap(
            0x80_0000_0000,
            0x4_0000_0000,
            0x8000,  // cap = 8 KiB
            backend,
            mapper.clone(),
        );
        let mk = |moff: u64| SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: moff,
        };
        session.setup(7, &mk(0)).unwrap();
        session.setup(7, &mk(0x4000)).unwrap();
        // Third one would push current_mapped to 12 KiB; cap rejects.
        let err = session.setup(7, &mk(0x8000)).unwrap_err();
        assert_eq!(err, super::super::backend::ENOSPC);
        assert_eq!(session.current_mapped_bytes(), 0x8000);
    }

    #[test]
    fn remove_releases_cap_reservation() {
        let backend = Arc::new(StubDaxBackend::new());
        let mapper = Arc::new(MockHvfMapper::new());
        let session = DaxSession::with_cap(
            0x80_0000_0000,
            0x4_0000_0000,
            0x4000, // cap = 4 KiB; only one slot at a time
            backend,
            mapper.clone(),
        );
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0,
        };
        session.setup(7, &req).unwrap();
        // Second mapping is rejected because the cap is now full —
        // the cap check fires before the overlap check, so ENOSPC.
        assert_eq!(
            session.setup(7, &req.clone()).unwrap_err(),
            super::super::backend::ENOSPC
        );
        // Remove the first; cap is freed.
        let rem = RemoveMappingOne { moffset: 0, len: 0x4000 };
        session.remove(7, &rem).unwrap();
        assert_eq!(session.current_mapped_bytes(), 0);
        // Now a fresh mapping fits.
        session.setup(7, &req).unwrap();
    }

    #[test]
    fn snapshot_restore_round_trip_preserves_entries() {
        let (sess, _m) = make_session();
        // Three slots with different perms / offsets — exercises
        // the BTreeMap key ordering, the prot bits, and the
        // moffset bookkeeping all at once.
        let req = |moff: u64, foff: u64, len: u64, flags: u64| SetupMappingIn {
            fh: 1, foffset: foff, len, flags, moffset: moff,
        };
        sess.setup(7, &req(0x0_0000, 0, 0x4000, FUSE_SETUPMAPPING_FLAG_READ)).unwrap();
        sess.setup(11, &req(0x8000, 0x1_0000, 0x4000,
            FUSE_SETUPMAPPING_FLAG_READ | FUSE_SETUPMAPPING_FLAG_WRITE)).unwrap();
        sess.setup(13, &req(0x10_0000, 0x4000, 0x8000,
            FUSE_SETUPMAPPING_FLAG_WRITE)).unwrap();
        assert_eq!(sess.active_slot_count(), 3);
        let pre_bytes = sess.current_mapped_bytes();

        let blob = sess.snapshot_state();
        // Header is 4 (magic) + 4 (version) + 8*5 (base, window,
        // cap, current_mapped, entry_count) = 48 bytes.
        // Each entry is 4*u64 + 2*u32 = 40 bytes.
        assert!(blob.len() >= 48 + 3 * 40,
            "blob.len()={} expected >= {}", blob.len(), 48 + 3 * 40);
        assert_eq!(&blob[..4], b"DAXS");

        // Fresh session with the same window layout — restore
        // populates it with metadata-only entries.
        let (sess2, m2) = make_session();
        sess2.restore_state(&blob).expect("restore should succeed");
        assert_eq!(sess2.active_slot_count(), 3);
        assert_eq!(sess2.current_mapped_bytes(), pre_bytes);

        // No mapper.map calls during restore — we don't re-bind
        // eagerly. The pure-metadata restore is the whole point.
        assert!(m2.calls().is_empty(), "restore must not call mapper.map");

        // Every restored entry should be metadata-only.
        let active = sess2.active.lock().unwrap();
        for (moff, slot) in active.iter() {
            assert!(slot.host_va.is_none(),
                "restored slot at moffset={moff:#x} must start with host_va=None");
        }
    }

    #[test]
    fn restored_entries_have_no_host_va_until_fault() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1, foffset: 0x1000, len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ, moffset: 0x2_0000,
        };
        sess.setup(42, &req).unwrap();
        let blob = sess.snapshot_state();

        let (sess2, _) = make_session();
        sess2.restore_state(&blob).unwrap();
        let active = sess2.active.lock().unwrap();
        let slot = active.get(&0x2_0000).expect("entry should be restored");
        assert!(slot.host_va.is_none());
        assert_eq!(slot.len, 0x4000);
        assert_eq!(slot.nodeid, 42);
        assert_eq!(slot.foffset, 0x1000);
        assert_eq!(slot.prot, DAX_PROT_READ);
    }

    #[test]
    fn handle_stage2_fault_rehydrates_restored_slot() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1, foffset: 0x5000, len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ | FUSE_SETUPMAPPING_FLAG_WRITE,
            moffset: 0x4_0000,
        };
        sess.setup(99, &req).unwrap();
        let blob = sess.snapshot_state();

        // Fresh session + fresh mapper to assert the post-restore
        // re-bind calls dax_map AND mapper.map.
        let (sess2, m2) = make_session();
        sess2.restore_state(&blob).unwrap();
        assert!(m2.calls().is_empty());

        // Fault somewhere inside the slot's GPA range — slot starts
        // at moffset 0x4_0000 (gpa 0x80_0004_0000) and is 0x4000
        // long; fault at gpa + 0x1234 lands inside.
        let fault_gpa = 0x80_0000_0000 + 0x4_0000 + 0x1234;
        let handled = sess2.handle_stage2_fault(fault_gpa).unwrap();
        assert!(handled, "fault inside a restored slot must be claimed");

        // mapper.map should now have been called exactly once with
        // the slot's start gpa + length + original prot.
        let calls = m2.calls();
        assert_eq!(calls.len(), 1);
        match &calls[0] {
            MockCall::Map { gpa, len, prot, .. } => {
                assert_eq!(*gpa, 0x80_0000_0000 + 0x4_0000);
                assert_eq!(*len, 0x4000);
                assert_eq!(*prot, DAX_PROT_READ | DAX_PROT_WRITE);
            }
            _ => panic!("expected Map"),
        }

        // The slot should now report host_va = Some.
        let active = sess2.active.lock().unwrap();
        assert!(active.get(&0x4_0000).unwrap().host_va.is_some());
    }

    #[test]
    fn handle_stage2_fault_out_of_window_returns_false_without_lock() {
        let (sess, _) = make_session();
        // dax_base_gpa = 0x80_0000_0000, window = 0x4_0000_0000;
        // fault at 0x10000 is way below the window.
        assert_eq!(sess.handle_stage2_fault(0x10000).unwrap(), false);
        // Also above the window.
        assert_eq!(
            sess.handle_stage2_fault(0x80_0000_0000 + 0x4_0000_0000 + 0x1000).unwrap(),
            false
        );
    }

    #[test]
    fn handle_stage2_fault_in_window_but_no_slot_returns_false() {
        let (sess, m) = make_session();
        // Set up a slot at moffset 0x4000-0x8000; fault outside it
        // but inside the window (at moffset 0x10_0000).
        let req = SetupMappingIn {
            fh: 1, foffset: 0, len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ, moffset: 0x4000,
        };
        sess.setup(7, &req).unwrap();
        // Snapshot+restore to clear the active host_va (which
        // would mask the no-slot path otherwise).
        let blob = sess.snapshot_state();
        let (sess2, m2) = make_session();
        sess2.restore_state(&blob).unwrap();

        // Fault at moffset 0x10_0000 — no slot covers that.
        let fault_gpa = 0x80_0000_0000 + 0x10_0000;
        assert_eq!(sess2.handle_stage2_fault(fault_gpa).unwrap(), false);
        // The mapper was never touched (no slot to bind).
        assert!(m2.calls().is_empty());

        // Sanity: the first session's mapper was hit during setup.
        assert!(m.calls().iter().any(|c| matches!(c, MockCall::Map { .. })));
    }

    #[test]
    fn handle_stage2_fault_on_already_bound_slot_is_noop() {
        // Bind a slot, fault into it — handler returns Ok(true)
        // without re-calling the backend or mapper.
        let (sess, m) = make_session();
        let req = SetupMappingIn {
            fh: 1, foffset: 0, len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ, moffset: 0,
        };
        sess.setup(7, &req).unwrap();
        let calls_before = m.calls().len();
        // Slot is host_va = Some after setup.
        assert!(sess.handle_stage2_fault(0x80_0000_0000 + 0x1234).unwrap());
        let calls_after = m.calls().len();
        assert_eq!(calls_before, calls_after,
            "already-bound slot must NOT trigger extra mapper calls");
    }

    #[test]
    fn restore_state_rejects_malformed_blobs() {
        let (sess, _) = make_session();

        // Truncated header.
        assert!(sess.restore_state(&[1, 2, 3]).is_err());

        // Wrong magic.
        let mut bad = Vec::from(*b"WRONG_MAGIC");
        bad.extend_from_slice(&1u32.to_le_bytes());
        assert!(sess.restore_state(&bad).is_err());

        // Right magic, wrong version.
        let mut bad = Vec::from(*b"DAXS");
        bad.extend_from_slice(&999u32.to_le_bytes());
        bad.extend_from_slice(&0u64.to_le_bytes());
        bad.extend_from_slice(&0u64.to_le_bytes());
        bad.extend_from_slice(&0u64.to_le_bytes());
        bad.extend_from_slice(&0u64.to_le_bytes());
        bad.extend_from_slice(&0u64.to_le_bytes());
        assert!(sess.restore_state(&bad).is_err());

        // Window mismatch — wrong dax_base_gpa.
        let mut wrong_window = Vec::from(*b"DAXS");
        wrong_window.extend_from_slice(&1u32.to_le_bytes());
        wrong_window.extend_from_slice(&0xDEAD_BEEFu64.to_le_bytes()); // wrong base
        wrong_window.extend_from_slice(&sess.dax_window_len.to_le_bytes());
        wrong_window.extend_from_slice(&0u64.to_le_bytes()); // cap_bytes
        wrong_window.extend_from_slice(&0u64.to_le_bytes()); // current_mapped
        wrong_window.extend_from_slice(&0u64.to_le_bytes()); // entry_count
        assert!(sess.restore_state(&wrong_window).is_err());

        // Truncated after entry header.
        let mut trunc = Vec::from(*b"DAXS");
        trunc.extend_from_slice(&1u32.to_le_bytes());
        trunc.extend_from_slice(&sess.dax_base_gpa.to_le_bytes());
        trunc.extend_from_slice(&sess.dax_window_len.to_le_bytes());
        trunc.extend_from_slice(&0u64.to_le_bytes());
        trunc.extend_from_slice(&0u64.to_le_bytes());
        trunc.extend_from_slice(&3u64.to_le_bytes()); // claims 3 entries
        // ... but provides zero entry bytes.
        assert!(sess.restore_state(&trunc).is_err());
    }

    #[test]
    fn restore_state_is_idempotent() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1, foffset: 0, len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ, moffset: 0,
        };
        sess.setup(7, &req).unwrap();
        let blob = sess.snapshot_state();

        let (sess2, _) = make_session();
        sess2.restore_state(&blob).unwrap();
        let after_first = sess2.active_slot_count();
        sess2.restore_state(&blob).unwrap();
        let after_second = sess2.active_slot_count();
        assert_eq!(after_first, after_second);
        assert_eq!(after_second, 1);
        assert_eq!(sess2.current_mapped_bytes(), 0x4000);
    }

    #[test]
    fn parse_remove_payload_round_trip() {
        let entries = vec![
            RemoveMappingOne { moffset: 0x4000, len: 0x4000 },
            RemoveMappingOne { moffset: 0x8000, len: 0x8000 },
        ];
        let mut buf = Vec::new();
        let hdr = RemoveMappingIn { count: entries.len() as u32, _pad: 0 };
        buf.extend_from_slice(unsafe {
            std::slice::from_raw_parts(
                &hdr as *const RemoveMappingIn as *const u8,
                core::mem::size_of::<RemoveMappingIn>(),
            )
        });
        for e in &entries {
            buf.extend_from_slice(unsafe {
                std::slice::from_raw_parts(
                    e as *const RemoveMappingOne as *const u8,
                    core::mem::size_of::<RemoveMappingOne>(),
                )
            });
        }
        let parsed = parse_remove_payload(&buf).unwrap();
        assert_eq!(parsed.len(), 2);
        assert_eq!(parsed[0].moffset, 0x4000);
        assert_eq!(parsed[1].moffset, 0x8000);
    }
}