ziggy 1.6.0

mod add_seeds;
mod build;
mod clean;
mod coverage;
mod fuzz;
mod minimize;
mod plot;
mod run;
mod triage;
mod util;

use crate::fuzz::FuzzingConfig;
use anyhow::{Context, Result, anyhow};
use clap::{Args, Parser, Subcommand, ValueEnum};
use std::{
    fs,
    path::PathBuf,
    sync::{Arc, atomic::AtomicBool},
};

pub const DEFAULT_UNMODIFIED_TARGET: &str = "automatically guessed";

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
pub enum FuzzingEngines {
    All,
    AFLPlusPlus,
    Honggfuzz,
}

pub const DEFAULT_OUTPUT_DIR: &str = "./output";

pub const DEFAULT_CORPUS_DIR: &str = "{ziggy_output}/{target_name}/corpus/";

pub const DEFAULT_COVERAGE_DIR: &str = "{ziggy_output}/{target_name}/coverage/";

pub const DEFAULT_MINIMIZATION_DIR: &str = "{ziggy_output}/{target_name}/corpus_minimized/";

pub const DEFAULT_PLOT_DIR: &str = "{ziggy_output}/{target_name}/plot/";

pub const DEFAULT_CRASHES_DIR: &str = "{ziggy_output}/{target_name}/crashes/";

pub const DEFAULT_TRIAGE_DIR: &str = "{ziggy_output}/{target_name}/triage/";

#[derive(Parser)]
#[clap(name = "cargo")]
#[clap(bin_name = "cargo")]
pub enum Cargo {
    #[clap(subcommand)]
    Ziggy(Ziggy),
}

#[derive(Subcommand)]
#[clap(
    author,
    version,
    about = "A multi-fuzzer management utility for all of your Rust fuzzing needs 🧑‍🎤"
)]
pub enum Ziggy {
    /// Build the fuzzer and the runner binaries
    Build(Build),

    /// Fuzz targets using different fuzzers in parallel
    Fuzz(Fuzz),

    /// Run a specific input or a directory of inputs to analyze backtrace
    Run(Run),

    /// Minimize the input corpus using the given fuzzing target
    Minimize(Minimize),

    /// Generate code coverage information using the existing corpus
    Cover(Cover),

    /// Plot AFL++ data using afl-plot
    Plot(Plot),

    /// Add seeds to the running AFL++ fuzzers
    AddSeeds(AddSeeds),

    /// Triage crashes found with CASR - currently only works for AFL++
    Triage(Triage),

    /// Remove generated artifacts from the target directory
    Clean(Clean),
}

#[derive(Args)]
pub struct Build {
    /// No AFL++ (Fuzz only with honggfuzz)
    #[clap(long = "no-afl", action)]
    no_afl: bool,

    /// No honggfuzz (Fuzz only with AFL++)
    #[clap(long = "no-honggfuzz", action)]
    no_honggfuzz: bool,

    /// Compile in release mode (--release)
    #[clap(long = "release", action)]
    release: bool,

    /// Build with ASAN (nightly only)
    #[clap(long = "asan", action)]
    asan: bool,
}

#[derive(Args)]
pub struct Fuzz {
    /// Target to fuzz
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Shared corpus directory
    #[clap(short, long, value_parser, value_name = "DIR", default_value = DEFAULT_CORPUS_DIR)]
    corpus: PathBuf,

    /// Initial corpus directory (will only be read)
    #[clap(short, long, value_parser, value_name = "DIR")]
    initial_corpus: Option<PathBuf>,

    /// Compile in release mode (--release)
    #[clap(long = "release", action)]
    release: bool,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,

    /// Number of concurrent fuzzing jobs
    #[clap(short, long, value_name = "NUM", default_value_t = 1)]
    jobs: u32,

    /// Timeout for a single run
    #[clap(short, long, value_name = "SECS")]
    timeout: Option<u32>,

    /// Memory limit for the fuzz target. (If fuzzing with honggfuzz, a numeric value in MiB must be specified)
    #[clap(short, long, value_name = "STRING")]
    memory_limit: Option<String>,

    /// Perform initial minimization
    #[clap(short = 'M', long, action, default_value_t = false)]
    minimize: bool,

    /// Dictionary file (format:<http://llvm.org/docs/LibFuzzer.html#dictionaries>)
    #[clap(short = 'x', long = "dict", value_name = "FILE")]
    dictionary: Option<PathBuf>,

    /// Maximum length of input
    #[clap(short = 'G', long = "maxlength", default_value_t = 1048576)]
    max_length: u64,

    /// Minimum length of input (AFL++ only)
    #[clap(short = 'g', long = "minlength", default_value_t = 1)]
    min_length: u64,

    /// No AFL++ (Fuzz only with honggfuzz)
    #[clap(long = "no-afl", action)]
    no_afl: bool,

