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