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