spacewasm/

memory.rs

use crate::Rc;
use crate::{AllocError, MemType};
use core::alloc::Layout;
use core::fmt::{Debug, Formatter};
use core::ptr::NonNull;
use core::slice;

/// An allocator for allocating Wasm pages
pub trait WasmMemoryAllocator {
    /// Allocate a new memory region for linear memory
    fn allocate(&self, layout: Layout) -> Result<NonNull<u8>, AllocError>;

    /// Reallocate a memory region, moving data if needed.
    ///
    /// On success, returns a pointer to a block of `layout` whose leading
    /// `old_layout.size()` bytes hold the previous contents; the old allocation
    /// (`ptr`/`old_layout`) has been consumed and must not be used again.
    ///
    /// # Failure ownership contract
    ///
    /// On `Err`, the implementation MUST leave the original allocation
    /// (`ptr` with `old_layout`) valid, unmoved, and still owned by the caller.
    fn reallocate(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        layout: Layout,
    ) -> Result<NonNull<u8>, AllocError>;

    /// Deallocate memory that has been allocated
    fn deallocate(&self, ptr: NonNull<u8>, layout: Layout);
}

impl<T: WasmMemoryAllocator> Rc<T> {
    pub fn into_wasm_memory_allocator(self) -> Rc<dyn WasmMemoryAllocator>
    where
        T: WasmMemoryAllocator + 'static,
    {
        unsafe { self.into_dyn(|x| x as &dyn WasmMemoryAllocator) }
    }
}

/// A WebAssembly linear memory
///
/// Linear memory implements inner mutability so it is not thread-safe.
/// `Memory` holds a raw pointer, so it is automatically `!Send + !Sync`.
///
/// Needs to have inner mutability because multiple modules can point
/// to it (via imports) and we use `Rc<Memory>` for reference counting these ownerships.
///
/// # Aliasing model
///
/// The `store_*` mutators take `&self`, so the borrow checker will not stop you
/// from writing these bytes while a shared reference into them is alive. Drop
/// any [`Memory::load`] / [`Memory::get_slice`] slice before any write —
/// `store_*`, [`Memory::grow`] (which may reallocate), a host callback, or guest
/// re-entry. Holding one across a write is UB under Stacked Borrows.
/// [`Memory::as_shared_cells`] is the one aliasing-tolerant accessor.
pub struct Memory {
    ptr: *mut u8,
    size: usize,
    ty: MemType,
    allocator: Option<Rc<dyn WasmMemoryAllocator>>,
}

impl Debug for Memory {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Memory").finish()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MemoryError {
    OutOfBounds,
    OutOfMemory,
    AllocationFailed,
    PageTooSmall,
}

impl From<AllocError> for MemoryError {
    fn from(e: AllocError) -> MemoryError {
        match e {
            AllocError::AllocationFailed => MemoryError::AllocationFailed,
            AllocError::OutOfMemory => MemoryError::OutOfMemory,
            AllocError::PageTooSmall => MemoryError::PageTooSmall,
        }
    }
}

impl Default for Memory {
    fn default() -> Self {
        Memory::zero()
    }
}

impl Memory {
    pub fn zero() -> Memory {
        Memory {
            ptr: core::ptr::null_mut(),
            size: 0,
            ty: MemType::zero(),
            allocator: None,
        }
    }

    pub fn new(ty: MemType, allocator: Rc<dyn WasmMemoryAllocator>) -> Result<Memory, AllocError> {
        let size = if let Some(size) = (ty.min() as u64).checked_mul(ty.page_size() as u64) {
            if size > usize::MAX as u64 {
                // Make sure this platform can actually do this allocation
                return Err(AllocError::AllocationFailed);
            } else {
                size as usize
            }
        } else {
            return Err(AllocError::AllocationFailed);
        };

        let ptr = if size == 0 {
            core::ptr::null_mut()
        } else {
            let ptr = allocator
                .allocate(
                    Layout::from_size_align(size, ty.page_alignment())
                        .map_err(|_| AllocError::AllocationFailed)?,
                )?
                .as_ptr();

            // Clear the pages
            unsafe {
                ptr.write_bytes(0, size);
            }
            ptr
        };

        Ok(Memory {
            ptr,
            size,
            ty,
            allocator: Some(allocator),
        })
    }
}

