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