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