impl Memory {
    #[inline]
    fn check_in_bounds(&self, addr: usize, size: usize) -> Result<(), MemoryError> {
        if size > self.size || addr > self.size - size {
            Err(MemoryError::OutOfBounds)
        } else {
            Ok(())
        }
    }

    /// Resolve a guest memory access `base + offset` to a byte address.
    ///
    /// The sum is evaluated in `u64` and narrowed back to `usize`, so it can
    /// never wrap on 32-bit targets, where `usize` is 32-bit and `base + offset`
    /// may exceed `u32::MAX`. An address that does not fit `usize` is rejected as
    /// out of bounds rather than aliasing a valid cell. The resulting address is
    /// still bounds-checked against the memory size by the individual load/store
    /// operations.
    #[inline]
    pub fn effective_address(base: u32, offset: u32) -> Result<usize, MemoryError> {
        usize::try_from(base as u64 + offset as u64).map_err(|_| MemoryError::OutOfBounds)
    }

    pub fn store_u8(&self, addr: usize, i: u8) -> Result<(), MemoryError> {
        self.check_in_bounds(addr, 1)?;
        unsafe {
            self.ptr.add(addr).write(i);
        }
        Ok(())
    }

    pub fn store_u16(&self, addr: usize, i: u16) -> Result<(), MemoryError> {
        self.check_in_bounds(addr, 2)?;
        unsafe {
            self.ptr.add(addr).cast::<u16>().write_unaligned(i.to_le());
        }
        Ok(())
    }

    pub fn store_u32(&self, addr: usize, i: u32) -> Result<(), MemoryError> {
        self.check_in_bounds(addr, 4)?;
        unsafe {
            self.ptr.add(addr).cast::<u32>().write_unaligned(i.to_le());
        }
        Ok(())
    }

    pub fn store_u64(&self, addr: usize, i: u64) -> Result<(), MemoryError> {
        self.check_in_bounds(addr, 8)?;
        unsafe {
            self.ptr.add(addr).cast::<u64>().write_unaligned(i.to_le());
        }
        Ok(())
    }

    pub fn store(&self, addr: usize, data: &[u8]) -> Result<(), MemoryError> {
        self.check_in_bounds(addr, data.len())?;
        if data.is_empty() {
            return Ok(());
        }

        unsafe {
            data.as_ptr().copy_to(self.ptr.add(addr), data.len());
        }
        Ok(())
    }

    pub fn load_u8(&self, addr: usize) -> Result<u8, MemoryError> {
        self.check_in_bounds(addr, 1)?;
        unsafe { Ok(self.ptr.add(addr).read()) }
    }

    pub fn load_u16(&self, addr: usize) -> Result<u16, MemoryError> {
        self.check_in_bounds(addr, 2)?;
        unsafe {
            Ok(u16::from_le(
                self.ptr.add(addr).cast::<u16>().read_unaligned(),
            ))
        }
    }

    pub fn load_u32(&self, addr: usize) -> Result<u32, MemoryError> {
        self.check_in_bounds(addr, 4)?;
        unsafe {
            Ok(u32::from_le(
                self.ptr.add(addr).cast::<u32>().read_unaligned(),
            ))
        }
    }

    pub fn load_u64(&self, addr: usize) -> Result<u64, MemoryError> {
        self.check_in_bounds(addr, 8)?;
        unsafe {
            Ok(u64::from_le(
                self.ptr.add(addr).cast::<u64>().read_unaligned(),
            ))
        }
    }

    pub fn load(&self, addr: usize, len: usize) -> Result<&[u8], MemoryError> {
        self.check_in_bounds(addr, len)?;
        if len == 0 {
            return Ok(&[]);
        }
        Ok(unsafe { core::slice::from_raw_parts(self.ptr.add(addr), len) })
    }

