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