Quick and Dirty crash minimizer

minimizer.rs

// Quick and Dirty Crash Minimizer.
// Usage: minimize input output -- program_invocation program_flags FILE
// The program has to crash with either SIGSEGV or SIGABRT.
use log::{debug, info};
use nix::sys::ptrace;
use nix::sys::signal::Signal;
use nix::sys::wait::{waitpid, WaitStatus};
use nix::unistd::Pid;
use rand::prelude::SliceRandom;
use std::collections::HashSet;
use std::env;
use std::fs;
use std::io;
use std::io::BufWriter;
use std::io::Read;
use std::io::Write;
use std::os::unix::process::CommandExt;
use std::process::Command;
use std::process::Stdio;

fn print_usage() {
    println!("Usage: ");
    let program = env::args().nth(0).unwrap();
    println!(
        "{} input_file minimized_file -- program_invocation program_flags FILE",
        program
    );
}

fn run_program(program: &String, params: &Vec<String>) -> std::result::Result<bool, io::Error> {
    let mut cmd = Command::new(&program);
    // TODO: The setup + crash detection could be abstracted away.
    cmd.args(params.iter())
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .env_clear()
        .env(
            "ASAN_OPTIONS",
            vec![
                "abort_on_error=1",
                "exitcode=101",
                "detect_leaks=0",
                "symbolize=0",
                "disable_coredump=1",
            ]
            .join(":"),
        );

    let child = unsafe {
        cmd.pre_exec(|| {
            ptrace::traceme().map_err(|e| match e {
                nix::Error::Sys(e) => io::Error::from_raw_os_error(e as i32),
                _ => io::Error::new(io::ErrorKind::Other, "Unknown PTRACE_TRACEME error"),
            })
        })
        .spawn()?
    };

    let pid = Pid::from_raw(child.id() as i32);

    debug!("Process {} launched. Waiting for attachment.", pid);
    // TODO: Timeout for this waitpid?
    match waitpid(Some(pid), None) {
        Ok(WaitStatus::Stopped(_, Signal::SIGTRAP)) => Ok(child),
        _ => Err(io::Error::new(io::ErrorKind::Other, "Invalid child state")),
    }?;

    ptrace::cont(pid, None).unwrap();
    debug!("Waiting for child process to crash/finish");

    // TODO: Timeout for this waitpid?
    match waitpid(Some(pid), None) {
        Ok(WaitStatus::Stopped(_, Signal::SIGSEGV)) => {
            info!("SIGSEGV");
            ptrace::cont(pid, Signal::SIGKILL).unwrap();
            Ok(true)
        }
        Ok(WaitStatus::Stopped(_, Signal::SIGABRT)) => {
            info!("SIGABRT");
            ptrace::cont(pid, Signal::SIGKILL).unwrap();
            Ok(true)
        }
        Ok(WaitStatus::Exited(_, _)) => {
            info!("No crash");
            Ok(false)
        }
        Ok(w) => {
            info!("Unrecognized signal! {:?}", w);
            ptrace::cont(pid, Signal::SIGKILL).unwrap();
            Ok(false)
        }
        Err(nix::Error::Sys(e)) => Err(io::Error::from_raw_os_error(e as i32)),
        _ => Err(io::Error::new(
            io::ErrorKind::Other,
            "Unknown waitpid error",
        )),
    }
}

fn minimize(
    input_file: &String,
    output_file: &String,
    program: &String,
    params: &Vec<String>,
    splitter: fn(&Vec<u8>) -> Vec<&[u8]>,
    iterations: usize,
) {
    let mut rng = rand::thread_rng();
    let mut contents = Vec::new();
    {
        let mut file = fs::File::open(input_file).expect("failed to open file");
        file.read_to_end(&mut contents)
            .expect("unable to read file");
    }

    let mut chunks = splitter(&contents);
    info!("input file has {} chunks", chunks.len());

    for _ in 0..iterations {
        // TODO: it would be nice to allow for a way to group chunks.
        // Possible use cases:
        //  * groups of 5 consecutive lines/bytes.
        //  * pairs of lines.
        // TODO: Maybe consider using a bitset? We are only using skipped
        // to check whether a small integer belongs to it.
        let mut skipped: HashSet<usize> = HashSet::new();
        let mut nums: Vec<usize> = (0..chunks.len()).collect();
        nums.shuffle(&mut rng);

        for idx in nums.iter() {
            let to_skip = *idx;

            debug!("Skipping chunk {}", to_skip);

            {
                let mut file =
                    BufWriter::new(fs::File::create("FILE").expect("failed to create test file"));
                for i in 0..chunks.len() {
                    if i == to_skip || skipped.contains(&i) {
                        continue;
                    }
                    file.write_all(&chunks[i])
                        .expect("failed to write test file");
                }
            }

            if let Ok(true) = run_program(program, params) {
                skipped.insert(to_skip);
                info!("Skipped chunks: {}", skipped.len());
            }
        }

        chunks = chunks
            .into_iter()
            .enumerate()
            .filter(|&(i, _)| !skipped.contains(&i))
            .map(|(_, e)| e)
            .collect();

        if skipped.len() == 0 {
            break;
        }
    }

    {
        info!("Final file has {} chunks", chunks.len());
        let mut file =
            BufWriter::new(fs::File::create(output_file).expect("failed to create test file"));

        for i in 0..chunks.len() {
            file.write_all(&chunks[i])
                .expect("failed to write test file");
        }
    }
}

// split_lines takes a vector of bytes and returns a vector of slices to each of the lines.
fn split_lines(contents: &Vec<u8>) -> Vec<&[u8]> {
    let mut lines = Vec::new();
    let mut left = 0;
    for i in 0..contents.len() {
        if contents[i] == '\n' as u8 {
            lines.push(&contents[left..i + 1]);
            left = i + 1;
        }
    }
    if left < contents.len() {
        lines.push(&contents[left..contents.len()]);
    }
    return lines;
}

// split_bytes takes a vector of bytes and returns a vector of slices to each of the bytes.
fn split_bytes(contents: &Vec<u8>) -> Vec<&[u8]> {
    let mut bytes = Vec::new();
    for i in 0..contents.len() {
        bytes.push(&contents[i..i + 1]);
    }
    return bytes;
}

fn main() {
    env_logger::init();

    let args: Vec<_> = env::args().into_iter().collect();

    if args.len() < 5 || args[3] != "--" {
        print_usage();
        return;
    }

    let input_file = &args[1];
    let output_file = &args[2];

    fs::copy(input_file, "FILE").expect("failed to create temp file");

    if let Ok(false) = run_program(&args[4], &args[5..].to_vec()) {
        println!("Program Does not crash. Exiting");
        return;
    }

    info!("Minimizing Lines");
    minimize(
        input_file,
        &String::from("minimized_lines"),
        &args[4],
        &args[5..].to_vec(),
        split_lines,
        10,
    );

    info!("Minimizing Bytes");
    minimize(
        &String::from("minimized_lines"),
        output_file,
        &args[4],
        &args[5..].to_vec(),
        split_bytes,
        10,
    );

    fs::remove_file("FILE").expect("failed to remove temporary file");
}