1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
use super::*;
use anyhow::{Result, bail};
impl Module {
/// Ensure that all of this Wasm module's functions will terminate when
/// executed.
///
/// This adds a new mutable, exported global to the module to keep track of
/// how much "fuel" is left. Fuel is decremented at the head of each loop
/// and function. When fuel reaches zero, a trap is raised.
///
/// The index of the fuel global is returned, so that you may control how
/// much fuel the module is given.
///
/// # Errors
///
/// Returns an error if any function body was generated with
/// possibly-invalid bytes rather than being generated by wasm-smith. In
/// such a situation this pass does not parse the input bytes and inject
/// instructions, instead it returns an error.
pub fn ensure_termination(&mut self, default_fuel: u32) -> Result<u32> {
let fuel_global = self.globals.len() as u32;
self.globals.push(GlobalType {
val_type: ValType::I32,
mutable: true,
shared: false,
});
self.defined_globals
.push((fuel_global, ConstExpr::i32_const(default_fuel as i32)));
let defined_funcs = &self.funcs[self.funcs.len() - self.code.len()..];
for ((_, ty), code) in defined_funcs.iter().zip(self.code.iter_mut()) {
let params = ty.params.len();
let mut temp_i32_local = None;
let mut temp_i32_local = |locals: &mut Vec<_>| {
if let Some(l) = temp_i32_local {
return l;
}
let l = u32::try_from(locals.len() + params).unwrap();
locals.push(ValType::I32);
temp_i32_local = Some(l);
l
};
let mut temp_i64_local = None;
let mut temp_i64_local = |locals: &mut Vec<_>| {
if let Some(l) = temp_i64_local {
return l;
}
let l = u32::try_from(locals.len() + params).unwrap();
locals.push(ValType::I64);
temp_i64_local = Some(l);
l
};
let check_fuel_1 = |insts: &mut Vec<Instruction>| {
// if fuel == 0 { trap }
insts.push(Instruction::GlobalGet(fuel_global));
insts.push(Instruction::I32Eqz);
insts.push(Instruction::If(BlockType::Empty));
insts.push(Instruction::Unreachable);
insts.push(Instruction::End);
// fuel -= 1
insts.push(Instruction::GlobalGet(fuel_global));
insts.push(Instruction::I32Const(1));
insts.push(Instruction::I32Sub);
insts.push(Instruction::GlobalSet(fuel_global));
};
let check_fuel_n = |insts: &mut Vec<Instruction>, local: u32| {
// if local >= fuel { trap }
insts.push(Instruction::LocalGet(local));
insts.push(Instruction::GlobalGet(fuel_global));
insts.push(Instruction::I32GeU);
insts.push(Instruction::If(BlockType::Empty));
insts.push(Instruction::Unreachable);
insts.push(Instruction::End);
// fuel -= local
insts.push(Instruction::GlobalGet(fuel_global));
insts.push(Instruction::LocalGet(local));
insts.push(Instruction::I32Sub);
insts.push(Instruction::GlobalSet(fuel_global));
};
let check_fuel_32_or_64 = |locals: &mut Vec<ValType>,
temp_i32_local: &mut dyn FnMut(&mut Vec<ValType>) -> u32,
temp_i64_local: &mut dyn FnMut(&mut Vec<ValType>) -> u32,
new_insts: &mut Vec<Instruction>,
inst: Instruction,
is_64: bool| {
if is_64 {
let local64 = temp_i64_local(locals);
let local32 = temp_i32_local(locals);
new_insts.push(Instruction::LocalTee(local64));
new_insts.push(Instruction::I32WrapI64);
new_insts.push(Instruction::LocalSet(local32));
check_fuel_n(new_insts, local32);
new_insts.push(Instruction::LocalGet(local64));
new_insts.push(inst);
} else {
let local = temp_i32_local(locals);
new_insts.push(Instruction::LocalTee(local));
check_fuel_n(new_insts, local);
new_insts.push(inst);
}
};
let instrs = match &mut code.instructions {
Instructions::Generated(list) => list,
Instructions::Arbitrary(_) => {
bail!(
"failed to ensure that a function generated due to it \
containing arbitrary instructions"
)
}
};
let mut new_insts = Vec::with_capacity(instrs.len() * 2);
// Check fuel at the start of functions to deal with infinite
// recursion.
check_fuel_1(&mut new_insts);
for inst in mem::replace(instrs, vec![]) {
match &inst {
// Check fuel at loop heads to deal with infinite loops.
Instruction::Loop(_) => {
new_insts.push(inst);
check_fuel_1(&mut new_insts);
}
// Check fuel on instructions that imply loops and decrement
// fuel accordingly and always have `len: i32` on top of the
// stack.
Instruction::ArrayCopy { .. }
| Instruction::ArrayFill(_)
| Instruction::ArrayInitData { .. }
| Instruction::ArrayInitElem { .. }
| Instruction::ArrayNew(_)
| Instruction::ArrayNewDefault(_)
| Instruction::ArrayNewData { .. }
| Instruction::ArrayNewElem { .. } => {
let local = temp_i32_local(&mut code.locals);
new_insts.push(Instruction::LocalTee(local));
check_fuel_n(&mut new_insts, local);
new_insts.push(inst);
}
// Check fuel on `table.init`, whose `len` operand is always
// `i32`, even for `table64`.
Instruction::TableInit { .. } => {
let local = temp_i32_local(&mut code.locals);
new_insts.push(Instruction::LocalTee(local));
check_fuel_n(&mut new_insts, local);
new_insts.push(inst);
}
// Check fuel on `table.fill`, whose `len` operand has the
// table's index type.
Instruction::TableFill(table) => {
let table = usize::try_from(*table).unwrap();
let is_64 = self.tables[table].table64;
check_fuel_32_or_64(
&mut code.locals,
&mut temp_i32_local,
&mut temp_i64_local,
&mut new_insts,
inst,
is_64,
);
}
// Check fuel on `table.copy`, whose `len` operand has the
// smaller of the source and destination tables' index types.
Instruction::TableCopy {
dst_table,
src_table,
} => {
let dst_table = usize::try_from(*dst_table).unwrap();
let src_table = usize::try_from(*src_table).unwrap();
check_fuel_32_or_64(
&mut code.locals,
&mut temp_i32_local,
&mut temp_i64_local,
&mut new_insts,
inst,
self.tables[dst_table].table64 && self.tables[src_table].table64,
);
}
// Check fuel on `memory.init`, whose `len` operand is
// always `i32`, even for `memory64`.
Instruction::MemoryInit { .. } => {
let local = temp_i32_local(&mut code.locals);
new_insts.push(Instruction::LocalTee(local));
check_fuel_n(&mut new_insts, local);
new_insts.push(inst);
}
// Check fuel on `memory.fill`, whose `len` operand has the
// memory's index type.
Instruction::MemoryFill(mem) => {
let mem = usize::try_from(*mem).unwrap();
check_fuel_32_or_64(
&mut code.locals,
&mut temp_i32_local,
&mut temp_i64_local,
&mut new_insts,
inst,
self.memories[mem].memory64,
);
}
// Check fuel on `memory.copy`, whose `len` operand has the
// smaller of the source and destination memories' index
// types.
Instruction::MemoryCopy { dst_mem, src_mem } => {
let dst_mem = usize::try_from(*dst_mem).unwrap();
let src_mem = usize::try_from(*src_mem).unwrap();
check_fuel_32_or_64(
&mut code.locals,
&mut temp_i32_local,
&mut temp_i64_local,
&mut new_insts,
inst,
self.memories[dst_mem].memory64 && self.memories[src_mem].memory64,
);
}
// Otherwise, just keep the instruction.
_ => new_insts.push(inst),
}
}
*instrs = new_insts;
}
Ok(fuel_global)
}
}