feat: Implement parts of the extism runtime in WebAssembly (#384)

This PR adds the `kernel` directory which contains a port of the Extism
memory allocator compiled to WebAssembly and removes
`runtime/src/memory.rs` completely.

Being able to re-use memory functions as a WASM module allows us to
begin to experiment with porting Extism to new runtimes!

This is in a draft state while I'm verifying some of these changes.

Cargo.toml

@@ -5,3 +5,4 @@ members = [
   "rust",
   "libextism",
 ]
+exclude = ["kernel"]

Makefile

@@ -21,6 +21,10 @@ endif
 build:
 	cargo build --release $(FEATURE_FLAGS) --manifest-path libextism/Cargo.toml
 
+.PHONY: kernel
+kernel:
+	cd kernel && bash build.sh
+
 lint:
 	cargo clippy --release --no-deps --manifest-path runtime/Cargo.toml
 

kernel/.cargo/config

@@ -0,0 +1,2 @@
+[build]
+target = "wasm32-unknown-unknown"

kernel/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "extism-runtime-kernel"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+
+[workspace]
+members = [
+  "."
+]

kernel/README.md

@@ -0,0 +1,20 @@
+# Extism kernel
+
+The Extism kernel implements core parts of the Extism runtime in Rust compiled to WebAssembly. This code is a conceptual
+re-write of [memory.rs][] with the goal of making core parts of the Extism implementation more portable across WebAssembly
+runtimes.
+
+See [lib.rs][] for more details about the implementation itself.
+
+## Building
+
+Because this crate is built using the `wasm32-unknown-unknown` target, it is a separate build process from the `extism-runtime` crate. 
+
+To build `extism-runtime.wasm`, strip it and copy it to the proper location in the `extism-runtime` tree you can run:
+
+```shell
+$ sh build.sh
+```
+
+[memory.rs]: https://github.com/extism/extism/blob/f4aa139eced4a74eb4a103f78222ba503e146109/runtime/src/memory.rs
+[lib.rs]: ./src/lib.rs

kernel/build.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+
+cargo build --release --target wasm32-unknown-unknown --package extism-runtime-kernel --bin extism-runtime
+cp target/wasm32-unknown-unknown/release/extism-runtime.wasm .
+wasm-strip extism-runtime.wasm || :
+mv extism-runtime.wasm ../runtime/src/extism-runtime.wasm
+

kernel/src/bin/extism-runtime.rs

@@ -0,0 +1,10 @@
+#![no_main]
+#![no_std]
+
+pub use extism_runtime_kernel::*;
+
+#[cfg(target_arch = "wasm32")]
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    core::arch::wasm32::unreachable()
+}

kernel/src/lib.rs

@@ -0,0 +1,446 @@
+//! # Extism kernel
+//!
+//! - Isolated memory from both host and plugin
+//! - An allocator for managing that memory
+//! - Input/output handling
+//! - Error message handling
+//! - Backward compatible `extism_*` functions
+//!
+//! ## Allocator
+//!
+//! The Extism allocator is a bump allocator that tracks the `length` of the total number of bytes
+//! available to the allocator and `position` to track how much of the data has been used. Things like memory
+//! have not really been optimized at all. When a new allocation that is larger than the remaning size is made,
+//! the allocator attempts to call `memory.grow` if that fails a `0` offset is returned, which should be interpreted
+//! as a failed allocation.
+//!
+//! ## Input/Output
+//!
+//! Input and output are just allocated blocks of memory that are marked as either input or output using
+//! the `extism_input_set` or `extism_output_set` functions. The global variables `INPUT_OFFSET` contains
+//! the offset in memory to the input data and `INPUT_LENGTH` contains the size of the input data. `OUTPUT_OFFSET`
+//! and `OUTPUT_LENGTH` are used for the output data.
+//!
+//! ## Error handling
+//!
+//! The `ERROR` global is used to track the current error message. If it is set to `0` then there is no error.
+//! The length of the error message can be retreived using `extism_length`.
+//!
+//! ## Memory offsets
+//! An offset of `0` is similar to a `NULL` pointer in C - it implies an allocation failure or memory error
+//! of some kind
+//!
+//! ## Extism functions
+//!
+//! These functions are backward compatible with the pre-kernel runtime, but a few new functions are added to
+//! give runtimes more access to the internals necesarry to load data in and out of a plugin.
+#![no_std]
+#![allow(clippy::missing_safety_doc)]
+
+use core::sync::atomic::*;
+
+pub type Pointer = u64;
+pub type Length = u64;
+
+/// WebAssembly page size
+const PAGE_SIZE: usize = 65536;
+
+/// Determines the amount of bytes that can be wasted by re-using a block. If more than this number is wasted by re-using
+/// a block then it will be split into two smaller blocks.
+const BLOCK_SPLIT_SIZE: usize = 128;
+
+/// Offset to the input data
+static mut INPUT_OFFSET: Pointer = 0;
+
+/// Length of the input data
+static mut INPUT_LENGTH: Length = 0;
+
+/// Offset to the output data
+static mut OUTPUT_OFFSET: Pointer = 0;
+
+/// Offset to the input data
+static mut OUTPUT_LENGTH: Length = 0;
+
+/// Current error message
+static mut ERROR: AtomicU64 = AtomicU64::new(0);
+
+/// Determines if the kernel has been initialized already
+static mut INITIALIZED: AtomicBool = AtomicBool::new(false);
+
+/// A pointer to the first page that will be managed by Extism, this is set during initialization
+static mut START_PAGE: usize = 0;
+
+/// Provides information about the usage status of a `MemoryBlock`
+#[repr(u8)]
+#[derive(PartialEq)]
+pub enum MemoryStatus {
+    /// Unused memory that is available b
+    Unused = 0,
+    /// In-use memory
+    Active = 1,
+    /// Free memory that is available for re-use
+    Free = 2,
+}
+
+/// A single `MemoryRoot` exists at the start of the memory to track information about the total
+/// size of the allocated memory and the position of the bump allocator.
+///
+/// The overall layout of the Extism-manged memory is organized like this:
+
+/// |------|-------|---------|-------|--------------|
+/// | Root | Block |  Data   | Block |     Data     | ...
+/// |------|-------|---------|-------|--------------|
+///
+/// Where `Root` and `Block` are fixed to the size of the `MemoryRoot` and `MemoryBlock` structs. But
+/// the size of `Data` is dependent on the allocation size.
+#[repr(C)]
+pub struct MemoryRoot {
+    /// Position of the bump allocator, relative to `START_PAGE`
+    pub position: AtomicU64,
+    /// The total size of all data allocated using this allocator
+    pub length: AtomicU64,
+    /// A pointer to where the blocks begin
+    pub blocks: [MemoryBlock; 0],
+}
+
+/// A `MemoryBlock` contains some metadata about a single allocation
+#[repr(C)]
+pub struct MemoryBlock {
+    /// The usage status of the block, `Unused` or `Free` blocks can be re-used.
+    pub status: AtomicU8,
+    /// The total size of the allocation
+    pub size: usize,
+    /// The number of bytes currently being used. If this block is a fresh allocation then `size` and `used` will
+    /// always be the same. If a block is re-used then these numbers may differ.
+    pub used: usize,
+    /// A pointer to the block data
+    pub data: [u8; 0],
+}
+
+/// Returns the number of pages needed for the given number of bytes
+pub fn num_pages(nbytes: u64) -> usize {
+    let nbytes = nbytes as f64;
+    let page = PAGE_SIZE as f64;
+    ((nbytes / page) + 0.5) as usize
+}
+
+// Get the `MemoryRoot` at the correct offset in memory
+#[inline]
+unsafe fn memory_root() -> &'static mut MemoryRoot {
+    &mut *((START_PAGE * PAGE_SIZE) as *mut MemoryRoot)
+}
+
+impl MemoryRoot {
+    /// Initialize or load the `MemoryRoot` from the correct position in memory
+    pub unsafe fn new() -> &'static mut MemoryRoot {
+        // If this fails then `INITIALIZED` is already `true` and we can just return the
+        // already initialized `MemoryRoot`
+        if INITIALIZED
+            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+            .is_err()
+        {
+            return memory_root();
+        }
+
+        // Ensure that at least one page is allocated to store the `MemoryRoot` data
+        START_PAGE = core::arch::wasm32::memory_grow(0, 1);
+        if START_PAGE == usize::MAX {
+            panic!("Out of memory");
+        }
+
+        // Initialize the `MemoryRoot` length, position and data
+        let root = memory_root();
+        root.length.store(
+            PAGE_SIZE as u64 - core::mem::size_of::<MemoryRoot>() as u64,
+            Ordering::Release,
+        );
+        root.position.store(0, Ordering::Release);
+
+        // Ensure the first block is marked as `Unused`
+        #[allow(clippy::size_of_in_element_count)]
+        core::ptr::write_bytes(
+            root.blocks.as_mut_ptr() as *mut _,
+            MemoryStatus::Unused as u8,
+            core::mem::size_of::<MemoryBlock>(),
+        );
+        root
+    }
+
+    /// Resets the position of the allocator and zeroes out all allocations
+    pub unsafe fn reset(&mut self) {
+        core::ptr::write_bytes(
+            self.blocks.as_mut_ptr() as *mut u8,
+            0,
+            self.length.load(Ordering::Acquire) as usize,
+        );
+        self.position.store(0, Ordering::Release);
+    }
+
+    // Find a block that is free to use, this can be a new block or an existing freed block. The `self_position` argument
+    // is used to avoid loading the allocators position more than once when performing an allocation.
+    unsafe fn find_free_block(
+        &mut self,
+        length: Length,
+        self_position: u64,
+    ) -> Option<&'static mut MemoryBlock> {
+        // Get the first block
+        let mut block = self.blocks.as_mut_ptr();
+
+        // Only loop while the block pointer is less then the current position
+        while (block as u64) < self.blocks.as_ptr() as u64 + self_position {
+            let b = &mut *block;
+
+            let status = b.status.load(Ordering::Acquire);
+
+            // An unused block is safe to use
+            if status == MemoryStatus::Unused as u8 {
+                return Some(b);
+            }
+
+            // Re-use freed blocks when they're large enough
+            if status == MemoryStatus::Free as u8 && b.size >= length as usize {
+                // Split block if there is too much excess
+                if b.size - length as usize >= BLOCK_SPLIT_SIZE {
+                    b.size -= length as usize;
+                    b.used = 0;
+
+                    let block1 = b.data.as_mut_ptr().add(b.size) as *mut MemoryBlock;
+                    let b1 = &mut *block1;
+                    b1.size = length as usize;
+                    b1.used = 0;
+                    b1.status.store(MemoryStatus::Free as u8, Ordering::Release);
+                    return Some(b1);
+                }
+
+                // Otherwise return the whole block
+                return Some(b);
+            }
+
+            // Get the next block
+            block = b.next_ptr();
+        }
+
+        None
+    }
+
+    /// Create a new `MemoryBlock`, when `Some(block)` is returned, `block` will contain at least enough room for `length` bytes
+    /// but may be as large as `length` + `BLOCK_SPLIT_SIZE` bytes. When `None` is returned the allocation has failed.
+    pub unsafe fn alloc(&mut self, length: Length) -> Option<&'static mut MemoryBlock> {
+        let self_position = self.position.load(Ordering::Acquire);
+        let self_length = self.length.load(Ordering::Acquire);
+        let b = self.find_free_block(length, self_position);
+
+        // If there's a free block then re-use it
+        if let Some(b) = b {
+            b.used = length as usize;
+            b.status
+                .store(MemoryStatus::Active as u8, Ordering::Release);
+            return Some(b);
+        }
+
+        // Get the current index for a new block
+        let curr = self.blocks.as_ptr() as u64 + self_position;
+
+        // Get the number of bytes available
+        let mem_left = self_length - self_position;
+
+        // When the allocation is larger than the number of bytes available
+        // we will need to try to grow the memory
+        if length >= mem_left {
+            // Calculate the number of pages needed to cover the remaining bytes
+            let npages = num_pages(length);
+            let x = core::arch::wasm32::memory_grow(0, npages);
+            if x == usize::MAX {
+                return None;
+            }
+            self.length
+                .fetch_add(npages as u64 * PAGE_SIZE as u64, Ordering::SeqCst);
+        }
+
+        // Bump the position by the size of the actual data + the size of the MemoryBlock structure
+        self.position.fetch_add(
+            length + core::mem::size_of::<MemoryBlock>() as u64,
+            Ordering::SeqCst,
+        );
+
+        // Initialize a new block at the current position
+        let ptr = curr as *mut MemoryBlock;
+        let block = &mut *ptr;
+        block
+            .status
+            .store(MemoryStatus::Active as u8, Ordering::Release);
+        block.size = length as usize;
+        block.used = length as usize;
+        Some(block)
+    }
+
+    /// Finds the block at an offset in memory
+    pub unsafe fn find_block(&mut self, offs: Pointer) -> Option<&mut MemoryBlock> {
+        if offs >= self.blocks.as_ptr() as Pointer + self.length.load(Ordering::Acquire) as Pointer
+        {
+            return None;
+        }
+        let ptr = offs - core::mem::size_of::<MemoryBlock>() as u64;
+        let ptr = ptr as *mut MemoryBlock;
+        Some(&mut *ptr)
+    }
+}
+
+impl MemoryBlock {
+    /// Get a pointer to the next block
+    ///
+    /// NOTE: This does no checking to ensure the resulting pointer is valid, the offset
+    /// is calculated based on metadata provided by the current block
+    #[inline]
+    pub unsafe fn next_ptr(&mut self) -> *mut MemoryBlock {
+        self.data
+            .as_mut_ptr()
+            .add(self.size + core::mem::size_of::<MemoryBlock>()) as *mut MemoryBlock
+    }
+
+    /// Mark a block as free
+    pub fn free(&mut self) {
+        self.status
+            .store(MemoryStatus::Free as u8, Ordering::Release);
+    }
+}
+
+// Extism functions - these functions should be
+
+/// Allocate a block of memory and return the offset
+#[no_mangle]
+pub unsafe fn extism_alloc(n: Length) -> Pointer {
+    let region = MemoryRoot::new();
+    let block = region.alloc(n);
+    match block {
+        Some(block) => block.data.as_mut_ptr() as Pointer,
+        None => 0,
+    }
+}
+
+/// Free allocated memory
+#[no_mangle]
+pub unsafe fn extism_free(p: Pointer) {
+    let block = MemoryRoot::new().find_block(p);
+    if let Some(block) = block {
+        block.free();
+    }
+}
+
+/// Get the length of an allocated memory block
+#[no_mangle]
+pub unsafe fn extism_length(p: Pointer) -> Length {
+    if p == 0 {
+        return 0;
+    }
+    if let Some(block) = MemoryRoot::new().find_block(p) {
+        block.used as Length
+    } else {
+        0
+    }
+}
+
+/// Load a byte from Extism-managed memory
+#[no_mangle]
+pub unsafe fn extism_load_u8(p: Pointer) -> u8 {
+    *(p as *mut u8)
+}
+
+/// Load a u64 from Extism-managed memory
+#[no_mangle]
+pub unsafe fn extism_load_u64(p: Pointer) -> u64 {
+    *(p as *mut u64)
+}
+
+/// Load a byte from the input data
+#[no_mangle]
+pub unsafe fn extism_input_load_u8(p: Pointer) -> u8 {
+    *((INPUT_OFFSET + p) as *mut u8)
+}
+
+/// Load a u64 from the input data
+#[no_mangle]
+pub unsafe fn extism_input_load_u64(p: Pointer) -> u64 {
+    *((INPUT_OFFSET + p) as *mut u64)
+}
+
+/// Write a byte in Extism-managed memory
+#[no_mangle]
+pub unsafe fn extism_store_u8(p: Pointer, x: u8) {
+    *(p as *mut u8) = x;
+}
+
+/// Write a u64 in Extism-managed memory
+#[no_mangle]
+pub unsafe fn extism_store_u64(p: Pointer, x: u64) {
+    unsafe {
+        *(p as *mut u64) = x;
+    }
+}
+
+/// Set the range of the input data in memory
+#[no_mangle]
+pub fn extism_input_set(p: Pointer, len: Length) {
+    unsafe {
+        INPUT_OFFSET = p;
+        INPUT_LENGTH = len;
+    }
+}
+
+/// Set the range of the output data in memory
+#[no_mangle]
+pub fn extism_output_set(p: Pointer, len: Length) {
+    unsafe {
+        OUTPUT_OFFSET = p;
+        OUTPUT_LENGTH = len;
+    }
+}
+
+/// Get the input length
+#[no_mangle]
+pub fn extism_input_length() -> Length {
+    unsafe { INPUT_LENGTH }
+}
+
+/// Get the input offset in Exitsm-managed memory
+#[no_mangle]
+pub fn extism_input_offset() -> Length {
+    unsafe { INPUT_OFFSET }
+}
+
+/// Get the output length
+#[no_mangle]
+pub fn extism_output_length() -> Length {
+    unsafe { OUTPUT_LENGTH }
+}
+
+/// Get the output offset in Extism-managed memory
+#[no_mangle]
+pub fn extism_output_offset() -> Length {
+    unsafe { OUTPUT_OFFSET }
+}
+
+/// Reset the allocator
+#[no_mangle]
+pub unsafe fn extism_reset() {
+    ERROR.store(0, Ordering::SeqCst);
+    MemoryRoot::new().reset()
+}
+
+/// Set the error message offset
+#[no_mangle]
+pub unsafe fn extism_error_set(ptr: Pointer) {
+    ERROR.store(ptr, Ordering::SeqCst);
+}
+
+/// Get the error message offset, if it's `0` then no error has been set
+#[no_mangle]
+pub unsafe fn extism_error_get() -> Pointer {
+    ERROR.load(Ordering::SeqCst)
+}
+
+/// Get the position of the allocator, this can be used as an indication of how many bytes are currently in-use
+#[no_mangle]
+pub unsafe fn extism_memory_bytes() -> Length {
+    MemoryRoot::new().position.load(Ordering::Acquire)
+}