    /// Grow the memory by n pages
    /// If the memory growth succeeds, return the old number of pages
    pub fn grow(&mut self, n: u32) -> Result<u32, MemoryError> {
        let old_pages = self.size();
        let total_pages = old_pages.checked_add(n).ok_or(MemoryError::OutOfMemory)?;

        if !self.ty.can_hold(total_pages) {
            return Err(MemoryError::OutOfMemory);
        }

        let allocator = self.allocator.as_ref().ok_or(MemoryError::OutOfMemory)?;
        if n == 0 {
            return Ok(old_pages);
        }

        let old_size = self.size;
        let new_size = usize::try_from(u64::from(total_pages) * self.ty.page_size() as u64)
            .map_err(|_| MemoryError::OutOfMemory)?;
        let new_layout = Layout::from_size_align(new_size, self.ty.page_alignment())
            .map_err(|_| MemoryError::AllocationFailed)?;

        let ptr = if old_size == 0 {
            allocator.allocate(new_layout)?
        } else {
            allocator.reallocate(
                NonNull::new(self.ptr).ok_or(MemoryError::OutOfMemory)?,
                Layout::from_size_align(old_size, self.ty.page_alignment())
                    .map_err(|_| MemoryError::AllocationFailed)?,
                new_layout,
            )?
        };

        // Clear the new memory before committing the larger size.
        unsafe {
            ptr.as_ptr()
                .add(old_size)
                .write_bytes(0, new_size - old_size);
        }

        self.ptr = ptr.as_ptr();
        self.size = new_size;

        Ok(old_pages)
    }

    pub fn mem_type(&self) -> MemType {
        self.ty
    }

    pub fn size(&self) -> u32 {
        (self.size / self.ty.page_size()) as u32
    }

    /// Returns whether this is the absent-memory sentinel created by [`Memory::zero`].
    pub fn is_zero(&self) -> bool {
        self.ptr.is_null() && self.allocator.is_none()
    }

    pub fn get_slice(&self) -> &[u8] {
        if self.ptr.is_null() {
            &[]
        } else {
            unsafe { slice::from_raw_parts(self.ptr, self.size) }
        }
    }

    /// Expose linear memory as `&[UnsafeCell<u8>]`.
    ///
    /// Unlike [`get_slice`], this view stays valid across writes, which is what
    /// makes it usable as a `wiggle` `GuestMemory::Shared` backing store. Still
    /// invalidated by [`Memory::grow`]. Empty for the [`Memory::zero`] sentinel.
    ///
    /// [`get_slice`]: Memory::get_slice
    pub fn as_shared_cells(&self) -> &[core::cell::UnsafeCell<u8>] {
        if self.ptr.is_null() {
            &[]
        } else {
            unsafe {
                core::slice::from_raw_parts(
                    self.ptr as *const core::cell::UnsafeCell<u8>,
                    self.size,
                )
            }
        }
    }
}

impl Drop for Memory {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            let allocator = self.allocator.take().unwrap();
            allocator.deallocate(
                NonNull::new(self.ptr).unwrap(),
                Layout::from_size_align(self.size, self.ty.page_alignment()).unwrap(),
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::MemPageSize;
    extern crate std;

    /// Test allocator backed by the system heap, mirroring the one used in
    /// `interpreter_tests`.
    struct TestAllocator;
    impl WasmMemoryAllocator for TestAllocator {
        fn allocate(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
            assert_ne!(layout.size(), 0);
            unsafe { NonNull::new(std::alloc::alloc(layout)).ok_or(AllocError::AllocationFailed) }
        }

        fn reallocate(
            &self,
            ptr: NonNull<u8>,
            old_layout: Layout,
            layout: Layout,
        ) -> Result<NonNull<u8>, AllocError> {
            assert_ne!(old_layout.size(), 0);
            assert_ne!(layout.size(), 0);
            unsafe {
                let new_ptr = std::alloc::realloc(ptr.as_ptr(), old_layout, layout.size());
                NonNull::new(new_ptr).ok_or(AllocError::AllocationFailed)
            }
        }

        fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
            assert_ne!(layout.size(), 0);
            unsafe { std::alloc::dealloc(ptr.as_ptr(), layout) }
        }
    }

    fn allocator() -> Rc<dyn WasmMemoryAllocator> {
        Rc::new(TestAllocator).unwrap().into_wasm_memory_allocator()
    }

