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