libextism/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "libextism"
-version = "0.3.0"
+version = "0.4.0"
 edition = "2021"
 authors = ["The Extism Authors", "oss@extism.org"]
 license = "BSD-3-Clause"

manifest/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "extism-manifest"
-version = "0.3.0"
+version = "0.4.0"
 edition = "2021"
 authors = ["The Extism Authors", "oss@extism.org"]
 license = "BSD-3-Clause"

manifest/src/lib.rs

@@ -4,7 +4,7 @@ use std::path::{Path, PathBuf};
 #[deprecated]
 pub type ManifestMemory = MemoryOptions;
 
-#[derive(Default, serde::Serialize, serde::Deserialize)]
+#[derive(Default, Clone, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "json_schema", derive(schemars::JsonSchema))]
 #[serde(deny_unknown_fields)]
 pub struct MemoryOptions {
@@ -12,7 +12,7 @@ pub struct MemoryOptions {
     pub max_pages: Option<u32>,
 }
 
-#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Clone, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "json_schema", derive(schemars::JsonSchema))]
 #[serde(deny_unknown_fields)]
 pub struct HttpRequest {
@@ -43,7 +43,7 @@ impl HttpRequest {
     }
 }
 
-#[derive(Default, serde::Serialize, serde::Deserialize)]
+#[derive(Default, Clone, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "json_schema", derive(schemars::JsonSchema))]
 #[serde(deny_unknown_fields)]
 pub struct WasmMetadata {
@@ -81,7 +81,7 @@ impl From<Vec<u8>> for Wasm {
 #[deprecated]
 pub type ManifestWasm = Wasm;
 
-#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Clone, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "json_schema", derive(schemars::JsonSchema))]
 #[serde(untagged)]
 #[serde(deny_unknown_fields)]
@@ -153,7 +153,7 @@ fn base64_schema(gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::
     schema.into()
 }
 
-#[derive(Default, serde::Serialize, serde::Deserialize)]
+#[derive(Default, Clone, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "json_schema", derive(schemars::JsonSchema))]
 #[serde(deny_unknown_fields)]
 pub struct Manifest {

python/example.py

@@ -5,7 +5,9 @@ import hashlib
 import pathlib
 
 sys.path.append(".")
-from extism import Function, host_fn, ValType, Plugin
+from extism import Function, host_fn, ValType, Plugin, set_log_file
+
+set_log_file("stderr", "trace")
 
 
 @host_fn
@@ -26,7 +28,7 @@ def main(args):
     if len(args) > 1:
         data = args[1].encode()
     else:
-        data = b"some data from python!"
+        data = b"a" * 1024
 
     wasm_file_path = (
         pathlib.Path(__file__).parent.parent / "wasm" / "code-functions.wasm"
@@ -46,11 +48,13 @@ def main(args):
     ]
     plugin = Plugin(manifest, wasi=True, functions=functions)
     # Call `count_vowels`
-    wasm_vowel_count = json.loads(plugin.call("count_vowels", data))
+    wasm_vowel_count = plugin.call("count_vowels", data)
+    print(wasm_vowel_count)
+    j = json.loads(wasm_vowel_count)
 
-    print("Number of vowels:", wasm_vowel_count["count"])
+    print("Number of vowels:", j["count"])
 
-    assert wasm_vowel_count["count"] == count_vowels(data)
+    assert j["count"] == count_vowels(data)
 
 
 if __name__ == "__main__":

runtime/Cargo.toml

@@ -22,8 +22,7 @@ log4rs = "1.1"
 url = "2"
 glob = "0.3"
 ureq = {version = "2.5", optional=true}
-extism-manifest = { version = "0.3.0", path = "../manifest" }
+extism-manifest = { version = "0.4.0", path = "../manifest" }
-pretty-hex = { version = "0.3" }
 uuid = { version = "1", features = ["v4"] }
 libc = "0.2"
 

runtime/build.rs

@@ -1,4 +1,16 @@
 fn main() {
+    println!("cargo:rerun-if-changed=src/extism-runtime.wasm");
+    let dir = std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
+
+    // Attempt to build the kernel, this is only done as a convenience when developing the
+    // kernel an should not be relied on. When changes are made to the kernel run
+    // `sh build.sh` in the `kernel/` directory to ensure it run successfully.
+    let _ = std::process::Command::new("bash")
+        .args(&["build.sh"])
+        .current_dir(dir.join("../kernel"))
+        .status()
+        .unwrap();
+
     let fn_macro = "
 #define EXTISM_FUNCTION(N) extern void N(ExtismCurrentPlugin*, const ExtismVal*, ExtismSize, ExtismVal*, ExtismSize, void*)
 #define EXTISM_GO_FUNCTION(N) extern void N(void*, ExtismVal*, ExtismSize, ExtismVal*, ExtismSize, uintptr_t)

runtime/src/internal.rs

@@ -2,60 +2,6 @@ use std::collections::BTreeMap;
 
 use crate::*;
 
-/// Internal stores data that is available to the caller in PDK functions
-pub struct Internal {
-    /// Call input length
-    pub input_length: usize,
-
-    /// Pointer to call input
-    pub input: *const u8,
-
-    /// Memory offset that points to the output
-    pub output_offset: usize,
-
-    /// Length of output in memory
-    pub output_length: usize,
-
-    /// WASI context
-    pub wasi: Option<Wasi>,
-
-    /// Keep track of the status from the last HTTP request
-    pub http_status: u16,
-
-    /// Store plugin-specific error messages
-    pub last_error: std::cell::RefCell<Option<std::ffi::CString>>,
-
-    /// Plugin variables
-    pub vars: BTreeMap<String, Vec<u8>>,
-
-    /// A pointer to the plugin memory, this should mostly be used from the PDK
-    pub memory: *mut PluginMemory,
-
-    pub(crate) available_pages: Option<u32>,
-}
-
-/// InternalExt provides a unified way of acessing `memory`, `store` and `internal` values
-pub trait InternalExt {
-    fn memory(&self) -> &PluginMemory;
-    fn memory_mut(&mut self) -> &mut PluginMemory;
-
-    fn store(&self) -> &Store<Internal> {
-        self.memory().store()
-    }
-
-    fn store_mut(&mut self) -> &mut Store<Internal> {
-        self.memory_mut().store_mut()
-    }
-
-    fn internal(&self) -> &Internal {
-        self.store().data()
-    }
-
-    fn internal_mut(&mut self) -> &mut Internal {
-        self.store_mut().data_mut()
-    }
-}
-
 /// WASI context
 pub struct Wasi {
     /// wasi
@@ -64,13 +10,211 @@ pub struct Wasi {
     /// wasi-nn
     #[cfg(feature = "nn")]
     pub nn: wasmtime_wasi_nn::WasiNnCtx,
-    #[cfg(not(feature = "nn"))]
+}
-    pub nn: (),
+
+/// Internal stores data that is available to the caller in PDK functions
+pub struct Internal {
+    /// Store
+    pub store: *mut Store<Internal>,
+
+    /// Linker
+    pub linker: *mut wasmtime::Linker<Internal>,
+
+    /// WASI context
+    pub wasi: Option<Wasi>,
+
+    /// Keep track of the status from the last HTTP request
+    pub http_status: u16,
+
+    /// Plugin variables
+    pub vars: BTreeMap<String, Vec<u8>>,
+
+    pub manifest: Manifest,
+
+    pub available_pages: Option<u32>,
+
+    pub(crate) memory_limiter: Option<MemoryLimiter>,
+}
+
+/// InternalExt provides a unified way of acessing `memory`, `store` and `internal` values
+pub trait InternalExt {
+    fn store(&self) -> &Store<Internal>;
+
+    fn store_mut(&mut self) -> &mut Store<Internal>;
+
+    fn linker(&self) -> &Linker<Internal>;
+
+    fn linker_mut(&mut self) -> &mut Linker<Internal>;
+
+    fn linker_and_store(&mut self) -> (&mut Linker<Internal>, &mut Store<Internal>);
+
+    fn internal(&self) -> &Internal {
+        self.store().data()
+    }
+
+    fn internal_mut(&mut self) -> &mut Internal {
+        self.store_mut().data_mut()
+    }
+
+    fn memory_ptr(&mut self) -> *mut u8 {
+        let (linker, mut store) = self.linker_and_store();
+        if let Some(mem) = linker.get(&mut store, "env", "memory") {
+            if let Some(mem) = mem.into_memory() {
+                return mem.data_ptr(&mut store);
+            }
+        }
+
+        std::ptr::null_mut()
+    }
+
+    fn memory(&mut self) -> &mut [u8] {
+        let (linker, mut store) = self.linker_and_store();
+        let mem = linker
+            .get(&mut store, "env", "memory")
+            .unwrap()
+            .into_memory()
+            .unwrap();
+        let ptr = mem.data_ptr(&store);
+        if ptr.is_null() {
+            return &mut [];
+        }
+        let size = mem.data_size(&store);
+        unsafe { std::slice::from_raw_parts_mut(ptr, size) }
+    }
+
+    fn memory_read(&mut self, offs: u64, len: Size) -> &[u8] {
+        let offs = offs as usize;
+        let len = len as usize;
+        let mem = self.memory();
+        &mem[offs..offs + len]
+    }
+
+    fn memory_read_str(&mut self, offs: u64) -> Result<&str, std::str::Utf8Error> {
+        let len = self.memory_length(offs);
+        std::str::from_utf8(self.memory_read(offs, len))
+    }
+
+    fn memory_write(&mut self, offs: u64, bytes: impl AsRef<[u8]>) {
+        let b = bytes.as_ref();
+        let offs = offs as usize;
+        let len = b.len();
+        self.memory()[offs..offs + len].copy_from_slice(bytes.as_ref());
+    }
+
+    fn memory_alloc(&mut self, n: Size) -> Result<u64, Error> {
+        let (linker, mut store) = self.linker_and_store();
+        let output = &mut [Val::I64(0)];
+        linker
+            .get(&mut store, "env", "extism_alloc")
+            .unwrap()
+            .into_func()
+            .unwrap()
+            .call(&mut store, &[Val::I64(n as i64)], output)?;
+        let offs = output[0].unwrap_i64() as u64;
+        if offs == 0 {
+            anyhow::bail!("out of memory")
+        }
+        Ok(offs)
+    }
+
+    fn memory_alloc_bytes(&mut self, bytes: impl AsRef<[u8]>) -> Result<u64, Error> {
+        let b = bytes.as_ref();
+        let offs = self.memory_alloc(b.len() as Size)?;
+        self.memory_write(offs, b);
+        Ok(offs)
+    }
+
+    fn memory_free(&mut self, offs: u64) {
+        let (linker, mut store) = self.linker_and_store();
+        linker
+            .get(&mut store, "env", "extism_free")
+            .unwrap()
+            .into_func()
+            .unwrap()
+            .call(&mut store, &[Val::I64(offs as i64)], &mut [])
+            .unwrap();
+    }
+
+    fn memory_length(&mut self, offs: u64) -> u64 {
+        let (linker, mut store) = self.linker_and_store();
+        let output = &mut [Val::I64(0)];
+        linker
+            .get(&mut store, "env", "extism_length")
+            .unwrap()
+            .into_func()
+            .unwrap()
+            .call(&mut store, &[Val::I64(offs as i64)], output)
+            .unwrap();
+        output[0].unwrap_i64() as u64
+    }
+
+    // A convenience method to set the plugin error and return a value
+    fn error<E>(&mut self, e: impl std::fmt::Debug, x: E) -> E {
+        let s = format!("{e:?}");
+        debug!("Set error: {:?}", s);
+        if let Ok(offs) = self.memory_alloc_bytes(&s) {
+            let (linker, mut store) = self.linker_and_store();
+            if let Some(f) = linker.get(&mut store, "env", "extism_error_set") {
+                f.into_func()
+                    .unwrap()
+                    .call(&mut store, &[Val::I64(offs as i64)], &mut [])
+                    .unwrap();
+            }
+        }
+        x
+    }
+
+    fn clear_error(&mut self) {
+        let (linker, mut store) = self.linker_and_store();
+        if let Some(f) = linker.get(&mut store, "env", "extism_error_set") {
+            f.into_func()
+                .unwrap()
+                .call(&mut store, &[Val::I64(0)], &mut [])
+                .unwrap();
+        }
+    }
+
+    fn has_error(&mut self) -> bool {
+        let (linker, mut store) = self.linker_and_store();
+        let output = &mut [Val::I64(0)];
+        linker
+            .get(&mut store, "env", "extism_error_get")
+            .unwrap()
+            .into_func()
+            .unwrap()
+            .call(&mut store, &[], output)
+            .unwrap();
+        output[0].unwrap_i64() != 0
+    }
+
+    fn get_error(&mut self) -> Option<&str> {
+        let (linker, mut store) = self.linker_and_store();
+        let output = &mut [Val::I64(0)];
+        linker
+            .get(&mut store, "env", "extism_error_get")
+            .unwrap()
+            .into_func()
+            .unwrap()
+            .call(&mut store, &[], output)
+            .unwrap();
+        let offs = output[0].unwrap_i64() as u64;
+        if offs == 0 {
+            return None;
+        }
+
+        let length = self.memory_length(offs);
+        let data = self.memory_read(offs, length);
+        let s = std::str::from_utf8(data);
+        match s {
+            Ok(s) => Some(s),
+            Err(_) => None,
+        }
+    }
 }
 
 impl Internal {
     pub(crate) fn new(
-        manifest: &Manifest,
+        manifest: Manifest,
         wasi: bool,
         available_pages: Option<u32>,
     ) -> Result<Self, Error> {
@@ -91,49 +235,106 @@ impl Internal {
             #[cfg(feature = "nn")]
             let nn = wasmtime_wasi_nn::WasiNnCtx::new()?;
 
-            #[cfg(not(feature = "nn"))]
-            #[allow(clippy::let_unit_value)]
-            let nn = ();
-
             Some(Wasi {
                 ctx: ctx.build(),
+                #[cfg(feature = "nn")]
                 nn,
             })
         } else {
             None
         };
 
+        let memory_limiter = if let Some(pgs) = available_pages {
+            let n = pgs as usize * 65536;
+            Some(MemoryLimiter {
+                max_bytes: n,
+                bytes_left: n,
+            })
+        } else {
+            None
+        };
+
         Ok(Internal {
-            input_length: 0,
-            output_offset: 0,
-            output_length: 0,
-            input: std::ptr::null(),
             wasi,
-            memory: std::ptr::null_mut(),
+            manifest,
             http_status: 0,
-            last_error: std::cell::RefCell::new(None),
             vars: BTreeMap::new(),
+            linker: std::ptr::null_mut(),
+            store: std::ptr::null_mut(),
             available_pages,
+            memory_limiter,
         })
     }
 
-    pub fn set_error(&self, e: impl std::fmt::Debug) {
+    pub fn linker(&self) -> &wasmtime::Linker<Internal> {
-        debug!("Set error: {:?}", e);
+        unsafe { &*self.linker }
-        *self.last_error.borrow_mut() = Some(error_string(e));
     }
 
-    /// Unset `last_error` field
+    pub fn linker_mut(&mut self) -> &mut wasmtime::Linker<Internal> {
-    pub fn clear_error(&self) {
+        unsafe { &mut *self.linker }
-        *self.last_error.borrow_mut() = None;
     }
 }
 
 impl InternalExt for Internal {
-    fn memory(&self) -> &PluginMemory {
+    fn store(&self) -> &Store<Internal> {
-        unsafe { &*self.memory }
+        unsafe { &*self.store }
     }
 
-    fn memory_mut(&mut self) -> &mut PluginMemory {
+    fn store_mut(&mut self) -> &mut Store<Internal> {
-        unsafe { &mut *self.memory }
+        unsafe { &mut *self.store }
+    }
+
+    fn linker(&self) -> &Linker<Internal> {
+        unsafe { &*self.linker }
+    }
+
+    fn linker_mut(&mut self) -> &mut Linker<Internal> {
+        unsafe { &mut *self.linker }
+    }
+
+    fn linker_and_store(&mut self) -> (&mut Linker<Internal>, &mut Store<Internal>) {
+        unsafe { (&mut *self.linker, &mut *self.store) }
+    }
+}
+
+pub(crate) struct MemoryLimiter {
+    bytes_left: usize,
+    max_bytes: usize,
+}
+
+impl MemoryLimiter {
+    pub(crate) fn reset(&mut self) {
+        self.bytes_left = self.max_bytes;
+    }
+}
+
+impl wasmtime::ResourceLimiter for MemoryLimiter {
+    fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        maximum: Option<usize>,
+    ) -> Result<bool> {
+        if let Some(max) = maximum {
+            if desired > max {
+                return Ok(false);
+            }
+        }
+
+        let d = desired - current;
+        if d > self.bytes_left {
+            return Ok(false);
+        }
+
+        self.bytes_left -= d;
+        Ok(true)
+    }
+
+    fn table_growing(&mut self, _current: u32, desired: u32, maximum: Option<u32>) -> Result<bool> {
+        if let Some(max) = maximum {
+            return Ok(desired <= max);
+        }
+
+        Ok(true)
     }
 }

runtime/src/lib.rs

@@ -5,7 +5,6 @@ mod context;
 mod function;
 mod internal;
 pub mod manifest;
-mod memory;
 pub(crate) mod pdk;
 mod plugin;
 mod plugin_ref;
@@ -16,7 +15,6 @@ pub use context::Context;
 pub use function::{Function, UserData, Val, ValType};
 pub use internal::{Internal, InternalExt, Wasi};
 pub use manifest::Manifest;
-pub use memory::{MemoryBlock, PluginMemory, ToMemoryBlock};
 pub use plugin::Plugin;
 pub use plugin_ref::PluginRef;
 pub(crate) use timer::{Timer, TimerAction};

runtime/src/manifest.rs

@@ -7,7 +7,7 @@ use sha2::Digest;
 use crate::*;
 
 /// Manifest wraps the manifest exported by `extism_manifest`
-#[derive(Default, serde::Serialize, serde::Deserialize)]
+#[derive(Default, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(transparent)]
 pub struct Manifest(extism_manifest::Manifest);
 
@@ -60,6 +60,8 @@ fn check_hash(hash: &Option<String>, data: &[u8]) -> Result<(), Error> {
     }
 }
 
+const WASM: &[u8] = include_bytes!("extism-runtime.wasm");
+
 /// Convert from manifest to a wasmtime Module
 fn to_module(engine: &Engine, wasm: &extism_manifest::Wasm) -> Result<(String, Module), Error> {
     match wasm {
@@ -167,6 +169,7 @@ const WASM_MAGIC: [u8; 4] = [0x00, 0x61, 0x73, 0x6d];
 impl Manifest {
     /// Create a new Manifest, returns the manifest and a map of modules
     pub fn new(engine: &Engine, data: &[u8]) -> Result<(Self, BTreeMap<String, Module>), Error> {
+        let extism_module = Module::new(engine, WASM)?;
         let has_magic = data.len() >= 4 && data[0..4] == WASM_MAGIC;
         let is_wast = data.starts_with(b"(module") || data.starts_with(b";;");
         if !has_magic && !is_wast {
@@ -178,12 +181,14 @@ impl Manifest {
             }
 
             let t = serde_json::from_slice::<Self>(data)?;
-            let m = t.modules(engine)?;
+            let mut m = t.modules(engine)?;
+            m.insert("env".to_string(), extism_module);
             return Ok((t, m));
         }
 
         let m = Module::new(engine, data)?;
         let mut modules = BTreeMap::new();
+        modules.insert("env".to_string(), extism_module);
         modules.insert("main".to_string(), m);
         Ok((Manifest::default(), modules))
     }

runtime/src/memory.rs

@@ -1,376 +0,0 @@
-use std::collections::BTreeMap;
-
-use crate::*;
-
-use pretty_hex::PrettyHex;
-
-/// Handles memory for plugins
-pub struct PluginMemory {
-    /// wasmtime Store
-    pub store: Option<Store<Internal>>,
-
-    /// WASM memory
-    pub memory: Memory,
-
-    /// Tracks allocated blocks
-    pub live_blocks: BTreeMap<usize, usize>,
-
-    /// Tracks free blocks
-    pub free: Vec<MemoryBlock>,
-
-    /// Tracks current offset in memory
-    pub position: usize,
-
-    /// Extism manifest
-    pub manifest: Manifest,
-}
-
-/// `ToMemoryBlock` is used to convert from Rust values to blocks of WASM memory
-pub trait ToMemoryBlock {
-    fn to_memory_block(&self, mem: &PluginMemory) -> Result<MemoryBlock, Error>;
-}
-
-impl ToMemoryBlock for MemoryBlock {
-    fn to_memory_block(&self, _mem: &PluginMemory) -> Result<MemoryBlock, Error> {
-        Ok(*self)
-    }
-}
-
-impl ToMemoryBlock for (usize, usize) {
-    fn to_memory_block(&self, _mem: &PluginMemory) -> Result<MemoryBlock, Error> {
-        Ok(MemoryBlock {
-            offset: self.0,
-            length: self.1,
-        })
-    }
-}
-
-impl ToMemoryBlock for usize {
-    fn to_memory_block(&self, mem: &PluginMemory) -> Result<MemoryBlock, Error> {
-        match mem.at_offset(*self) {
-            Some(x) => Ok(x),
-            None => Err(Error::msg(format!("Invalid memory offset: {}", self))),
-        }
-    }
-}
-
-const PAGE_SIZE: u32 = 65536;
-
-// BLOCK_SIZE_THRESHOLD exists to ensure that free blocks are never split up any
-// smaller than this value
-const BLOCK_SIZE_THRESHOLD: usize = 32;
-
-impl PluginMemory {
-    /// Create memory for a plugin
-    pub fn new(store: Store<Internal>, memory: Memory, manifest: Manifest) -> Self {
-        PluginMemory {
-            free: Vec::new(),
-            live_blocks: BTreeMap::new(),
-            store: Some(store),
-            memory,
-            position: 1,
-            manifest,
-        }
-    }
-
-    pub fn store(&self) -> &Store<Internal> {
-        self.store.as_ref().unwrap()
-    }
-
-    pub fn store_mut(&mut self) -> &mut Store<Internal> {
-        self.store.as_mut().unwrap()
-    }
-
-    /// Moves module to a new store
-    pub fn reinstantiate(&mut self) -> Result<(), Error> {
-        if let Some(store) = self.store.take() {
-            let engine = store.engine().clone();
-            let internal = store.into_data();
-            let pages = internal.available_pages;
-            let mut store = Store::new(&engine, internal);
-            store.epoch_deadline_callback(|_internal| Err(Error::msg("timeout")));
-            self.memory = Memory::new(&mut store, MemoryType::new(2, pages))?;
-            self.store = Some(store);
-        }
-
-        self.reset();
-
-        Ok(())
-    }
-
-    /// Write byte to memory
-    pub(crate) fn store_u8(&mut self, offs: usize, data: u8) -> Result<(), MemoryAccessError> {
-        trace!("store_u8: offset={offs} data={data:#04x}");
-        if offs >= self.size() {
-            // This should raise MemoryAccessError
-            let buf = &mut [0];
-            self.memory.read(&self.store.as_ref().unwrap(), offs, buf)?;
-            return Ok(());
-        }
-        self.memory.data_mut(&mut self.store.as_mut().unwrap())[offs] = data;
-        Ok(())
-    }
-
-    /// Read byte from memory
-    pub(crate) fn load_u8(&self, offs: usize) -> Result<u8, MemoryAccessError> {
-        trace!("load_u8: offset={offs}");
-        if offs >= self.size() {
-            // This should raise MemoryAccessError
-            let buf = &mut [0];
-            self.memory.read(&self.store.as_ref().unwrap(), offs, buf)?;
-            return Ok(0);
-        }
-        Ok(self.memory.data(&self.store.as_ref().unwrap())[offs])
-    }
-
-    /// Write u64 to memory
-    pub(crate) fn store_u64(&mut self, offs: usize, data: u64) -> Result<(), Error> {
-        trace!("store_u64: offset={offs} data={data:#18x}");
-        let handle = MemoryBlock {
-            offset: offs,
-            length: 8,
-        };
-        self.write(handle, data.to_ne_bytes())?;
-        Ok(())
-    }
-
-    /// Read u64 from memory
-    pub(crate) fn load_u64(&self, offs: usize) -> Result<u64, Error> {
-        trace!("load_u64: offset={offs}");
-        let mut buf = [0; 8];
-        let handle = MemoryBlock {
-            offset: offs,
-            length: 8,
-        };
-        self.read(handle, &mut buf)?;
-        Ok(u64::from_ne_bytes(buf))
-    }
-
-    /// Write slice to memory
-    pub fn write(&mut self, pos: impl ToMemoryBlock, data: impl AsRef<[u8]>) -> Result<(), Error> {
-        let pos = pos.to_memory_block(self)?;
-        assert!(data.as_ref().len() <= pos.length);
-        self.memory
-            .write(&mut self.store.as_mut().unwrap(), pos.offset, data.as_ref())?;
-        Ok(())
-    }
-
-    /// Read slice from memory
-    pub fn read(&self, pos: impl ToMemoryBlock, mut data: impl AsMut<[u8]>) -> Result<(), Error> {
-        let pos = pos.to_memory_block(self)?;
-        assert!(data.as_mut().len() <= pos.length);
-        self.memory
-            .read(&self.store.as_ref().unwrap(), pos.offset, data.as_mut())?;
-        Ok(())
-    }
-
-    /// Size of memory in bytes
-    pub fn size(&self) -> usize {
-        self.memory.data_size(&self.store.as_ref().unwrap())
-    }
-
-    /// Size of memory in pages
-    pub fn pages(&self) -> u32 {
-        self.memory.size(&self.store.as_ref().unwrap()) as u32
-    }
-
-    /// Reserve `n` bytes of memory
-    pub fn alloc(&mut self, n: usize) -> Result<MemoryBlock, Error> {
-        debug!("Allocating {n} bytes");
-
-        for (i, block) in self.free.iter_mut().enumerate() {
-            if block.length == n {
-                let block = self.free.swap_remove(i);
-                self.live_blocks.insert(block.offset, block.length);
-                debug!("Found block with exact size at offset {}", block.offset);
-                return Ok(block);
-            } else if block.length.saturating_sub(n) >= BLOCK_SIZE_THRESHOLD {
-                let handle = MemoryBlock {
-                    offset: block.offset,
-                    length: n,
-                };
-                debug!(
-                    "Using block with size {} at offset {}",
-                    block.length, block.offset
-                );
-
-                block.offset += n;
-                block.length -= n;
-                self.live_blocks.insert(handle.offset, handle.length);
-                return Ok(handle);
-            }
-        }
-
-        let new_offset = self.position.saturating_add(n);
-
-        // If there aren't enough bytes, try to grow the memory size
-        if new_offset >= self.size() {
-            debug!("Need more memory");
-
-            let bytes_needed = (new_offset as f64 - self.size() as f64) / PAGE_SIZE as f64;
-            let mut pages_needed = bytes_needed.ceil() as u64;
-            if pages_needed == 0 {
-                pages_needed = 1
-            }
-
-            debug!("Requesting {pages_needed} more pages");
-            // This will fail if we've already allocated the maximum amount of memory allowed
-            self.memory
-                .grow(&mut self.store.as_mut().unwrap(), pages_needed)?;
-        }
-
-        let mem = MemoryBlock {
-            offset: self.position,
-            length: n,
-        };
-
-        debug!(
-            "Allocated new block: {} bytes at offset {}",
-            mem.length, mem.offset
-        );
-
-        self.live_blocks.insert(mem.offset, mem.length);
-        self.position += n;
-        Ok(mem)
-    }
-
-    /// Allocate and copy `data` into the wasm memory
-    pub fn alloc_bytes(&mut self, data: impl AsRef<[u8]>) -> Result<MemoryBlock, Error> {
-        let handle = self.alloc(data.as_ref().len())?;
-        self.write(handle, data)?;
-        Ok(handle)
-    }
-
-    /// Free the block allocated at `offset`
-    pub fn free(&mut self, offset: usize) {
-        debug!("Freeing block at {offset}");
-        if let Some(length) = self.live_blocks.remove(&offset) {
-            self.free.push(MemoryBlock { offset, length });
-        } else {
-            return;
-        }
-
-        let free_size: usize = self.free.iter().map(|x| x.length).sum();
-
-        // Perform compaction if there is at least 1kb of free memory available
-        if free_size >= 1024 {
-            let mut last: Option<MemoryBlock> = None;
-            let mut free = Vec::new();
-            for block in self.free.iter() {
-                match last {
-                    None => {
-                        free.push(*block);
-                    }
-                    Some(last) => {
-                        if last.offset + last.length == block.offset {
-                            free.push(MemoryBlock {
-                                offset: last.offset,
-                                length: last.length + block.length,
-                            });
-                        }
-                    }
-                }
-                last = Some(*block);
-            }
-            self.free = free;
-        }
-    }
-
-    /// Log entire memory as hexdump using the `trace` log level
-    pub fn dump(&self) {
-        let data = self.memory.data(self.store.as_ref().unwrap());
-
-        trace!("{:?}", data[..self.position].hex_dump());
-    }
-
-    /// Reset memory - clears free-list and live blocks and resets position
-    pub fn reset(&mut self) {
-        self.free.clear();
-        self.live_blocks.clear();
-        self.position = 1;
-    }
-
-    /// Get memory as a slice of bytes
-    pub fn data(&self) -> &[u8] {
-        self.memory.data(self.store.as_ref().unwrap())
-    }
-
-    /// Get memory as a mutable slice of bytes
-    pub fn data_mut(&mut self) -> &mut [u8] {
-        self.memory.data_mut(self.store.as_mut().unwrap())
-    }
-
-    /// Get bytes occupied by the provided memory handle
-    pub fn get(&self, handle: impl ToMemoryBlock) -> Result<&[u8], Error> {
-        let handle = handle.to_memory_block(self)?;
-        Ok(&self.memory.data(self.store.as_ref().unwrap())
-            [handle.offset..handle.offset + handle.length])
-    }
-
-    /// Get mutable bytes occupied by the provided memory handle
-    pub fn get_mut(&mut self, handle: impl ToMemoryBlock) -> Result<&mut [u8], Error> {
-        let handle = handle.to_memory_block(self)?;
-        Ok(&mut self.memory.data_mut(self.store.as_mut().unwrap())
-            [handle.offset..handle.offset + handle.length])
-    }
-
-    /// Get str occupied by the provided memory handle
-    pub fn get_str(&self, handle: impl ToMemoryBlock) -> Result<&str, Error> {
-        let handle = handle.to_memory_block(self)?;
-        Ok(std::str::from_utf8(
-            &self.memory.data(self.store.as_ref().unwrap())
-                [handle.offset..handle.offset + handle.length],
-        )?)
-    }
-
-    /// Get mutable str occupied by the provided memory handle
-    pub fn get_mut_str(&mut self, handle: impl ToMemoryBlock) -> Result<&mut str, Error> {
-        let handle = handle.to_memory_block(self)?;
-        Ok(std::str::from_utf8_mut(
-            &mut self.memory.data_mut(self.store.as_mut().unwrap())
-                [handle.offset..handle.offset + handle.length],
-        )?)
-    }
-
-    /// Pointer to the provided memory handle
-    pub fn ptr(&self, handle: impl ToMemoryBlock) -> Result<*mut u8, Error> {
-        let handle = handle.to_memory_block(self)?;
-        Ok(unsafe {
-            self.memory
-                .data_ptr(&self.store.as_ref().unwrap())
-                .add(handle.offset)
-        })
-    }
-
-    /// Get the length of the block starting at `offs`
-    pub fn block_length(&self, offs: usize) -> Option<usize> {
-        self.live_blocks.get(&offs).cloned()
-    }
-
-    /// Get the block at the specified offset
-    pub fn at_offset(&self, offset: usize) -> Option<MemoryBlock> {
-        let block_length = self.block_length(offset);
-        block_length.map(|length| MemoryBlock { offset, length })
-    }
-}
-
-#[derive(Clone, Copy)]
-pub struct MemoryBlock {
-    pub offset: usize,
-    pub length: usize,
-}
-
-impl From<(usize, usize)> for MemoryBlock {
-    fn from(x: (usize, usize)) -> Self {
-        MemoryBlock {
-            offset: x.0,
-            length: x.1,
-        }
-    }
-}
-
-impl MemoryBlock {
-    pub fn new(offset: usize, length: usize) -> Self {
-        MemoryBlock { offset, length }
-    }
-}

runtime/src/pdk.rs

@@ -18,177 +18,6 @@ macro_rules! args {
     };
 }
 
-/// Get the input length
-/// Params: none
-/// Returns: i64 (length)
-pub(crate) fn input_length(
-    caller: Caller<Internal>,
-    _input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &Internal = caller.data();
-    output[0] = Val::I64(data.input_length as i64);
-    Ok(())
-}
-
-/// Load a byte from input
-/// Params: i64 (offset)
-/// Returns: i32 (byte)
-pub(crate) fn input_load_u8(
-    caller: Caller<Internal>,
-    input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &Internal = caller.data();
-    if data.input.is_null() {
-        return Ok(());
-    }
-    output[0] = unsafe { Val::I32(*data.input.add(input[0].unwrap_i64() as usize) as i32) };
-    Ok(())
-}
-
-/// Load an unsigned 64 bit integer from input
-/// Params: i64 (offset)
-/// Returns: i64 (int)
-pub(crate) fn input_load_u64(
-    caller: Caller<Internal>,
-    input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &Internal = caller.data();
-    if data.input.is_null() {
-        return Ok(());
-    }
-    let offs = args!(input, 0, i64) as usize;
-    let slice = unsafe { std::slice::from_raw_parts(data.input.add(offs), 8) };
-    let byte = u64::from_ne_bytes(slice.try_into().unwrap());
-    output[0] = Val::I64(byte as i64);
-    Ok(())
-}
-
-/// Store a byte in memory
-/// Params: i64 (offset), i32 (byte)
-/// Returns: none
-pub(crate) fn store_u8(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    _output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let (offset, byte) = args!(input, (0, i64), (1, i32));
-    data.memory_mut().store_u8(offset as usize, byte as u8)?;
-    Ok(())
-}
-
-/// Load a byte from memory
-/// Params: i64 (offset)
-/// Returns: i32 (byte)
-pub(crate) fn load_u8(
-    caller: Caller<Internal>,
-    input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &Internal = caller.data();
-    let offset = args!(input, 0, i64) as usize;
-    let byte = data.memory().load_u8(offset)?;
-    output[0] = Val::I32(byte as i32);
-    Ok(())
-}
-
-/// Store an unsigned 64 bit integer in memory
-/// Params: i64 (offset), i64 (int)
-/// Returns: none
-pub(crate) fn store_u64(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    _output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let (offset, b) = args!(input, (0, i64), (1, i64));
-    data.memory_mut().store_u64(offset as usize, b as u64)?;
-    Ok(())
-}
-
-/// Load an unsigned 64 bit integer from memory
-/// Params: i64 (offset)
-/// Returns: i64 (int)
-pub(crate) fn load_u64(
-    caller: Caller<Internal>,
-    input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &Internal = caller.data();
-    let offset = args!(input, 0, i64) as usize;
-    let byte = data.memory().load_u64(offset)?;
-    output[0] = Val::I64(byte as i64);
-    Ok(())
-}
-
-/// Set output offset and length
-/// Params: i64 (offset), i64 (length)
-/// Returns: none
-pub(crate) fn output_set(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    _output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let (offset, length) = args!(input, (0, i64), (1, i64));
-    data.output_offset = offset as usize;
-    data.output_length = length as usize;
-    Ok(())
-}
-
-/// Allocate bytes
-/// Params: i64 (length)
-/// Returns: i64 (offset)
-pub(crate) fn alloc(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let offs = data.memory_mut().alloc(input[0].unwrap_i64() as _)?;
-    output[0] = Val::I64(offs.offset as i64);
-
-    Ok(())
-}
-
-/// Free memory
-/// Params: i64 (offset)
-/// Returns: none
-pub(crate) fn free(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    _output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let offset = args!(input, 0, i64) as usize;
-    data.memory_mut().free(offset);
-    Ok(())
-}
-
-/// Set the error message, this can be checked by the host program
-/// Params: i64 (offset)
-/// Returns: none
-pub(crate) fn error_set(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    _output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let offset = args!(input, 0, i64) as usize;
-
-    if offset == 0 {
-        *data.last_error.borrow_mut() = None;
-        return Ok(());
-    }
-
-    let s = data.memory().get_str(offset)?;
-    data.set_error(s);
-    Ok(())
-}
-
 /// Get a configuration value
 /// Params: i64 (offset)
 /// Returns: i64 (offset)
@@ -199,21 +28,24 @@ pub(crate) fn config_get(
 ) -> Result<(), Error> {
     let data: &mut Internal = caller.data_mut();
 
-    let offset = args!(input, 0, i64) as usize;
+    let offset = args!(input, 0, i64) as u64;
-    let key = data.memory().get_str(offset)?;
+    let key = data.memory_read_str(offset)?;
-    let val = data.memory().manifest.as_ref().config.get(key);
+    let key = unsafe {
+        std::str::from_utf8_unchecked(std::slice::from_raw_parts(key.as_ptr(), key.len()))
+    };
+    let val = data.internal().manifest.as_ref().config.get(key);
     let ptr = val.map(|x| (x.len(), x.as_ptr()));
     let mem = match ptr {
         Some((len, ptr)) => {
             let bytes = unsafe { std::slice::from_raw_parts(ptr, len) };
-            data.memory_mut().alloc_bytes(bytes)?
+            data.memory_alloc_bytes(bytes)?
         }
         None => {
             output[0] = Val::I64(0);
             return Ok(());
         }
     };
-    output[0] = Val::I64(mem.offset as i64);
+    output[0] = Val::I64(mem as i64);
     Ok(())
 }
 
@@ -227,23 +59,24 @@ pub(crate) fn var_get(
 ) -> Result<(), Error> {
     let data: &mut Internal = caller.data_mut();
 
-    let offset = args!(input, 0, i64) as usize;
+    let offset = args!(input, 0, i64) as u64;
-    let key = data.memory().get_str(offset)?;
+    let key = data.memory_read_str(offset)?;
-    let val = data.vars.get(key);
+    let key = unsafe {
+        std::str::from_utf8_unchecked(std::slice::from_raw_parts(key.as_ptr(), key.len()))
+    };
+    let val = data.internal().vars.get(key);
     let ptr = val.map(|x| (x.len(), x.as_ptr()));
-
     let mem = match ptr {
         Some((len, ptr)) => {
             let bytes = unsafe { std::slice::from_raw_parts(ptr, len) };
-            data.memory_mut().alloc_bytes(bytes)?
+            data.memory_alloc_bytes(bytes)?
         }
         None => {
             output[0] = Val::I64(0);
             return Ok(());
         }
     };
-
+    output[0] = Val::I64(mem as i64);
-    output[0] = Val::I64(mem.offset as i64);
     Ok(())
 }
 
@@ -262,16 +95,16 @@ pub(crate) fn var_set(
         size += v.len();
     }
 
-    let voffset = args!(input, 1, i64) as usize;
+    let voffset = args!(input, 1, i64) as u64;
 
     // If the store is larger than 100MB then stop adding things
     if size > 1024 * 1024 * 100 && voffset != 0 {
         return Err(Error::msg("Variable store is full"));
     }
 
-    let key_offs = args!(input, 0, i64) as usize;
+    let key_offs = args!(input, 0, i64) as u64;
     let key = {
-        let key = data.memory().get_str(key_offs)?;
+        let key = data.memory_read_str(key_offs)?;
         let key_len = key.len();
         let key_ptr = key.as_ptr();
         unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(key_ptr, key_len)) }
@@ -283,10 +116,11 @@ pub(crate) fn var_set(
         return Ok(());
     }
 
-    let value = data.memory().get(voffset)?;
+    let vlen = data.memory_length(voffset);
+    let value = data.memory_read(voffset, vlen).to_vec();
 
     // Insert the value from memory into the `vars` map
-    data.vars.insert(key.to_string(), value.to_vec());
+    data.vars.insert(key.to_string(), value);
 
     Ok(())
 }
@@ -312,18 +146,19 @@ pub(crate) fn http_request(
     {
         use std::io::Read;
         let data: &mut Internal = caller.data_mut();
-        let http_req_offset = args!(input, 0, i64) as usize;
+        let http_req_offset = args!(input, 0, i64) as u64;
 
+        let http_req_len = data.memory_length(http_req_offset);
         let req: extism_manifest::HttpRequest =
-            serde_json::from_slice(data.memory().get(http_req_offset)?)?;
+            serde_json::from_slice(data.memory_read(http_req_offset, http_req_len))?;
 
-        let body_offset = args!(input, 1, i64) as usize;
+        let body_offset = args!(input, 1, i64) as u64;
 
         let url = match url::Url::parse(&req.url) {
             Ok(u) => u,
             Err(e) => return Err(Error::msg(format!("Invalid URL: {e:?}"))),
         };
-        let allowed_hosts = &data.memory().manifest.as_ref().allowed_hosts;
+        let allowed_hosts = &data.internal().manifest.as_ref().allowed_hosts;
         let host_str = url.host_str().unwrap_or_default();
         let host_matches = if let Some(allowed_hosts) = allowed_hosts {
             allowed_hosts.iter().any(|url| {
@@ -352,7 +187,8 @@ pub(crate) fn http_request(
         }
 
         let res = if body_offset > 0 {
-            let buf = data.memory().get(body_offset)?;
+            let len = data.memory_length(body_offset);
+            let buf = data.memory_read(body_offset, len);
             r.send_bytes(buf)
         } else {
             r.call()
@@ -379,9 +215,8 @@ pub(crate) fn http_request(
                 .take(1024 * 1024 * 50) // TODO: make this limit configurable
                 .read_to_end(&mut buf)?;
 
-            let mem = data.memory_mut().alloc_bytes(buf)?;
+            let mem = data.memory_alloc_bytes(buf)?;
-
+            output[0] = Val::I64(mem as i64);
-            output[0] = Val::I64(mem.offset as i64);
         } else {
             output[0] = Val::I64(0);
         }
@@ -403,39 +238,17 @@ pub(crate) fn http_status_code(
     Ok(())
 }
 
-/// Get the length of an allocated block given the offset
-/// Params: i64 (offset)
-/// Returns: i64 (length or 0)
-pub(crate) fn length(
-    mut caller: Caller<Internal>,
-    input: &[Val],
-    output: &mut [Val],
-) -> Result<(), Error> {
-    let data: &mut Internal = caller.data_mut();
-    let offset = args!(input, 0, i64) as usize;
-    if offset == 0 {
-        output[0] = Val::I64(0);
-        return Ok(());
-    }
-    let length = match data.memory().block_length(offset) {
-        Some(x) => x,
-        None => return Err(Error::msg("Unable to find length for offset")),
-    };
-    output[0] = Val::I64(length as i64);
-    Ok(())
-}
-
 pub fn log(
     level: log::Level,
-    caller: Caller<Internal>,
+    mut caller: Caller<Internal>,
     input: &[Val],
     _output: &mut [Val],
 ) -> Result<(), Error> {
-    let data: &Internal = caller.data();
+    let data: &mut Internal = caller.data_mut();
-    let offset = args!(input, 0, i64) as usize;
+    let offset = args!(input, 0, i64) as u64;
-    let buf = data.memory().get(offset)?;
+    let buf = data.memory_read_str(offset);
 
-    match std::str::from_utf8(buf) {
+    match buf {
         Ok(buf) => log::log!(level, "{}", buf),
         Err(_) => log::log!(level, "{:?}", buf),
     }

runtime/src/plugin.rs

@@ -9,6 +9,7 @@ pub struct Plugin {
 
     /// Used to define functions and create new instances
     pub linker: Linker<Internal>,
+    pub store: Store<Internal>,
 
     /// Instance provides the ability to call functions in a module
     pub instance: Instance,
@@ -18,9 +19,6 @@ pub struct Plugin {
     /// actually cleaned up along with a `Store`
     pub instantiations: usize,
 
-    /// Handles interactions with WASM memory
-    pub memory: std::cell::UnsafeCell<PluginMemory>,
-
     /// The ID used to identify this plugin with the `Timer`
     pub timer_id: uuid::Uuid,
 
@@ -33,12 +31,24 @@ pub struct Plugin {
 }
 
 impl InternalExt for Plugin {
-    fn memory(&self) -> &PluginMemory {
+    fn store(&self) -> &Store<Internal> {
-        unsafe { &*self.memory.get() }
+        &self.store
     }
 
-    fn memory_mut(&mut self) -> &mut PluginMemory {
+    fn store_mut(&mut self) -> &mut Store<Internal> {
-        self.memory.get_mut()
+        &mut self.store
+    }
+
+    fn linker(&self) -> &Linker<Internal> {
+        &self.linker
+    }
+
+    fn linker_mut(&mut self) -> &mut Linker<Internal> {
+        &mut self.linker
+    }
+
+    fn linker_and_store(&mut self) -> (&mut Linker<Internal>, &mut Store<Internal>) {
+        (&mut self.linker, &mut self.store)
     }
 }
 
@@ -68,9 +78,13 @@ fn calculate_available_memory(
     let mut fail_memory_check = false;
     let mut total_memory_needed = 0;
     for (name, module) in modules.iter() {
+        if name == "env" {
+            continue;
+        }
         let mut memories = 0;
         for export in module.exports() {
             if let Some(memory) = export.ty().memory() {
+                memories += 1;
                 let memory_max = memory.maximum();
                 match memory_max {
                     None => anyhow::bail!("Unbounded memory in module {name}, when `memory.max_pages` is set in the manifest all modules \
@@ -78,16 +92,15 @@ fn calculate_available_memory(
                     Some(m) => {
                         total_memory_needed += m;
                         if !fail_memory_check {
-                            continue
+                            continue;
                         }
 
                         *available_pages = available_pages.saturating_sub(m as u32);
                         if *available_pages == 0 {
                             fail_memory_check = true;
                         }
-                    },
                     }
-                memories += 1;
+                }
             }
         }
 
@@ -132,14 +145,13 @@ impl Plugin {
 
         let mut store = Store::new(
             &engine,
-            Internal::new(&manifest, with_wasi, available_pages)?,
+            Internal::new(manifest, with_wasi, available_pages)?,
         );
         store.epoch_deadline_callback(|_internal| Err(Error::msg("timeout")));
 
-        // Create memory
+        if available_pages.is_some() {
-        let memory = Memory::new(&mut store, MemoryType::new(2, available_pages))?;
+            store.limiter(|internal| internal.memory_limiter.as_mut().unwrap());
-        let mut memory = PluginMemory::new(store, memory, manifest);
+        }
-
         let mut linker = Linker::new(&engine);
         linker.allow_shadowing(true);
 
@@ -177,7 +189,10 @@ impl Plugin {
         }
 
         // Add builtins
-        for (_name, module) in modules.iter() {
+        for (name, module) in modules.iter() {
+            if name != main_name {
+                linker.module(&mut store, name, module)?;
+            }
             for import in module.imports() {
                 let module_name = import.module();
                 let name = import.name();
@@ -185,23 +200,11 @@ impl Plugin {
 
                 if module_name == EXPORT_MODULE_NAME {
                     define_funcs!(name,  {
-                        alloc(I64) -> I64;
-                        free(I64);
-                        load_u8(I64) -> I32;
-                        load_u64(I64) -> I64;
-                        store_u8(I64, I32);
-                        store_u64(I64, I64);
-                        input_length() -> I64;
-                        input_load_u8(I64) -> I32;
-                        input_load_u64(I64) -> I64;
-                        output_set(I64, I64);
-                        error_set(I64);
                         config_get(I64) -> I64;
                         var_get(I64) -> I64;
                         var_set(I64, I64);
                         http_request(I64, I64) -> I64;
                         http_status_code() -> I32;
-                        length(I64) -> I64;
                         log_warn(I64);
                         log_info(I64);
                         log_debug(I64);
@@ -219,20 +222,13 @@ impl Plugin {
             }
         }
 
-        // Add modules to linker
+        let instance = linker.instantiate(&mut store, main)?;
-        for (name, module) in modules.iter() {
-            if name != main_name {
-                linker.module(&mut memory.store_mut(), name, module)?;
-            }
-        }
-
-        let instance = linker.instantiate(&mut memory.store_mut(), main)?;
         let timer_id = uuid::Uuid::new_v4();
         let mut plugin = Plugin {
             modules,
             linker,
-            memory: std::cell::UnsafeCell::new(memory),
             instance,
+            store,
             instantiations: 1,
             runtime: None,
             timer_id,
@@ -242,8 +238,8 @@ impl Plugin {
             },
         };
 
-        // Make sure `Internal::memory` is initialized
+        plugin.internal_mut().store = &mut plugin.store;
-        plugin.internal_mut().memory = plugin.memory.get();
+        plugin.internal_mut().linker = &mut plugin.linker;
 
         // Then detect runtime before returning the new plugin
         plugin.detect_runtime();
@@ -252,35 +248,55 @@ impl Plugin {
 
     /// Get a function by name
     pub fn get_func(&mut self, function: impl AsRef<str>) -> Option<Func> {
-        self.instance
+        self.instance.get_func(&mut self.store, function.as_ref())
-            .get_func(&mut self.memory.get_mut().store_mut(), function.as_ref())
-    }
-
-    // A convenience method to set the plugin error and return a value
-    pub fn error<E>(&self, e: impl std::fmt::Debug, x: E) -> E {
-        self.store().data().set_error(e);
-        x
     }
 
     /// Store input in memory and initialize `Internal` pointer
-    pub fn set_input(&mut self, input: *const u8, mut len: usize) {
+    pub(crate) fn set_input(
+        &mut self,
+        input: *const u8,
+        mut len: usize,
+        tx: std::sync::mpsc::SyncSender<TimerAction>,
+    ) -> Result<(), Error> {
         if input.is_null() {
             len = 0;
         }
-        let ptr = self.memory.get();
+
+        {
+            let store = &mut self.store as *mut _;
+            let linker = &mut self.linker as *mut _;
             let internal = self.internal_mut();
-        internal.input = input;
+            internal.store = store;
-        internal.input_length = len;
+            internal.linker = linker;
-        internal.memory = ptr
         }
 
-    /// Dump memory using trace! logging
+        let bytes = unsafe { std::slice::from_raw_parts(input, len) };
-    pub fn dump_memory(&self) {
+        trace!("Input size: {}", bytes.len());
-        self.memory().dump();
+
+        self.start_timer(&tx)?;
+        if let Some(f) = self.linker.get(&mut self.store, "env", "extism_reset") {
+            f.into_func().unwrap().call(&mut self.store, &[], &mut [])?;
+        }
+
+        let offs = self.memory_alloc_bytes(bytes)?;
+
+        if let Some(f) = self.linker.get(&mut self.store, "env", "extism_input_set") {
+            f.into_func().unwrap().call(
+                &mut self.store,
+                &[Val::I64(offs as i64), Val::I64(len as i64)],
+                &mut [],
+            )?;
+        }
+
+        Ok(())
     }
 
     /// Create a new instance from the same modules
     pub fn reinstantiate(&mut self) -> Result<(), Error> {
+        if let Some(limiter) = self.internal_mut().memory_limiter.as_mut() {
+            limiter.reset();
+        }
+
         let (main_name, main) = self
             .modules
             .get("main")
@@ -290,24 +306,34 @@ impl Plugin {
                 (entry.0.as_str(), entry.1)
             });
 
-        // Avoid running into resource limits, after 5 instantiations reset the store. This will
-        // release any old `Instance` objects
         if self.instantiations > 5 {
-            self.memory.get_mut().reinstantiate()?;
+            let engine = self.store.engine().clone();
+            let internal = self.internal();
+            self.store = Store::new(
+                &engine,
+                Internal::new(
+                    internal.manifest.clone(),
+                    internal.wasi.is_some(),
+                    internal.available_pages,
+                )?,
+            );
+            self.store
+                .epoch_deadline_callback(|_internal| Err(Error::msg("timeout")));
+
+            if self.internal().available_pages.is_some() {
+                self.store
+                    .limiter(|internal| internal.memory_limiter.as_mut().unwrap());
+            }
 
-            // Get the `main` module, or the last one if `main` doesn't exist
             for (name, module) in self.modules.iter() {
                 if name != main_name {
-                    self.linker
+                    self.linker.module(&mut self.store, name, module)?;
-                        .module(&mut self.memory.get_mut().store_mut(), name, module)?;
                 }
             }
             self.instantiations = 0;
         }
 
-        let instance = self
+        let instance = self.linker.instantiate(&mut self.store, main)?;
-            .linker
-            .instantiate(&mut self.memory.get_mut().store_mut(), &main)?;
         self.instance = instance;
         self.detect_runtime();
         self.instantiations += 1;
@@ -316,7 +342,7 @@ impl Plugin {
 
     /// Determine if wasi is enabled
     pub fn has_wasi(&self) -> bool {
-        self.memory().store().data().wasi.is_some()
+        self.internal().wasi.is_some()
     }
 
     fn detect_runtime(&mut self) {
@@ -325,13 +351,10 @@ impl Plugin {
         // by calling the `hs_init` export
         if let Some(init) = self.get_func("hs_init") {
             if let Some(cleanup) = self.get_func("hs_exit") {
-                if init
+                if init.typed::<(i32, i32), ()>(&self.store()).is_err() {
-                    .typed::<(i32, i32), ()>(&self.memory().store())
-                    .is_err()
-                {
                     trace!(
                         "hs_init function found with type {:?}",
-                        init.ty(&self.memory().store())
+                        init.ty(self.store())
                     );
                 }
                 self.runtime = Some(Runtime::Haskell { init, cleanup });
@@ -339,38 +362,48 @@ impl Plugin {
             return;
         }
 
-        // Check for `__wasm__call_ctors` and `__wasm_call_dtors`, this is used by WASI to
+        // Check for `__wasm_call_ctors` and `__wasm_call_dtors`, this is used by WASI to
         // initialize certain interfaces.
         if self.has_wasi() {
-            if let Some(init) = self.get_func("__wasm_call_ctors") {
+            let init = if let Some(init) = self.get_func("__wasm_call_ctors") {
-                if init.typed::<(), ()>(&self.memory().store()).is_err() {
+                if init.typed::<(), ()>(&self.store()).is_err() {
                     trace!(
                         "__wasm_call_ctors function found with type {:?}",
-                        init.ty(&self.memory().store())
+                        init.ty(self.store())
                     );
                     return;
                 }
                 trace!("WASI runtime detected");
-                if let Some(cleanup) = self.get_func("__wasm_call_dtors") {
+                init
-                    if cleanup.typed::<(), ()>(&self.memory().store()).is_err() {
+            } else if let Some(init) = self.get_func("_initialize") {
+                if init.typed::<(), ()>(&self.store()).is_err() {
                     trace!(
-                            "__wasm_call_dtors function found with type {:?}",
+                        "_initialize function found with type {:?}",
-                            cleanup.ty(&self.memory().store())
+                        init.ty(self.store())
                     );
                     return;
                 }
-                    self.runtime = Some(Runtime::Wasi {
+                trace!("WASI reactor module detected");
-                        init,
+                init
-                        cleanup: Some(cleanup),
+            } else {
-                    });
                 return;
-                }
+            };
 
-                self.runtime = Some(Runtime::Wasi {
+            let cleanup = if let Some(cleanup) = self.get_func("__wasm_call_dtors") {
-                    init,
+                if cleanup.typed::<(), ()>(&self.store()).is_err() {
-                    cleanup: None,
+                    trace!(
-                });
+                        "__wasm_call_dtors function found with type {:?}",
+                        cleanup.ty(self.store())
+                    );
+                    None
+                } else {
+                    Some(cleanup)
                 }
+            } else {
+                None
+            };
+
+            self.runtime = Some(Runtime::Wasi { init, cleanup });
             return;
         }
 
@@ -378,22 +411,22 @@ impl Plugin {
     }
 
     pub(crate) fn initialize_runtime(&mut self) -> Result<(), Error> {
+        let mut store = &mut self.store;
         if let Some(runtime) = &self.runtime {
             trace!("Plugin::initialize_runtime");
             match runtime {
                 Runtime::Haskell { init, cleanup: _ } => {
-                    let mut results =
+                    let mut results = vec![Val::null(); init.ty(&store).results().len()];
-                        vec![Val::null(); init.ty(&self.memory().store()).results().len()];
                     init.call(
-                        &mut self.memory.get_mut().store_mut(),
+                        &mut store,
                         &[Val::I32(0), Val::I32(0)],
                         results.as_mut_slice(),
                     )?;
                     debug!("Initialized Haskell language runtime");
                 }
                 Runtime::Wasi { init, cleanup: _ } => {
-                    debug!("Calling __wasm_call_ctors");
+                    init.call(&mut store, &[], &mut [])?;
-                    init.call(&mut self.memory.get_mut().store_mut(), &[], &mut [])?;
+                    debug!("Initialied WASI runtime");
                 }
             }
         }
@@ -411,7 +444,7 @@ impl Plugin {
                     cleanup: Some(cleanup),
                 } => {
                     debug!("Calling __wasm_call_dtors");
-                    cleanup.call(&mut self.memory_mut().store_mut(), &[], &mut [])?;
+                    cleanup.call(self.store_mut(), &[], &mut [])?;
                 }
                 Runtime::Wasi {
                     init: _,
@@ -420,13 +453,8 @@ impl Plugin {
                 // Cleanup Haskell runtime if `hs_exit` and `hs_exit` are present,
                 // by calling the `hs_exit` export
                 Runtime::Haskell { init: _, cleanup } => {
-                    let mut results =
+                    let mut results = vec![Val::null(); cleanup.ty(self.store()).results().len()];
-                        vec![Val::null(); cleanup.ty(&self.memory().store()).results().len()];
+                    cleanup.call(self.store_mut(), &[], results.as_mut_slice())?;
-                    cleanup.call(
-                        &mut self.memory_mut().store_mut(),
-                        &[],
-                        results.as_mut_slice(),
-                    )?;
                     debug!("Cleaned up Haskell language runtime");
                 }
             }
@@ -442,14 +470,14 @@ impl Plugin {
         tx: &std::sync::mpsc::SyncSender<TimerAction>,
     ) -> Result<(), Error> {
         let duration = self
-            .memory()
+            .internal()
             .manifest
             .as_ref()
             .timeout_ms
             .map(std::time::Duration::from_millis);
         self.cancel_handle.epoch_timer_tx = Some(tx.clone());
-        self.memory_mut().store_mut().set_epoch_deadline(1);
+        self.store_mut().set_epoch_deadline(1);
-        let engine: Engine = self.memory().store().engine().clone();
+        let engine: Engine = self.store().engine().clone();
         tx.send(TimerAction::Start {
             id: self.timer_id,
             duration,

runtime/src/plugin_ref.rs

@@ -3,6 +3,7 @@ use crate::*;
 // PluginRef is used to access a plugin from a context-scoped plugin registry
 pub struct PluginRef<'a> {
     pub id: PluginIndex,
+    running: bool,
     pub(crate) epoch_timer_tx: std::sync::mpsc::SyncSender<TimerAction>,
     plugin: *mut Plugin,
     _t: std::marker::PhantomData<&'a ()>,
@@ -10,23 +11,18 @@ pub struct PluginRef<'a> {
 
 impl<'a> PluginRef<'a> {
     /// Initialize the plugin for a new call
-    ///
+    pub fn start_call(mut self) -> Self {
-    /// - Resets memory offsets
+        trace!("PluginRef::start_call: {}", self.id,);
-    /// - Updates `input` pointer
-    pub fn init(mut self, data: *const u8, data_len: usize) -> Self {
-        trace!("PluginRef::init: {}", self.id,);
         let plugin = self.as_mut();
-        plugin.memory_mut().reset();
+
         if plugin.has_wasi() || plugin.runtime.is_some() {
             if let Err(e) = plugin.reinstantiate() {
                 error!("Failed to reinstantiate: {e:?}");
-                plugin
+                plugin.error(format!("Failed to reinstantiate: {e:?}"), ());
-                    .internal()
-                    .set_error(format!("Failed to reinstantiate: {e:?}"));
                 return self;
             }
         }
-        plugin.set_input(data, data_len);
+        self.running = true;
         self
     }
 
@@ -45,12 +41,25 @@ impl<'a> PluginRef<'a> {
             return ctx.error(format!("Plugin does not exist: {plugin_id}"), None);
         };
 
+        {
+            let plugin = unsafe { &mut *plugin };
+            // Start timer
+            if let Err(e) = plugin.start_timer(&epoch_timer_tx) {
+                let id = plugin.timer_id;
+                plugin.error(
+                    format!("Unable to start timeout manager for {id}: {e:?}"),
+                    (),
+                );
+                return None;
+            }
+        }
+
         if clear_error {
             trace!("Clearing context error");
             ctx.error = None;
             trace!("Clearing plugin error: {plugin_id}");
             unsafe {
-                (&*plugin).internal().clear_error();
+                (*plugin).clear_error();
             }
         }
 
@@ -59,6 +68,7 @@ impl<'a> PluginRef<'a> {
             plugin,
             epoch_timer_tx,
             _t: std::marker::PhantomData,
+            running: false,
         })
     }
 }
@@ -78,6 +88,14 @@ impl<'a> AsMut<Plugin> for PluginRef<'a> {
 impl<'a> Drop for PluginRef<'a> {
     fn drop(&mut self) {
         trace!("Dropping PluginRef {}", self.id);
-        // Cleanup?
+        if self.running {
+            let plugin = self.as_mut();
+
+            // Stop timer
+            if let Err(e) = plugin.stop_timer() {
+                let id = plugin.timer_id;
+                error!("Failed to stop timeout manager for {id}: {e:?}");
+            }
+        }
     }
 }

runtime/src/sdk.rs

@@ -99,7 +99,7 @@ pub unsafe extern "C" fn extism_current_plugin_memory(plugin: *mut Internal) ->
     }
 
     let plugin = &mut *plugin;
-    plugin.memory_mut().data_mut().as_mut_ptr()
+    plugin.memory_ptr()
 }
 
 /// Allocate a memory block in the currently running plugin
@@ -111,16 +111,7 @@ pub unsafe extern "C" fn extism_current_plugin_memory_alloc(plugin: *mut Interna
     }
 
     let plugin = &mut *plugin;
-
+    plugin.memory_alloc(n as u64).unwrap_or_default()
-    let mem = match plugin.memory_mut().alloc(n as usize) {
-        Ok(x) => x,
-        Err(e) => {
-            plugin.set_error(e);
-            return 0;
-        }
-    };
-
-    mem.offset as u64
 }
 
 /// Get the length of an allocated block
@@ -135,11 +126,7 @@ pub unsafe extern "C" fn extism_current_plugin_memory_length(
     }
 
     let plugin = &mut *plugin;
-
+    plugin.memory_length(n)
-    match plugin.memory().block_length(n as usize) {
-        Some(x) => x as Size,
-        None => 0,
-    }
 }
 
 /// Free an allocated memory block
@@ -151,7 +138,7 @@ pub unsafe extern "C" fn extism_current_plugin_memory_free(plugin: *mut Internal
     }
 
     let plugin = &mut *plugin;
-    plugin.memory_mut().free(ptr as usize);
+    plugin.memory_free(ptr);
 }
 
 /// Create a new host function
@@ -436,21 +423,21 @@ pub unsafe extern "C" fn extism_plugin_config(
             }
         };
 
-    let wasi = &mut plugin.memory.get_mut().store_mut().data_mut().wasi;
+    let wasi = &mut plugin.internal_mut().wasi;
     if let Some(Wasi { ctx, .. }) = wasi {
         for (k, v) in json.iter() {
             match v {
                 Some(v) => {
-                    let _ = ctx.push_env(&k, &v);
+                    let _ = ctx.push_env(k, v);
                 }
                 None => {
-                    let _ = ctx.push_env(&k, "");
+                    let _ = ctx.push_env(k, "");
                 }
             }
         }
     }
 
-    let config = &mut plugin.memory.get_mut().manifest.as_mut().config;
+    let config = &mut plugin.internal_mut().manifest.as_mut().config;
     for (k, v) in json.into_iter() {
         match v {
             Some(v) => {
@@ -512,15 +499,11 @@ pub unsafe extern "C" fn extism_plugin_call(
     // needed before a new call
     let mut plugin_ref = match PluginRef::new(ctx, plugin_id, true) {
         None => return -1,
-        Some(p) => p.init(data, data_len as usize),
+        Some(p) => p.start_call(),
     };
     let tx = plugin_ref.epoch_timer_tx.clone();
     let plugin = plugin_ref.as_mut();
 
-    if plugin.internal().last_error.borrow().is_some() {
-        return -1;
-    }
-
     // Find function
     let name = std::ffi::CStr::from_ptr(func_name);
     let name = match name.to_str() {
@@ -534,15 +517,6 @@ pub unsafe extern "C" fn extism_plugin_call(
         None => return plugin.error(format!("Function not found: {name}"), -1),
     };
 
-    // Start timer
-    if let Err(e) = plugin.start_timer(&tx) {
-        let id = plugin.timer_id;
-        return plugin.error(
-            format!("Unable to start timeout manager for {id}: {e:?}"),
-            -1,
-        );
-    }
-
     // Check the number of results, reject functions with more than 1 result
     let n_results = func.ty(plugin.store()).results().len();
     if n_results > 1 {
@@ -559,13 +533,19 @@ pub unsafe extern "C" fn extism_plugin_call(
         }
     }
 
+    if let Err(e) = plugin.set_input(data, data_len as usize, tx) {
+        return plugin.error(e, -1);
+    }
+
+    if plugin.has_error() {
+        return -1;
+    }
+
     debug!("Calling function: {name} in plugin {plugin_id}");
 
     // Call the function
     let mut results = vec![wasmtime::Val::null(); n_results];
-    let res = func.call(&mut plugin.store_mut(), &[], results.as_mut_slice());
+    let res = func.call(plugin.store_mut(), &[], results.as_mut_slice());
-
-    plugin.dump_memory();
 
     // Cleanup runtime
     if !is_start {
@@ -574,18 +554,10 @@ pub unsafe extern "C" fn extism_plugin_call(
         }
     }
 
-    // Stop timer
-    if let Err(e) = plugin.stop_timer() {
-        let id = plugin.timer_id;
-        return plugin.error(
-            format!("Failed to stop timeout manager for {id}: {e:?}"),
-            -1,
-        );
-    }
-
     match res {
         Ok(()) => (),
         Err(e) => {
+            plugin.store.set_epoch_deadline(1);
             if let Some(exit) = e.downcast_ref::<wasmtime_wasi::I32Exit>() {
                 trace!("WASI return code: {}", exit.0);
                 if exit.0 != 0 {
@@ -635,20 +607,27 @@ pub unsafe extern "C" fn extism_error(ctx: *mut Context, plugin: PluginIndex) ->
         return get_context_error(ctx);
     }
 
-    let plugin_ref = match PluginRef::new(ctx, plugin, false) {
+    let mut plugin_ref = match PluginRef::new(ctx, plugin, false) {
         None => return std::ptr::null(),
         Some(p) => p,
     };
-    let plugin = plugin_ref.as_ref();
+    let plugin = plugin_ref.as_mut();
-
+    let output = &mut [Val::I64(0)];
-    let err = plugin.internal().last_error.borrow();
+    if let Some(f) = plugin
-    match err.as_ref() {
+        .linker
-        Some(e) => e.as_ptr() as *const _,
+        .get(&mut plugin.store, "env", "extism_error_get")
-        None => {
+    {
+        f.into_func()
+            .unwrap()
+            .call(&mut plugin.store, &[], output)
+            .unwrap();
+    }
+    if output[0].unwrap_i64() == 0 {
         trace!("Error is NULL");
-            std::ptr::null()
+        return std::ptr::null();
-        }
     }
+
+    plugin.memory_ptr().add(output[0].unwrap_i64() as usize) as *const _
 }
 
 /// Get the length of a plugin's output data
@@ -660,13 +639,20 @@ pub unsafe extern "C" fn extism_plugin_output_length(
     trace!("Call to extism_plugin_output_length for plugin {plugin}");
 
     let ctx = &mut *ctx;
-    let plugin_ref = match PluginRef::new(ctx, plugin, true) {
+    let mut plugin_ref = match PluginRef::new(ctx, plugin, true) {
         None => return 0,
         Some(p) => p,
     };
-    let plugin = plugin_ref.as_ref();
+    let plugin = plugin_ref.as_mut();
-
+    let out = &mut [Val::I64(0)];
-    let len = plugin.internal().output_length as Size;
+    let _ = plugin
+        .linker
+        .get(&mut plugin.store, "env", "extism_output_length")
+        .unwrap()
+        .into_func()
+        .unwrap()
+        .call(&mut plugin.store_mut(), &[], out);
+    let len = out[0].unwrap_i64() as Size;
     trace!("Output length: {len}");
     len
 }
@@ -680,20 +666,23 @@ pub unsafe extern "C" fn extism_plugin_output_data(
     trace!("Call to extism_plugin_output_data for plugin {plugin}");
 
     let ctx = &mut *ctx;
-    let plugin_ref = match PluginRef::new(ctx, plugin, true) {
+    let mut plugin_ref = match PluginRef::new(ctx, plugin, true) {
         None => return std::ptr::null(),
         Some(p) => p,
     };
-    let plugin = plugin_ref.as_ref();
+    let plugin = plugin_ref.as_mut();
-    let internal = plugin.internal();
+    let ptr = plugin.memory_ptr();
+    let out = &mut [Val::I64(0)];
+    let mut store = &mut *(plugin.store_mut() as *mut Store<_>);
     plugin
-        .memory()
+        .linker
-        .ptr(MemoryBlock::new(
+        .get(&mut store, "env", "extism_output_offset")
-            internal.output_offset,
+        .unwrap()
-            internal.output_length,
+        .into_func()
-        ))
+        .unwrap()
-        .map(|x| x as *const _)
+        .call(&mut store, &[], out)
-        .unwrap_or(std::ptr::null())
+        .unwrap();
+    ptr.add(out[0].unwrap_i64() as usize)
 }
 
 /// Set log file and level

rust/Cargo.toml

@@ -9,7 +9,7 @@ repository = "https://github.com/extism/extism"
 description = "Extism Host SDK for Rust"
 
 [dependencies]
-extism-manifest = { version = "0.3.0", path = "../manifest" }
+extism-manifest = { version = "0.4.0", path = "../manifest" }
 extism-runtime = { version = "0.4.0", path = "../runtime"}
 serde_json = "1"
 log = "0.4"

rust/src/lib.rs

@@ -1,6 +1,6 @@
 pub use extism_manifest::{self as manifest, Manifest};
 pub use extism_runtime::{
-    sdk as bindings, Function, Internal as CurrentPlugin, MemoryBlock, UserData, Val, ValType,
+    sdk as bindings, Function, Internal as CurrentPlugin, UserData, Val, ValType,
 };
 
 mod context;

zig/src/current_plugin.zig

@@ -13,7 +13,7 @@ pub fn getMemory(self: Self, offset: u64) []const u8 {
     const len = c.extism_current_plugin_memory_length(self.c_currplugin, offset);
     const c_data = c.extism_current_plugin_memory(self.c_currplugin);
     const data: [*:0]u8 = std.mem.span(c_data);
-    return data[offset .. len + 1];
+    return data[offset .. offset + len];
 }
 
 pub fn alloc(self: *Self, n: u64) u64 {

zig/src/utils.zig

@@ -81,7 +81,7 @@ fn stringify(
             try out_stream.writeByte('{');
             var field_output = false;
             var child_options = options;
-            child_options.whitespace.indent_level += 1;
+            child_options.whitespace = .indent_2;
             inline for (S.fields) |Field| {
                 // don't include void fields
                 if (Field.type == void) continue;
@@ -103,18 +103,14 @@ fn stringify(
                     } else {
                         try out_stream.writeByte(',');
                     }
-                    try child_options.whitespace.outputIndent(out_stream);
                     try json.encodeJsonString(Field.name, options, out_stream);
                     try out_stream.writeByte(':');
-                    if (child_options.whitespace.separator) {
+                    if (child_options.whitespace != .minified) {
                         try out_stream.writeByte(' ');
                     }
                     try stringify(@field(value, Field.name), child_options, out_stream);
                 }
             }
-            if (field_output) {
-                try options.whitespace.outputIndent(out_stream);
-            }
             try out_stream.writeByte('}');
             return;
         },
@@ -133,24 +129,19 @@ fn stringify(
             },
             // TODO: .Many when there is a sentinel (waiting for https://github.com/ziglang/zig/pull/3972)
             .Slice => {
-                if (ptr_info.child == u8 and options.string == .String and std.unicode.utf8ValidateSlice(value)) {
+                if (ptr_info.child == u8 and std.unicode.utf8ValidateSlice(value)) {
                     try json.encodeJsonString(value, options, out_stream);
                     return;
                 }
 
                 try out_stream.writeByte('[');
                 var child_options = options;
-                child_options.whitespace.indent_level += 1;
                 for (value, 0..) |x, i| {
                     if (i != 0) {
                         try out_stream.writeByte(',');
                     }
-                    try child_options.whitespace.outputIndent(out_stream);
                     try stringify(x, child_options, out_stream);
                 }
-                if (value.len != 0) {
-                    try options.whitespace.outputIndent(out_stream);
-                }
                 try out_stream.writeByte(']');
                 return;
             },