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