1use std::{
2 collections::HashMap, fs::File, future::Future, marker::PhantomData, path::Path, pin::Pin,
3};
4
5use anyhow::{anyhow, Error};
6use serde::{de::DeserializeOwned, Serialize};
7
8use wasi_common::{dir, file};
9use wasmtime::IntoFunc;
10use wasmtime::{Caller, Config, Engine, Instance, Linker, Module, Store, TypedFunc};
11use wasmtime_wasi::{Dir, WasiCtx, WasiCtxBuilder};
12
13pub struct WasiResource(u32);
14
15pub struct WasiFn<A: Serialize, R: DeserializeOwned> {
16 function: TypedFunc<(u32, u32), u32>,
17 _function_type: PhantomData<fn(A) -> R>,
18}
19
20impl<A: Serialize, R: DeserializeOwned> Copy for WasiFn<A, R> {}
21
22impl<A: Serialize, R: DeserializeOwned> Clone for WasiFn<A, R> {
23 fn clone(&self) -> Self {
24 Self {
25 function: self.function,
26 _function_type: PhantomData,
27 }
28 }
29}
30
31// impl<A: Serialize, R: DeserializeOwned> WasiFn<A, R> {
32// #[inline(always)]
33// pub async fn call(&self, runtime: &mut Wasi, arg: A) -> Result<R, Error> {
34// runtime.call(self, arg).await
35// }
36// }
37
38pub struct Wasi {
39 engine: Engine,
40 module: Module,
41 store: Store<WasiCtx>,
42 instance: Instance,
43 alloc_buffer: TypedFunc<u32, u32>,
44 // free_buffer: TypedFunc<(u32, u32), ()>,
45}
46
47// type signature derived from:
48// https://docs.rs/wasmtime/latest/wasmtime/struct.Linker.html#method.func_wrap2_async
49// macro_rules! dynHostFunction {
50// () => {
51// Box<
52// dyn for<'a> Fn(Caller<'a, WasiCtx>, u32, u32)
53// -> Box<dyn Future<Output = u32> + Send + 'a>
54// + Send
55// + Sync
56// + 'static
57// >
58// };
59// }
60
61// macro_rules! implHostFunction {
62// () => {
63// impl for<'a> Fn(Caller<'a, WasiCtx>, u32, u32)
64// -> Box<dyn Future<Output = u32> + Send + 'a>
65// + Send
66// + Sync
67// + 'static
68// };
69// }
70
71// This type signature goodness gracious
72pub type HostFunction = Box<dyn IntoFunc<WasiCtx, (u32, u32), u32>>;
73
74pub struct WasiPluginBuilder {
75 host_functions: HashMap<String, Box<dyn Fn(&str, &mut Linker<WasiCtx>) -> Result<(), Error>>>,
76 wasi_ctx_builder: WasiCtxBuilder,
77}
78
79impl WasiPluginBuilder {
80 pub fn new() -> Self {
81 WasiPluginBuilder {
82 host_functions: HashMap::new(),
83 wasi_ctx_builder: WasiCtxBuilder::new(),
84 }
85 }
86
87 pub fn new_with_default_ctx() -> WasiPluginBuilder {
88 let mut this = Self::new();
89 this.wasi_ctx_builder = this.wasi_ctx_builder.inherit_stdin().inherit_stderr();
90 this
91 }
92
93 pub fn host_function<A: Serialize, R: DeserializeOwned>(
94 mut self,
95 name: &str,
96 function: &dyn Fn(A) -> R + Send + Sync + 'static,
97 ) -> Self {
98 let name = name.to_string();
99 self.host_functions.insert(
100 name,
101 Box::new(move |name: &str, linker: &mut Linker<WasiCtx>| {
102 linker.func_wrap("env", name, |ptr: u32, len: u32| {
103 function(todo!());
104 7u32
105 })?;
106 Ok(())
107 }),
108 );
109 self
110 }
111
112 pub fn wasi_ctx(mut self, config: impl FnOnce(WasiCtxBuilder) -> WasiCtxBuilder) -> Self {
113 self.wasi_ctx_builder = config(self.wasi_ctx_builder);
114 self
115 }
116
117 pub async fn init<T: AsRef<[u8]>>(self, module: T) -> Result<Wasi, Error> {
118 let plugin = WasiPlugin {
119 module: module.as_ref().to_vec(),
120 wasi_ctx: self.wasi_ctx_builder.build(),
121 host_functions: self.host_functions,
122 };
123
124 Wasi::init(plugin).await
125 }
126}
127
128/// Represents a to-be-initialized plugin.
129/// Please use [`WasiPluginBuilder`], don't use this directly.
130pub struct WasiPlugin {
131 pub module: Vec<u8>,
132 pub wasi_ctx: WasiCtx,
133 pub host_functions:
134 HashMap<String, Box<dyn Fn(&str, &mut Linker<WasiCtx>) -> Result<(), Error>>>,
135}
136
137impl Wasi {
138 pub fn dump_memory(data: &[u8]) {
139 for (i, byte) in data.iter().enumerate() {
140 if i % 32 == 0 {
141 println!();
142 }
143 if i % 4 == 0 {
144 print!("|");
145 }
146 if *byte == 0 {
147 print!("__")
148 } else {
149 print!("{:02x}", byte);
150 }
151 }
152 println!();
153 }
154}
155
156impl Wasi {
157 async fn init(plugin: WasiPlugin) -> Result<Self, Error> {
158 let mut config = Config::default();
159 config.async_support(true);
160 let engine = Engine::new(&config)?;
161 let mut linker = Linker::new(&engine);
162
163 for (name, add_to_linker) in plugin.host_functions.into_iter() {
164 add_to_linker(&name, &mut linker)?;
165 }
166
167 linker
168 .func_wrap("env", "__command", |x: u32, y: u32| x + y)
169 .unwrap();
170 linker.func_wrap("env", "__hello", |x: u32| x * 2).unwrap();
171 linker.func_wrap("env", "__bye", |x: u32| x / 2).unwrap();
172
173 wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
174
175 let mut store: Store<_> = Store::new(&engine, plugin.wasi_ctx);
176 let module = Module::new(&engine, plugin.module)?;
177
178 linker.module_async(&mut store, "", &module).await?;
179 let instance = linker.instantiate_async(&mut store, &module).await?;
180
181 let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
182 // let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
183
184 Ok(Wasi {
185 engine,
186 module,
187 store,
188 instance,
189 alloc_buffer,
190 // free_buffer,
191 })
192 }
193
194 /// Attaches a file or directory the the given system path to the runtime.
195 /// Note that the resource must be freed by calling `remove_resource` afterwards.
196 pub fn attach_path<T: AsRef<Path>>(&mut self, path: T) -> Result<WasiResource, Error> {
197 // grab the WASI context
198 let ctx = self.store.data_mut();
199
200 // open the file we want, and convert it into the right type
201 // this is a footgun and a half
202 let file = File::open(&path).unwrap();
203 let dir = Dir::from_std_file(file);
204 let dir = Box::new(wasmtime_wasi::dir::Dir::from_cap_std(dir));
205
206 // grab an empty file descriptor, specify capabilities
207 let fd = ctx.table().push(Box::new(()))?;
208 let caps = dir::DirCaps::all();
209 let file_caps = file::FileCaps::all();
210
211 // insert the directory at the given fd,
212 // return a handle to the resource
213 ctx.insert_dir(fd, dir, caps, file_caps, path.as_ref().to_path_buf());
214 Ok(WasiResource(fd))
215 }
216
217 /// Returns `true` if the resource existed and was removed.
218 pub fn remove_resource(&mut self, resource: WasiResource) -> Result<(), Error> {
219 self.store
220 .data_mut()
221 .table()
222 .delete(resource.0)
223 .ok_or_else(|| anyhow!("Resource did not exist, but a valid handle was passed in"))?;
224 Ok(())
225 }
226
227 // pub fn with_resource<T>(
228 // &mut self,
229 // resource: WasiResource,
230 // callback: fn(&mut Self) -> Result<T, Error>,
231 // ) -> Result<T, Error> {
232 // let result = callback(self);
233 // self.remove_resource(resource)?;
234 // return result;
235 // }
236
237 // So this call function is kinda a dance, I figured it'd be a good idea to document it.
238 // the high level is we take a serde type, serialize it to a byte array,
239 // (we're doing this using bincode for now)
240 // then toss that byte array into webassembly.
241 // webassembly grabs that byte array, does some magic,
242 // and serializes the result into yet another byte array.
243 // we then grab *that* result byte array and deserialize it into a result.
244 //
245 // phew...
246 //
247 // now the problem is, webassambly doesn't support buffers.
248 // only really like i32s, that's it (yeah, it's sad. Not even unsigned!)
249 // (ok, I'm exaggerating a bit).
250 //
251 // the Wasm function that this calls must have a very specific signature:
252 //
253 // fn(pointer to byte array: i32, length of byte array: i32)
254 // -> pointer to (
255 // pointer to byte_array: i32,
256 // length of byte array: i32,
257 // ): i32
258 //
259 // This pair `(pointer to byte array, length of byte array)` is called a `Buffer`
260 // and can be found in the cargo_test plugin.
261 //
262 // so on the wasm side, we grab the two parameters to the function,
263 // stuff them into a `Buffer`,
264 // and then pray to the `unsafe` Rust gods above that a valid byte array pops out.
265 //
266 // On the flip side, when returning from a wasm function,
267 // we convert whatever serialized result we get into byte array,
268 // which we stuff into a Buffer and allocate on the heap,
269 // which pointer to we then return.
270 // Note the double indirection!
271 //
272 // So when returning from a function, we actually leak memory *twice*:
273 //
274 // 1) once when we leak the byte array
275 // 2) again when we leak the allocated `Buffer`
276 //
277 // This isn't a problem because Wasm stops executing after the function returns,
278 // so the heap is still valid for our inspection when we want to pull things out.
279
280 /// Takes an item, allocates a buffer, serializes the argument to that buffer,
281 /// and returns a (ptr, len) pair to that buffer.
282 async fn serialize_to_buffer<T: Serialize>(&mut self, item: T) -> Result<(u32, u32), Error> {
283 // serialize the argument using bincode
284 let item = bincode::serialize(&item)?;
285 let buffer_len = item.len() as u32;
286
287 // allocate a buffer and write the argument to that buffer
288 let buffer_ptr = self
289 .alloc_buffer
290 .call_async(&mut self.store, buffer_len)
291 .await?;
292 let plugin_memory = self
293 .instance
294 .get_memory(&mut self.store, "memory")
295 .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
296 plugin_memory.write(&mut self.store, buffer_ptr as usize, &item)?;
297 Ok((buffer_ptr, buffer_len))
298 }
299
300 /// Takes a ptr to a (ptr, len) pair and returns the corresponding deserialized buffer
301 fn deserialize_from_buffer<R: DeserializeOwned>(&mut self, buffer: u32) -> Result<R, Error> {
302 // create a buffer to read the (ptr, length) pair into
303 // this is a total of 4 + 4 = 8 bytes.
304 let raw_buffer = &mut [0; 8];
305 let plugin_memory = self
306 .instance
307 .get_memory(&mut self.store, "memory")
308 .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
309 plugin_memory.read(&mut self.store, buffer as usize, raw_buffer)?;
310
311 // use these bytes (wasm stores things little-endian)
312 // to get a pointer to the buffer and its length
313 let b = raw_buffer;
314 let buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
315 let buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
316 let buffer_end = buffer_ptr + buffer_len;
317
318 // read the buffer at this point into a byte array
319 // deserialize the byte array into the provided serde type
320 let result = &plugin_memory.data(&mut self.store)[buffer_ptr..buffer_end];
321 let result = bincode::deserialize(result)?;
322
323 // TODO: this is handled wasm-side, but I'd like to double-check
324 // // deallocate the argument buffer
325 // self.free_buffer.call(&mut self.store, arg_buffer);
326
327 Ok(result)
328 }
329
330 pub fn function<A: Serialize, R: DeserializeOwned, T: AsRef<str>>(
331 &mut self,
332 name: T,
333 ) -> Result<WasiFn<A, R>, Error> {
334 let fun_name = format!("__{}", name.as_ref());
335 let fun = self
336 .instance
337 .get_typed_func::<(u32, u32), u32, _>(&mut self.store, &fun_name)?;
338 Ok(WasiFn {
339 function: fun,
340 _function_type: PhantomData,
341 })
342 }
343
344 // TODO: dont' use as for conversions
345 pub async fn call<A: Serialize, R: DeserializeOwned>(
346 &mut self,
347 handle: &WasiFn<A, R>,
348 arg: A,
349 ) -> Result<R, Error> {
350 // dbg!(&handle.name);
351 // dbg!(serde_json::to_string(&arg)).unwrap();
352
353 // write the argument to linear memory
354 // this returns a (ptr, lentgh) pair
355 let arg_buffer = self.serialize_to_buffer(arg).await?;
356
357 // get the webassembly function we want to actually call
358 // TODO: precompute handle
359 // let fun_name = format!("__{}", handle);
360 // let fun = self
361 // .instance
362 // .get_typed_func::<(u32, u32), u32, _>(&mut self.store, &fun_name)?;
363 let fun = handle.function;
364
365 // call the function, passing in the buffer and its length
366 // this returns a ptr to a (ptr, lentgh) pair
367 let result_buffer = fun.call_async(&mut self.store, arg_buffer).await?;
368
369 self.deserialize_from_buffer(result_buffer)
370 }
371}