    fn mem_type(min: u32, max: Option<u32>) -> MemType {
        MemType {
            initial_pages: min,
            max_pages: max,
            page_size: MemPageSize::_65536,
        }
    }

    #[test]
    fn zero_memory_is_zero() {
        let mem = Memory::zero();
        assert!(mem.is_zero());
        assert_eq!(mem.size(), 0);
    }

    #[test]
    fn new_allocates_and_reports_size() {
        let mem = Memory::new(mem_type(2, Some(4)), allocator()).unwrap();
        assert!(!mem.is_zero());
        assert_eq!(mem.size(), 2);
        assert_eq!(mem.mem_type().min(), 2);
        // The freshly allocated pages are zeroed.
        assert!(mem.get_slice().iter().all(|&b| b == 0));
        assert_eq!(mem.get_slice().len(), 2 * MemPageSize::_65536.size());
    }

    #[test]
    fn zero_page_memory_allocates_lazily() {
        let mut mem = Memory::new(mem_type(0, None), allocator()).unwrap();

        assert!(!mem.is_zero());
        assert_eq!(mem.size(), 0);
        assert!(mem.get_slice().is_empty());
        assert_eq!(mem.load(0, 0), Ok(&[][..]));
        assert_eq!(mem.store(0, &[]), Ok(()));
        assert_eq!(mem.load(1, 0), Err(MemoryError::OutOfBounds));
        assert_eq!(mem.store(1, &[]), Err(MemoryError::OutOfBounds));

        assert_eq!(mem.grow(0), Ok(0));
        assert_eq!(mem.grow(1), Ok(0));
        assert_eq!(mem.size(), 1);
        assert_eq!(mem.grow(0), Ok(1));
        assert_eq!(mem.get_slice().len(), MemPageSize::_65536.size());
        assert!(mem.get_slice().iter().all(|&b| b == 0));
        assert_eq!(mem.load(0, 1), Ok(&[0][..]));
        assert_eq!(mem.store(0, &[0xA5]), Ok(()));
        assert_eq!(mem.load(0, 1), Ok(&[0xA5][..]));
    }

    #[test]
    fn stores_are_little_endian() {
        // Wasm mandates little-endian memory regardless of host byte order.
        // Assert the actual byte layout, which round-trip tests cannot catch.
        let mem = Memory::new(mem_type(1, None), allocator()).unwrap();

        mem.store_u16(0, 0x1122).unwrap();
        assert_eq!(mem.load(0, 2).unwrap(), &[0x22, 0x11][..]);

        mem.store_u32(8, 0x1122_3344).unwrap();
        assert_eq!(mem.load(8, 4).unwrap(), &[0x44, 0x33, 0x22, 0x11][..]);

        mem.store_u64(16, 0x1122_3344_5566_7788).unwrap();
        assert_eq!(
            mem.load(16, 8).unwrap(),
            &[0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11][..]
        );
    }

    #[test]
    fn loads_are_little_endian() {
        // A little-endian byte pattern must decode to the expected value.
        let mem = Memory::new(mem_type(1, None), allocator()).unwrap();

        mem.store(0, &[0x22, 0x11]).unwrap();
        assert_eq!(mem.load_u16(0).unwrap(), 0x1122);

        mem.store(8, &[0x44, 0x33, 0x22, 0x11]).unwrap();
        assert_eq!(mem.load_u32(8).unwrap(), 0x1122_3344);

        mem.store(16, &[0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11])
            .unwrap();
        assert_eq!(mem.load_u64(16).unwrap(), 0x1122_3344_5566_7788);
    }

    #[test]
    fn grow_extends_and_zeroes() {
        let mut mem = Memory::new(mem_type(1, Some(4)), allocator()).unwrap();
        // Write a marker into the first page.
        mem.store_u32(0, 0xAABB_CCDD).unwrap();

        let old_pages = mem.grow(2).unwrap();
        assert_eq!(old_pages, 1);
        assert_eq!(mem.size(), 3);

        // Existing data survives the grow...
        assert_eq!(mem.load_u32(0).unwrap(), 0xAABB_CCDD);
        // ...and the newly added region is zeroed.
        let first_new_byte = MemPageSize::_65536.size();
        assert_eq!(mem.load_u8(first_new_byte).unwrap(), 0);
    }

