wasm-mutate: Add a mutator to modify data segments

crates/wasm-mutate/src/info.rs

@@ -250,6 +250,10 @@ impl<'a> ModuleInfo<'a> {
         self.raw_sections[self.exports.unwrap()]
     }
 
+    pub fn get_data_section(&self) -> RawSection<'a> {
+        self.raw_sections[self.data.unwrap()]
+    }
+
     pub fn has_exports(&self) -> bool {
         self.exports != None
     }

crates/wasm-mutate/src/lib.rs

@@ -17,14 +17,10 @@ mod mutators;
 pub use error::*;
 
 use crate::mutators::{
-    codemotion::CodemotionMutator,
-    custom::RemoveCustomSection,
-    function_body_unreachable::FunctionBodyUnreachable,
-    peephole::PeepholeMutator,
-    remove_export::RemoveExportMutator,
-    remove_item::{Item, RemoveItemMutator},
-    rename_export::RenameExportMutator,
-    snip_function::SnipMutator,
+    codemotion::CodemotionMutator, custom::RemoveCustomSection,
+    function_body_unreachable::FunctionBodyUnreachable, modify_data::ModifyDataMutator,
+    peephole::PeepholeMutator, remove_export::RemoveExportMutator, remove_item::RemoveItemMutator,
+    rename_export::RenameExportMutator, snip_function::SnipMutator, Item,
 };
 use info::ModuleInfo;
 use mutators::Mutator;
@@ -173,7 +169,7 @@ pub struct WasmMutate<'wasm> {
     // Note: this is only exposed via the programmatic interface, not via the
     // CLI.
     #[cfg_attr(feature = "clap", clap(skip = None))]
-    raw_mutate_func: Option<Arc<dyn Fn(&mut Vec<u8>) -> Result<()>>>,
+    raw_mutate_func: Option<Arc<dyn Fn(&mut Vec<u8>, usize) -> Result<()>>>,
 
     #[cfg_attr(feature = "clap", clap(skip = None))]
     rng: Option<SmallRng>,
@@ -239,7 +235,7 @@ impl<'wasm> WasmMutate<'wasm> {
     /// to get raw bytes from `libFuzzer`, for example.
     pub fn raw_mutate_func(
         &mut self,
-        raw_mutate_func: Option<Arc<dyn Fn(&mut Vec<u8>) -> Result<()>>>,
+        raw_mutate_func: Option<Arc<dyn Fn(&mut Vec<u8>, usize) -> Result<()>>>,
     ) -> &mut Self {
         self.raw_mutate_func = raw_mutate_func;
         self
@@ -259,8 +255,7 @@ impl<'wasm> WasmMutate<'wasm> {
         &'a mut self,
         input_wasm: &'wasm [u8],
     ) -> Result<Box<dyn Iterator<Item = Result<Vec<u8>>> + 'a>> {
-        self.info = Some(ModuleInfo::new(input_wasm)?);
-        self.rng = Some(SmallRng::seed_from_u64(self.seed));
+        self.setup(input_wasm)?;
 
         // This macro just expands the logic to return an iterator form the
         // mutators
@@ -287,19 +282,68 @@ impl<'wasm> WasmMutate<'wasm> {
                 RemoveItemMutator(Item::Data),
                 RemoveItemMutator(Item::Element),
                 RemoveItemMutator(Item::Tag),
+                ModifyDataMutator {
+                    max_data_size: 10 << 20, // 10MB
+                },
             )
         );
 
         Err(Error::no_mutations_applicable())
     }
 
+    fn setup(&mut self, input_wasm: &'wasm [u8]) -> Result<()> {
+        self.info = Some(ModuleInfo::new(input_wasm)?);
+        self.rng = Some(SmallRng::seed_from_u64(self.seed));
+        Ok(())
+    }
+
     pub(crate) fn rng(&mut self) -> &mut SmallRng {
         self.rng.as_mut().unwrap()
     }
 
     pub(crate) fn info(&self) -> &ModuleInfo<'wasm> {
         self.info.as_ref().unwrap()
     }
