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