    #[test]
    fn grow_past_max_fails() {
        let mut mem = Memory::new(mem_type(1, Some(2)), allocator()).unwrap();
        // Growing to 3 pages exceeds the declared maximum of 2.
        assert_eq!(mem.grow(2), Err(MemoryError::OutOfMemory));
        // The memory is left untouched.
        assert_eq!(mem.size(), 1);
    }

    #[test]
    fn grow_overflow_page_count_fails() {
        let mut mem = Memory::new(mem_type(1, None), allocator()).unwrap();
        // `size() + n` overflows u32.
        assert_eq!(mem.grow(u32::MAX), Err(MemoryError::OutOfMemory));
    }

    #[test]
    fn grow_zero_memory_fails() {
        // The absent-memory sentinel has no backing allocator.
        let mut mem = Memory::zero();
        assert_eq!(mem.grow(0), Err(MemoryError::OutOfMemory));
        assert_eq!(mem.grow(1), Err(MemoryError::OutOfMemory));
    }

    #[test]
    fn load_store_out_of_bounds() {
        let mem = Memory::new(mem_type(1, None), allocator()).unwrap();
        let size = MemPageSize::_65536.size();
        assert_eq!(mem.load_u32(size - 3), Err(MemoryError::OutOfBounds));
        assert_eq!(mem.store_u64(size - 4, 0), Err(MemoryError::OutOfBounds));
        assert_eq!(mem.load(size - 1, 4), Err(MemoryError::OutOfBounds));
    }

    #[test]
    fn effective_address_rejects_overflow_on_32bit() {
        // On 64-bit `usize`, base+offset always fits; the sum is exact.
        assert_eq!(Memory::effective_address(10, 20).unwrap(), 30);
    }

    #[test]
    fn as_shared_cells_aliases_and_mutates() {
        let mem = Memory::new(mem_type(1, None), allocator()).unwrap();

        // The cell view covers exactly the linear memory bytes.
        let cells = mem.as_shared_cells();
        assert_eq!(cells.len(), mem.get_slice().len());

        // A store is observable through the cell view...
        mem.store_u8(3, 0x5A).unwrap();
        assert_eq!(unsafe { *cells[3].get() }, 0x5A);

        // ...and a write through a cell is observable through a load, which is
        // the shared-but-mutable access wiggle relies on.
        unsafe {
            *cells[7].get() = 0xC3;
        }
        assert_eq!(mem.load_u8(7).unwrap(), 0xC3);
    }

    #[test]
    fn as_shared_cells_zero_memory_is_empty() {
        assert!(Memory::zero().as_shared_cells().is_empty());
    }

    #[test]
    fn memory_error_from_alloc_error() {
        assert_eq!(
            MemoryError::from(AllocError::AllocationFailed),
            MemoryError::AllocationFailed
        );
        assert_eq!(
            MemoryError::from(AllocError::OutOfMemory),
            MemoryError::OutOfMemory
        );
        assert_eq!(
            MemoryError::from(AllocError::PageTooSmall),
            MemoryError::PageTooSmall
        );
    }
}

#[cfg(kani)]
mod kani_proofs {
    use super::*;
    use crate::Allocator;
    use crate::MemPageSize;
    extern crate std;

    use crate::test_support::RustSystemAllocator;