+
+    fn raw_mutate(&mut self, data: &mut Vec<u8>, max_size: usize) -> Result<()> {
+        // If a raw mutation function is configured then that's prioritized.
+        if let Some(mutate) = &self.raw_mutate_func {
+            return mutate(data, max_size);
+        }
+
+        // If no raw mutation function is configured then we apply a naive
+        // default heuristic. For now that heuristic is to simply replace a
+        // subslice of data with a random slice of other data.
+        //
+        // First up start/end indices are picked.
+        let a = self.rng().gen_range(0, data.len() + 1);
+        let b = self.rng().gen_range(0, data.len() + 1);
+        let start = a.min(b);
+        let end = a.max(b);
+
+        // Next a length of the replacement is chosen. Note that the replacement
+        // is always smaller than the input if reduction is requested, otherwise
+        // we choose some arbitrary length of bytes to insert.
+        let max_size = if self.reduce || self.rng().gen() {
+            0
+        } else {
+            max_size
+        };
+        let len = self
+            .rng()
+            .gen_range(0, end - start + max_size.saturating_sub(data.len()) + 1);
+
+        // With parameters chosen the `Vec::splice` method is used to replace
+        // the data in the input.
+        data.splice(
+            start..end,
+            self.rng()
+                .sample_iter(rand::distributions::Standard)
+                .take(len),
+        );
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]

crates/wasm-mutate/src/mutators.rs

@@ -23,13 +23,18 @@
 pub mod codemotion;
 pub mod custom;
 pub mod function_body_unreachable;
+pub mod modify_data;
 pub mod peephole;
 pub mod remove_export;
 pub mod remove_item;
 pub mod rename_export;
 pub mod snip_function;
 pub mod start;
 
+mod translate;
+pub use self::translate::Item;
+use self::translate::{DefaultTranslator, Translator};
+
 use std::borrow::Cow;
 
 use super::Result;
@@ -100,60 +105,65 @@ pub trait Mutator {
 pub type OperatorAndByteOffset<'a> = (Operator<'a>, usize);
 
 #[cfg(test)]
-pub(crate) fn match_mutation<T>(original: &str, mutator: T, expected: &str)
+fn match_mutation<T>(original: &str, mutator: T, expected: &str)
 where
     T: Mutator + Copy,
 {
-    use crate::info::ModuleInfo;
-    use crate::ErrorKind;
-    use rand::{prelude::SmallRng, SeedableRng};
-    use wasmparser::WasmFeatures;
-
-    let mut wasmmutate = WasmMutate::default();
-    let original = &wat::parse_str(original).unwrap();
-
-    let expected = &wat::parse_str(expected).unwrap();
-    let expected_text = wasmprinter::print_bytes(expected).unwrap();
-
-    let info = ModuleInfo::new(original).unwrap();
-    wasmmutate.info = Some(info);
-    let rnd = SmallRng::seed_from_u64(0);
-    wasmmutate.rng = Some(rnd);
-
-    let can_mutate = mutator.can_mutate(&wasmmutate);
-
-    assert!(can_mutate);
-
-    let attempts = 100;
-
-    for _ in 0..attempts {
-        let mutation = match mutator
-            .mutate(&mut wasmmutate)
-            .and_then(|mut mutation| mutation.next().unwrap())
-        {
-            Ok(mutation) => mutation,
-            Err(e) if matches!(e.kind(), ErrorKind::NoMutationsApplicable) => continue,
-            Err(e) => panic!("mutation error: {}", e),
-        };
-
-        let mutation_bytes = mutation.finish();
-
-        let mut validator = wasmparser::Validator::new();
-        validator.wasm_features(WasmFeatures {
-            multi_memory: true,
-            ..WasmFeatures::default()
-        });
-        crate::validate(&mut validator, &mutation_bytes);
-
-        // If it fails, it is probably an invalid
-        // reformatting expected
-        let text = wasmprinter::print_bytes(mutation_bytes).unwrap();
-        assert_eq!(text.trim(), expected_text.trim());
-        return;
-    }
+    WasmMutate::default().match_mutation(original, mutator, expected)
+}
 
