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( 193 wasm_encoder::ValType::FuncRef, 194 wasm_encoder::Elements::Functions( 195 &(0..num_imported_funcs + u32::try_from(self.funcs.len()).unwrap()) 196 .collect::<Vec<_>>(), 197 ), 198 ); 199 section(&mut module, elems); 200 201 let check_fuel = |body: &mut wasm_encoder::Function| { 202 // Trap if we are out of fuel. 203 body.instruction(&Instruction::GlobalGet(fuel_global)) 204 .instruction(&Instruction::I32Eqz) 205 .instruction(&Instruction::If(wasm_encoder::BlockType::Empty)) 206 .instruction(&Instruction::Unreachable) 207 .instruction(&Instruction::End); 208 209 // Decrement fuel. 210 body.instruction(&Instruction::GlobalGet(fuel_global)) 211 .instruction(&Instruction::I32Const(1)) 212 .instruction(&Instruction::I32Sub) 213 .instruction(&Instruction::GlobalSet(fuel_global)); 214 }; 215 216 let push_func_to_stack = |body: &mut wasm_encoder::Function, func: u32| { 217 // Add this function to our internal stack. 218 // 219 // Note that we know our `stack_len_global` can't go beyond memory 220 // bounds because we limit fuel to at most `u8::MAX` and each stack 221 // entry is an `i32` and `u8::MAX * size_of(i32)` still fits in one 222 // Wasm page. 223 body.instruction(&Instruction::GlobalGet(stack_len_global)) 224 .instruction(&Instruction::I32Const(func as i32)) 225 .instruction(&Instruction::I32Store(wasm_encoder::MemArg { 226 offset: 0, 227 align: 0, 228 memory_index: memory, 229 })) 230 .instruction(&Instruction::GlobalGet(stack_len_global)) 231 .instruction(&Instruction::I32Const(mem::size_of::<i32>() as i32)) 232 .instruction(&Instruction::I32Add) 233 .instruction(&Instruction::GlobalSet(stack_len_global)); 234 }; 235 236 let pop_func_from_stack = |body: &mut wasm_encoder::Function| { 237 // Remove this function from our internal stack. 238 body.instruction(&Instruction::GlobalGet(stack_len_global)) 239 .instruction(&Instruction::I32Const(mem::size_of::<i32>() as i32)) 240 .instruction(&Instruction::I32Sub) 241 .instruction(&Instruction::GlobalSet(stack_len_global)); 242 }; 243 244 let mut code = wasm_encoder::CodeSection::new(); 245 for (func_index, func) in self.funcs.iter().enumerate() { 246 let mut body = wasm_encoder::Function::new(vec![]); 247 248 push_func_to_stack( 249 &mut body, 250 num_imported_funcs + u32::try_from(func_index).unwrap(), 251 ); 252 check_fuel(&mut body); 253 254 // Perform our specified operations. 255 for op in &func.ops { 256 match op { 257 Op::CheckStackInHost => { 258 body.instruction(&Instruction::Call(check_stack_func)); 259 } 260 Op::Call(f) => { 261 body.instruction(&Instruction::Call(f + num_imported_funcs)); 262 } 263 Op::CallThroughHost(f) => { 264 body.instruction(&Instruction::RefFunc(f + num_imported_funcs)) 265 .instruction(&Instruction::Call(call_func_func)); 266 } 267 } 268 } 269 270 // Potentially trap at the end of our function as well, so that we 271 // exercise the scenario where the Wasm-to-host trampoline 272 // initialized `last_wasm_exit_sp` et al when calling out to a host 273 // function, but then we returned back to Wasm and then trapped 274 // while `last_wasm_exit_sp` et al are still initialized from that 275 // previous host call. 276 check_fuel(&mut body); 277 278 pop_func_from_stack(&mut body); 279 280 function(&mut code, body); 281 } 282 283 let mut run_body = wasm_encoder::Function::new(vec![]); 284 285 // Reset the bump pointer for the internal stack (this allows us to 286 // reuse an instance in the oracle, rather than re-instantiate). 287 run_body 288 .instruction(&Instruction::I32Const(0)) 289 .instruction(&Instruction::GlobalSet(stack_len_global)); 290 291 // Initialize the fuel global. 292 run_body 293 .instruction(&Instruction::LocalGet(0)) 294 .instruction(&Instruction::GlobalSet(fuel_global)); 295 296 push_func_to_stack(&mut run_body, run_func); 297 298 // Make sure to check for out-of-fuel in the `run` function as well, so 299 // that we also capture stack traces with only one frame, not just `run` 300 // followed by the first locally-defined function and then zero or more 301 // extra frames. 302 check_fuel(&mut run_body); 303 304 // Call the first locally defined function. 305 run_body.instruction(&Instruction::Call(num_imported_funcs)); 306 307 check_fuel(&mut run_body); 308 pop_func_from_stack(&mut run_body); 309 310 function(&mut code, run_body); 311 312 let mut get_stack_body = wasm_encoder::Function::new(vec![]); 313 get_stack_body 314 .instruction(&Instruction::I32Const(0)) 315 .instruction(&Instruction::GlobalGet(stack_len_global)); 316 function(&mut code, get_stack_body); 317 318 section(&mut module, code); 319 320 return module.finish(); 321 322 // Helper that defines a section in the module and takes ownership of it 323 // so that it is dropped and its memory reclaimed after adding it to the 324 // module. 325 fn section(module: &mut wasm_encoder::Module, section: impl wasm_encoder::Section) { 326 module.section(§ion); 327 } 328 329 // Helper that defines a function body in the code section and takes 330 // ownership of it so that it is dropped and its memory reclaimed after 331 // adding it to the module. 332 fn function(code: &mut wasm_encoder::CodeSection, mut func: wasm_encoder::Function) { 333 func.instruction(&Instruction::End); 334 code.function(&func); 335 } 336 } 337 } 338 339 #[cfg(test)] 340 mod tests { 341 use super::*; 342 use rand::prelude::*; 343 use wasmparser::Validator; 344 345 #[test] 346 fn stacks_generates_valid_wasm_modules() { 347 let mut rng = SmallRng::seed_from_u64(0); 348 let mut buf = vec![0; 2048]; 349 for _ in 0..1024 { 350 rng.fill_bytes(&mut buf); 351 let u = Unstructured::new(&buf); 352 if let Ok(stacks) = Stacks::arbitrary_take_rest(u) { 353 let wasm = stacks.wasm(); 354 validate(&wasm); 355 } 356 } 357 } 358 359 fn validate(wasm: &[u8]) { 360 let mut validator = Validator::new(); 361 let err = match validator.validate_all(wasm) { 362 Ok(_) => return, 363 Err(e) => e, 364 }; 365 drop(std::fs::write("test.wasm", wasm)); 366 if let Ok(text) = wasmprinter::print_bytes(wasm) { 367 drop(std::fs::write("test.wat", &text)); 368 } 369 panic!("wasm failed to validate: {}", err); 370 } 371 } 372