/src/wasm-tools/crates/wasm-mutate/src/mutators/codemotion.rs
Line | Count | Source |
1 | | //! This mutator applies a random mutation over the control flow AST of an input |
2 | | //! binary |
3 | | //! |
4 | | //! To extend `wasm-mutate` with another code motion mutator, the new mutator |
5 | | //! struct should implement the [AstMutator] trait and we strongly recommend the |
6 | | //! usage of the [ir::AstWriter] to define how the mutator writes the new AST back |
7 | | //! to Wasm. |
8 | | //! |
9 | | //! For and example take a look at [IfComplementMutator][IfComplementMutator] |
10 | | //! |
11 | | //! Register the new mutator then in the meta [CodemotionMutator] logic. |
12 | | //! ```ignore |
13 | | //! let mutators: Vec<Box<dyn AstMutator>> = vec![ |
14 | | //! Box::new(IfComplementMutator), |
15 | | //! ]; |
16 | | //! ``` |
17 | | |
18 | | pub mod if_complement; |
19 | | pub mod ir; |
20 | | pub mod loop_unrolling; |
21 | | |
22 | | use self::ir::parse_context::Ast; |
23 | | use super::Mutator; |
24 | | use crate::{ |
25 | | Error, Result, WasmMutate, |
26 | | module::map_type, |
27 | | mutators::{ |
28 | | OperatorAndByteOffset, |
29 | | codemotion::{ |
30 | | if_complement::IfComplementMutator, ir::AstBuilder, loop_unrolling::LoopUnrollMutator, |
31 | | }, |
32 | | }, |
33 | | }; |
34 | | use rand::{RngExt, prelude::*}; |
35 | | use wasm_encoder::{CodeSection, Function, Module, ValType}; |
36 | | use wasmparser::{CodeSectionReader, FunctionBody}; |
37 | | |
38 | | /// Code motion meta mutator, it groups all code motion mutators and select a |
39 | | /// valid random one when an input Wasm binary is passed to it. |
40 | | #[derive(Clone, Copy)] |
41 | | pub struct CodemotionMutator; |
42 | | |
43 | | impl CodemotionMutator { |
44 | 157 | fn copy_locals(&self, reader: FunctionBody) -> Result<Vec<(u32, ValType)>> { |
45 | | // Create the new function |
46 | 157 | let mut localreader = reader.get_locals_reader()?; |
47 | | // Get current locals and map to encoder types |
48 | 157 | let mut local_count = 0; |
49 | 157 | let current_locals = (0..localreader.get_count()) |
50 | 2.30k | .map(|_| { |
51 | 2.30k | let (count, ty) = localreader.read().unwrap(); |
52 | 2.30k | local_count += count; |
53 | 2.30k | (count, map_type(ty).unwrap()) |
54 | 2.30k | }) |
55 | 157 | .collect::<Vec<(u32, ValType)>>(); |
56 | | |
57 | 157 | Ok(current_locals) |
58 | 157 | } |
59 | | |
60 | 350 | fn random_mutate( |
61 | 350 | &self, |
62 | 350 | config: &mut WasmMutate, |
63 | 350 | mutators: &[Box<dyn AstMutator>], |
64 | 350 | ) -> crate::Result<(Function, u32)> { |
65 | 350 | let original_code_section = config.info().code.unwrap(); |
66 | 350 | let reader = config.info().get_binary_reader(original_code_section); |
67 | 350 | let input_code_section = config.info().get_code_section().data; |
68 | 350 | let sectionreader = CodeSectionReader::new(reader)?; |
69 | 350 | let function_count = sectionreader.count(); |
70 | 350 | let function_to_mutate = config.rng().random_range(0..function_count); |
71 | | |
72 | | // This split strategy will avoid very often mutating the first function |
73 | | // and very rarely mutating the last function |
74 | 350 | let all_readers = sectionreader.into_iter().collect::<Result<Vec<_>, _>>()?; |
75 | | |
76 | 2.78k | for fidx in (function_to_mutate..function_count).chain(0..function_to_mutate) { |
77 | 2.78k | config.consume_fuel(1)?; |
78 | 2.78k | let reader = all_readers[fidx as usize].clone(); |
79 | 2.78k | let operatorreader = reader.get_operators_reader()?; |
80 | | |
81 | 2.78k | let operators = operatorreader |
82 | 2.78k | .into_iter_with_offsets() |
83 | 2.78k | .collect::<wasmparser::Result<Vec<OperatorAndByteOffset>>>()?; |
84 | | |
85 | | // build Ast |
86 | 2.78k | let ast = AstBuilder.build_ast(&operators)?; |
87 | | // filter mutators by those applicable |
88 | 2.74k | let filtered = mutators |
89 | 2.74k | .iter() |
90 | 5.49k | .filter(|m| m.can_mutate(config, &ast)) |
91 | 2.74k | .collect::<Vec<_>>(); |
92 | | // If no mutator, just continue to the next function |
93 | 2.74k | if filtered.is_empty() { |
94 | 2.58k | continue; |
95 | 157 | } |
96 | | |
97 | 157 | match filtered.choose(config.rng()) { |
98 | 157 | Some(chosen_mutator) => { |
99 | 157 | let newfunc = chosen_mutator.mutate( |
100 | 157 | config, |
101 | 157 | &ast, |
102 | 157 | &self.copy_locals(reader)?, |
103 | 157 | &operators, |
104 | 157 | input_code_section, |
105 | 0 | )?; |
106 | 157 | return Ok((newfunc, fidx)); |
107 | | } |
108 | 0 | None => continue, |
109 | | } |
110 | | } |
111 | | |
112 | 149 | Err(Error::no_mutations_applicable()) |
113 | 350 | } |
114 | | } |
115 | | /// Trait to be implemented by all code motion mutators |
116 | | pub trait AstMutator { |
117 | | /// Transform the function AST in order to generate a new Wasm module |
118 | | fn mutate<'a>( |
119 | | &self, |
120 | | config: &'a mut WasmMutate, |
121 | | ast: &Ast, |
122 | | locals: &[(u32, ValType)], |
123 | | operators: &Vec<OperatorAndByteOffset>, |
124 | | input_wasm: &'a [u8], |
125 | | ) -> Result<Function>; |
126 | | |
127 | | /// Checks if this mutator can be applied to the passed `ast` |
128 | | fn can_mutate<'a>(&self, config: &'a crate::WasmMutate, ast: &Ast) -> bool; |
129 | | } |
130 | | |
131 | | /// Meta mutator for peephole |
132 | | impl Mutator for CodemotionMutator { |
133 | 350 | fn mutate<'a>( |
134 | 350 | &self, |
135 | 350 | config: &mut WasmMutate<'a>, |
136 | 350 | ) -> Result<Box<dyn Iterator<Item = Result<Module>> + 'a>> { |
137 | | // Initialize mutators |
138 | 350 | let mutators: Vec<Box<dyn AstMutator>> = vec![ |
139 | 350 | Box::new(IfComplementMutator), |
140 | 350 | Box::new(LoopUnrollMutator), // Add the other here |
141 | | ]; |
142 | | |
143 | 350 | let (newfunc, function_to_mutate) = self.random_mutate(config, &mutators)?; |
144 | | |
145 | 157 | let mut codes = CodeSection::new(); |
146 | 157 | let code_section = config.info().get_binary_reader(config.info().code.unwrap()); |
147 | 157 | let sectionreader = CodeSectionReader::new(code_section)?; |
148 | | |
149 | 2.24k | for (fidx, reader) in sectionreader.into_iter().enumerate() { |
150 | 2.24k | let reader = reader?; |
151 | 2.24k | if fidx as u32 == function_to_mutate { |
152 | 157 | log::trace!("Mutating function {fidx}"); |
153 | 157 | codes.function(&newfunc); |
154 | 2.08k | } else { |
155 | 2.08k | codes.raw(reader.as_bytes()); |
156 | 2.08k | } |
157 | | } |
158 | 157 | let module = config |
159 | 157 | .info() |
160 | 157 | .replace_section(config.info().code.unwrap(), &codes); |
161 | 157 | Ok(Box::new(std::iter::once(Ok(module)))) |
162 | 350 | } |
163 | | |
164 | 538 | fn can_mutate<'a>(&self, config: &'a WasmMutate) -> bool { |
165 | 538 | config.info().has_code() && config.info().num_local_functions() > 0 |
166 | 538 | } |
167 | | } |
168 | | |
169 | | #[cfg(test)] |
170 | | mod tests { |
171 | | use crate::{WasmMutate, mutators::codemotion::CodemotionMutator}; |
172 | | |
173 | | fn test_motion_mutator(original: &str, expected: &str, seed: u64) { |
174 | | let mut config = WasmMutate::default(); |
175 | | config.seed(seed); |
176 | | config.match_mutation(original, CodemotionMutator, expected); |
177 | | } |
178 | | |
179 | | #[test] |
180 | | fn test_if_swap() { |
181 | | test_motion_mutator( |
182 | | r#" |
183 | | (module |
184 | | (memory 1) |
185 | | (func (export "exported_func") (param i32) (result i32) |
186 | | (local i32 i32) |
187 | | local.get 0 |
188 | | if (result i32) |
189 | | i32.const 50 |
190 | | else |
191 | | i32.const 41 |
192 | | end |
193 | | ) |
194 | | ) |
195 | | "#, |
196 | | r#" |
197 | | (module |
198 | | (type (;0;) (func (param i32) (result i32))) |
199 | | (func (;0;) (type 0) (param i32) (result i32) |
200 | | (local i32 i32) |
201 | | local.get 0 |
202 | | i32.eqz |
203 | | if (result i32) ;; label = @1 |
204 | | i32.const 41 |
205 | | else |
206 | | i32.const 50 |
207 | | end) |
208 | | (memory (;0;) 1) |
209 | | (export "exported_func" (func 0))) |
210 | | "#, |
211 | | 0, |
212 | | ); |
213 | | } |
214 | | |
215 | | #[test] |
216 | | fn test_if_swap2() { |
217 | | test_motion_mutator( |
218 | | r#" |
219 | | (module |
220 | | (memory 1) |
221 | | (func (export "exported_func") (param i32) (result i32) |
222 | | local.get 0 |
223 | | local.get 0 |
224 | | i32.add |
225 | | local.get 0 |
226 | | if |
227 | | i32.const 150 |
228 | | drop |
229 | | else |
230 | | i32.const 200 |
231 | | drop |
232 | | end |
233 | | if (result i32) |
234 | | i32.const 50 |
235 | | else |
236 | | i32.const 41 |
237 | | end |
238 | | ) |
239 | | ) |
240 | | "#, |
241 | | r#" |
242 | | (module |
243 | | (type (;0;) (func (param i32) (result i32))) |
244 | | (func (;0;) (type 0) (param i32) (result i32) |
245 | | local.get 0 |
246 | | local.get 0 |
247 | | i32.add |
248 | | local.get 0 |
249 | | if ;; label = @1 |
250 | | i32.const 150 |
251 | | drop |
252 | | else |
253 | | i32.const 200 |
254 | | drop |
255 | | end |
256 | | i32.eqz |
257 | | if (result i32) ;; label = @1 |
258 | | i32.const 41 |
259 | | else |
260 | | i32.const 50 |
261 | | end) |
262 | | (memory (;0;) 1) |
263 | | (export "exported_func" (func 0))) |
264 | | "#, |
265 | | 1, |
266 | | ); |
267 | | } |
268 | | |
269 | | #[test] |
270 | | fn test_if_swap3() { |
271 | | test_motion_mutator( |
272 | | r#" |
273 | | (module |
274 | | (memory 1) |
275 | | (func (export "exported_func") (param i32) (result i32) |
276 | | local.get 0 |
277 | | local.get 0 |
278 | | i32.add |
279 | | local.get 0 |
280 | | if |
281 | | i32.const 150 |
282 | | drop |
283 | | else |
284 | | i32.const 200 |
285 | | drop |
286 | | end |
287 | | if (result i32) |
288 | | i32.const 50 |
289 | | else |
290 | | unreachable |
291 | | end |
292 | | ) |
293 | | ) |
294 | | "#, |
295 | | r#" |
296 | | (module |
297 | | (type (;0;) (func (param i32) (result i32))) |
298 | | (func (;0;) (type 0) (param i32) (result i32) |
299 | | local.get 0 |
300 | | local.get 0 |
301 | | i32.add |
302 | | local.get 0 |
303 | | if ;; label = @1 |
304 | | i32.const 150 |
305 | | drop |
306 | | else |
307 | | i32.const 200 |
308 | | drop |
309 | | end |
310 | | i32.eqz |
311 | | if (result i32) ;; label = @1 |
312 | | unreachable |
313 | | else |
314 | | i32.const 50 |
315 | | end) |
316 | | (memory (;0;) 1) |
317 | | (export "exported_func" (func 0))) |
318 | | "#, |
319 | | 1, |
320 | | ); |
321 | | } |
322 | | |
323 | | #[test] |
324 | | fn test_unrolling1() { |
325 | | test_motion_mutator( |
326 | | r#" |
327 | | (module |
328 | | (memory 1) |
329 | | (func (export "exported_func") (param i32) (result i32) |
330 | | local.get 0 |
331 | | local.get 0 |
332 | | i32.add |
333 | | drop |
334 | | loop |
335 | | i32.const 1 |
336 | | local.get 0 |
337 | | i32.add |
338 | | local.tee 0 |
339 | | i32.const 100 |
340 | | i32.le_u |
341 | | br_if 0 |
342 | | end |
343 | | local.get 0 |
344 | | ) |
345 | | ) |
346 | | "#, |
347 | | r#" |
348 | | (module |
349 | | (type (;0;) (func (param i32) (result i32))) |
350 | | (func (;0;) (type 0) (param i32) (result i32) |
351 | | local.get 0 |
352 | | local.get 0 |
353 | | i32.add |
354 | | drop |
355 | | block ;; label = @1 |
356 | | block ;; label = @2 |
357 | | i32.const 1 |
358 | | local.get 0 |
359 | | i32.add |
360 | | local.tee 0 |
361 | | i32.const 100 |
362 | | i32.le_u |
363 | | br_if 0 (;@2;) |
364 | | br 1 (;@1;) |
365 | | end |
366 | | loop ;; label = @2 |
367 | | i32.const 1 |
368 | | local.get 0 |
369 | | i32.add |
370 | | local.tee 0 |
371 | | i32.const 100 |
372 | | i32.le_u |
373 | | br_if 0 (;@2;) |
374 | | end |
375 | | end |
376 | | local.get 0) |
377 | | (memory (;0;) 1) |
378 | | (export "exported_func" (func 0))) |
379 | | |
380 | | "#, |
381 | | 1, |
382 | | ); |
383 | | } |
384 | | |
385 | | #[test] |
386 | | fn test_unrolling2() { |
387 | | test_motion_mutator( |
388 | | r#" |
389 | | (module |
390 | | (memory 1) |
391 | | (func (export "exported_func") (param i32) (result i32) |
392 | | local.get 0 |
393 | | local.get 0 |
394 | | i32.add |
395 | | drop |
396 | | loop |
397 | | i32.const 1 |
398 | | local.get 0 |
399 | | i32.add |
400 | | local.tee 0 |
401 | | i32.const 100 |
402 | | i32.le_u |
403 | | br_if 0 |
404 | | local.get 0 |
405 | | if |
406 | | i32.const 200 |
407 | | local.get 0 |
408 | | i32.add |
409 | | local.set 0 |
410 | | else |
411 | | i32.const 300 |
412 | | local.get 0 |
413 | | i32.add |
414 | | local.set 0 |
415 | | end |
416 | | end |
417 | | local.get 0 |
418 | | ) |
419 | | ) |
420 | | "#, |
421 | | r#" |
422 | | (module |
423 | | (type (;0;) (func (param i32) (result i32))) |
424 | | (func (;0;) (type 0) (param i32) (result i32) |
425 | | local.get 0 |
426 | | local.get 0 |
427 | | i32.add |
428 | | drop |
429 | | block ;; label = @1 |
430 | | block ;; label = @2 |
431 | | i32.const 1 |
432 | | local.get 0 |
433 | | i32.add |
434 | | local.tee 0 |
435 | | i32.const 100 |
436 | | i32.le_u |
437 | | br_if 0 (;@2;) |
438 | | local.get 0 |
439 | | if ;; label = @3 |
440 | | i32.const 200 |
441 | | local.get 0 |
442 | | i32.add |
443 | | local.set 0 |
444 | | else |
445 | | i32.const 300 |
446 | | local.get 0 |
447 | | i32.add |
448 | | local.set 0 |
449 | | end |
450 | | br 1 (;@1;) |
451 | | end |
452 | | loop ;; label = @2 |
453 | | i32.const 1 |
454 | | local.get 0 |
455 | | i32.add |
456 | | local.tee 0 |
457 | | i32.const 100 |
458 | | i32.le_u |
459 | | br_if 0 (;@2;) |
460 | | local.get 0 |
461 | | if ;; label = @3 |
462 | | i32.const 200 |
463 | | local.get 0 |
464 | | i32.add |
465 | | local.set 0 |
466 | | else |
467 | | i32.const 300 |
468 | | local.get 0 |
469 | | i32.add |
470 | | local.set 0 |
471 | | end |
472 | | end |
473 | | end |
474 | | local.get 0) |
475 | | (memory (;0;) 1) |
476 | | (export "exported_func" (func 0))) |
477 | | |
478 | | "#, |
479 | | 1, |
480 | | ); |
481 | | } |
482 | | |
483 | | #[test] |
484 | | fn test_unrolling3() { |
485 | | test_motion_mutator( |
486 | | r#" |
487 | | (module |
488 | | (memory 1) |
489 | | (func (export "exported_func") (param i32) (result i32) |
490 | | local.get 0 |
491 | | local.get 0 |
492 | | i32.add |
493 | | drop |
494 | | loop |
495 | | i32.const 1 |
496 | | local.get 0 |
497 | | i32.add |
498 | | local.tee 0 |
499 | | if |
500 | | i32.const 200 |
501 | | local.get 0 |
502 | | i32.add |
503 | | local.tee 0 |
504 | | br_if 1 |
505 | | else |
506 | | local.get 0 |
507 | | br_if 1 |
508 | | end |
509 | | local.get 0 |
510 | | br_if 0 |
511 | | end |
512 | | local.get 0 |
513 | | ) |
514 | | ) |
515 | | "#, |
516 | | r#" |
517 | | (module |
518 | | (type (;0;) (func (param i32) (result i32))) |
519 | | (func (;0;) (type 0) (param i32) (result i32) |
520 | | local.get 0 |
521 | | local.get 0 |
522 | | i32.add |
523 | | drop |
524 | | block ;; label = @1 |
525 | | block ;; label = @2 |
526 | | i32.const 1 |
527 | | local.get 0 |
528 | | i32.add |
529 | | local.tee 0 |
530 | | if ;; label = @3 |
531 | | i32.const 200 |
532 | | local.get 0 |
533 | | i32.add |
534 | | local.tee 0 |
535 | | br_if 1 (;@2;) |
536 | | else |
537 | | local.get 0 |
538 | | br_if 1 (;@2;) |
539 | | end |
540 | | local.get 0 |
541 | | br_if 0 (;@2;) |
542 | | br 1 (;@1;) |
543 | | end |
544 | | loop ;; label = @2 |
545 | | i32.const 1 |
546 | | local.get 0 |
547 | | i32.add |
548 | | local.tee 0 |
549 | | if ;; label = @3 |
550 | | i32.const 200 |
551 | | local.get 0 |
552 | | i32.add |
553 | | local.tee 0 |
554 | | br_if 1 (;@2;) |
555 | | else |
556 | | local.get 0 |
557 | | br_if 1 (;@2;) |
558 | | end |
559 | | local.get 0 |
560 | | br_if 0 (;@2;) |
561 | | end |
562 | | end |
563 | | local.get 0) |
564 | | (memory (;0;) 1) |
565 | | (export "exported_func" (func 0))) |
566 | | |
567 | | "#, |
568 | | 1, |
569 | | ); |
570 | | } |
571 | | |
572 | | #[test] |
573 | | fn test_unrolling4() { |
574 | | test_motion_mutator( |
575 | | r#" |
576 | | (module |
577 | | (memory 1) |
578 | | (func (export "exported_func") (param i32) (result i32) |
579 | | block |
580 | | loop |
581 | | loop |
582 | | local.get 0 |
583 | | i32.const 100 |
584 | | i32.ge_s |
585 | | br_if 2 |
586 | | end |
587 | | local.get 0 |
588 | | i32.const 200 |
589 | | i32.le_s |
590 | | br_if 0 |
591 | | end |
592 | | end |
593 | | local.get 0 |
594 | | ) |
595 | | ) |
596 | | "#, |
597 | | r#" |
598 | | (module |
599 | | (type (;0;) (func (param i32) (result i32))) |
600 | | (func (;0;) (type 0) (param i32) (result i32) |
601 | | block ;; label = @1 |
602 | | block ;; label = @2 |
603 | | block ;; label = @3 |
604 | | loop ;; label = @4 |
605 | | local.get 0 |
606 | | i32.const 100 |
607 | | i32.ge_s |
608 | | br_if 3 (;@1;) |
609 | | end |
610 | | local.get 0 |
611 | | i32.const 200 |
612 | | i32.le_s |
613 | | br_if 0 (;@3;) |
614 | | br 1 (;@2;) |
615 | | end |
616 | | loop ;; label = @3 |
617 | | loop ;; label = @4 |
618 | | local.get 0 |
619 | | i32.const 100 |
620 | | i32.ge_s |
621 | | br_if 3 (;@1;) |
622 | | end |
623 | | local.get 0 |
624 | | i32.const 200 |
625 | | i32.le_s |
626 | | br_if 0 (;@3;) |
627 | | end |
628 | | end |
629 | | end |
630 | | local.get 0) |
631 | | (memory (;0;) 1) |
632 | | (export "exported_func" (func 0))) |
633 | | |
634 | | "#, |
635 | | 1, |
636 | | ); |
637 | | } |
638 | | |
639 | | #[test] |
640 | | fn test_unrolling5() { |
641 | | test_motion_mutator( |
642 | | r#" |
643 | | (module |
644 | | (memory 1) |
645 | | (func (export "exported_func") (param i32) (result i32) |
646 | | block |
647 | | loop |
648 | | loop |
649 | | local.get 0 |
650 | | br_table 1 2 2 2 2 |
651 | | end |
652 | | local.get 0 |
653 | | i32.const 200 |
654 | | i32.le_s |
655 | | br_if 0 |
656 | | end |
657 | | end |
658 | | local.get 0 |
659 | | ) |
660 | | ) |
661 | | "#, |
662 | | r#" |
663 | | (module |
664 | | (type (;0;) (func (param i32) (result i32))) |
665 | | (func (;0;) (type 0) (param i32) (result i32) |
666 | | block ;; label = @1 |
667 | | block ;; label = @2 |
668 | | block ;; label = @3 |
669 | | loop ;; label = @4 |
670 | | local.get 0 |
671 | | br_table 1 (;@3;) 3 (;@1;) 3 (;@1;) 3 (;@1;) 3 (;@1;) |
672 | | end |
673 | | local.get 0 |
674 | | i32.const 200 |
675 | | i32.le_s |
676 | | br_if 0 (;@3;) |
677 | | br 1 (;@2;) |
678 | | end |
679 | | loop ;; label = @3 |
680 | | loop ;; label = @4 |
681 | | local.get 0 |
682 | | br_table 1 (;@3;) 3 (;@1;) 3 (;@1;) 3 (;@1;) 3 (;@1;) |
683 | | end |
684 | | local.get 0 |
685 | | i32.const 200 |
686 | | i32.le_s |
687 | | br_if 0 (;@3;) |
688 | | end |
689 | | end |
690 | | end |
691 | | local.get 0) |
692 | | (memory (;0;) 1) |
693 | | (export "exported_func" (func 0))) |
694 | | "#, |
695 | | 1, |
696 | | ); |
697 | | } |
698 | | } |