1use std::{collections::HashMap, fs::File, path::Path};
2
3use anyhow::{anyhow, Error};
4use serde::{de::DeserializeOwned, Serialize};
5
6use wasmtime::{Engine, Func, Instance, Linker, Memory, MemoryType, Module, Store, TypedFunc};
7use wasmtime_wasi::{dir, Dir, WasiCtx, WasiCtxBuilder};
8
9pub struct Wasi {
10 engine: Engine,
11 module: Module,
12 store: Store<WasiCtx>,
13 instance: Instance,
14 alloc_buffer: TypedFunc<i32, i32>,
15 // free_buffer: TypedFunc<(i32, i32), ()>,
16}
17
18pub struct WasiPlugin {
19 pub module: Vec<u8>,
20 pub wasi_ctx: WasiCtx,
21}
22
23impl Wasi {
24 pub fn dump_memory(data: &[u8]) {
25 for (i, byte) in data.iter().enumerate() {
26 if i % 32 == 0 {
27 println!();
28 }
29 if i % 4 == 0 {
30 print!("|");
31 }
32 if *byte == 0 {
33 print!("__")
34 } else {
35 print!("{:02x}", byte);
36 }
37 }
38 println!();
39 }
40}
41
42impl Wasi {
43 pub fn default_ctx() -> WasiCtx {
44 WasiCtxBuilder::new()
45 .inherit_stdout()
46 .inherit_stderr()
47 .build()
48 }
49
50 pub fn init(plugin: WasiPlugin) -> Result<Self, Error> {
51 let engine = Engine::default();
52 let mut linker = Linker::new(&engine);
53 println!("linking");
54 wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
55 println!("linked");
56 let mut store: Store<_> = Store::new(&engine, plugin.wasi_ctx);
57 println!("moduling");
58 let module = Module::new(&engine, plugin.module)?;
59 println!("moduled");
60
61 linker.module(&mut store, "", &module)?;
62 println!("linked again");
63 let instance = linker.instantiate(&mut store, &module)?;
64 println!("instantiated");
65
66 let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
67 // let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
68 println!("can alloc");
69
70 Ok(Wasi {
71 engine,
72 module,
73 store,
74 instance,
75 alloc_buffer,
76 // free_buffer,
77 })
78 }
79
80 pub fn attach_file<T: AsRef<Path>>(&mut self, path: T) -> Result<(), Error> {
81 let ctx = self.store.data_mut();
82 let file = File::open(&path).unwrap();
83 let dir = Dir::from_std_file(file);
84 // this is a footgun and a half.
85 let dir = dir::Dir::from_cap_std(dir);
86 ctx.push_preopened_dir(Box::new(dir), path)?;
87 Ok(())
88 }
89
90 // pub fn remove_file<T: AsRef<Path>>(&mut self, path: T) -> Result<(), Error> {
91 // let ctx = self.store.data_mut();
92 // ctx.remove
93 // Ok(())
94 // }
95
96 // So this call function is kinda a dance, I figured it'd be a good idea to document it.
97 // the high level is we take a serde type, serialize it to a byte array,
98 // (we're doing this using bincode for now)
99 // then toss that byte array into webassembly.
100 // webassembly grabs that byte array, does some magic,
101 // and serializes the result into yet another byte array.
102 // we then grab *that* result byte array and deserialize it into a result.
103 //
104 // phew...
105 //
106 // now the problem is, webassambly doesn't support buffers.
107 // only really like i32s, that's it (yeah, it's sad. Not even unsigned!)
108 // (ok, I'm exaggerating a bit).
109 //
110 // the Wasm function that this calls must have a very specific signature:
111 //
112 // fn(pointer to byte array: i32, length of byte array: i32)
113 // -> pointer to (
114 // pointer to byte_array: i32,
115 // length of byte array: i32,
116 // ): i32
117 //
118 // This pair `(pointer to byte array, length of byte array)` is called a `Buffer`
119 // and can be found in the cargo_test plugin.
120 //
121 // so on the wasm side, we grab the two parameters to the function,
122 // stuff them into a `Buffer`,
123 // and then pray to the `unsafe` Rust gods above that a valid byte array pops out.
124 //
125 // On the flip side, when returning from a wasm function,
126 // we convert whatever serialized result we get into byte array,
127 // which we stuff into a Buffer and allocate on the heap,
128 // which pointer to we then return.
129 // Note the double indirection!
130 //
131 // So when returning from a function, we actually leak memory *twice*:
132 //
133 // 1) once when we leak the byte array
134 // 2) again when we leak the allocated `Buffer`
135 //
136 // This isn't a problem because Wasm stops executing after the function returns,
137 // so the heap is still valid for our inspection when we want to pull things out.
138
139 // TODO: dont' use as for conversions
140 pub fn call<A: Serialize, R: DeserializeOwned>(
141 &mut self,
142 handle: &str,
143 arg: A,
144 ) -> Result<R, Error> {
145 // serialize the argument using bincode
146 let arg = bincode::serialize(&arg)?;
147 let arg_buffer_len = arg.len();
148
149 // allocate a buffer and write the argument to that buffer
150 let arg_buffer_ptr = self
151 .alloc_buffer
152 .call(&mut self.store, arg_buffer_len as i32)?;
153 let plugin_memory = self
154 .instance
155 .get_memory(&mut self.store, "memory")
156 .ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
157 plugin_memory.write(&mut self.store, arg_buffer_ptr as usize, &arg)?;
158
159 // get the webassembly function we want to actually call
160 // TODO: precompute handle
161 let fun_name = format!("__{}", handle);
162 let fun = self
163 .instance
164 .get_typed_func::<(i32, i32), i32, _>(&mut self.store, &fun_name)?;
165
166 // call the function, passing in the buffer and its length
167 // this should return a pointer to a (ptr, lentgh) pair
168 let arg_buffer = (arg_buffer_ptr, arg_buffer_len as i32);
169 let result_buffer = fun.call(&mut self.store, arg_buffer)?;
170
171 // create a buffer to read the (ptr, length) pair into
172 // this is a total of 4 + 4 = 8 bytes.
173 let buffer = &mut [0; 8];
174 plugin_memory.read(&mut self.store, result_buffer as usize, buffer)?;
175
176 // use these bytes (wasm stores things little-endian)
177 // to get a pointer to the buffer and its length
178 let b = buffer;
179 let result_buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
180 let result_buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
181 let result_buffer_end = result_buffer_ptr + result_buffer_len;
182
183 // read the buffer at this point into a byte array
184 // deserialize the byte array into the provided serde type
185 let result = &plugin_memory.data(&mut self.store)[result_buffer_ptr..result_buffer_end];
186 let result = bincode::deserialize(result)?;
187
188 // TODO: this is handled wasm-side, but I'd like to double-check
189 // // deallocate the argument buffer
190 // self.free_buffer.call(&mut self.store, arg_buffer);
191
192 return Ok(result);
193 }
194}