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}