canfuzz 0.1.0

//! This module orchestrates the fuzzing process using the `libafl` fuzzing framework.
//!
//! It defines the `FuzzerOrchestrator` trait, which provides a generic interface
//! for setting up and running fuzz tests against IC canisters. The main `run` function
//! configures and starts the `libafl` fuzzing loop, while the `test_one_input` function
//! provides a convenient way to debug specific inputs.

use candid::Principal;
use chrono::Local;
use ic_management_canister_types::CanisterId;
use libafl::feedback_or;
use libafl::feedbacks::{ExitKindFeedback, TimeoutFeedback};
use pocket_ic::PocketIc;
use std::fs::{self, File};
use std::io::Read;
use std::path::PathBuf;
use std::sync::Arc;

use crate::custom::oom_exit_kind::OomLogic;
use crate::libafl::{
    Evaluator,
    corpus::CachedOnDiskCorpus,
    events::SimpleEventManager,
    executors::{ExitKind, inprocess::InProcessExecutor},
    feedbacks::{CrashFeedback, map::AflMapFeedback},
    fuzzer::{Fuzzer, StdFuzzer},
    inputs::BytesInput,
    mutators::{HavocScheduledMutator, havoc_mutations},
    observers::map::{StdMapObserver, hitcount_map::HitcountsMapObserver},
    schedulers::QueueScheduler,
    stages::{AflStatsStage, CalibrationStage, mutational::StdMutationalStage},
    state::StdState,
};

use crate::libafl::monitors::SimpleMonitor;
// use libafl::monitors::tui::{ui::TuiUI, TuiMonitor};
use crate::libafl_bolts::{current_nanos, rands::StdRand, tuples::tuple_list};

use crate::constants::COVERAGE_FN_EXPORT_NAME;
use crate::fuzzer::FuzzerState;

/// A trait for types that can provide access to the global `FuzzerState`.
pub trait FuzzerStateProvider {
    /// Returns a reference to the `FuzzerState`.
    fn get_fuzzer_state(&self) -> &FuzzerState;
}

/// A trait that defines the necessary components for a canister fuzzing target.
///
/// Implementors of this trait provide the specific logic for setting up the environment,
/// executing a test case against one or more canisters, and cleaning up afterwards.
pub trait FuzzerOrchestrator: FuzzerStateProvider {
    /// Performs one-time initialization at the start of the fuzzing campaign.
    /// This is where canisters are typically installed.
    fn init(&mut self);

    /// Sets up the environment before each execution of a test case.
    /// This could involve resetting canister state to a clean snapshot.
    fn setup(&self) {}

    /// Executes a single fuzzing input against the target canister(s).
    ///
    /// # Arguments
    ///
    /// * `input` - The `BytesInput` generated by the fuzzer.
    ///
    /// # Returns
    ///
    /// * `ExitKind` - Indicates the outcome of the execution (e.g., `Ok`, `Crash`).
    fn execute(&self, input: BytesInput) -> ExitKind;

    /// Returns a thread-safe reference to the `PocketIc` instance.
    fn get_state_machine(&self) -> Arc<PocketIc> {
        self.get_fuzzer_state().get_state_machine()
    }

    /// Returns the `CanisterId` of the canister that has been instrumented for coverage.
    fn get_coverage_canister_id(&self) -> CanisterId {
        self.get_fuzzer_state().get_coverage_canister_id()
    }

    /// Creates and returns the path to a new timestamped directory for storing input items.
    ///
    /// The directory is structured as `$OUT_DIR/artifacts/<fuzzer_name>/<timestamp>/input`,
    /// where `$OUT_DIR` is the build script output directory (e.g., `target/debug/build/.../out`),
    /// `<fuzzer_name>` is the name provided to `FuzzerState::new`, and `<timestamp>` is based
    /// on the current time.
    ///
    /// # Panics
    ///
    /// Panics if the `OUT_DIR` environment variable is not set or if the directory cannot be created.
    fn input_dir(&self) -> PathBuf {
        let out_dir = std::env::var("OUT_DIR").expect("OUT_DIR is not set");
        let input_dir = PathBuf::from(out_dir)
            .join("artifacts")
            .join(self.get_fuzzer_state().name())
            .join(Local::now().format("%Y%m%d_%H%M").to_string())
            .join("input");
        fs::create_dir_all(&input_dir)
            .unwrap_or_else(|e| panic!("Failed to create input directory {input_dir:?}: {e}"));
        println!("Input directory: {input_dir:?}");
        input_dir
    }

    /// Creates and returns the path to a new timestamped directory for storing crashes.
    ///
    /// The directory is structured as `$OUT_DIR/artifacts/<fuzzer_name>/<timestamp>/crashes`,
    /// where `$OUT_DIR` is the build script output directory (e.g., `target/debug/build/.../out`),
    /// `<fuzzer_name>` is the name provided to `FuzzerState::new`, and `<timestamp>` is based
    /// on the current time.
    ///
    /// # Panics
    ///
    /// Panics if the `OUT_DIR` environment variable is not set or if the directory cannot be created.
    fn crashes_dir(&self) -> PathBuf {
        let out_dir = std::env::var("OUT_DIR").expect("OUT_DIR is not set");
        let crashes_dir = PathBuf::from(out_dir)
            .join("artifacts")
            .join(self.get_fuzzer_state().name())
            .join(Local::now().format("%Y%m%d_%H%M").to_string())
            .join("crashes");
        fs::create_dir_all(&crashes_dir)
            .unwrap_or_else(|e| panic!("Failed to create crashes directory {crashes_dir:?}: {e}"));
        println!("Crashes directory: {crashes_dir:?}");
        crashes_dir
    }

