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