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