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