wasm.rs

  1use std::collections::HashMap;
  2
  3use anyhow::anyhow;
  4
  5use wasmtime::{Engine, Func, Instance, Memory, MemoryType, Module, Store, TypedFunc};
  6
  7use crate::*;
  8
  9pub struct Wasm {
 10    engine: Engine,
 11    module: Module,
 12    store: Store<()>,
 13    instance: Instance,
 14    alloc_buffer: TypedFunc<i32, i32>,
 15    // free_buffer: TypedFunc<(i32, i32), ()>,
 16}
 17
 18pub struct WasmPlugin {
 19    pub source_bytes: Vec<u8>,
 20}
 21
 22impl Wasm {
 23    pub fn dump_memory(data: &[u8]) {
 24        for (i, byte) in data.iter().enumerate() {
 25            if i % 32 == 0 {
 26                println!();
 27            }
 28            if i % 4 == 0 {
 29                print!("|");
 30            }
 31            if *byte == 0 {
 32                print!("__")
 33            } else {
 34                print!("{:02x}", byte);
 35            }
 36        }
 37        println!();
 38    }
 39}
 40
 41impl Runtime for Wasm {
 42    type Plugin = WasmPlugin;
 43    type Error = anyhow::Error;
 44
 45    fn init(plugin: WasmPlugin) -> Result<Self, Self::Error> {
 46        let engine = Engine::default();
 47        let module = Module::new(&engine, plugin.source_bytes)?;
 48        let mut store: Store<()> = Store::new(&engine, ());
 49        let instance = Instance::new(&mut store, &module, &[])?;
 50
 51        let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
 52        // let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
 53
 54        Ok(Wasm {
 55            engine,
 56            module,
 57            store,
 58            instance,
 59            alloc_buffer,
 60            // free_buffer,
 61        })
 62    }
 63
 64    // fn constant<T: DeserializeOwned>(&mut self, handle: &Handle) -> Result<T, Self::Error> {
 65    //     let export = self
 66    //         .instance
 67    //         .get_export(&mut self.store, handle.inner())
 68    //         .ok_or_else(|| anyhow!("Could not get export"))?;
 69
 70    //     todo!()
 71    // }
 72
 73    // So this call function is kinda a dance, I figured it'd be a good idea to document it.
 74    // the high level is we take a serde type, serialize it to a byte array,
 75    // (we're doing this using bincode for now)
 76    // then toss that byte array into webassembly.
 77    // webassembly grabs that byte array, does some magic,
 78    // and serializes the result into yet another byte array.
 79    // we then grab *that* result byte array and deserialize it into a result.
 80    //
 81    // phew...
 82    //
 83    // now the problem is, webassambly doesn't support buffers.
 84    // only really like i32s, that's it (yeah, it's sad. Not even unsigned!)
 85    // (ok, I'm exaggerating a bit).
 86    //
 87    // the Wasm function that this calls must have a very specific signature:
 88    //
 89    // fn(pointer to byte array: i32, length of byte array: i32)
 90    //     -> pointer to (
 91    //            pointer to byte_array: i32,
 92    //            length of byte array: i32,
 93    //     ): i32
 94    //
 95    // This pair `(pointer to byte array, length of byte array)` is called a `Buffer`
 96    // and can be found in the cargo_test plugin.
 97    //
 98    // so on the wasm side, we grab the two parameters to the function,
 99    // stuff them into a `Buffer`,
100    // and then pray to the `unsafe` Rust gods above that a valid byte array pops out.
101    //
102    // On the flip side, when returning from a wasm function,
103    // we convert whatever serialized result we get into byte array,
104    // which we stuff into a Buffer and allocate on the heap,
105    // which pointer to we then return.
106    // Note the double indirection!
107    //
108    // So when returning from a function, we actually leak memory *twice*:
109    //
110    // 1) once when we leak the byte array
111    // 2) again when we leak the allocated `Buffer`
112    //
113    // This isn't a problem because Wasm stops executing after the function returns,
114    // so the heap is still valid for our inspection when we want to pull things out.
115
116    // TODO: dont' use as for conversions
117    fn call<A: Serialize, R: DeserializeOwned>(
118        &mut self,
119        handle: &str,
120        arg: A,
121    ) -> Result<R, Self::Error> {
122        // serialize the argument using bincode
123        let arg = bincode::serialize(&arg)?;
124        let arg_buffer_len = arg.len();
125
126        // allocate a buffer and write the argument to that buffer
127        let arg_buffer_ptr = self
128            .alloc_buffer
129            .call(&mut self.store, arg_buffer_len as i32)?;
130        let plugin_memory = self
131            .instance
132            .get_memory(&mut self.store, "memory")
133            .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
134        plugin_memory.write(&mut self.store, arg_buffer_ptr as usize, &arg)?;
135
136        // get the webassembly function we want to actually call
137        // TODO: precompute handle
138        let fun_name = format!("__{}", handle);
139        let fun = self
140            .instance
141            .get_typed_func::<(i32, i32), i32, _>(&mut self.store, &fun_name)?;
142
143        // call the function, passing in the buffer and its length
144        // this should return a pointer to a (ptr, lentgh) pair
145        let arg_buffer = (arg_buffer_ptr, arg_buffer_len as i32);
146        let result_buffer = fun.call(&mut self.store, arg_buffer)?;
147
148        // create a buffer to read the (ptr, length) pair into
149        // this is a total of 4 + 4 = 8 bytes.
150        let buffer = &mut [0; 8];
151        plugin_memory.read(&mut self.store, result_buffer as usize, buffer)?;
152
153        // use these bytes (wasm stores things little-endian)
154        // to get a pointer to the buffer and its length
155        let b = buffer;
156        let result_buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
157        let result_buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
158        let result_buffer_end = result_buffer_ptr + result_buffer_len;
159
160        // read the buffer at this point into a byte array
161        // deserialize the byte array into the provided serde type
162        let result = &plugin_memory.data(&mut self.store)[result_buffer_ptr..result_buffer_end];
163        let result = bincode::deserialize(result)?;
164
165        // TODO: this is handled wasm-side, but I'd like to double-check
166        // // deallocate the argument buffer
167        // self.free_buffer.call(&mut self.store, arg_buffer);
168
169        return Ok(result);
170    }
171
172    // fn register_handle<T: AsRef<str>>(&mut self, name: T) -> bool {
173    //     self.instance
174    //         .get_export(&mut self.store, name.as_ref())
175    //         .is_some()
176    // }
177}