nub_host_kvm/mem/

shared_mem.rs

/*
Copyright 2025  The Hyperlight Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

use std::any::type_name;
use std::ffi::c_void;
use std::io::Error;
use std::mem::{align_of, size_of};
use std::ptr::null_mut;
use std::sync::{Arc, RwLock};

use nub_host_common::mem::PAGE_SIZE_USIZE;
use tracing::{Span, instrument};

use super::memory_region::{
    HostGuestMemoryRegion, MemoryRegion, MemoryRegionFlags, MemoryRegionKind, MemoryRegionType,
};
use crate::{HyperlightError, Result, log_then_return, new_error};

/// Makes sure that the given `offset` and `size` are within the bounds of the memory with size `mem_size`.
macro_rules! bounds_check {
    ($offset:expr, $size:expr, $mem_size:expr) => {
        if $offset.checked_add($size).is_none_or(|end| end > $mem_size) {
            return Err(new_error!(
                "Cannot read value from offset {} with size {} in memory of size {}",
                $offset,
                $size,
                $mem_size
            ));
        }
    };
}

/// generates a reader function for the given type
macro_rules! generate_reader {
    ($fname:ident, $ty:ty) => {
        /// Read a value of type `$ty` from the memory at the given offset.
        #[allow(dead_code)]
        #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
        pub(crate) fn $fname(&self, offset: usize) -> Result<$ty> {
            let data = self.as_slice();
            bounds_check!(offset, std::mem::size_of::<$ty>(), data.len());
            Ok(<$ty>::from_le_bytes(
                data[offset..offset + std::mem::size_of::<$ty>()].try_into()?,
            ))
        }
    };
}

/// generates a writer function for the given type
macro_rules! generate_writer {
    ($fname:ident, $ty:ty) => {
        /// Write a value of type `$ty` to the memory at the given offset.
        #[allow(dead_code)]
        pub(crate) fn $fname(&mut self, offset: usize, value: $ty) -> Result<()> {
            let data = self.as_mut_slice();
            bounds_check!(offset, std::mem::size_of::<$ty>(), data.len());
            data[offset..offset + std::mem::size_of::<$ty>()].copy_from_slice(&value.to_le_bytes());
            Ok(())
        }
    };
}

/// A representation of a host mapping of a shared memory region,
/// which will be released when this structure is Drop'd. This is not
/// individually Clone (since it holds ownership of the mapping), or
/// Send or Sync, since it doesn't ensure any particular synchronization.
#[derive(Debug)]
pub struct HostMapping {
    ptr: *mut u8,
    size: usize,
}

impl Drop for HostMapping {
    fn drop(&mut self) {
        use libc::munmap;

        unsafe {
            munmap(self.ptr as *mut c_void, self.size);
        }
    }
}

/// These three structures represent various phases of the lifecycle of
/// a memory buffer that is shared with the guest. An
/// ExclusiveSharedMemory is used for certain operations that
/// unrestrictedly write to the shared memory, including setting it up
/// and taking snapshots.
#[derive(Debug)]
pub struct ExclusiveSharedMemory {
    region: Arc<HostMapping>,
}
unsafe impl Send for ExclusiveSharedMemory {}

/// A GuestSharedMemory is used to represent
/// the reference to all-of-memory that is taken by the virtual cpu.
/// Because of the memory model limitations that affect
/// HostSharedMemory, it is likely fairly important (to ensure that
/// our UB remains limited to interaction with an external compilation
/// unit that likely can't be discovered by the compiler) that _rust_
/// users do not perform racy accesses to the guest communication
/// buffers that are also accessed by HostSharedMemory.
#[derive(Debug)]
pub struct GuestSharedMemory {
    region: Arc<HostMapping>,
    /// The lock that indicates this shared memory is being used by non-Rust code
    ///
    /// This lock _must_ be held whenever the guest is executing,
    /// because it prevents the host from converting its
    /// HostSharedMemory to an ExclusiveSharedMemory. Since the guest
    /// may arbitrarily mutate the shared memory, only synchronized
    /// accesses from Rust should be allowed!
    ///
    /// We cannot enforce this in the type system, because the memory
    /// is mapped in to the VM at VM creation time.
    pub lock: Arc<RwLock<()>>,
}
// SAFETY: `GuestSharedMemory` is the KVM-facing handle for a mapping whose
// lifetime is owned by `Arc<HostMapping>`. The wrapper exposes no direct Rust
// slice access except through `SharedMemory::with_exclusivity`, which takes the
// shared RwLock's write side and therefore excludes host-side volatile access
// while exclusive remapping/snapshotting is in progress. Sharing this handle
// between host worker threads only shares the registered mapping/lifetime
// handle; synchronized host reads/writes still go through `HostSharedMemory`.
unsafe impl Send for GuestSharedMemory {}
unsafe impl Sync for GuestSharedMemory {}

/// A HostSharedMemory allows synchronized accesses to guest
/// communication buffers, allowing it to be used concurrently with a
/// GuestSharedMemory.
///
/// # Concurrency model
///
/// Given future requirements for asynchronous I/O with a minimum
/// amount of copying (e.g. WASIp3 streams), we would like it to be
/// possible to safely access these buffers concurrently with the
/// guest, ensuring that (1) data is read appropriately if the guest
/// is well-behaved; and (2) the host's behaviour is defined
/// regardless of whether or not the guest is well-behaved.
///
/// The ideal (future) flow for a guest->host message is something like
///   - Guest writes (unordered) bytes describing a work item into a buffer
///   - Guest reveals buffer via a release-store of a pointer into an
///     MMIO ring-buffer
///   - Host acquire-loads the buffer pointer from the "MMIO" ring
///     buffer
///   - Host (unordered) reads the bytes from the buffer
///   - Host performs validation of those bytes and uses them
///
/// Unfortunately, there appears to be no way to do this with defined
/// behaviour in present Rust (see
/// e.g. <https://github.com/rust-lang/unsafe-code-guidelines/issues/152>).
/// Rust does not yet have its own defined memory model, but in the
/// interim, it is widely treated as inheriting the current C/C++
/// memory models.  The most immediate problem is that regardless of
/// anything else, under those memory models \[1, p. 17-18; 2, p. 88\],
///
///   > The execution of a program contains a _data race_ if it
///   > contains two [C++23: "potentially concurrent"] conflicting
///   > actions [C23: "in different threads"], at least one of which
///   > is not atomic, and neither happens before the other [C++23: ",
///   > except for the special case for signal handlers described
///   > below"].  Any such data race results in undefined behavior.
///
/// Consequently, if a misbehaving guest fails to correctly
/// synchronize its stores with the host, the host's innocent loads
/// will trigger undefined behaviour for the entire program, including
/// the host.  Note that this also applies if the guest makes an
/// unsynchronized read of a location that the host is writing!
///
/// Despite Rust's de jure inheritance of the C memory model at the
/// present time, the compiler in many cases de facto adheres to LLVM
/// semantics, so it is worthwhile to consider what LLVM does in this
/// case as well.  According to the the LangRef \[3\] memory model,
/// loads which are involved in a race that includes at least one
/// non-atomic access (whether the load or a store) return `undef`,
/// making them roughly equivalent to reading uninitialized
/// memory. While this is much better, it is still bad.
///
/// Considering a different direction, recent C++ papers have seemed
/// to lean towards using `volatile` for similar use cases. For
/// example, in P1152R0 \[4\], JF Bastien notes that
///
///   > We’ve shown that volatile is purposely defined to denote
///   > external modifications. This happens for:
///   >   - Shared memory with untrusted code, where volatile is the
///   >     right way to avoid time-of-check time-of-use (ToCToU)
///   >     races which lead to security bugs such as \[PWN2OWN\] and
///   >     \[XENXSA155\].
///
/// Unfortunately, although this paper was adopted for C++20 (and,
/// sadly, mostly un-adopted for C++23, although that does not concern
/// us), the paper did not actually redefine volatile accesses or data
/// races to prevent volatile accesses from racing with other accesses
/// and causing undefined behaviour.  P1382R1 \[5\] would have amended
/// the wording of the data race definition to specifically exclude
/// volatile, but, unfortunately, despite receiving a
/// generally-positive reception at its first WG21 meeting more than
/// five years ago, it has not progressed.
///
/// Separately from the data race issue, there is also a concern that
/// according to the various memory models in use, there may be ways
/// in which the guest can semantically obtain uninitialized memory
/// and write it into the shared buffer, which may also result in
/// undefined behaviour on reads.  The degree to which this is a
/// concern is unclear, however, since it is unclear to what degree
/// the Rust abstract machine's conception of uninitialized memory
/// applies to the sandbox.  Returning briefly to the LLVM level,
/// rather than the Rust level, this, combined with the fact that
/// racing loads in LLVM return `undef`, as discussed above, we would
/// ideally `llvm.freeze` the result of any load out of the sandbox.
///
/// It would furthermore be ideal if we could run the flatbuffers
/// parsing code directly on the guest memory, in order to avoid
/// unnecessary copies.  That is unfortunately probably not viable at
/// the present time: because the generated flatbuffers parsing code
/// doesn't use atomic or volatile accesses, it is likely to introduce
/// double-read vulnerabilities.
///
/// In short, none of the Rust-level operations available to us do the
/// right thing, at the Rust spec level or the LLVM spec level. Our
/// major remaining options are therefore:
///   - Choose one of the options that is available to us, and accept
///     that we are doing something unsound according to the spec, but
///     hope that no reasonable compiler could possibly notice.
///   - Use inline assembly per architecture, for which we would only
///     need to worry about the _architecture_'s memory model (which
///     is far less demanding).
///
/// The leading candidate for the first option would seem to be to
/// simply use volatile accesses; there seems to be wide agreement
/// that this _should_ be a valid use case for them (even if it isn't
/// now), and projects like Linux and rust-vmm already use C11
/// `volatile` for this purpose.  It is also worth noting that because
/// we still do need to synchronize with the guest when it _is_ being
/// well-behaved, we would ideally use volatile acquire loads and
/// volatile release stores for interacting with the stack pointer in
/// the guest in this case.  Unfortunately, while those operations are
/// defined in LLVM, they are not presently exposed to Rust. While
/// atomic fences that are not associated with memory accesses
/// ([`std::sync::atomic::fence`]) might at first glance seem to help with
/// this problem, they unfortunately do not \[6\]:
///
///    > A fence ‘A’ which has (at least) Release ordering semantics,
///    > synchronizes with a fence ‘B’ with (at least) Acquire
///    > semantics, if and only if there exist operations X and Y,
///    > both operating on some atomic object ‘M’ such that A is
///    > sequenced before X, Y is sequenced before B and Y observes
///    > the change to M. This provides a happens-before dependence
///    > between A and B.
///
/// Note that the X and Y must be to an _atomic_ object.
///
/// We consequently assume that there has been a strong architectural
/// fence on a vmenter/vmexit between data being read and written.
/// This is unsafe (not guaranteed in the type system)!
///
/// \[1\] N3047 C23 Working Draft. <https://www.open-std.org/jtc1/sc22/wg14/www/docs/n3047.pdf>
/// \[2\] N4950 C++23 Working Draft. <https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2023/n4950.pdf>
/// \[3\] LLVM Language Reference Manual, Memory Model for Concurrent Operations. <https://llvm.org/docs/LangRef.html#memmodel>
/// \[4\] P1152R0: Deprecating `volatile`. JF Bastien. <https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2018/p1152r0.html>
/// \[5\] P1382R1: `volatile_load<T>` and `volatile_store<T>`. JF Bastien, Paul McKenney, Jeffrey Yasskin, and the indefatigable TBD. <https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p1382r1.pdf>
/// \[6\] Documentation for std::sync::atomic::fence. <https://doc.rust-lang.org/std/sync/atomic/fn.fence.html>
#[derive(Clone, Debug)]
pub struct HostSharedMemory {
    region: Arc<HostMapping>,
    lock: Arc<RwLock<()>>,
}
unsafe impl Send for HostSharedMemory {}

impl ExclusiveSharedMemory {
    /// Create a new region of shared memory with the given minimum
    /// size in bytes. The region will be surrounded by guard pages.
    ///
    /// Return `Err` if shared memory could not be allocated.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub fn new(min_size_bytes: usize) -> Result<Self> {
        use libc::{
            MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, PROT_READ, PROT_WRITE, c_int, mmap, off_t,
            size_t,
        };
        #[cfg(not(miri))]
        use libc::{MAP_NORESERVE, PROT_NONE, mprotect};

        if min_size_bytes == 0 {
            return Err(new_error!("Cannot create shared memory with size 0"));
        }

        let total_size = min_size_bytes
            .checked_add(2 * PAGE_SIZE_USIZE) // guard page around the memory
            .ok_or_else(|| new_error!("Memory required for sandbox exceeded usize::MAX"))?;

        if total_size % PAGE_SIZE_USIZE != 0 {
            return Err(new_error!(
                "shared memory must be a multiple of {}",
                PAGE_SIZE_USIZE
            ));
        }

        // usize and isize are guaranteed to be the same size, and
        // isize::MAX should be positive, so this cast should be safe.
        if total_size > isize::MAX as usize {
            return Err(HyperlightError::MemoryRequestTooBig(
                total_size,
                isize::MAX as usize,
            ));
        }

        // allocate the memory
        #[cfg(not(miri))]
        let flags = MAP_ANONYMOUS | MAP_PRIVATE | MAP_NORESERVE;
        #[cfg(miri)]
        let flags = MAP_ANONYMOUS | MAP_PRIVATE;

        let addr = unsafe {
            mmap(
                null_mut(),
                total_size as size_t,
                PROT_READ | PROT_WRITE,
                flags,
                -1 as c_int,
                0 as off_t,
            )
        };
        if addr == MAP_FAILED {
            log_then_return!(HyperlightError::MmapFailed(
                Error::last_os_error().raw_os_error()
            ));
        }

        // protect the guard pages
        #[cfg(not(miri))]
        {
            let res = unsafe { mprotect(addr, PAGE_SIZE_USIZE, PROT_NONE) };
            if res != 0 {
                return Err(HyperlightError::MprotectFailed(
                    Error::last_os_error().raw_os_error(),
                ));
            }
            let res = unsafe {
                mprotect(
                    (addr as *const u8).add(total_size - PAGE_SIZE_USIZE) as *mut c_void,
                    PAGE_SIZE_USIZE,
                    PROT_NONE,
                )
            };
            if res != 0 {
                return Err(HyperlightError::MprotectFailed(
                    Error::last_os_error().raw_os_error(),
                ));
            }
        }

        Ok(Self {
            // HostMapping is only non-Send/Sync because raw pointers
            // are not ("as a lint", as the Rust docs say). We don't
            // want to mark HostMapping Send/Sync immediately, because
            // that could socially imply that it's "safe" to use
            // unsafe accesses from multiple threads at once. Instead, we
            // directly impl Send and Sync on this type. Since this
            // type does have Send and Sync manually impl'd, the Arc
            // is not pointless as the lint suggests.
            #[allow(clippy::arc_with_non_send_sync)]
            region: Arc::new(HostMapping {
                ptr: addr as *mut u8,
                size: total_size,
            }),
        })
    }

    /// Internal helper method to get the backing memory as a mutable slice.
    ///
    /// # Safety
    /// As per std::slice::from_raw_parts_mut:
    /// - self.base_addr() must be valid for both reads and writes for
    ///   self.mem_size() * mem::size_of::<u8>() many bytes, and it
    ///   must be properly aligned.
    ///
    ///   The rules on validity are still somewhat unspecified, but we
    ///   assume that the result of our calls to mmap/CreateFileMappings may
    ///   be considered a single "allocated object". The use of
    ///   non-atomic accesses is alright from a Safe Rust standpoint,
    ///   because SharedMemoryBuilder is  not Sync.
    /// - self.base_addr() must point to self.mem_size() consecutive
    ///   properly initialized values of type u8
    ///
    ///   Again, the exact provenance restrictions on what is
    ///   considered to be initialized values are unclear, but we make
    ///   sure to use mmap(MAP_ANONYMOUS) and
    ///   CreateFileMapping(SEC_COMMIT), so the pages in question are
    ///   zero-initialized, which we hope counts for u8.
    /// - The memory referenced by the returned slice must not be
    ///   accessed through any other pointer (not derived from the
    ///   return value) for the duration of the lifetime 'a. Both read
    ///   and write accesses are forbidden.
    ///
    ///   Accesses from Safe Rust necessarily follow this rule,
    ///   because the returned slice's lifetime is the same as that of
    ///   a mutable borrow of self.
    /// - The total size self.mem_size() * mem::size_of::<u8>() of the
    ///   slice must be no larger than isize::MAX, and adding that
    ///   size to data must not "wrap around" the address space. See
    ///   the safety documentation of pointer::offset.
    ///
    ///   This is ensured by a check in ::new()
    pub(super) fn as_mut_slice(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.base_ptr(), self.mem_size()) }
    }

    /// Internal helper method to get the backing memory as a slice.
    ///
    /// # Safety
    /// See the discussion on as_mut_slice, with the third point
    /// replaced by:
    /// - The memory referenced by the returned slice must not be
    ///   mutated for the duration of lifetime 'a, except inside an
    ///   UnsafeCell.
    ///
    ///   Host accesses from Safe Rust necessarily follow this rule,
    ///   because the returned slice's lifetime is the same as that of
    ///   a borrow of self, preventing mutations via other methods.
    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
    pub fn as_slice<'a>(&'a self) -> &'a [u8] {
        unsafe { std::slice::from_raw_parts(self.base_ptr(), self.mem_size()) }
    }

    /// Copies all bytes from `src` to `self` starting at offset
    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
    pub fn copy_from_slice(&mut self, src: &[u8], offset: usize) -> Result<()> {
        let data = self.as_mut_slice();
        bounds_check!(offset, src.len(), data.len());
        data[offset..offset + src.len()].copy_from_slice(src);
        Ok(())
    }

    generate_reader!(read_u8, u8);
    generate_reader!(read_i8, i8);
    generate_reader!(read_u16, u16);
    generate_reader!(read_i16, i16);
    generate_reader!(read_u32, u32);
    generate_reader!(read_i32, i32);
    generate_reader!(read_u64, u64);
    generate_reader!(read_i64, i64);
    generate_reader!(read_usize, usize);
    generate_reader!(read_isize, isize);

    generate_writer!(write_u8, u8);
    generate_writer!(write_i8, i8);
    generate_writer!(write_u16, u16);
    generate_writer!(write_i16, i16);
    generate_writer!(write_u32, u32);
    generate_writer!(write_i32, i32);
    generate_writer!(write_u64, u64);
    generate_writer!(write_i64, i64);
    generate_writer!(write_usize, usize);
    generate_writer!(write_isize, isize);

    /// Convert the ExclusiveSharedMemory, which may be freely
    /// modified, into a GuestSharedMemory, which may be somewhat
    /// freely modified (mostly by the guest), and a HostSharedMemory,
    /// which may only make certain kinds of accesses that do not race
    /// in the presence of malicious code inside the guest mutating
    /// the GuestSharedMemory.
    pub fn build(self) -> (HostSharedMemory, GuestSharedMemory) {
        let lock = Arc::new(RwLock::new(()));
        let hshm = HostSharedMemory {
            region: self.region.clone(),
            lock: lock.clone(),
        };
        (
            hshm,
            GuestSharedMemory {
                region: self.region.clone(),
                lock,
            },
        )
    }
}

fn mapping_at(
    s: &impl SharedMemory,
    gpa: u64,
    size: usize,
    region_type: MemoryRegionType,
    flags: MemoryRegionFlags,
) -> MemoryRegion {
    let guest_base = gpa as usize;

    MemoryRegion {
        guest_region: guest_base..(guest_base + size),
        host_region: s.host_region_base()
            ..<HostGuestMemoryRegion as MemoryRegionKind>::add(s.host_region_base(), size),
        region_type,
        flags,
    }
}

impl GuestSharedMemory {
    /// Create a [`super::memory_region::MemoryRegion`] structure
    /// suitable for mapping this region into a VM
    pub(crate) fn mapping_at(
        &self,
        guest_base: u64,
        region_type: MemoryRegionType,
    ) -> MemoryRegion {
        let flags = match region_type {
            MemoryRegionType::Scratch => {
                MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE
            }
            #[allow(clippy::panic)]
            // This will not ever actually panic: the only places this
            // is called are HyperlightVm::update_snapshot_mapping and
            // HyperlightVm::update_scratch_mapping. The latter
            // statically uses the Scratch region type, and the former
            // does not use this at all — the snapshot scratch mapping is
            // a ReadonlySharedMemory, not a GuestSharedMemory.
            _ => panic!(
                "GuestSharedMemory::mapping_at should only be used for Scratch or Snapshot regions"
            ),
        };
        mapping_at(self, guest_base, self.mem_size(), region_type, flags)
    }
}

/// A trait that abstracts over the particular kind of SharedMemory,
/// used when invoking operations from Rust that absolutely must have
/// exclusive control over the shared memory for correctness +
/// performance, like snapshotting.
pub trait SharedMemory {
    /// Return a readonly reference to the host mapping backing this SharedMemory
    fn region(&self) -> &HostMapping;

    /// Return the base address of the host mapping of this
    /// region. Following the general Rust philosophy, this does not
    /// need to be marked as `unsafe` because doing anything with this
    /// pointer itself requires `unsafe`.
    fn base_addr(&self) -> usize {
        self.region().ptr as usize + PAGE_SIZE_USIZE
    }

    /// Return the base address of the host mapping of this region as
    /// a pointer. Following the general Rust philosophy, this does
    /// not need to be marked as `unsafe` because doing anything with
    /// this pointer itself requires `unsafe`.
    fn base_ptr(&self) -> *mut u8 {
        self.region().ptr.wrapping_add(PAGE_SIZE_USIZE)
    }

    /// Return the length of usable memory contained in `self`.
    /// The returned size does not include the size of the surrounding
    /// guard pages.
    fn mem_size(&self) -> usize {
        self.region().size - 2 * PAGE_SIZE_USIZE
    }

    /// Return the raw base address of the host mapping, including the
    /// guard pages.
    fn raw_ptr(&self) -> *mut u8 {
        self.region().ptr
    }

    /// Return the raw size of the host mapping, including the guard
    /// pages.
    fn raw_mem_size(&self) -> usize {
        self.region().size
    }

    /// Extract a base address that can be mapped into a VM for this
    /// SharedMemory.
    ///
    /// Returns the host base address as a raw `usize` pointer.
    fn host_region_base(&self) -> <HostGuestMemoryRegion as MemoryRegionKind>::HostBaseType {
        self.base_addr()
    }

    /// Return the end address of the host region (base + usable size).
    fn host_region_end(&self) -> <HostGuestMemoryRegion as MemoryRegionKind>::HostBaseType {
        <HostGuestMemoryRegion as MemoryRegionKind>::add(self.host_region_base(), self.mem_size())
    }

    /// Run some code with exclusive access to the SharedMemory
    /// underlying this.  If the SharedMemory is not an
    /// ExclusiveSharedMemory, any concurrent accesses to the relevant
    /// HostSharedMemory/GuestSharedMemory may make this fail, or be
    /// made to fail by this, and should be avoided.
    fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
        &mut self,
        f: F,
    ) -> Result<T>;

    /// Run some code that is allowed to access the contents of the
    /// SharedMemory as if it is a normal slice.  By default, this is
    /// implemented via [`SharedMemory::with_exclusivity`], which is
    /// the correct implementation for a memory that can be mutated,
    /// but a [`ReadonlySharedMemory`], can support this.
    fn with_contents<T, F: FnOnce(&[u8]) -> T>(&mut self, f: F) -> Result<T> {
        self.with_exclusivity(|m| f(m.as_slice()))
    }

    /// Zero a shared memory region
    fn zero(&mut self) -> Result<()> {
        self.with_exclusivity(|e| {
            #[allow(unused_mut)] // unused on some platforms, although not others
            let mut do_copy = true;
            // TODO: Compare & add heuristic thresholds: mmap, MADV_DONTNEED, MADV_REMOVE, MADV_FREE (?)
            #[cfg(feature = "kvm")]
            unsafe {
                let ret = libc::madvise(
                    e.region.ptr as *mut libc::c_void,
                    e.region.size,
                    libc::MADV_DONTNEED,
                );
                if ret == 0 {
                    do_copy = false;
                }
            }
            if do_copy {
                e.as_mut_slice().fill(0);
            }
        })
    }
}

impl SharedMemory for ExclusiveSharedMemory {
    fn region(&self) -> &HostMapping {
        &self.region
    }
    fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
        &mut self,
        f: F,
    ) -> Result<T> {
        Ok(f(self))
    }
}

impl SharedMemory for GuestSharedMemory {
    fn region(&self) -> &HostMapping {
        &self.region
    }
    fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
        &mut self,
        f: F,
    ) -> Result<T> {
        let guard = self
            .lock
            .try_write()
            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
        let mut excl = ExclusiveSharedMemory {
            region: self.region.clone(),
        };
        let ret = f(&mut excl);
        drop(excl);
        drop(guard);
        Ok(ret)
    }
}

/// An unsafe marker trait for types for which all bit patterns are valid.
/// This is required in order for it to be safe to read a value of a particular
/// type out of the sandbox from the HostSharedMemory.
///
/// # Safety
/// This must only be implemented for types for which all bit patterns
/// are valid. It requires that any (non-undef/poison) value of the
/// correct size can be transmuted to the type.
pub unsafe trait AllValid {}
unsafe impl AllValid for u8 {}
unsafe impl AllValid for u16 {}
unsafe impl AllValid for u32 {}
unsafe impl AllValid for u64 {}
unsafe impl AllValid for i8 {}
unsafe impl AllValid for i16 {}
unsafe impl AllValid for i32 {}
unsafe impl AllValid for i64 {}
unsafe impl AllValid for [u8; 16] {}

impl HostSharedMemory {
    /// Read a value of type T, whose representation is the same
    /// between the sandbox and the host, and which has no invalid bit
    /// patterns
    pub fn read<T: AllValid>(&self, offset: usize) -> Result<T> {
        bounds_check!(offset, std::mem::size_of::<T>(), self.mem_size());
        unsafe {
            let mut ret: core::mem::MaybeUninit<T> = core::mem::MaybeUninit::uninit();
            {
                let slice: &mut [u8] = core::slice::from_raw_parts_mut(
                    ret.as_mut_ptr() as *mut u8,
                    std::mem::size_of::<T>(),
                );
                self.copy_to_slice(slice, offset)?;
            }
            Ok(ret.assume_init())
        }
    }

    /// Write a value of type T, whose representation is the same
    /// between the sandbox and the host, and which has no invalid bit
    /// patterns
    pub fn write<T: AllValid>(&self, offset: usize, data: T) -> Result<()> {
        bounds_check!(offset, std::mem::size_of::<T>(), self.mem_size());
        unsafe {
            let slice: &[u8] = core::slice::from_raw_parts(
                core::ptr::addr_of!(data) as *const u8,
                std::mem::size_of::<T>(),
            );
            self.copy_from_slice(slice, offset)?;
        }
        Ok(())
    }

    /// Copy the contents of the slice into the sandbox at the
    /// specified offset
    pub fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
        bounds_check!(offset, slice.len(), self.mem_size());
        let base = self.base_ptr().wrapping_add(offset);
        let guard = self
            .lock
            .try_read()
            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;

        const CHUNK: usize = size_of::<u128>();
        let len = slice.len();
        let mut i = 0;

        // Handle unaligned head bytes until we reach u128 alignment.
        // Note: align_offset can return usize::MAX if alignment is impossible.
        // In that case, head_len = len via .min(), so we fall back to byte-by-byte
        // operations for the entire slice.
        let align_offset = base.align_offset(align_of::<u128>());
        let head_len = align_offset.min(len);
        while i < head_len {
            unsafe {
                slice[i] = base.add(i).read_volatile();
            }
            i += 1;
        }

        // Read aligned u128 chunks
        // SAFETY: After processing head_len bytes, base.add(i) is u128-aligned.
        // We use write_unaligned for the destination since the slice may not be u128-aligned.
        let dst = slice.as_mut_ptr();
        while i + CHUNK <= len {
            unsafe {
                let value = (base.add(i) as *const u128).read_volatile();
                std::ptr::write_unaligned(dst.add(i) as *mut u128, value);
            }
            i += CHUNK;
        }

        // Handle remaining tail bytes
        while i < len {
            unsafe {
                slice[i] = base.add(i).read_volatile();
            }
            i += 1;
        }

        drop(guard);
        Ok(())
    }

    /// Copy the contents of the sandbox at the specified offset into
    /// the slice
    pub fn copy_from_slice(&self, slice: &[u8], offset: usize) -> Result<()> {
        bounds_check!(offset, slice.len(), self.mem_size());
        let base = self.base_ptr().wrapping_add(offset);
        let guard = self
            .lock
            .try_read()
            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;

        const CHUNK: usize = size_of::<u128>();
        let len = slice.len();
        let mut i = 0;

        // Handle unaligned head bytes until we reach u128 alignment.
        // Note: align_offset can return usize::MAX if alignment is impossible.
        // In that case, head_len = len via .min(), so we fall back to byte-by-byte
        // operations for the entire slice.
        let align_offset = base.align_offset(align_of::<u128>());
        let head_len = align_offset.min(len);
        while i < head_len {
            unsafe {
                base.add(i).write_volatile(slice[i]);
            }
            i += 1;
        }

        // Write aligned u128 chunks
        // SAFETY: After processing head_len bytes, base.add(i) is u128-aligned.
        // We use read_unaligned for the source since the slice may not be u128-aligned.
        let src = slice.as_ptr();
        while i + CHUNK <= len {
            unsafe {
                let value = std::ptr::read_unaligned(src.add(i) as *const u128);
                (base.add(i) as *mut u128).write_volatile(value);
            }
            i += CHUNK;
        }

        // Handle remaining tail bytes
        while i < len {
            unsafe {
                base.add(i).write_volatile(slice[i]);
            }
            i += 1;
        }

        drop(guard);
        Ok(())
    }

    /// Fill the memory in the range `[offset, offset + len)` with `value`
    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
    pub fn fill(&mut self, value: u8, offset: usize, len: usize) -> Result<()> {
        bounds_check!(offset, len, self.mem_size());
        let base = self.base_ptr().wrapping_add(offset);
        let guard = self
            .lock
            .try_read()
            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;

        const CHUNK: usize = size_of::<u128>();
        let value_u128 = u128::from_ne_bytes([value; CHUNK]);
        let mut i = 0;

        // Handle unaligned head bytes until we reach u128 alignment.
        // Note: align_offset can return usize::MAX if alignment is impossible.
        // In that case, head_len = len via .min(), so we fall back to byte-by-byte
        // operations for the entire slice.
        let align_offset = base.align_offset(align_of::<u128>());
        let head_len = align_offset.min(len);
        while i < head_len {
            unsafe {
                base.add(i).write_volatile(value);
            }
            i += 1;
        }

        // Write aligned u128 chunks
        // SAFETY: After processing head_len bytes, base.add(i) is u128-aligned
        while i + CHUNK <= len {
            unsafe {
                (base.add(i) as *mut u128).write_volatile(value_u128);
            }
            i += CHUNK;
        }

        // Handle remaining tail bytes
        while i < len {
            unsafe {
                base.add(i).write_volatile(value);
            }
            i += 1;
        }

        drop(guard);
        Ok(())
    }

    /// Pushes the given data onto shared memory to the buffer at the given offset.
    /// NOTE! buffer_start_offset must point to the beginning of the buffer
    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
    pub fn push_buffer(
        &mut self,
        buffer_start_offset: usize,
        buffer_size: usize,
        data: &[u8],
    ) -> Result<()> {
        let stack_pointer_rel = self.read::<u64>(buffer_start_offset)? as usize;
        let buffer_size_u64: u64 = buffer_size.try_into()?;

        if stack_pointer_rel > buffer_size || stack_pointer_rel < 8 {
            return Err(new_error!(
                "Unable to push data to buffer: Stack pointer is out of bounds. Stack pointer: {}, Buffer size: {}",
                stack_pointer_rel,
                buffer_size_u64
            ));
        }

        let size_required = data.len() + 8;
        let size_available = buffer_size - stack_pointer_rel;

        if size_required > size_available {
            return Err(new_error!(
                "Not enough space in buffer to push data. Required: {}, Available: {}",
                size_required,
                size_available
            ));
        }

        // get absolute
        let stack_pointer_abs = stack_pointer_rel + buffer_start_offset;

        // write the actual data to the top of stack
        self.copy_from_slice(data, stack_pointer_abs)?;

        // write the offset to the newly written data, to the top of stack.
        // this is used when popping the stack, to know how far back to jump
        self.write::<u64>(stack_pointer_abs + data.len(), stack_pointer_rel as u64)?;

        // update stack pointer to point to the next free address
        self.write::<u64>(
            buffer_start_offset,
            (stack_pointer_rel + data.len() + 8) as u64,
        )?;
        Ok(())
    }

    /// Pop the top element of the ring as raw bytes. Unlike
    /// [`Self::try_pop_buffer_into`], this doesn't peek at the element's
    /// contents — the element size is recovered from the trailing
    /// back-pointer that [`Self::push_buffer`] wrote.
    pub fn try_pop_buffer_raw(
        &mut self,
        buffer_start_offset: usize,
        buffer_size: usize,
    ) -> Result<Vec<u8>> {
        let stack_pointer_rel = self.read::<u64>(buffer_start_offset)? as usize;

        if stack_pointer_rel > buffer_size || stack_pointer_rel < 16 {
            return Err(new_error!(
                "try_pop_buffer_raw: stack pointer {} out of bounds (size {})",
                stack_pointer_rel,
                buffer_size
            ));
        }

        let back_ptr_abs = stack_pointer_rel + buffer_start_offset - 8;
        let element_offset_rel = self.read::<u64>(back_ptr_abs)? as usize;

        if element_offset_rel < 8 || element_offset_rel > stack_pointer_rel.saturating_sub(8) {
            return Err(new_error!(
                "try_pop_buffer_raw: back-pointer {} outside [8, {}]",
                element_offset_rel,
                stack_pointer_rel.saturating_sub(8)
            ));
        }

        let element_size = stack_pointer_rel - element_offset_rel - 8;
        let element_abs = element_offset_rel + buffer_start_offset;
        let mut out = vec![0u8; element_size];
        self.copy_to_slice(&mut out, element_abs)?;

        // Pop: rewind stack pointer.
        self.write::<u64>(buffer_start_offset, element_offset_rel as u64)?;
        // Zero out the popped slot + its back-pointer.
        self.fill(0, element_abs, stack_pointer_rel - element_offset_rel)?;

        Ok(out)
    }

    /// Pops the given given buffer into a `T` and returns it.
    /// NOTE! the data must be a size-prefixed flatbuffer, and
    /// buffer_start_offset must point to the beginning of the buffer
    pub fn try_pop_buffer_into<T>(
        &mut self,
        buffer_start_offset: usize,
        buffer_size: usize,
    ) -> Result<T>
    where
        T: for<'b> TryFrom<&'b [u8]>,
    {
        // get the stackpointer
        let stack_pointer_rel = self.read::<u64>(buffer_start_offset)? as usize;

        if stack_pointer_rel > buffer_size || stack_pointer_rel < 16 {
            return Err(new_error!(
                "Unable to pop data from buffer: Stack pointer is out of bounds. Stack pointer: {}, Buffer size: {}",
                stack_pointer_rel,
                buffer_size
            ));
        }

        // make it absolute
        let last_element_offset_abs = stack_pointer_rel + buffer_start_offset;

        // go back 8 bytes to get offset to element on top of stack
        let last_element_offset_rel: usize =
            self.read::<u64>(last_element_offset_abs - 8)? as usize;

        // Validate element offset (guest-writable): must be in [8, stack_pointer_rel - 16]
        // to leave room for the 8-byte back-pointer plus at least 8 bytes of element data
        // (the minimum for a size-prefixed flatbuffer: 4-byte prefix + 4-byte root offset).
        if last_element_offset_rel > stack_pointer_rel.saturating_sub(16)
            || last_element_offset_rel < 8
        {
            return Err(new_error!(
                "Corrupt buffer back-pointer: element offset {} is outside valid range [8, {}].",
                last_element_offset_rel,
                stack_pointer_rel.saturating_sub(16),
            ));
        }

        // make it absolute
        let last_element_offset_abs = last_element_offset_rel + buffer_start_offset;

        // Max bytes the element can span (excluding the 8-byte back-pointer).
        let max_element_size = stack_pointer_rel - last_element_offset_rel - 8;

        // Get the size of the flatbuffer buffer from memory
        let fb_buffer_size = {
            let raw_prefix = self.read::<u32>(last_element_offset_abs)?;
            // flatbuffer byte arrays are prefixed by 4 bytes indicating
            // the remaining size; add 4 for the prefix itself.
            let total = raw_prefix.checked_add(4).ok_or_else(|| {
                new_error!(
                    "Corrupt buffer size prefix: value {} overflows when adding 4-byte header.",
                    raw_prefix
                )
            })?;
            usize::try_from(total)
        }?;

        if fb_buffer_size > max_element_size {
            return Err(new_error!(
                "Corrupt buffer size prefix: flatbuffer claims {} bytes but the element slot is only {} bytes.",
                fb_buffer_size,
                max_element_size
            ));
        }

        let mut result_buffer = vec![0; fb_buffer_size];

        self.copy_to_slice(&mut result_buffer, last_element_offset_abs)?;
        let to_return = T::try_from(result_buffer.as_slice()).map_err(|_e| {
            new_error!(
                "pop_buffer_into: failed to convert buffer to {}",
                type_name::<T>()
            )
        })?;

        // update the stack pointer to point to the element we just popped off since that is now free
        self.write::<u64>(buffer_start_offset, last_element_offset_rel as u64)?;

        // zero out the memory we just popped off
        let num_bytes_to_zero = stack_pointer_rel - last_element_offset_rel;
        self.fill(0, last_element_offset_abs, num_bytes_to_zero)?;

        Ok(to_return)
    }
}

impl SharedMemory for HostSharedMemory {
    fn region(&self) -> &HostMapping {
        &self.region
    }
    fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
        &mut self,
        f: F,
    ) -> Result<T> {
        let guard = self
            .lock
            .try_write()
            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
        let mut excl = ExclusiveSharedMemory {
            region: self.region.clone(),
        };
        let ret = f(&mut excl);
        drop(excl);
        drop(guard);
        Ok(ret)
    }
}

/// A ReadonlySharedMemory is a different kind of shared memory,
/// separate from the exclusive/host/guest lifecycle, used to
/// represent read-only mappings of snapshot pages into the guest
/// efficiently.
#[derive(Clone, Debug)]
pub struct ReadonlySharedMemory {
    region: Arc<HostMapping>,
    /// If `Some`, only this many bytes are mapped into guest PA space
    /// by `mapping_at`. If `None`, the full `mem_size()` is mapped.
    guest_mapped_size: Option<usize>,
}
// Safety: HostMapping is only non-Send/Sync (causing
// ReadonlySharedMemory to not be automatically Send/Sync) because raw
// pointers are not ("as a lint", as the Rust docs say). We don't want
// to mark HostMapping Send/Sync immediately, because that could
// socially imply that it's "safe" to use unsafe accesses from
// multiple threads at once in more cases, including ones that don't
// actually ensure immutability/synchronisation. Since
// ReadonlySharedMemory can only be accessed by reading, and reading
// concurrently from multiple threads is not racy,
// ReadonlySharedMemory can be Send and Sync.
unsafe impl Send for ReadonlySharedMemory {}
unsafe impl Sync for ReadonlySharedMemory {}

impl ReadonlySharedMemory {
    pub(crate) fn from_bytes(contents: &[u8]) -> Result<Self> {
        let mut anon = ExclusiveSharedMemory::new(contents.len())?;
        anon.copy_from_slice(contents, 0)?;
        Ok(ReadonlySharedMemory {
            region: anon.region,
            guest_mapped_size: None,
        })
    }

    /// The number of bytes that should be mapped into guest PA space.
    /// Returns `guest_mapped_size` if set, otherwise `mem_size()`.
    pub(crate) fn guest_mapped_size(&self) -> usize {
        self.guest_mapped_size.unwrap_or_else(|| self.mem_size())
    }

    pub(crate) fn as_slice(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.base_ptr(), self.mem_size()) }
    }

    pub(crate) fn build(self) -> (Self, Self) {
        (self.clone(), self)
    }

    pub(crate) fn mapping_at(
        &self,
        guest_base: u64,
        region_type: MemoryRegionType,
    ) -> MemoryRegion {
        #[allow(clippy::panic)]
        // This will not ever actually panic: the only place this is
        // called is HyperlightVm::update_snapshot_mapping, which
        // always calls it with the Snapshot region type.
        if region_type != MemoryRegionType::Snapshot {
            panic!("ReadonlySharedMemory::mapping_at should only be used for Snapshot regions");
        }
        // Register snapshot mem RWX at the KVM level. Upstream marked
        // this RX-only and relied on guest-PT CoW for write semantics,
        // which trapped first writes and resolved them to scratch frames
        // — driving a slow leak via prim_alloc. The underlying mmap is
        // already PROT_READ|PROT_WRITE; `ReadonlySharedMemory` is a
        // host-side Rust-API artifact, not a KVM-level constraint.
        mapping_at(
            self,
            guest_base,
            self.guest_mapped_size(),
            region_type,
            MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
        )
    }
}

impl SharedMemory for ReadonlySharedMemory {
    fn region(&self) -> &HostMapping {
        &self.region
    }
    // There's no way to get exclusive (and therefore writable) access
    // to a ReadonlySharedMemory.
    fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
        &mut self,
        _: F,
    ) -> Result<T> {
        Err(new_error!(
            "Cannot take exclusive access to a ReadonlySharedMemory"
        ))
    }
    // However, just access to the contents as a slice is doable
    fn with_contents<T, F: FnOnce(&[u8]) -> T>(&mut self, f: F) -> Result<T> {
        Ok(f(self.as_slice()))
    }
}

impl<S: SharedMemory> PartialEq<S> for ReadonlySharedMemory {
    fn eq(&self, other: &S) -> bool {
        self.raw_ptr() == other.raw_ptr()
    }
}