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