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