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