1use std::future::Future;
2
3use std::{fs::File, marker::PhantomData, path::Path};
4
5use anyhow::{anyhow, Error};
6use serde::{de::DeserializeOwned, Serialize};
7
8use wasi_common::{dir, file};
9use wasmtime::Memory;
10use wasmtime::{
11 AsContext, AsContextMut, Caller, Config, Engine, Extern, Instance, Linker, Module, Store, Trap,
12 TypedFunc,
13};
14use wasmtime_wasi::{Dir, WasiCtx, WasiCtxBuilder};
15
16/// Represents a resource currently managed by the plugin, like a file descriptor.
17pub struct PluginResource(u32);
18
19#[repr(C)]
20struct WasiBuffer {
21 ptr: u32,
22 len: u32,
23}
24
25impl WasiBuffer {
26 pub fn into_u64(self) -> u64 {
27 ((self.ptr as u64) << 32) | (self.len as u64)
28 }
29
30 pub fn from_u64(packed: u64) -> Self {
31 WasiBuffer {
32 ptr: (packed >> 32) as u32,
33 len: packed as u32,
34 }
35 }
36}
37
38/// Represents a typed WebAssembly function.
39pub struct WasiFn<A: Serialize, R: DeserializeOwned> {
40 function: TypedFunc<u64, u64>,
41 _function_type: PhantomData<fn(A) -> R>,
42}
43
44impl<A: Serialize, R: DeserializeOwned> Copy for WasiFn<A, R> {}
45
46impl<A: Serialize, R: DeserializeOwned> Clone for WasiFn<A, R> {
47 fn clone(&self) -> Self {
48 Self {
49 function: self.function,
50 _function_type: PhantomData,
51 }
52 }
53}
54
55/// This struct is used to build a new [`Plugin`], using the builder pattern.
56/// Create a new default plugin with `PluginBuilder::new_with_default_ctx`,
57/// and add host-side exported functions using `host_function` and `host_function_async`.
58/// Finalize the plugin by calling [`init`].
59pub struct PluginBuilder {
60 wasi_ctx: WasiCtx,
61 engine: Engine,
62 linker: Linker<WasiCtxAlloc>,
63}
64
65impl PluginBuilder {
66 /// Create a new [`PluginBuilder`] with the given WASI context.
67 /// Using the default context is a safe bet, see [`new_with_default_context`].
68 pub fn new(wasi_ctx: WasiCtx) -> Result<Self, Error> {
69 let mut config = Config::default();
70 config.async_support(true);
71 // config.epoch_interruption(true);
72 let engine = Engine::new(&config)?;
73 let linker = Linker::new(&engine);
74
75 Ok(PluginBuilder {
76 // host_functions: HashMap::new(),
77 wasi_ctx,
78 engine,
79 linker,
80 })
81 }
82
83 /// Create a new `PluginBuilder` that inherits the
84 /// host processes' access to `stdout` and `stderr`.
85 pub fn new_with_default_ctx() -> Result<Self, Error> {
86 let wasi_ctx = WasiCtxBuilder::new()
87 .inherit_stdout()
88 .inherit_stderr()
89 .build();
90 Self::new(wasi_ctx)
91 }
92
93 /// Add an `async` host function. See [`host_function`] for details.
94 pub fn host_function_async<F, A, R, Fut>(
95 mut self,
96 name: &str,
97 function: F,
98 ) -> Result<Self, Error>
99 where
100 F: Fn(A) -> Fut + Send + Sync + 'static,
101 Fut: Future<Output = R> + Send + 'static,
102 A: DeserializeOwned + Send + 'static,
103 R: Serialize + Send + Sync + 'static,
104 {
105 self.linker.func_wrap1_async(
106 "env",
107 &format!("__{}", name),
108 move |mut caller: Caller<'_, WasiCtxAlloc>, packed_buffer: u64| {
109 // TODO: use try block once avaliable
110 let result: Result<(WasiBuffer, Memory, _), Trap> = (|| {
111 // grab a handle to the memory
112 let mut plugin_memory = match caller.get_export("memory") {
113 Some(Extern::Memory(mem)) => mem,
114 _ => return Err(Trap::new("Could not grab slice of plugin memory"))?,
115 };
116
117 let buffer = WasiBuffer::from_u64(packed_buffer);
118
119 // get the args passed from Guest
120 let args = Plugin::buffer_to_bytes(&mut plugin_memory, &mut caller, &buffer)?;
121
122 let args: A = Plugin::deserialize_to_type(&args)?;
123
124 // Call the Host-side function
125 let result = function(args);
126
127 Ok((buffer, plugin_memory, result))
128 })();
129
130 Box::new(async move {
131 let (buffer, mut plugin_memory, future) = result?;
132
133 let result: R = future.await;
134 let result: Result<Vec<u8>, Error> = Plugin::serialize_to_bytes(result)
135 .map_err(|_| {
136 Trap::new("Could not serialize value returned from function").into()
137 });
138 let result = result?;
139
140 Plugin::buffer_to_free(caller.data().free_buffer(), &mut caller, buffer)
141 .await?;
142
143 let buffer = Plugin::bytes_to_buffer(
144 caller.data().alloc_buffer(),
145 &mut plugin_memory,
146 &mut caller,
147 result,
148 )
149 .await?;
150
151 Ok(buffer.into_u64())
152 })
153 },
154 )?;
155 Ok(self)
156 }
157
158 /// Add a new host function to the given `PluginBuilder`.
159 /// A host function is a function defined host-side, in Rust,
160 /// that is accessible guest-side, in WebAssembly.
161 /// You can specify host-side functions to import using
162 /// the `#[input]` macro attribute:
163 /// ```ignore
164 /// #[input]
165 /// fn total(counts: Vec<f64>) -> f64;
166 /// ```
167 /// When loading a plugin, you need to provide all host functions the plugin imports:
168 /// ```ignore
169 /// let plugin = PluginBuilder::new_with_default_context()
170 /// .host_function("total", |counts| counts.iter().fold(0.0, |tot, n| tot + n))
171 /// // and so on...
172 /// ```
173 /// And that's a wrap!
174 pub fn host_function<A, R>(
175 mut self,
176 name: &str,
177 function: impl Fn(A) -> R + Send + Sync + 'static,
178 ) -> Result<Self, Error>
179 where
180 A: DeserializeOwned + Send,
181 R: Serialize + Send + Sync,
182 {
183 self.linker.func_wrap1_async(
184 "env",
185 &format!("__{}", name),
186 move |mut caller: Caller<'_, WasiCtxAlloc>, packed_buffer: u64| {
187 // TODO: use try block once avaliable
188 let result: Result<(WasiBuffer, Memory, Vec<u8>), Trap> = (|| {
189 // grab a handle to the memory
190 let mut plugin_memory = match caller.get_export("memory") {
191 Some(Extern::Memory(mem)) => mem,
192 _ => return Err(Trap::new("Could not grab slice of plugin memory"))?,
193 };
194
195 let buffer = WasiBuffer::from_u64(packed_buffer);
196
197 // get the args passed from Guest
198 let args = Plugin::buffer_to_type(&mut plugin_memory, &mut caller, &buffer)?;
199
200 // Call the Host-side function
201 let result: R = function(args);
202
203 // Serialize the result back to guest
204 let result = Plugin::serialize_to_bytes(result).map_err(|_| {
205 Trap::new("Could not serialize value returned from function")
206 })?;
207
208 Ok((buffer, plugin_memory, result))
209 })();
210
211 Box::new(async move {
212 let (buffer, mut plugin_memory, result) = result?;
213
214 Plugin::buffer_to_free(caller.data().free_buffer(), &mut caller, buffer)
215 .await?;
216
217 let buffer = Plugin::bytes_to_buffer(
218 caller.data().alloc_buffer(),
219 &mut plugin_memory,
220 &mut caller,
221 result,
222 )
223 .await?;
224
225 Ok(buffer.into_u64())
226 })
227 },
228 )?;
229 Ok(self)
230 }
231
232 /// Initializes a [`Plugin`] from a given compiled Wasm module.
233 /// Both binary (`.wasm`) and text (`.wat`) module formats are supported.
234 pub async fn init<T: AsRef<[u8]>>(self, precompiled: bool, module: T) -> Result<Plugin, Error> {
235 Plugin::init(precompiled, module.as_ref().to_vec(), self).await
236 }
237}
238
239#[derive(Copy, Clone)]
240struct WasiAlloc {
241 alloc_buffer: TypedFunc<u32, u32>,
242 free_buffer: TypedFunc<u64, ()>,
243}
244
245struct WasiCtxAlloc {
246 wasi_ctx: WasiCtx,
247 alloc: Option<WasiAlloc>,
248}
249
250impl WasiCtxAlloc {
251 fn alloc_buffer(&self) -> TypedFunc<u32, u32> {
252 self.alloc
253 .expect("allocator has been not initialized, cannot allocate buffer!")
254 .alloc_buffer
255 }
256
257 fn free_buffer(&self) -> TypedFunc<u64, ()> {
258 self.alloc
259 .expect("allocator has been not initialized, cannot free buffer!")
260 .free_buffer
261 }
262
263 fn init_alloc(&mut self, alloc: WasiAlloc) {
264 self.alloc = Some(alloc)
265 }
266}
267
268/// Represents a WebAssembly plugin, with access to the WebAssembly System Inferface.
269/// Build a new plugin using [`PluginBuilder`].
270pub struct Plugin {
271 engine: Engine,
272 module: Module,
273 store: Store<WasiCtxAlloc>,
274 instance: Instance,
275}
276
277impl Plugin {
278 /// Dumps the *entirety* of Wasm linear memory to `stdout`.
279 /// Don't call this unless you're debugging a memory issue!
280 pub fn dump_memory(data: &[u8]) {
281 for (i, byte) in data.iter().enumerate() {
282 if i % 32 == 0 {
283 println!();
284 }
285 if i % 4 == 0 {
286 print!("|");
287 }
288 if *byte == 0 {
289 print!("__")
290 } else {
291 print!("{:02x}", byte);
292 }
293 }
294 println!();
295 }
296}
297
298impl Plugin {
299 async fn init(
300 precompiled: bool,
301 module: Vec<u8>,
302 plugin: PluginBuilder,
303 ) -> Result<Self, Error> {
304 // initialize the WebAssembly System Interface context
305 let engine = plugin.engine;
306 let mut linker = plugin.linker;
307 wasmtime_wasi::add_to_linker(&mut linker, |s| &mut s.wasi_ctx)?;
308
309 // create a store, note that we can't initialize the allocator,
310 // because we can't grab the functions until initialized.
311 let mut store: Store<WasiCtxAlloc> = Store::new(
312 &engine,
313 WasiCtxAlloc {
314 wasi_ctx: plugin.wasi_ctx,
315 alloc: None,
316 },
317 );
318 // store.epoch_deadline_async_yield_and_update(todo!());
319 let module = if precompiled {
320 unsafe { Module::deserialize(&engine, module)? }
321 } else {
322 Module::new(&engine, module)?
323 };
324
325 // load the provided module into the asynchronous runtime
326 linker.module_async(&mut store, "", &module).await?;
327 let instance = linker.instantiate_async(&mut store, &module).await?;
328
329 // now that the module is initialized,
330 // we can initialize the store's allocator
331 let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
332 let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
333 store.data_mut().init_alloc(WasiAlloc {
334 alloc_buffer,
335 free_buffer,
336 });
337
338 Ok(Plugin {
339 engine,
340 module,
341 store,
342 instance,
343 })
344 }
345
346 /// Attaches a file or directory the the given system path to the runtime.
347 /// Note that the resource must be freed by calling `remove_resource` afterwards.
348 pub fn attach_path<T: AsRef<Path>>(&mut self, path: T) -> Result<PluginResource, Error> {
349 // grab the WASI context
350 let ctx = self.store.data_mut();
351
352 // open the file we want, and convert it into the right type
353 // this is a footgun and a half
354 let file = File::open(&path).unwrap();
355 let dir = Dir::from_std_file(file);
356 let dir = Box::new(wasmtime_wasi::dir::Dir::from_cap_std(dir));
357
358 // grab an empty file descriptor, specify capabilities
359 let fd = ctx.wasi_ctx.table().push(Box::new(()))?;
360 let caps = dir::DirCaps::all();
361 let file_caps = file::FileCaps::all();
362
363 // insert the directory at the given fd,
364 // return a handle to the resource
365 ctx.wasi_ctx
366 .insert_dir(fd, dir, caps, file_caps, path.as_ref().to_path_buf());
367 Ok(PluginResource(fd))
368 }
369
370 /// Returns `true` if the resource existed and was removed.
371 /// Currently the only resource we support is adding scoped paths (e.g. folders and files)
372 /// to plugins using [`attach_path`].
373 pub fn remove_resource(&mut self, resource: PluginResource) -> Result<(), Error> {
374 self.store
375 .data_mut()
376 .wasi_ctx
377 .table()
378 .delete(resource.0)
379 .ok_or_else(|| anyhow!("Resource did not exist, but a valid handle was passed in"))?;
380 Ok(())
381 }
382
383 // So this call function is kinda a dance, I figured it'd be a good idea to document it.
384 // the high level is we take a serde type, serialize it to a byte array,
385 // (we're doing this using bincode for now)
386 // then toss that byte array into webassembly.
387 // webassembly grabs that byte array, does some magic,
388 // and serializes the result into yet another byte array.
389 // we then grab *that* result byte array and deserialize it into a result.
390 //
391 // phew...
392 //
393 // now the problem is, webassambly doesn't support buffers.
394 // only really like i32s, that's it (yeah, it's sad. Not even unsigned!)
395 // (ok, I'm exaggerating a bit).
396 //
397 // the Wasm function that this calls must have a very specific signature:
398 //
399 // fn(pointer to byte array: i32, length of byte array: i32)
400 // -> pointer to (
401 // pointer to byte_array: i32,
402 // length of byte array: i32,
403 // ): i32
404 //
405 // This pair `(pointer to byte array, length of byte array)` is called a `Buffer`
406 // and can be found in the cargo_test plugin.
407 //
408 // so on the wasm side, we grab the two parameters to the function,
409 // stuff them into a `Buffer`,
410 // and then pray to the `unsafe` Rust gods above that a valid byte array pops out.
411 //
412 // On the flip side, when returning from a wasm function,
413 // we convert whatever serialized result we get into byte array,
414 // which we stuff into a Buffer and allocate on the heap,
415 // which pointer to we then return.
416 // Note the double indirection!
417 //
418 // So when returning from a function, we actually leak memory *twice*:
419 //
420 // 1) once when we leak the byte array
421 // 2) again when we leak the allocated `Buffer`
422 //
423 // This isn't a problem because Wasm stops executing after the function returns,
424 // so the heap is still valid for our inspection when we want to pull things out.
425
426 /// Serializes a given type to bytes.
427 fn serialize_to_bytes<A: Serialize>(item: A) -> Result<Vec<u8>, Error> {
428 // serialize the argument using bincode
429 let bytes = bincode::serialize(&item)?;
430 Ok(bytes)
431 }
432
433 /// Deserializes a given type from bytes.
434 fn deserialize_to_type<R: DeserializeOwned>(bytes: &[u8]) -> Result<R, Error> {
435 // serialize the argument using bincode
436 let bytes = bincode::deserialize(bytes)?;
437 Ok(bytes)
438 }
439
440 // fn deserialize<R: DeserializeOwned>(
441 // plugin_memory: &mut Memory,
442 // mut store: impl AsContextMut<Data = WasiCtxAlloc>,
443 // buffer: WasiBuffer,
444 // ) -> Result<R, Error> {
445 // let buffer_start = buffer.ptr as usize;
446 // let buffer_end = buffer_start + buffer.len as usize;
447
448 // // read the buffer at this point into a byte array
449 // // deserialize the byte array into the provided serde type
450 // let item = &plugin_memory.data(store.as_context())[buffer_start..buffer_end];
451 // let item = bincode::deserialize(bytes)?;
452 // Ok(item)
453 // }
454
455 /// Takes an item, allocates a buffer, serializes the argument to that buffer,
456 /// and returns a (ptr, len) pair to that buffer.
457 async fn bytes_to_buffer(
458 alloc_buffer: TypedFunc<u32, u32>,
459 plugin_memory: &mut Memory,
460 mut store: impl AsContextMut<Data = WasiCtxAlloc>,
461 item: Vec<u8>,
462 ) -> Result<WasiBuffer, Error> {
463 // allocate a buffer and write the argument to that buffer
464 let len = item.len() as u32;
465 let ptr = alloc_buffer.call_async(&mut store, len).await?;
466 plugin_memory.write(&mut store, ptr as usize, &item)?;
467 Ok(WasiBuffer { ptr, len })
468 }
469
470 /// Takes a `(ptr, len)` pair and returns the corresponding deserialized buffer.
471 fn buffer_to_type<R: DeserializeOwned>(
472 plugin_memory: &Memory,
473 store: impl AsContext<Data = WasiCtxAlloc>,
474 buffer: &WasiBuffer,
475 ) -> Result<R, Error> {
476 let buffer_start = buffer.ptr as usize;
477 let buffer_end = buffer_start + buffer.len as usize;
478
479 // read the buffer at this point into a byte array
480 // deserialize the byte array into the provided serde type
481 let result = &plugin_memory.data(store.as_context())[buffer_start..buffer_end];
482 let result = bincode::deserialize(result)?;
483
484 Ok(result)
485 }
486
487 // TODO: don't allocate a new `Vec`!
488 /// Takes a `(ptr, len)` pair and returns the corresponding deserialized buffer.
489 fn buffer_to_bytes<'a>(
490 plugin_memory: &'a Memory,
491 store: impl AsContext<Data = WasiCtxAlloc> + 'a,
492 buffer: &WasiBuffer,
493 ) -> Result<Vec<u8>, Error> {
494 let buffer_start = buffer.ptr as usize;
495 let buffer_end = buffer_start + buffer.len as usize;
496
497 // read the buffer at this point into a byte array
498 // deserialize the byte array into the provided serde type
499 let result = plugin_memory.data(store.as_context())[buffer_start..buffer_end].to_vec();
500 Ok(result)
501 }
502
503 async fn buffer_to_free(
504 free_buffer: TypedFunc<u64, ()>,
505 mut store: impl AsContextMut<Data = WasiCtxAlloc>,
506 buffer: WasiBuffer,
507 ) -> Result<(), Error> {
508 // deallocate the argument buffer
509 Ok(free_buffer
510 .call_async(&mut store, buffer.into_u64())
511 .await?)
512 }
513
514 /// Retrieves the handle to a function of a given type.
515 pub fn function<A: Serialize, R: DeserializeOwned, T: AsRef<str>>(
516 &mut self,
517 name: T,
518 ) -> Result<WasiFn<A, R>, Error> {
519 let fun_name = format!("__{}", name.as_ref());
520 let fun = self
521 .instance
522 .get_typed_func::<u64, u64, _>(&mut self.store, &fun_name)?;
523 Ok(WasiFn {
524 function: fun,
525 _function_type: PhantomData,
526 })
527 }
528
529 // TODO: dont' use as for conversions
530 /// Asynchronously calls a function defined Guest-side.
531 pub async fn call<A: Serialize, R: DeserializeOwned>(
532 &mut self,
533 handle: &WasiFn<A, R>,
534 arg: A,
535 ) -> Result<R, Error> {
536 let mut plugin_memory = self
537 .instance
538 .get_memory(&mut self.store, "memory")
539 .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
540
541 // write the argument to linear memory
542 // this returns a (ptr, lentgh) pair
543 let arg_buffer = Self::bytes_to_buffer(
544 self.store.data().alloc_buffer(),
545 &mut plugin_memory,
546 &mut self.store,
547 Self::serialize_to_bytes(arg)?,
548 )
549 .await?;
550
551 // call the function, passing in the buffer and its length
552 // this returns a ptr to a (ptr, lentgh) pair
553 let result_buffer = handle
554 .function
555 .call_async(&mut self.store, arg_buffer.into_u64())
556 .await?;
557
558 Self::buffer_to_type(
559 &mut plugin_memory,
560 &mut self.store,
561 &WasiBuffer::from_u64(result_buffer),
562 )
563 }
564}