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