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