flaron_sdk/

mem.rs

//! Guest-side memory helpers and the per-invocation bump arena.
//!
//! Every host function that returns data does so by writing into the guest's
//! linear memory at a slot the SDK allocates with [`guest_alloc`]. The host
//! returns a packed `(ptr, len)` value that callers decode with the helpers
//! in this module.
//!
//! ## Wasm vs host builds
//!
//! On `wasm32-unknown-unknown` the bump arena is a static byte array and a
//! "pointer" returned from [`guest_alloc`] is a real linear-memory address -
//! the wasm host writes/reads directly through it.
//!
//! On native targets (used by `cargo test`) the arena is a `thread_local` Box
//! and a "pointer" is an OFFSET into that buffer. Real 64-bit Rust pointers
//! cannot fit in the i32 ABI slots the wasm host expects, so the test mock
//! and the SDK both route through the offset model. The cfg-gates in this
//! module are the only place that distinction shows up - every other module
//! is target-agnostic.
//!
//! This module is `pub(crate)` aside from the two arena hooks
//! ([`guest_alloc`] and [`reset_arena`]) - it is plumbing, not API.

/// Decode a packed i64 returned from the host into `(ptr, len)`.
///
/// Host convention: high 32 bits = pointer into the guest's address space,
/// low 32 bits = length in bytes.
pub(crate) fn decode_ptr_len(packed: i64) -> (u32, u32) {
    let packed = packed as u64;
    ((packed >> 32) as u32, (packed & 0xFFFF_FFFF) as u32)
}

/// Pack a `(ptr, len)` pair into a host-style i64 return value.
///
/// Inverse of [`decode_ptr_len`]. Used by the test mock host to assemble
/// responses; the real wasm host packs values the same way.
#[allow(dead_code)]
pub(crate) fn encode_ptr_len(ptr: u32, len: u32) -> i64 {
    (((ptr as u64) << 32) | (len as u64)) as i64
}

/// Read raw bytes from the guest at the given pointer/offset and length.
///
/// # Safety
/// On wasm the address must lie within memory the host has just written into
/// the guest's linear memory via the bump arena. On non-wasm the value is
/// treated as an offset into the thread-local arena buffer instead.
pub(crate) unsafe fn read_bytes(ptr: u32, len: u32) -> Vec<u8> {
    if len == 0 {
        return Vec::new();
    }
    #[cfg(target_arch = "wasm32")]
    {
        core::slice::from_raw_parts(ptr as usize as *const u8, len as usize).to_vec()
    }
    #[cfg(not(target_arch = "wasm32"))]
    {
        arena::read_at(ptr as usize, len as usize)
    }
}

/// Read a UTF-8 string from guest memory.
///
/// # Safety
/// Same requirements as [`read_bytes`]. Bytes that are not valid UTF-8 are
/// replaced with the Unicode replacement character - we accept that loss to
/// avoid undefined behaviour from `from_utf8_unchecked`.
pub(crate) unsafe fn read_string(ptr: u32, len: u32) -> String {
    if len == 0 {
        return String::new();
    }
    let bytes = read_bytes(ptr, len);
    String::from_utf8_lossy(&bytes).into_owned()
}

/// Read a string from a packed `(ptr, len)` return value.
///
/// Returns `None` when the host returned `0` ("no result" / not found).
///
/// # Safety
/// The packed pointer must point at memory the host wrote into the guest's
/// address space within the current invocation. See [`read_string`].
pub(crate) unsafe fn read_packed_string(packed: i64) -> Option<String> {
    if packed == 0 {
        return None;
    }
    let (ptr, len) = decode_ptr_len(packed);
    Some(read_string(ptr, len))
}

/// Read raw bytes from a packed `(ptr, len)` return value.
///
/// Returns `None` when the host returned `0` ("no result" / not found).
///
/// # Safety
/// The packed pointer must point at memory the host wrote into the guest's
/// address space within the current invocation.
pub(crate) unsafe fn read_packed_bytes(packed: i64) -> Option<Vec<u8>> {
    if packed == 0 {
        return None;
    }
    let (ptr, len) = decode_ptr_len(packed);
    Some(read_bytes(ptr, len))
}

/// Decode a hex string into raw bytes. Returns `None` on the first invalid
/// character - callers MUST treat this as a hard failure (corrupted host
/// response, never silently zero-fill, especially for crypto material).
pub(crate) fn hex_decode(hex: &str) -> Option<Vec<u8>> {
    let hex = hex.as_bytes();
    if !hex.len().is_multiple_of(2) {
        return None;
    }
    let mut bytes = Vec::with_capacity(hex.len() / 2);
    for chunk in hex.chunks_exact(2) {
        let hi = hex_nibble(chunk[0])?;
        let lo = hex_nibble(chunk[1])?;
        bytes.push((hi << 4) | lo);
    }
    Some(bytes)
}

