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