ncp_runtime/
engine.rs

1use anyhow::{bail, Context, Result};
2use wasmtime::{Engine, Linker, Module, Store, StoreLimits, StoreLimitsBuilder};
3
4struct StoreState {
5    limits: StoreLimits,
6}
7
8pub struct CompiledBrick {
9    engine: Engine,
10    module: Module,
11}
12
13impl CompiledBrick {
14    pub fn new(wasm_bytes: &[u8]) -> Result<Self> {
15        // Keep this shape so you can later do:
16        // let mut cfg = wasmtime::Config::new();
17        // cfg.consume_fuel(true);
18        // let engine = Engine::new(&cfg)?;
19        let engine = Engine::default();
20        let module = Module::new(&engine, wasm_bytes)
21            .map_err(|e| anyhow::anyhow!("compiling WASM module: {e}"))?;
22        Ok(Self { engine, module })
23    }
24
25    /// Invoke using Model B: alloc → write → invoke → read → free
26    pub fn invoke(
27        &self,
28        envelope: &[u8],
29        max_mem_mb: u64,
30        max_output_bytes: u64,
31    ) -> Result<Vec<u8>> {
32        let mem_bytes = (max_mem_mb as usize)
33            .saturating_mul(1024)
34            .saturating_mul(1024);
35        if mem_bytes == 0 {
36            bail!("limits.max_mem_mb must be > 0");
37        }
38        let limits = StoreLimitsBuilder::new().memory_size(mem_bytes).build();
39        let mut store = Store::new(&self.engine, StoreState { limits });
40        store.limiter(|state| &mut state.limits);
41
42        let linker = Linker::new(&self.engine);
43        let instance = linker
44            .instantiate(&mut store, &self.module)
45            .map_err(|e| anyhow::anyhow!("instantiating WASM module: {e}"))?;
46
47        let memory = instance
48            .get_memory(&mut store, "memory")
49            .ok_or_else(|| anyhow::anyhow!("WASM module must export 'memory'"))?;
50
51        let alloc_fn = instance
52            .get_typed_func::<i32, i32>(&mut store, "alloc")
53            .map_err(|e| anyhow::anyhow!("WASM module must export 'alloc(i32)->i32': {e}"))?;
54
55        let free_fn = instance
56            .get_typed_func::<(i32, i32), ()>(&mut store, "free")
57            .map_err(|e| anyhow::anyhow!("WASM module must export 'free(i32,i32)': {e}"))?;
58
59        let invoke_fn = instance
60            .get_typed_func::<(i32, i32), i32>(&mut store, "invoke")
61            .map_err(|e| anyhow::anyhow!("WASM module must export 'invoke(i32,i32)->i32': {e}"))?;
62
63        let mut to_free: Vec<(i32, i32)> = Vec::new();
64
65        let envelope_len_i32: i32 = envelope
66            .len()
67            .try_into()
68            .context("envelope too large for i32 length")?;
69
70        // 1) alloc envelope
71        let envelope_ptr = alloc_fn
72            .call(&mut store, envelope_len_i32)
73            .map_err(|e| anyhow::anyhow!("calling alloc for envelope: {e}"))?;
74        if envelope_ptr == 0 {
75            bail!("alloc returned 0 (OOM) for envelope — map to Failure(RESOURCE_EXCEEDED)");
76        }
77        to_free.push((envelope_ptr, envelope_len_i32));
78
79        // Run invoke + read in a closure so we always free afterward
80        let res = (|| -> Result<Vec<u8>> {
81            // 2) write envelope
82            memory
83                .write(&mut store, envelope_ptr as usize, envelope)
84                .map_err(|e| anyhow::anyhow!("writing envelope to WASM memory: {e}"))?;
85
86            // 3) invoke
87            let result_ptr = invoke_fn
88                .call(&mut store, (envelope_ptr, envelope_len_i32))
89                .map_err(|e| anyhow::anyhow!("calling invoke: {e}"))?;
90            if result_ptr == 0 {
91                bail!("invoke returned 0 result_ptr");
92            }
93
94            // 4) read len prefix
95            let mut len_buf = [0u8; 4];
96            memory
97                .read(&store, result_ptr as usize, &mut len_buf)
98                .map_err(|e| anyhow::anyhow!("reading result length prefix: {e}"))?;
99            let result_len = u32::from_le_bytes(len_buf) as u64;
100
101            if result_len == 0 {
102                bail!("invoke returned zero-length result");
103            }
104            if result_len > max_output_bytes {
105                bail!(
106                    "result too large: {} bytes > limits.max_output_bytes {}",
107                    result_len,
108                    max_output_bytes
109                );
110            }
111
112            // schedule free of result buffer: prefix(4) + payload
113            let total = 4u64 + result_len;
114            let total_i32: i32 = total
115                .try_into()
116                .context("result buffer too large for i32 length")?;
117            to_free.push((result_ptr, total_i32));
118
119            // 5) read payload
120            let mut result_bytes = vec![0u8; result_len as usize];
121            memory
122                .read(&store, result_ptr as usize + 4, &mut result_bytes)
123                .map_err(|e| anyhow::anyhow!("reading result CBOR payload: {e}"))?;
124
125            Ok(result_bytes)
126        })();
127
128        // 6) free (best-effort, reverse order for allocator friendliness) — always runs
129        for (ptr, len) in to_free.into_iter().rev() {
130            let _ = free_fn.call(&mut store, (ptr, len));
131        }
132
133        res
134    }
135}
ncp_runtime/engine.rs

ncp_runtime/
engine.rs