fn hex_nibble(c: u8) -> Option<u8> {
    match c {
        b'0'..=b'9' => Some(c - b'0'),
        b'a'..=b'f' => Some(c - b'a' + 10),
        b'A'..=b'F' => Some(c - b'A' + 10),
        _ => None,
    }
}

/// Pack a borrowed byte slice as a `(ptr, len)` host-call argument.
///
/// On wasm this is a zero-cost cast: the i32 is the slice's pointer, which
/// the wasm host can dereference directly.
///
/// On non-wasm targets the slice is copied into the per-thread arena and
/// the offset is returned - real 64-bit Rust pointers don't fit in i32 ABI
/// slots, so we use offsets uniformly with the response side.
pub(crate) fn host_arg_bytes(data: &[u8]) -> (i32, i32) {
    #[cfg(target_arch = "wasm32")]
    {
        (data.as_ptr() as i32, data.len() as i32)
    }
    #[cfg(not(target_arch = "wasm32"))]
    {
        if data.is_empty() {
            return (0, 0);
        }
        let offset = guest_alloc(data.len() as i32);
        if offset == 0 {
            return (0, 0);
        }
        arena::write_at(offset as usize, data);
        (offset, data.len() as i32)
    }
}

/// Convenience wrapper for string args.
pub(crate) fn host_arg_str(s: &str) -> (i32, i32) {
    host_arg_bytes(s.as_bytes())
}

// ---------- Per-invocation bump allocator ----------
//
// The flaron host calls the guest's exported `alloc(size)` every time it
// needs to hand a value back (every spark_get, every header lookup, every WS
// event, etc.). A naive global allocator would leak that memory for the
// lifetime of the entire WASM instance - long-lived flares would steadily
// grow until OOM.
//
// Instead we use a per-invocation bump arena. The flare's WASM exports
// (`handle_request`, `ws_open`, `ws_message`, `ws_close`) reset the arena at
// the top of each invocation; everything the host allocates during that
// invocation is reclaimed when control returns to the host.

const ARENA_SIZE: usize = 256 * 1024;

#[cfg(target_arch = "wasm32")]
mod arena {
    use super::ARENA_SIZE;
    use core::cell::UnsafeCell;

    struct BumpArena {
        buf: UnsafeCell<[u8; ARENA_SIZE]>,
        offset: UnsafeCell<usize>,
    }

    // SAFETY: WASM guests are single-threaded; nothing else can race the arena.
    unsafe impl Sync for BumpArena {}

    static ARENA: BumpArena = BumpArena {
        buf: UnsafeCell::new([0; ARENA_SIZE]),
        offset: UnsafeCell::new(0),
    };

    pub fn reset() {
        // SAFETY: single-threaded guest; we own the arena.
        unsafe {
            *ARENA.offset.get() = 0;
        }
    }

