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