    #[kani::proof]
    fn proof_store_load_correctness() {
        let alloc = RustSystemAllocator;
        let size = 64;

        let ptr = unsafe {
            alloc
                .alloc(Layout::from_size_align(size, (MemPageSize::_65536).alignment()).unwrap())
                .unwrap()
        };
        let mem = Memory {
            ptr,
            size,
            ty: MemType::zero(),
            allocator: None,
        };

        // Test all integer sizes with symbolic values and addresses
        let addr: usize = kani::any();
        kani::assume(addr <= size - 8); // Reserve space for largest type (u64)

        // Test u8 store/load
        let val_u8: u8 = kani::any();
        mem.store_u8(addr, val_u8).unwrap();
        assert_eq!(mem.load_u8(addr).unwrap(), val_u8, "u8 round-trip failed");

        // Test u16 store/load
        let val_u16: u16 = kani::any();
        mem.store_u16(addr, val_u16).unwrap();
        assert_eq!(
            mem.load_u16(addr).unwrap(),
            val_u16,
            "u16 round-trip failed"
        );

        // Test u32 store/load
        let val_u32: u32 = kani::any();
        mem.store_u32(addr, val_u32).unwrap();
        assert_eq!(
            mem.load_u32(addr).unwrap(),
            val_u32,
            "u32 round-trip failed"
        );

        // Test u64 store/load
        let val_u64: u64 = kani::any();
        mem.store_u64(addr, val_u64).unwrap();
        assert_eq!(
            mem.load_u64(addr).unwrap(),
            val_u64,
            "u64 round-trip failed"
        );

        // Test out-of-bounds detection
        assert!(
            mem.store_u32(size - 2, 0).is_err(),
            "Out-of-bounds store must fail"
        );
        assert!(
            mem.store_u32(size, 0).is_err(),
            "Store at boundary must fail"
        );
        assert!(
            mem.load_u32(size - 2).is_err(),
            "Out-of-bounds load must fail"
        );
        assert!(mem.load_u32(size).is_err(), "Load at boundary must fail");

        // Test addresses that cause overflow in bounds check
        let overflow_addr = usize::MAX - 1;
        assert!(
            mem.store_u32(overflow_addr, 0).is_err(),
            "Overflow address must be rejected"
        );
        assert!(
            mem.load_u32(overflow_addr).is_err(),
            "Overflow address must be rejected"
        );

        // Verify safe addresses have lossless isize conversion
        if addr <= isize::MAX as usize {
            let offset = addr as isize;
            assert!(offset >= 0, "Valid addr converts to non-negative offset");
            assert!(offset as usize == addr, "isize conversion must be lossless");
        }

        // Prevent Drop from calling GlobalAllocator FFI
        core::mem::forget(mem);
        unsafe {
            alloc.dealloc(
                ptr,
                Layout::from_size_align(size, (MemPageSize::_65536).alignment()).unwrap(),
            )
        };
    }

    #[kani::proof]
    fn proof_byte_slice_operations() {
        let alloc = RustSystemAllocator;
        let size = 16;

        let ptr = unsafe {
            alloc
                .alloc(Layout::from_size_align(size, (MemPageSize::_65536).alignment()).unwrap())
                .unwrap()
        };
        let mem = Memory {
            ptr,
            size,
            ty: MemType::zero(),
            allocator: None,
        };

        // Test fixed-size byte slice (4 bytes) at symbolic address
        let addr: usize = kani::any();
        kani::assume(addr <= size - 4);

        // Use symbolic 4-byte array
        let data: [u8; 4] = kani::any();

        // Store and load back
        mem.store(addr, &data).unwrap();
        let loaded = mem.load(addr, 4).unwrap();
        assert_eq!(loaded, &data, "Byte slice round-trip failed");

        // Test empty slice
        mem.store(0, &[]).unwrap();
        assert_eq!(mem.load(0, 0).unwrap().len(), 0, "Empty slice must work");

        // Test out-of-bounds slice access
        assert!(
            mem.store(size - 2, &[1, 2, 3, 4]).is_err(),
            "Out-of-bounds store must fail"
        );
        assert!(
            mem.load(size - 2, 4).is_err(),
            "Out-of-bounds load must fail"
        );

        // Prevent Drop from calling GlobalAllocator FFI
        core::mem::forget(mem);
        unsafe {
            alloc.dealloc(
                ptr,
                Layout::from_size_align(size, (MemPageSize::_65536).alignment()).unwrap(),
            )
        };
    }

    /// The effective address is the exact sum of `base` and `offset`; it is
    /// never silently wrapped modulo the pointer width.
    #[kani::proof]
    fn proof_effective_address_no_wrap() {
        let base: u32 = kani::any();
        let offset: u32 = kani::any();
        let ea = Memory::effective_address(base, offset).unwrap();
        assert_eq!(ea as u64, base as u64 + offset as u64);
    }
}