1 //! Generate a Wasm program that keeps track of its current stack frames. 2 //! 3 //! We can then compare the stack trace we observe in Wasmtime to what the Wasm 4 //! program believes its stack should be. Any discrepencies between the two 5 //! points to a bug in either this test case generator or Wasmtime's stack 6 //! walker. 7 8 use std::mem; 9 10 use arbitrary::{Arbitrary, Result, Unstructured}; 11 use wasm_encoder::Instruction; 12 13 const MAX_FUNCS: usize = 20; 14 const MAX_OPS: usize = 1_000; 15 16 /// Generate a Wasm module that keeps track of its current call stack, to 17 /// compare to the host. 18 #[derive(Debug)] 19 pub struct Stacks { 20 funcs: Vec<Function>, 21 inputs: Vec<u8>, 22 } 23 24 #[derive(Debug, Default)] 25 struct Function { 26 ops: Vec<Op>, 27 } 28 29 #[derive(Arbitrary, Debug, Clone, Copy)] 30 enum Op { 31 CheckStackInHost, 32 Call(u32), 33 CallThroughHost(u32), 34 } 35 36 impl<'a> Arbitrary<'a> for Stacks { 37 fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> { 38 let funcs = Self::arbitrary_funcs(u)?; 39 let n = u.len().min(200); 40 let inputs = u.bytes(n)?.to_vec(); 41 Ok(Stacks { funcs, inputs }) 42 } 43 } 44 45 impl Stacks { 46 fn arbitrary_funcs(u: &mut Unstructured) -> Result<Vec<Function>> { 47 let mut funcs = vec![Function::default()]; 48 49 // The indices of functions within `funcs` that we still need to 50 // generate. 51 let mut work_list = vec![0]; 52 53 while let Some(f) = work_list.pop() { 54 let mut ops = Vec::with_capacity(u.arbitrary_len::<Op>()?.min(MAX_OPS)); 55 for _ in 0..ops.capacity() { 56 ops.push(u.arbitrary()?); 57 } 58 for op in &mut ops { 59 match op { 60 Op::CallThroughHost(idx) | Op::Call(idx) => { 61 if u.is_empty() || funcs.len() >= MAX_FUNCS || u.ratio(4, 5)? { 62 // Call an existing function. 63 *idx = *idx % u32::try_from(funcs.len()).unwrap(); 64 } else { 65 // Call a new function... 66 *idx = u32::try_from(funcs.len()).unwrap(); 67 // ...which means we also need to eventually define it. 68 work_list.push(funcs.len()); 69 funcs.push(Function::default()); 70 } 71 } 72 Op::CheckStackInHost => {} 73 } 74 } 75 funcs[f].ops = ops; 76 } 77 78 Ok(funcs) 79 } 80 81 /// Get the input values to run the Wasm module with. 82 pub fn inputs(&self) -> &[u8] { 83 &self.inputs 84 } 85 86 /// Get this test case's Wasm module. 87 /// 88 /// The Wasm module has the following imports: 89 /// 90 /// * `host.check_stack: [] -> []`: The host can check the Wasm's 91 /// understanding of its own stack against the host's understanding of the 92 /// Wasm stack to find discrepency bugs. 93 /// 94 /// * `host.call_func: [funcref] -> []`: The host should call the given 95 /// `funcref`, creating a call stack with multiple sequences of contiguous 96 /// Wasm frames on the stack like `[..., wasm, host, wasm]`. 97 /// 98 /// The Wasm module has the following exports: 99 /// 100 /// * `run: [i32] -> []`: This function should be called with each of the 101 /// input values to run this generated test case. 102 /// 103 /// * `get_stack: [] -> [i32 i32]`: Get the pointer and length of the `u32` 104 /// array of this Wasm's understanding of its stack. This is useful for 105 /// checking whether the host's view of the stack at a trap matches the 106 /// Wasm program's understanding. 107 pub fn wasm(&self) -> Vec<u8> { 108 let mut module = wasm_encoder::Module::new(); 109 110 let mut types = wasm_encoder::TypeSection::new(); 111 112 let run_type = types.len(); 113 types.function(vec![wasm_encoder::ValType::I32], vec![]); 114 115 let get_stack_type = types.len(); 116 types.function( 117 vec![], 118 vec![wasm_encoder::ValType::I32, wasm_encoder::ValType::I32], 119 ); 120 121 let null_type = types.len(); 122 types.function(vec![], vec![]); 123 124 let call_func_type = types.len(); 125 types.function(vec![wasm_encoder::ValType::FUNCREF], vec![]); 126 127 section(&mut module, types); 128 129 let mut imports = wasm_encoder::ImportSection::new(); 130 let check_stack_func = 0; 131 imports.import( 132 "host", 133 "check_stack", 134 wasm_encoder::EntityType::Function(null_type), 135 ); 136 let call_func_func = 1; 137 imports.import( 138 "host", 139 "call_func", 140 wasm_encoder::EntityType::Function(call_func_type), 141 ); 142 let num_imported_funcs = 2; 143 section(&mut module, imports); 144 145 let mut funcs = wasm_encoder::FunctionSection::new(); 146 for _ in &self.funcs { 147 funcs.function(null_type); 148 } 149 let run_func = funcs.len() + num_imported_funcs; 150 funcs.function(run_type); 151 let get_stack_func = funcs.len() + num_imported_funcs; 152 funcs.function(get_stack_type); 153 section(&mut module, funcs); 154 155 let mut mems = wasm_encoder::MemorySection::new(); 156 let memory = mems.len(); 157 mems.memory(wasm_encoder::MemoryType { 158 minimum: 1, 159 maximum: Some(1), 160 memory64: false, 161 shared: false, 162 }); 163 section(&mut module, mems); 164 165 let mut globals = wasm_encoder::GlobalSection::new(); 166 let fuel_global = globals.len(); 167 globals.global( 168 wasm_encoder::GlobalType { 169 val_type: wasm_encoder::ValType::I32, 170 mutable: true, 171 }, 172 &wasm_encoder::ConstExpr::i32_const(0), 173 ); 174 let stack_len_global = globals.len(); 175 globals.global( 176 wasm_encoder::GlobalType { 177 val_type: wasm_encoder::ValType::I32, 178 mutable: true, 179 }, 180 &wasm_encoder::ConstExpr::i32_const(0), 181 ); 182 section(&mut module, globals); 183 184 let mut exports = wasm_encoder::ExportSection::new(); 185 exports.export("run", wasm_encoder::ExportKind::Func, run_func); 186 exports.export("get_stack", wasm_encoder::ExportKind::Func, get_stack_func); 187 exports.export("memory", wasm_encoder::ExportKind::Memory, memory); 188 exports.export("fuel", wasm_encoder::ExportKind::Global, fuel_global); 189 section(&mut module, exports); 190 191 let mut elems = wasm_encoder::ElementSection::new(); 192 elems.declared(wasm_encoder::Elements::Functions( 193 &(0..num_imported_funcs + u32::try_from(self.funcs.len()).unwrap()).collect::<Vec<_>>(), 194 )); 195 section(&mut module, elems); 196 197 let check_fuel = |body: &mut wasm_encoder::Function| { 198 // Trap if we are out of fuel. 199 body.instruction(&Instruction::GlobalGet(fuel_global)) 200 .instruction(&Instruction::I32Eqz) 201 .instruction(&Instruction::If(wasm_encoder::BlockType::Empty)) 202 .instruction(&Instruction::Unreachable) 203 .instruction(&Instruction::End); 204 205 // Decrement fuel. 206 body.instruction(&Instruction::GlobalGet(fuel_global)) 207 .instruction(&Instruction::I32Const(1)) 208 .instruction(&Instruction::I32Sub) 209 .instruction(&Instruction::GlobalSet(fuel_global)); 210 }; 211 212 let push_func_to_stack = |body: &mut wasm_encoder::Function, func: u32| { 213 // Add this function to our internal stack. 214 // 215 // Note that we know our `stack_len_global` can't go beyond memory 216 // bounds because we limit fuel to at most `u8::MAX` and each stack 217 // entry is an `i32` and `u8::MAX * size_of(i32)` still fits in one 218 // Wasm page. 219 body.instruction(&Instruction::GlobalGet(stack_len_global)) 220 .instruction(&Instruction::I32Const(func as i32)) 221 .instruction(&Instruction::I32Store(wasm_encoder::MemArg { 222 offset: 0, 223 align: 0, 224 memory_index: memory, 225 })) 226 .instruction(&Instruction::GlobalGet(stack_len_global)) 227 .instruction(&Instruction::I32Const(mem::size_of::<i32>() as i32)) 228 .instruction(&Instruction::I32Add) 229 .instruction(&Instruction::GlobalSet(stack_len_global)); 230 }; 231 232 let pop_func_from_stack = |body: &mut wasm_encoder::Function| { 233 // Remove this function from our internal stack. 234 body.instruction(&Instruction::GlobalGet(stack_len_global)) 235 .instruction(&Instruction::I32Const(mem::size_of::<i32>() as i32)) 236 .instruction(&Instruction::I32Sub) 237 .instruction(&Instruction::GlobalSet(stack_len_global)); 238 }; 239 240 let mut code = wasm_encoder::CodeSection::new(); 241 for (func_index, func) in self.funcs.iter().enumerate() { 242 let mut body = wasm_encoder::Function::new(vec![]); 243 244 push_func_to_stack( 245 &mut body, 246 num_imported_funcs + u32::try_from(func_index).unwrap(), 247 ); 248 check_fuel(&mut body); 249 250 // Perform our specified operations. 251 for op in &func.ops { 252 match op { 253 Op::CheckStackInHost => { 254 body.instruction(&Instruction::Call(check_stack_func)); 255 } 256 Op::Call(f) => { 257 body.instruction(&Instruction::Call(f + num_imported_funcs)); 258 } 259 Op::CallThroughHost(f) => { 260 body.instruction(&Instruction::RefFunc(f + num_imported_funcs)) 261 .instruction(&Instruction::Call(call_func_func)); 262 } 263 } 264 } 265 266 // Potentially trap at the end of our function as well, so that we 267 // exercise the scenario where the Wasm-to-host trampoline 268 // initialized `last_wasm_exit_sp` et al when calling out to a host 269 // function, but then we returned back to Wasm and then trapped 270 // while `last_wasm_exit_sp` et al are still initialized from that 271 // previous host call. 272 check_fuel(&mut body); 273 274 pop_func_from_stack(&mut body); 275 276 function(&mut code, body); 277 } 278 279 let mut run_body = wasm_encoder::Function::new(vec![]); 280 281 // Reset the bump pointer for the internal stack (this allows us to 282 // reuse an instance in the oracle, rather than re-instantiate). 283 run_body 284 .instruction(&Instruction::I32Const(0)) 285 .instruction(&Instruction::GlobalSet(stack_len_global)); 286 287 // Initialize the fuel global. 288 run_body 289 .instruction(&Instruction::LocalGet(0)) 290 .instruction(&Instruction::GlobalSet(fuel_global)); 291 292 push_func_to_stack(&mut run_body, run_func); 293 294 // Make sure to check for out-of-fuel in the `run` function as well, so 295 // that we also capture stack traces with only one frame, not just `run` 296 // followed by the first locally-defined function and then zero or more 297 // extra frames. 298 check_fuel(&mut run_body); 299 300 // Call the first locally defined function. 301 run_body.instruction(&Instruction::Call(num_imported_funcs)); 302 303 check_fuel(&mut run_body); 304 pop_func_from_stack(&mut run_body); 305 306 function(&mut code, run_body); 307 308 let mut get_stack_body = wasm_encoder::Function::new(vec![]); 309 get_stack_body 310 .instruction(&Instruction::I32Const(0)) 311 .instruction(&Instruction::GlobalGet(stack_len_global)); 312 function(&mut code, get_stack_body); 313 314 section(&mut module, code); 315 316 return module.finish(); 317 318 // Helper that defines a section in the module and takes ownership of it 319 // so that it is dropped and its memory reclaimed after adding it to the 320 // module. 321 fn section(module: &mut wasm_encoder::Module, section: impl wasm_encoder::Section) { 322 module.section(§ion); 323 } 324 325 // Helper that defines a function body in the code section and takes 326 // ownership of it so that it is dropped and its memory reclaimed after 327 // adding it to the module. 328 fn function(code: &mut wasm_encoder::CodeSection, mut func: wasm_encoder::Function) { 329 func.instruction(&Instruction::End); 330 code.function(&func); 331 } 332 } 333 } 334 335 #[cfg(test)] 336 mod tests { 337 use super::*; 338 use rand::prelude::*; 339 use wasmparser::Validator; 340 341 #[test] 342 fn stacks_generates_valid_wasm_modules() { 343 let mut rng = SmallRng::seed_from_u64(0); 344 let mut buf = vec![0; 2048]; 345 for _ in 0..1024 { 346 rng.fill_bytes(&mut buf); 347 let u = Unstructured::new(&buf); 348 if let Ok(stacks) = Stacks::arbitrary_take_rest(u) { 349 let wasm = stacks.wasm(); 350 validate(&wasm); 351 } 352 } 353 } 354 355 fn validate(wasm: &[u8]) { 356 let mut validator = Validator::new(); 357 let err = match validator.validate_all(wasm) { 358 Ok(_) => return, 359 Err(e) => e, 360 }; 361 drop(std::fs::write("test.wasm", wasm)); 362 if let Ok(text) = wasmprinter::print_bytes(wasm) { 363 drop(std::fs::write("test.wat", &text)); 364 } 365 panic!("wasm failed to validate: {}", err); 366 } 367 } 368