Merge pull request #226 from ia0/v2_fuzz

Add debug helper for fuzzing

libraries/persistent_store/fuzz/examples/store.rs

@@ -0,0 +1,116 @@
+// Copyright 2019-2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use fuzz_store::{fuzz, StatKey, Stats};
+use std::io::Write;
+use std::io::{stdout, Read};
+use std::path::Path;
+
+fn usage(program: &str) {
+    println!(
+        r#"Usage: {} {{ [<artifact_file>] | <corpus_directory> <bucket_predicate>.. }}
+
+If <artifact_file> is not provided, it is read from standard input.
+
+When <bucket_predicate>.. are provided, only runs matching all predicates are shown. The format of
+each <bucket_predicate> is <bucket_key>=<bucket_value>."#,
+        program
+    );
+}
+
+fn debug(data: &[u8]) {
+    println!("{:02x?}", data);
+    fuzz(data, true, None);
+}
+
+/// Bucket predicate.
+struct Predicate {
+    /// Bucket key.
+    key: StatKey,
+
+    /// Bucket value.
+    value: usize,
+}
+
+impl std::str::FromStr for Predicate {
+    type Err = String;
+
+    fn from_str(input: &str) -> Result<Self, Self::Err> {
+        let predicate: Vec<&str> = input.split('=').collect();
+        if predicate.len() != 2 {
+            return Err("Predicate should have exactly one equal sign.".to_string());
+        }
+        let key = predicate[0]
+            .parse()
+            .map_err(|_| format!("Predicate key `{}` is not recognized.", predicate[0]))?;
+        let value: usize = predicate[1]
+            .parse()
+            .map_err(|_| format!("Predicate value `{}` is not a number.", predicate[1]))?;
+        if value != 0 && !value.is_power_of_two() {
+            return Err(format!(
+                "Predicate value `{}` is not a bucket.",
+                predicate[1]
+            ));
+        }
+        Ok(Predicate { key, value })
+    }
+}
+
+fn analyze(corpus: &Path, predicates: Vec<Predicate>) {
+    let mut stats = Stats::default();
+    let mut count = 0;
+    let total = std::fs::read_dir(corpus).unwrap().count();
+    for entry in std::fs::read_dir(corpus).unwrap() {
+        let data = std::fs::read(entry.unwrap().path()).unwrap();
+        let mut stat = Stats::default();
+        fuzz(&data, false, Some(&mut stat));
+        if predicates
+            .iter()
+            .all(|p| stat.get_count(p.key, p.value).is_some())
+        {
+            stats.merge(&stat);
+        }
+        count += 1;
+        print!("\u{1b}[K{} / {}\r", count, total);
+        stdout().flush().unwrap();
+    }
+    // NOTE: To avoid reloading the corpus each time we want to check a different filter, we can
+    // start an interactive loop here taking filters as input and printing the filtered stats. We
+    // would keep all individual stats for each run in a vector.
+    print!("{}", stats);
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    // No arguments reads from stdin.
+    if args.len() <= 1 {
+        let stdin = std::io::stdin();
+        let mut data = Vec::new();
+        stdin.lock().read_to_end(&mut data).unwrap();
+        return debug(&data);
+    }
+    let path = Path::new(&args[1]);
+    // File argument assumes artifact.
+    if path.is_file() && args.len() == 2 {
+        return debug(&std::fs::read(path).unwrap());
+    }
+    // Directory argument assumes corpus.
+    if path.is_dir() {
+        match args[2..].iter().map(|x| x.parse()).collect() {
+            Ok(predicates) => return analyze(path, predicates),
+            Err(error) => eprintln!("Error: {}", error),
+        }
+    }
+    usage(&args[0]);
+}