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