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