Skip to content

Commit

Permalink
Merge pull request google#226 from ia0/v2_fuzz
Browse files Browse the repository at this point in the history
Add debug helper for fuzzing
  • Loading branch information
ia0 authored Nov 30, 2020
2 parents 5f80a4b + a0e3048 commit 9d80b87
Showing 1 changed file with 116 additions and 0 deletions.
116 changes: 116 additions & 0 deletions libraries/persistent_store/fuzz/examples/store.rs
Original file line number Diff line number Diff line change
@@ -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]);
}

0 comments on commit 9d80b87

Please sign in to comment.