    /// Returns the path to the seed corpus directory.
    ///
    /// This directory should contain initial valid inputs to kickstart the fuzzing process.
    fn corpus_dir(&self) -> PathBuf;

    /// Fetches the coverage map from the instrumented canister and updates the global `COVERAGE_MAP`.
    ///
    /// It makes a query call to the `export_coverage` function on the coverage canister.
    /// If the query fails, the coverage map is not updated.
    #[allow(static_mut_refs)]
    fn set_coverage_map(&self) {
        let test = self.get_state_machine();
        let result = test.update_call(
            self.get_coverage_canister_id(),
            Principal::anonymous(),
            COVERAGE_FN_EXPORT_NAME,
            vec![],
        );
        if let Ok(result) = result {
            unsafe { crate::instrumentation::COVERAGE_MAP.copy_from_slice(&result) };
        }
    }

    /// Provides a mutable reference to the global `COVERAGE_MAP`.
    fn get_coverage_map(&self) -> &'static mut [u8] {
        unsafe { crate::instrumentation::COVERAGE_MAP }
    }

    /// The main entry point for running a fuzzing campaign.
    ///
    /// This function orchestrates the entire fuzzing process:
    /// 1. Calls `self.init()` for one-time setup.
    /// 2. Defines a `harness` closure that wraps `self.execute()` and updates the coverage map.
    /// 3. Sets up `libafl` components:
    ///    - A `HitcountsMapObserver` to monitor the `COVERAGE_MAP`.
    ///    - `AflMapFeedback` for coverage-guided feedback and `CrashFeedback` for finding crashes.
    ///    - A `StdState` to hold the fuzzer's state (corpus, solutions, etc.).
    ///    - A `SimpleEventManager` with a `SimpleMonitor` for logging.
    ///    - A `QueueScheduler` to decide which input to fuzz next.
    ///    - An `InProcessExecutor` to run the harness.
    /// 4. Loads the initial seed corpus from the directory provided by `corpus_dir()`.
    /// 5. Configures mutational stages, including a `HavocScheduledMutator`.
    /// 6. Starts the main fuzzing loop.
    fn run(&mut self) {
        self.init();

        // The harness is a closure that `libafl` will call for each fuzzed input.
        let mut harness = |input: &BytesInput| {
            self.setup();
            let result = self.execute(input.clone());
            self.set_coverage_map();
            result
        };

        let hitcount_map_observer = HitcountsMapObserver::new(unsafe {
            StdMapObserver::new("coverage_map", self.get_coverage_map())
        });

        // Feedback mechanisms tell the fuzzer if an input is "interesting"
        let afl_map_feedback = AflMapFeedback::new(&hitcount_map_observer);
        let mut feedback = afl_map_feedback;
        let calibration_stage = CalibrationStage::new(&feedback);

        // The objective is to find crashes, timeouts or oom
        let crash_feedback = CrashFeedback::new();
        let timeout_feedback = TimeoutFeedback::new();
        let oom_feedback: ExitKindFeedback<OomLogic> = ExitKindFeedback::new();
        let mut objective = feedback_or!(crash_feedback, timeout_feedback, oom_feedback);

        // A stats stage to print statistics about the fuzzing run.
        let stats_stage = AflStatsStage::builder()
            .map_observer(&hitcount_map_observer)
            .build()
            .unwrap();

        let mut state = StdState::new(
            StdRand::with_seed(current_nanos()),
            CachedOnDiskCorpus::new(self.input_dir(), 512).unwrap(),
            CachedOnDiskCorpus::new(self.crashes_dir(), 512).unwrap(),
            &mut feedback,
            &mut objective,
        )
        .unwrap();

        let mon = SimpleMonitor::new(|s| println!("{s}"));
        // A TUI monitor can be used for a more sophisticated display.
        // Example:
        // let ui = TuiUI::with_version(String::from("My Fuzzer"), String::from("0.1.0"), false);
        // let mon = TuiMonitor::new(ui);
        let mut mgr = SimpleEventManager::new(mon);
        let scheduler = QueueScheduler::new();
        let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);

        let mut executor = InProcessExecutor::new(
            &mut harness,
            tuple_list!(hitcount_map_observer),
            &mut fuzzer,
            &mut state,
            &mut mgr,
        )
        .expect("Failed to create the Executor");

        // Load initial inputs from the corpus directory
        let paths = fs::read_dir(self.corpus_dir()).unwrap();
        for path in paths {
            let p = path.unwrap().path();
            let mut f = File::open(p.clone()).unwrap();
            let mut buffer = Vec::new();
            f.read_to_end(&mut buffer).unwrap();
            fuzzer
                .evaluate_input(
                    &mut state,
                    &mut executor,
                    &mut mgr,
                    &BytesInput::new(buffer),
                )
                .unwrap();
        }
        // Standard mutational stage with a havoc mutator
        let mutator = HavocScheduledMutator::new(havoc_mutations());
        let mut stages = tuple_list!(
            calibration_stage,
            StdMutationalStage::new(mutator),
            stats_stage
        );

        // Start the fuzzing loop!
        fuzzer
            .fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)
            .expect("Error in the fuzzing loop");
    }

    /// Executes a single input against the orchestrator's harness.
    ///
    /// This function is useful for debugging specific inputs, such as those that
    /// have caused a crash, without running the full fuzzing loop. It calls
    /// `init`, `setup`, `execute` in sequence for the given input.
    ///
    /// # Arguments
    ///
    /// * `bytes` - The raw byte vector of the input to be tested.
    fn test_one_input(&mut self, bytes: Vec<u8>) {
        self.init();
        self.setup();
        let result = self.execute(BytesInput::new(bytes));
        println!("Execution result: {result:?}");
    }
}