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