-    panic!(
-        "never found any applicable mutations after {} attempts",
-        attempts
-    );
+impl WasmMutate<'_> {
+    #[cfg(test)]
+    fn match_mutation<T>(&mut self, original: &str, mutator: T, expected: &str)
+    where
+        T: Mutator + Copy,
+    {
+        use crate::ErrorKind;
+        use wasmparser::WasmFeatures;
+
+        let original = &wat::parse_str(original).unwrap();
+
+        let expected = &wat::parse_str(expected).unwrap();
+        let expected_text = wasmprinter::print_bytes(expected).unwrap();
+
+        let mut config = self.clone();
+        config.setup(&original).unwrap();
+
+        let can_mutate = mutator.can_mutate(&config);
+
+        assert!(can_mutate);
+
+        let attempts = 100;
+
+        for _ in 0..attempts {
+            let mutation = match mutator
+                .mutate(&mut config)
+                .and_then(|mut mutation| mutation.next().unwrap())
+            {
+                Ok(mutation) => mutation,
+                Err(e) if matches!(e.kind(), ErrorKind::NoMutationsApplicable) => continue,
+                Err(e) => panic!("mutation error: {}", e),
+            };
+
+            let mutation_bytes = mutation.finish();
+
+            let mut validator = wasmparser::Validator::new();
+            validator.wasm_features(WasmFeatures {
+                multi_memory: true,
+                ..WasmFeatures::default()
+            });
+            crate::validate(&mut validator, &mutation_bytes);
+
+            // If it fails, it is probably an invalid
+            // reformatting expected
+            let text = wasmprinter::print_bytes(mutation_bytes).unwrap();
+            assert_eq!(text.trim(), expected_text.trim());
+            return;
+        }
+
+        panic!(
+            "never found any applicable mutations after {} attempts",
+            attempts
+        );
+    }
 }

crates/wasm-mutate/src/mutators/modify_data.rs

@@ -0,0 +1,89 @@
+use super::Mutator;
+use crate::{Result, WasmMutate};
+
+use crate::mutators::{DefaultTranslator, Translator};
+use rand::Rng;
+use wasm_encoder::{DataSection, DataSegment, DataSegmentMode, Module};
+use wasmparser::{DataKind, DataSectionReader};
+
+/// Mutator that modifies a data segment, either adding or removing bytes.
+#[derive(Clone, Copy)]
+pub struct ModifyDataMutator {
+    pub max_data_size: usize,
+}
+
+impl Mutator for ModifyDataMutator {
+    fn mutate<'a>(
+        self,
+        config: &'a mut WasmMutate,
+    ) -> Result<Box<dyn Iterator<Item = Result<Module>> + 'a>> {
+        let mut new_section = DataSection::new();
+        let mut reader = DataSectionReader::new(config.info().get_data_section().data, 0)?;
+
+        // Select an arbitrary data segment to modify.
+        let data_to_modify = config.rng().gen_range(0, reader.get_count());
+
+        // Iterate over all data segments in the old data section and re-add
+        // them to the `new_section` one-by-one.
+        for i in 0..reader.get_count() {
+            let data = reader.read()?;
+            let offset;
+            // Preserve the mode of the data segment
+            let mode = match &data.kind {
+                DataKind::Active {
+                    memory_index,
+                    init_expr,
+                } => {
+                    offset = DefaultTranslator.translate_init_expr(init_expr)?;
+                    DataSegmentMode::Active {
+                        memory_index: *memory_index,
+                        offset: &offset,
+                    }
+                }
+                DataKind::Passive => DataSegmentMode::Passive,
+            };
+            // If this is the correct data segment apply the mutation,
+            // otherwise preserve the data.
+            let mut data = data.data.to_vec();
+            if i == data_to_modify {
+                config.raw_mutate(&mut data, self.max_data_size)?;
+            }
+            new_section.segment(DataSegment { mode, data });
+        }
+
+        Ok(Box::new(std::iter::once(Ok(config
+            .info()
+            .replace_section(
+                config.info().data.unwrap(),
+                &new_section,
+            )))))
+    }
+
+    fn can_mutate<'a>(&self, config: &'a WasmMutate) -> bool {
+        // Modifying a data segment doesn't preserve the semantics of the
+        // original module and also only works if there's actually some data.
+        !config.preserve_semantics && config.info().num_data() > 0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::ModifyDataMutator;
+    use crate::WasmMutate;
+    use std::sync::Arc;
+
+    #[test]
+    fn test_remove_export_mutator() {
+        let mut config = WasmMutate::default();
+        config.raw_mutate_func(Some(Arc::new(|data, _| {
+            assert_eq!(data, b"x");
+            *data = "y".to_string().into_bytes();
+            Ok(())
+        })));
+        config.match_mutation(
+            r#"(module (data "x"))"#,
+            ModifyDataMutator { max_data_size: 100 },
+            r#"(module (data "y"))"#,
+        );
+    }
+}