    /// No honggfuzz (Fuzz only with AFL++)
    #[clap(long = "no-honggfuzz", action)]
    no_honggfuzz: bool,

    // This value helps us create a global timer for our display
    #[clap(skip = std::time::Instant::now())]
    start_time: std::time::Instant,

    /// Pass flags to AFL++ directly
    #[clap(short, long)]
    afl_flags: Vec<String>,

    /// AFL++ configuration
    #[clap(short = 'C', long, default_value = "generic")]
    config: FuzzingConfig,

    /// With a coverage worker
    #[clap(long)]
    coverage_worker: bool,

    /// Coverage generation interval in minutes
    #[clap(long, default_value = "15")]
    coverage_interval: u64,

    /// Corpus sync interval in minutes
    #[clap(long, default_value = "10")]
    corpus_sync_interval: u64,

    /// Fuzz an already AFL++ instrumented binary; the ziggy way
    #[clap(short, long)]
    binary: Option<PathBuf>,

    /// Build with ASAN (nightly only)
    #[clap(long = "asan", action)]
    asan: bool,

    /// Foreign fuzzer directories to sync with (AFL++ -F option)
    #[clap(long = "foreign-sync", short = 'F', action)]
    foreign_sync_dirs: Vec<PathBuf>,
}

#[derive(Args)]
pub struct Run {
    /// Target to use
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Input directories and/or files to run
    #[clap(short, long, value_name = "DIR", default_value = DEFAULT_CORPUS_DIR)]
    inputs: Vec<PathBuf>,

    /// Recursively run nested directories for all input directories
    #[clap(short, long)]
    recursive: bool,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,

    /// Build with ASAN (nightly only)
    #[clap(long = "asan", action)]
    asan: bool,

    /// Activate these features on the target
    #[clap(short = 'F', long, num_args = 0..)]
    features: Vec<String>,

    /// Stop the run after the first crash is encountered
    #[clap(short = 'x', long)]
    stop_on_crash: bool,
}

#[derive(Args, Clone)]
pub struct Minimize {
    /// Target to use
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Corpus directory to minimize
    #[clap(short, long, default_value = DEFAULT_CORPUS_DIR)]
    input_corpus: PathBuf,

    /// Minimized corpus output directory
    #[clap(short, long, default_value = DEFAULT_MINIMIZATION_DIR)]
    output_corpus: PathBuf,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,

    /// Number of concurrent minimizing jobs (AFL++ only)
    #[clap(short, long, value_name = "NUM", default_value_t = 1)]
    jobs: u32,

    /// Timeout for a single run
    #[clap(short, long, value_name = "MILLI_SECS", default_value_t = 5000)]
    timeout: u32,

    #[clap(short, long, value_enum, default_value_t = FuzzingEngines::All)]
    engine: FuzzingEngines,
}

#[derive(Args)]
pub struct Cover {
    /// Target to generate coverage for
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Output directory for code coverage report
    #[clap(short, long, value_parser, value_name = "DIR", default_value = DEFAULT_COVERAGE_DIR)]
    output: PathBuf,

    /// Input corpus directory to run target on
    #[clap(short, long, value_parser, value_name = "DIR", default_value = DEFAULT_CORPUS_DIR)]
    input: PathBuf,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,

    /// Source directory of covered code
    #[clap(short, long, value_parser, value_name = "DIR")]
    source: Option<PathBuf>,

    /// Keep coverage data files (WARNING: Do not use if source code has changed)
    #[clap(short, long, default_value_t = false)]
    keep: bool,

    /// Comma separated list of output types. See grcov --help to see supported output types. Default: html
    #[clap(short = 't', long)]
    output_types: Option<String>,

    /// Number of concurrent jobs
    #[clap(short, long, value_name = "NUM")]
    jobs: Option<usize>,
}

#[derive(Args)]
pub struct Plot {
    /// Target to generate plot for
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Name of AFL++ fuzzer to use as data source
    #[clap(short, long, value_name = "NAME", default_value = "mainaflfuzzer")]
    input: String,

    /// Output directory for plot
    #[clap(short, long, value_parser, value_name = "DIR", default_value = DEFAULT_PLOT_DIR)]
    output: PathBuf,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,
}

#[derive(Args)]
pub struct Triage {
    /// Target to use
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Triage output directory to be written to (will be overwritten)
    #[clap(short, long, value_name = "DIR", default_value = DEFAULT_TRIAGE_DIR)]
    output: PathBuf,

    /// Number of concurrent fuzzing jobs
    #[clap(short, long, value_name = "NUM", default_value_t = 1)]
    jobs: u32,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,

    /// Terminate runner after x seconds
    #[clap(short, long, value_name = "SECS")]
    timeout: Option<u32>,
    /* future feature, wait for casr
    /// Crash directory to be sourced from
    #[clap(short, long, value_parser, value_name = "DIR", default_value = DEFAULT_CRASHES_DIR)]
    input: PathBuf,
    */
}

