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