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