#[derive(Args)]
pub struct AddSeeds {
    /// Target to use
    #[clap(value_name = "TARGET", default_value = DEFAULT_UNMODIFIED_TARGET)]
    target: String,

    /// Seeds directory to be added
    #[clap(short, long, value_parser, value_name = "DIR")]
    input: PathBuf,

    /// Fuzzers output directory
    #[clap(
        short, long, env = "ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value = DEFAULT_OUTPUT_DIR
    )]
    ziggy_output: PathBuf,
}

#[derive(Args)]
pub struct Clean {
    /// Arguments passed through to cargo clean, see cargo clean --help
    #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
    args: Vec<std::ffi::OsString>,
}

#[derive(Debug)]
pub struct Common {
    terminate: Arc<AtomicBool>,
    sigs_done: Option<()>,
    pub cargo_path: PathBuf,
}

impl Common {
    fn new() -> Self {
        Self {
            terminate: Arc::new(AtomicBool::new(false)),
            sigs_done: Some(()),
            cargo_path: std::env::var("CARGO")
                .unwrap_or_else(|_| String::from("cargo"))
                .into(),
        }
    }
    fn is_terminated(&self) -> bool {
        self.terminate.load(std::sync::atomic::Ordering::Acquire)
    }

    fn shutdown_deferred(&self) {
        self.terminate
            .store(false, std::sync::atomic::Ordering::Release);
    }

    fn shutdown_immediate(&self) {
        self.terminate
            .store(true, std::sync::atomic::Ordering::Release);
    }

    fn setup_signal_handling(&mut self) -> Result<(), anyhow::Error> {
        if self.sigs_done.take().is_some() {
            for signal in signal_hook::consts::TERM_SIGNALS {
                signal_hook::flag::register_conditional_shutdown(
                    *signal,
                    1,
                    Arc::clone(&self.terminate),
                )
                .context("Setting up signal handler")?;
                signal_hook::flag::register(*signal, Arc::clone(&self.terminate))
                    .context("Setting up signal handler")?;
            }
        }
        Ok(())
    }

    fn cargo(&self) -> std::process::Command {
        let mut cmd = std::process::Command::new(&self.cargo_path);
        cmd.stdin(std::process::Stdio::null());
        cmd
    }
}

fn main() -> Result<(), anyhow::Error> {
    let mut common = Common::new();
    common.shutdown_immediate();
    common.setup_signal_handling()?;

    let Cargo::Ziggy(command) = Cargo::parse();
    match command {
        Ziggy::Build(args) => args.build().context("Failed to build the fuzzers"),
        Ziggy::Fuzz(mut args) => args.fuzz(&common).context("Failure running fuzzers"),
        Ziggy::Run(mut args) => args.run().context("Failure running inputs"),
        Ziggy::Minimize(mut args) => args.minimize().context("Failure running minimization"),
        Ziggy::Cover(mut args) => args
            .generate_coverage()
            .context("Failure generating coverage"),
        Ziggy::Plot(mut args) => args.generate_plot().context("Failure generating plot"),
        Ziggy::AddSeeds(args) => args.add_seeds().context("Failure adding seeds to AFL"),
        Ziggy::Triage(args) => args.triage().context("Failure triaging with casr"),
        Ziggy::Clean(args) => args.clean().context("Failure cleaning build artifacts"),
    }
}

pub fn find_target(target: &String) -> Result<String, anyhow::Error> {
    // If the target is already set, we're done here
    if target != DEFAULT_UNMODIFIED_TARGET {
        return Ok(target.into());
    }

    let new_target_result = guess_target();

    new_target_result.context("Target is not obvious")
}

fn guess_target() -> Result<String> {
    let metadata = cargo_metadata::MetadataCommand::new().exec()?;
    let default_bin = metadata
        .workspace_default_members
        .iter()
        .find_map(|default_pkg| {
            metadata
                .packages
                .iter()
                .find(|p| p.id == *default_pkg)
                .and_then(|p| {
                    p.targets
                        .iter()
                        .find_map(|target| target.is_bin().then(|| target.name.clone()))
                })
        });

    default_bin.ok_or_else(|| anyhow!("Please specify a target"))
}

fn target_dir() -> &'static cargo_metadata::camino::Utf8PathBuf {
    use std::sync::LazyLock;

    static TARGET_DIR: LazyLock<cargo_metadata::camino::Utf8PathBuf> = LazyLock::new(|| {
        cargo_metadata::MetadataCommand::new().exec().map_or_else(
            |_| cargo_metadata::camino::Utf8PathBuf::from("target"),
            |metadata| metadata.target_directory,
        )
    });

    &TARGET_DIR
}