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