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