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 let engine = Engine::new(&config)?;
109 let linker = Linker::new(&engine);
110
111 match yield_when {
112 PluginYield::Epoch { .. } => {
113 config.epoch_interruption(true);
114 }
115 PluginYield::Fuel(_) => {
116 config.consume_fuel(true);
117 }
118 }
119
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<T: AsRef<[u8]>>(self, precompiled: bool, module: T) -> Result<Plugin, Error> {
349 Plugin::init(precompiled, module.as_ref(), 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
382/// Represents a WebAssembly plugin, with access to the WebAssembly System Inferface.
383/// Build a new plugin using [`PluginBuilder`].
384pub struct Plugin {
385 store: Store<WasiCtxAlloc>,
386 instance: Instance,
387}
388
389impl Plugin {
390 /// Dumps the *entirety* of Wasm linear memory to `stdout`.
391 /// Don't call this unless you're debugging a memory issue!
392 pub fn dump_memory(data: &[u8]) {
393 for (i, byte) in data.iter().enumerate() {
394 if i % 32 == 0 {
395 println!();
396 }
397 if i % 4 == 0 {
398 print!("|");
399 }
400 if *byte == 0 {
401 print!("__")
402 } else {
403 print!("{:02x}", byte);
404 }
405 }
406 println!();
407 }
408
409 async fn init(precompiled: bool, module: &[u8], plugin: PluginBuilder) -> Result<Self, Error> {
410 // initialize the WebAssembly System Interface context
411 let engine = plugin.engine;
412 let mut linker = plugin.linker;
413 wasmtime_wasi::add_to_linker(&mut linker, |s| &mut s.wasi_ctx)?;
414
415 // create a store, note that we can't initialize the allocator,
416 // because we can't grab the functions until initialized.
417 let mut store: Store<WasiCtxAlloc> = Store::new(
418 &engine,
419 WasiCtxAlloc {
420 wasi_ctx: plugin.wasi_ctx,
421 alloc: None,
422 },
423 );
424
425 let module = if precompiled {
426 unsafe { Module::deserialize(&engine, module)? }
427 } else {
428 Module::new(&engine, module)?
429 };
430
431 // set up automatic yielding based on configuration
432 match plugin.yield_when {
433 PluginYield::Epoch {
434 yield_epoch: PluginYieldEpoch { delta, .. },
435 initialize_incrementer,
436 } => {
437 store.epoch_deadline_async_yield_and_update(delta);
438 initialize_incrementer(engine);
439 }
440 PluginYield::Fuel(PluginYieldFuel { initial, refill }) => {
441 store.add_fuel(initial).unwrap();
442 store.out_of_fuel_async_yield(u64::MAX, refill);
443 }
444 }
445
446 // load the provided module into the asynchronous runtime
447 linker.module_async(&mut store, "", &module).await?;
448 let instance = linker.instantiate_async(&mut store, &module).await?;
449
450 // now that the module is initialized,
451 // we can initialize the store's allocator
452 let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
453 let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
454 store.data_mut().init_alloc(WasiAlloc {
455 alloc_buffer,
456 free_buffer,
457 });
458
459 Ok(Plugin { store, instance })
460 }
461
462 /// Attaches a file or directory the the given system path to the runtime.
463 /// Note that the resource must be freed by calling `remove_resource` afterwards.
464 pub fn attach_path<T: AsRef<Path>>(&mut self, path: T) -> Result<PluginResource, Error> {
465 // grab the WASI context
466 let ctx = self.store.data_mut();
467
468 // open the file we want, and convert it into the right type
469 // this is a footgun and a half
470 let file = File::open(&path).unwrap();
471 let dir = Dir::from_std_file(file);
472 let dir = Box::new(wasmtime_wasi::dir::Dir::from_cap_std(dir));
473
474 // grab an empty file descriptor, specify capabilities
475 let fd = ctx.wasi_ctx.table().push(Box::new(()))?;
476 let caps = dir::DirCaps::all();
477 let file_caps = file::FileCaps::all();
478
479 // insert the directory at the given fd,
480 // return a handle to the resource
481 ctx.wasi_ctx
482 .insert_dir(fd, dir, caps, file_caps, path.as_ref().to_path_buf());
483 Ok(PluginResource(fd))
484 }
485
486 /// Returns `true` if the resource existed and was removed.
487 /// Currently the only resource we support is adding scoped paths (e.g. folders and files)
488 /// to plugins using [`attach_path`].
489 pub fn remove_resource(&mut self, resource: PluginResource) -> Result<(), Error> {
490 self.store
491 .data_mut()
492 .wasi_ctx
493 .table()
494 .delete(resource.0)
495 .ok_or_else(|| anyhow!("Resource did not exist, but a valid handle was passed in"))?;
496 Ok(())
497 }
498
499 // So this call function is kinda a dance, I figured it'd be a good idea to document it.
500 // the high level is we take a serde type, serialize it to a byte array,
501 // (we're doing this using bincode for now)
502 // then toss that byte array into webassembly.
503 // webassembly grabs that byte array, does some magic,
504 // and serializes the result into yet another byte array.
505 // we then grab *that* result byte array and deserialize it into a result.
506 //
507 // phew...
508 //
509 // now the problem is, webassambly doesn't support buffers.
510 // only really like i32s, that's it (yeah, it's sad. Not even unsigned!)
511 // (ok, I'm exaggerating a bit).
512 //
513 // the Wasm function that this calls must have a very specific signature:
514 //
515 // fn(pointer to byte array: i32, length of byte array: i32)
516 // -> pointer to (
517 // pointer to byte_array: i32,
518 // length of byte array: i32,
519 // ): i32
520 //
521 // This pair `(pointer to byte array, length of byte array)` is called a `Buffer`
522 // and can be found in the cargo_test plugin.
523 //
524 // so on the wasm side, we grab the two parameters to the function,
525 // stuff them into a `Buffer`,
526 // and then pray to the `unsafe` Rust gods above that a valid byte array pops out.
527 //
528 // On the flip side, when returning from a wasm function,
529 // we convert whatever serialized result we get into byte array,
530 // which we stuff into a Buffer and allocate on the heap,
531 // which pointer to we then return.
532 // Note the double indirection!
533 //
534 // So when returning from a function, we actually leak memory *twice*:
535 //
536 // 1) once when we leak the byte array
537 // 2) again when we leak the allocated `Buffer`
538 //
539 // This isn't a problem because Wasm stops executing after the function returns,
540 // so the heap is still valid for our inspection when we want to pull things out.
541
542 /// Serializes a given type to bytes.
543 fn serialize_to_bytes<A: Serialize>(item: A) -> Result<Vec<u8>, Error> {
544 // serialize the argument using bincode
545 let bytes = bincode::serialize(&item)?;
546 Ok(bytes)
547 }
548
549 /// Deserializes a given type from bytes.
550 fn deserialize_to_type<R: DeserializeOwned>(bytes: &[u8]) -> Result<R, Error> {
551 // serialize the argument using bincode
552 let bytes = bincode::deserialize(bytes)?;
553 Ok(bytes)
554 }
555
556 // fn deserialize<R: DeserializeOwned>(
557 // plugin_memory: &mut Memory,
558 // mut store: impl AsContextMut<Data = WasiCtxAlloc>,
559 // buffer: WasiBuffer,
560 // ) -> Result<R, Error> {
561 // let buffer_start = buffer.ptr as usize;
562 // let buffer_end = buffer_start + buffer.len as usize;
563
564 // // read the buffer at this point into a byte array
565 // // deserialize the byte array into the provided serde type
566 // let item = &plugin_memory.data(store.as_context())[buffer_start..buffer_end];
567 // let item = bincode::deserialize(bytes)?;
568 // Ok(item)
569 // }
570
571 /// Takes an item, allocates a buffer, serializes the argument to that buffer,
572 /// and returns a (ptr, len) pair to that buffer.
573 async fn bytes_to_buffer(
574 alloc_buffer: TypedFunc<u32, u32>,
575 plugin_memory: &mut Memory,
576 mut store: impl AsContextMut<Data = WasiCtxAlloc>,
577 item: Vec<u8>,
578 ) -> Result<WasiBuffer, Error> {
579 // allocate a buffer and write the argument to that buffer
580 let len = item.len() as u32;
581 let ptr = alloc_buffer.call_async(&mut store, len).await?;
582 plugin_memory.write(&mut store, ptr as usize, &item)?;
583 Ok(WasiBuffer { ptr, len })
584 }
585
586 /// Takes a `(ptr, len)` pair and returns the corresponding deserialized buffer.
587 fn buffer_to_type<R: DeserializeOwned>(
588 plugin_memory: &Memory,
589 store: impl AsContext<Data = WasiCtxAlloc>,
590 buffer: &WasiBuffer,
591 ) -> Result<R, Error> {
592 let buffer_start = buffer.ptr as usize;
593 let buffer_end = buffer_start + buffer.len as usize;
594
595 // read the buffer at this point into a byte array
596 // deserialize the byte array into the provided serde type
597 let result = &plugin_memory.data(store.as_context())[buffer_start..buffer_end];
598 let result = bincode::deserialize(result)?;
599
600 Ok(result)
601 }
602
603 /// Takes a `(ptr, len)` pair and returns the corresponding deserialized buffer.
604 fn buffer_to_bytes<'a>(
605 plugin_memory: &'a Memory,
606 store: wasmtime::StoreContext<'a, WasiCtxAlloc>,
607 buffer: &'a WasiBuffer,
608 ) -> Result<&'a [u8], Error> {
609 let buffer_start = buffer.ptr as usize;
610 let buffer_end = buffer_start + buffer.len as usize;
611
612 // read the buffer at this point into a byte array
613 // deserialize the byte array into the provided serde type
614 let result = &plugin_memory.data(store)[buffer_start..buffer_end];
615 Ok(result)
616 }
617
618 async fn buffer_to_free(
619 free_buffer: TypedFunc<u64, ()>,
620 mut store: impl AsContextMut<Data = WasiCtxAlloc>,
621 buffer: WasiBuffer,
622 ) -> Result<(), Error> {
623 // deallocate the argument buffer
624 Ok(free_buffer
625 .call_async(&mut store, buffer.into_u64())
626 .await?)
627 }
628
629 /// Retrieves the handle to a function of a given type.
630 pub fn function<A: Serialize, R: DeserializeOwned, T: AsRef<str>>(
631 &mut self,
632 name: T,
633 ) -> Result<WasiFn<A, R>, Error> {
634 let fun_name = format!("__{}", name.as_ref());
635 let fun = self
636 .instance
637 .get_typed_func::<u64, u64, _>(&mut self.store, &fun_name)?;
638 Ok(WasiFn {
639 function: fun,
640 _function_type: PhantomData,
641 })
642 }
643
644 /// Asynchronously calls a function defined Guest-side.
645 pub async fn call<A: Serialize, R: DeserializeOwned>(
646 &mut self,
647 handle: &WasiFn<A, R>,
648 arg: A,
649 ) -> Result<R, Error> {
650 let mut plugin_memory = self
651 .instance
652 .get_memory(&mut self.store, "memory")
653 .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
654
655 // write the argument to linear memory
656 // this returns a (ptr, lentgh) pair
657 let arg_buffer = Self::bytes_to_buffer(
658 self.store.data().alloc_buffer(),
659 &mut plugin_memory,
660 &mut self.store,
661 Self::serialize_to_bytes(arg)?,
662 )
663 .await?;
664
665 // call the function, passing in the buffer and its length
666 // this returns a ptr to a (ptr, lentgh) pair
667 let result_buffer = handle
668 .function
669 .call_async(&mut self.store, arg_buffer.into_u64())
670 .await?;
671
672 Self::buffer_to_type(
673 &mut plugin_memory,
674 &mut self.store,
675 &WasiBuffer::from_u64(result_buffer),
676 )
677 }
678}