wasm_smith/core/terminate.rs
1use super::*;
2use anyhow::{bail, Result};
3
4impl Module {
5 /// Ensure that all of this Wasm module's functions will terminate when
6 /// executed.
7 ///
8 /// This adds a new mutable, exported global to the module to keep track of
9 /// how much "fuel" is left. Fuel is decremented at the head of each loop
10 /// and function. When fuel reaches zero, a trap is raised.
11 ///
12 /// The index of the fuel global is returned, so that you may control how
13 /// much fuel the module is given.
14 ///
15 /// # Errors
16 ///
17 /// Returns an error if any function body was generated with
18 /// possibly-invalid bytes rather than being generated by wasm-smith. In
19 /// such a situation this pass does not parse the input bytes and inject
20 /// instructions, instead it returns an error.
21 pub fn ensure_termination(&mut self, default_fuel: u32) -> Result<u32> {
22 let fuel_global = self.globals.len() as u32;
23 self.globals.push(GlobalType {
24 val_type: ValType::I32,
25 mutable: true,
26 shared: false,
27 });
28 self.defined_globals
29 .push((fuel_global, ConstExpr::i32_const(default_fuel as i32)));
30
31 for code in &mut self.code {
32 let check_fuel = |insts: &mut Vec<Instruction>| {
33 // if fuel == 0 { trap }
34 insts.push(Instruction::GlobalGet(fuel_global));
35 insts.push(Instruction::I32Eqz);
36 insts.push(Instruction::If(BlockType::Empty));
37 insts.push(Instruction::Unreachable);
38 insts.push(Instruction::End);
39
40 // fuel -= 1
41 insts.push(Instruction::GlobalGet(fuel_global));
42 insts.push(Instruction::I32Const(1));
43 insts.push(Instruction::I32Sub);
44 insts.push(Instruction::GlobalSet(fuel_global));
45 };
46
47 let instrs = match &mut code.instructions {
48 Instructions::Generated(list) => list,
49 Instructions::Arbitrary(_) => {
50 bail!(
51 "failed to ensure that a function generated due to it \
52 containing arbitrary instructions"
53 )
54 }
55 };
56 let mut new_insts = Vec::with_capacity(instrs.len() * 2);
57
58 // Check fuel at the start of functions to deal with infinite
59 // recursion.
60 check_fuel(&mut new_insts);
61
62 for inst in mem::replace(instrs, vec![]) {
63 let is_loop = matches!(&inst, Instruction::Loop(_));
64 new_insts.push(inst);
65
66 // Check fuel at loop heads to deal with infinite loops.
67 if is_loop {
68 check_fuel(&mut new_insts);
69 }
70 }
71
72 *instrs = new_insts;
73 }
74
75 Ok(fuel_global)
76 }
77}