Tiny Fuzzer

tinyfuzz.rs

pub fn target(a: i32, b: i32, c: i32) {
    if a > b {
        instrument::branch();

        // ..snip

        if b > c {
            instrument::branch();

            // ..snip
        } else {
            instrument::branch();

            // ..snip
        }
    } else if a > c {
        instrument::branch();

        // ..snip
    } else {
        instrument::branch();

        // ..snip
    }
}

pub mod instrument {
    /*
    Infrastructure for our fuzzer.
    */

    use crate::hash;
    use std::{collections::HashMap, mem, panic::Location, sync::Mutex};

    /*
    A global value representing the current path through an instrumented function.

    This value is mutated by the [`branch`] function.
    */
    static CURRENT_PATH: Mutex<u64> = Mutex::new(0);

    /*
    Track a branch in an instrumented function.

    This function needs to be inserted at each branch in an instrumented function.
    */
    #[track_caller]
    pub fn branch() {
        // Get a value representing the current point in the program
        let loc = Location::caller();

        // Compute the hash of the location
        let hash = hash::of(loc);

        // Add the hash to the current path
        let mut path = CURRENT_PATH.lock().unwrap();

        *path += 1;
        *path ^= hash;
    }

    /*
    Execute an instrumented function, returning an identifier for the path it took.
    */
    pub fn exec(target: impl Fn()) -> u64 {
        // Clear the current path in case it was left inconsistent
        *CURRENT_PATH.lock().unwrap() = 0;

        // Execute the instrumented function
        target();

        // Take the final value representing the path `target` took
        mem::take(&mut *CURRENT_PATH.lock().unwrap())
    }

    /*
    Fuzz a target function.

    The target needs to be instrumented using [`branch`]. To assist the fuzzer,
    the caller provides two functions that manipulate `target`'s inputs:

    - `reduce`: A function that aims to simplify arguments without changing the
      execution path taken. The fuzzer will check whether a reduction produces the
      same path or not.
    - `permute`: A function that aims to scramble arguments to find new execution paths.
    */
    pub fn fuzz<T>(
        iterations: usize,
        cases: Vec<T>,
        reduce: impl Fn(&T) -> Vec<T>,
        permute: impl Fn(&T) -> Vec<T>,
        target: impl Fn(&T),
    ) -> Vec<T> {
        // Start with an empty corpus and a default test case
        let mut corpus = HashMap::new();
        let mut queue = cases;

        for _ in 0..iterations {
            // Drain the queue, invoking the target function for each of its inputs
            while let Some(args) = queue.pop() {
                // Execute our instrumented function
                let path = exec(|| target(&args));

                // Record this path in our corpus, if it isn't present already
                let args = corpus.entry(path).or_insert(args);

                // Try reduce the arguments to a simpler form
                for reduction in reduce(args) {
                    if exec(|| target(&reduction)) == path {
                        *args = reduction;
                    }
                }
            }

            // Refill the queue with a new set of inputs
            for args in corpus.values() {
                for permutation in permute(args) {
                    queue.push(permutation);
                }
            }
        }

        // Return the final corpus of inputs we discovered
        corpus.into_values().collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn check() {
        let corpus = instrument::fuzz::<[i32; 3]>(
            100_000,
            // Start with a default input
            vec![[0, 0, 0]],
            // Reduce our variables together
            |args| vec![[args[0] / 2, args[1] / 2, args[2] / 2]],
            // Permute our variables independently
            |args| {
                (0..args.len())
                    .into_iter()
                    .map(|i| {
                        let mut args = *args;
                        args[i] = rand::random();

                        args
                    })
                    .collect()
            },
            // Fuzz our `target` function
            |args| target(args[0], args[1], args[2]),
        );

        for case in corpus {
            println!("{case:?}");
        }
    }
}