    pub fn alloc(size: i32) -> i32 {
        if size <= 0 {
            return 0;
        }
        let size = size as usize;
        // SAFETY: single-threaded guest; only this function writes `offset`.
        unsafe {
            let offset = &mut *ARENA.offset.get();
            let aligned = (*offset + 7) & !7;
            if aligned.checked_add(size).is_none_or(|end| end > ARENA_SIZE) {
                return 0;
            }
            *offset = aligned + size;
            let buf = ARENA.buf.get() as *mut u8;
            buf.add(aligned) as i32
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
mod arena {
    use super::ARENA_SIZE;
    use std::cell::RefCell;

    /// First valid offset. The SDK treats a packed return of `0` as "no
    /// result", so we never hand out offset 0 - the first 8 bytes of the
    /// buffer are reserved sentinel space.
    const FIRST_OFFSET: usize = 8;

    struct ThreadArena {
        buf: Box<[u8; ARENA_SIZE]>,
        offset: usize,
    }

    impl ThreadArena {
        fn new() -> Self {
            Self {
                buf: Box::new([0; ARENA_SIZE]),
                offset: FIRST_OFFSET,
            }
        }
    }

    thread_local! {
        static ARENA: RefCell<ThreadArena> = RefCell::new(ThreadArena::new());
    }

    pub fn reset() {
        ARENA.with(|a| a.borrow_mut().offset = FIRST_OFFSET);
    }

    /// Returns the OFFSET (not a real pointer) into the per-thread arena.
    /// On non-wasm targets the SDK and the test mock both interpret the
    /// returned i32 as `arena[offset..]`, never as a Rust pointer.
    pub fn alloc(size: i32) -> i32 {
        if size <= 0 {
            return 0;
        }
        let size = size as usize;
        ARENA.with(|a| {
            let mut a = a.borrow_mut();
            let aligned = (a.offset + 7) & !7;
            if aligned.checked_add(size).is_none_or(|end| end > ARENA_SIZE) {
                return 0;
            }
            a.offset = aligned + size;
            aligned as i32
        })
    }

    pub fn write_at(offset: usize, data: &[u8]) {
        ARENA.with(|a| {
            let mut a = a.borrow_mut();
            a.buf[offset..offset + data.len()].copy_from_slice(data);
        })
    }

    pub fn read_at(offset: usize, len: usize) -> Vec<u8> {
        ARENA.with(|a| a.borrow().buf[offset..offset + len].to_vec())
    }
}

/// Reset the bump arena. Call this at the top of every guest export so the
/// next host invocation starts with a fresh 256 KiB scratch space.
///
/// The flare entry-point macros ([`crate::handle_request!`],
/// [`crate::ws_handlers!`]) call this for you - only invoke it directly if
/// you are wiring up your own export functions.
pub fn reset_arena() {
    arena::reset();
}

/// Guest memory allocator the host calls (via the `alloc` export generated
/// by [`crate::export_alloc!`]) to write return values into the WASM linear
/// memory. Hands out 8-byte aligned slices from the bump arena.
///
/// Returns `0` on failure (size not positive, arena exhausted) - the host
/// treats `0` as "guest cannot accept this value" and propagates an error.
pub fn guest_alloc(size: i32) -> i32 {
    arena::alloc(size)
}

/// Test-only re-exports of the arena helpers. Used by `crate::ffi`'s mock
/// host on non-wasm targets to write canned responses.
#[cfg(not(target_arch = "wasm32"))]
pub(crate) use arena::{read_at as arena_read_at, write_at as arena_write_at};

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn decode_ptr_len_splits_high_low() {
        let packed = encode_ptr_len(0xDEAD_BEEF, 0x1234);
        let (ptr, len) = decode_ptr_len(packed);
        assert_eq!(ptr, 0xDEAD_BEEF);
        assert_eq!(len, 0x1234);
    }

    #[test]
    fn encode_decode_round_trip_zero() {
        assert_eq!(decode_ptr_len(0), (0, 0));
    }

    #[test]
    fn hex_decode_lowercase() {
        assert_eq!(hex_decode("deadbeef"), Some(vec![0xde, 0xad, 0xbe, 0xef]));
    }

    #[test]
    fn hex_decode_uppercase() {
        assert_eq!(hex_decode("DEADBEEF"), Some(vec![0xde, 0xad, 0xbe, 0xef]));
    }

    #[test]
    fn hex_decode_mixed_case() {
        assert_eq!(hex_decode("DeAdBeEf"), Some(vec![0xde, 0xad, 0xbe, 0xef]));
    }

    #[test]
    fn hex_decode_odd_length_fails() {
        assert!(hex_decode("abc").is_none());
    }

    #[test]
    fn hex_decode_invalid_char_fails() {
        assert!(hex_decode("zz").is_none());
    }

    #[test]
    fn hex_decode_empty_ok() {
        assert_eq!(hex_decode(""), Some(vec![]));
    }

    #[test]
    fn guest_alloc_returns_zero_for_non_positive() {
        reset_arena();
        assert_eq!(guest_alloc(0), 0);
        assert_eq!(guest_alloc(-1), 0);
    }

    #[test]
    fn guest_alloc_aligned_to_eight() {
        reset_arena();
        let p1 = guest_alloc(1);
        assert_ne!(p1, 0);
        let p2 = guest_alloc(1);
        assert_ne!(p2, 0);
        // p2 must be 8 bytes after p1 (alignment).
        assert_eq!(p2 - p1, 8);
    }

    #[test]
    fn guest_alloc_returns_zero_when_exhausted() {
        reset_arena();
        // Single allocation larger than the arena.
        assert_eq!(guest_alloc((ARENA_SIZE + 1) as i32), 0);
    }

    #[test]
    fn reset_arena_recycles_space() {
        reset_arena();
        let p1 = guest_alloc(64);
        reset_arena();
        let p2 = guest_alloc(64);
        assert_eq!(p1, p2);
    }

    #[test]
    fn read_packed_zero_is_none() {
        assert!(unsafe { read_packed_string(0) }.is_none());
        assert!(unsafe { read_packed_bytes(0) }.is_none());
    }

    #[test]
    fn read_bytes_zero_len_is_empty() {
        let bytes = unsafe { read_bytes(0, 0) };
        assert!(bytes.is_empty());
    }

    #[test]
    fn host_arg_bytes_round_trip() {
        reset_arena();
        let (ptr, len) = host_arg_bytes(b"hello");
        #[cfg(not(target_arch = "wasm32"))]
        {
            assert!(ptr > 0);
            assert_eq!(len, 5);
            let read = unsafe { read_bytes(ptr as u32, len as u32) };
            assert_eq!(read, b"hello");
        }
        #[cfg(target_arch = "wasm32")]
        {
            let _ = (ptr, len);
        }
    }
}