canfuzz 0.6.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 [`FuzzerOrchestrator::run`]
//! function configures and starts the `libafl` fuzzing loop, while
//! [`FuzzerOrchestrator::test_one_input`] provides a convenient way to debug specific inputs.
//!
//! When [`FuzzerOrchestrator::instruction_config`] returns an [`InstructionConfig`] with
//! `enabled: true`, the fuzzer additionally tracks instruction counts via
//! [`set_instruction_count`](FuzzerOrchestrator::set_instruction_count) and uses
//! [`InstructionCountFeedback`](crate::custom::feedback::instruction_count::InstructionCountFeedback)
//! to guide inputs toward higher instruction consumption. Each time the maximum instruction
//! count is broken, a detailed log line is printed and the input is saved to disk.
//! If [`InstructionConfig::max_instruction_count`] is set, inputs that exceed the threshold
//! are treated as crashes.

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, Write as IoWrite};
use std::path::PathBuf;
use std::sync::Arc;

/// Diagnostic information about the current fuzzing session.
/// Stored in a global [`OnceLock`] so that panic hooks and Ctrl+C handlers can print
/// a summary before the process exits.
struct SessionInfo {
    name: String,
    corpus_dir: PathBuf,
    input_dir: PathBuf,
    crashes_dir: PathBuf,
    rng_seed: u64,
}

static SESSION_INFO: std::sync::OnceLock<SessionInfo> = std::sync::OnceLock::new();

/// Print a diagnostic summary of the current fuzzing session to stderr.
fn print_session_info() {
    if let Some(info) = SESSION_INFO.get() {
        eprintln!("\n### Fuzzer session summary ###");
        eprintln!("  Fuzzer:       {}", info.name);
        eprintln!("  Seed corpus:  {}", info.corpus_dir.display());
        eprintln!("  Input dir:    {}", info.input_dir.display());
        eprintln!("  Crashes dir:  {}", info.crashes_dir.display());
        eprintln!("  RNG seed:     {}", info.rng_seed);
        eprintln!("#############################");
    }
}

use crate::custom::feedback::oom_exit_kind::OomLogic;
use crate::custom::mutator::candid::{CandidParserMutator, CandidTypeDefArgs};
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::{
        CanTrack,
        map::{StdMapObserver, hitcount_map::HitcountsMapObserver},
    },
    schedulers::{
        IndexesLenTimeMinimizerScheduler, StdWeightedScheduler, powersched::PowerSchedule,
    },
    stages::{AflStatsStage, CalibrationStage, StdPowerMutationalStage},
    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, INSTRUCTION_COUNT_FN_EXPORT_NAME};
use crate::custom::observer::instruction_count::INSTRUCTION_MAP;
use crate::fuzzer::FuzzerState;

/// Configuration for instruction count maximization.
///
/// Returned by [`FuzzerOrchestrator::instruction_config`]. When `enabled` is true,
/// the fuzzer tracks instruction counts and considers inputs that increase the maximum
/// as "interesting".
#[derive(Debug, Clone, Default)]
pub struct InstructionConfig {
    /// Enable instruction count maximization feedback.
    pub enabled: bool,
    /// If set, inputs whose instruction count exceeds this threshold are treated as crashes.
    pub max_instruction_count: Option<u64>,
}

/// 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: AsRef<FuzzerState> + AsMut<FuzzerState> {
    /// 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.as_ref().get_state_machine()
    }

    /// Returns the `CanisterId` of the canister that has been instrumented for coverage.
    fn get_coverage_canister_id(&self) -> CanisterId {
        self.as_ref().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.as_ref().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.as_ref().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 update call to the `__export_coverage_for_afl` function on the coverage canister.
    /// If the update 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 }
    }

    /// Provides configuration for the `CandidParserMutator`.
    ///
    /// By default, this returns `None`, which disables the Candid-aware mutator.
    /// To enable it, override this method in your fuzzer implementation to return
    /// `Some(CandidTypeDefArgs { ... })`. You will need to provide the path to the
    /// `.did` file and the name of the canister method you intend to fuzz.
    ///
    /// This allows the fuzzer to perform structure-aware mutations on Candid-encoded inputs.
    fn get_candid_args() -> Option<CandidTypeDefArgs> {
        None
    }

    /// Returns configuration for instruction count maximization.
    ///
    /// Override this to return an [`InstructionConfig`] with `enabled: true` to track
    /// instruction counts and guide the fuzzer toward inputs that consume more IC instructions.
    /// Optionally set `max_instruction_count` to a threshold — inputs that exceed it will be
    /// treated as crashes.
    /// Requires `instrument_instruction_count: true` in `InstrumentationArgs`.
    fn instruction_config() -> InstructionConfig {
        InstructionConfig::default()
    }

    /// Fetches the instruction count from the instrumented canister and updates the global `INSTRUCTION_MAP`.
    ///
    /// It makes a query call to the `__export_instruction_count_for_afl` function on the coverage canister.
    /// If the instruction count exceeds the previous maximum, the input is marked as interesting.
    /// Returns `true` if the instruction count exceeded the configured
    /// [`InstructionConfig::max_instruction_count`] threshold (i.e. should be treated as a crash).
    #[allow(static_mut_refs)]
    fn set_instruction_count(&self, input: &BytesInput) -> bool {
        let test = self.get_state_machine();
        let result = test.query_call(
            self.get_coverage_canister_id(),
            Principal::anonymous(),
            INSTRUCTION_COUNT_FN_EXPORT_NAME,
            vec![],
        );
        if let Ok(result) = result
            && result.len() >= 8
        {
            let instructions = u64::from_le_bytes(result[0..8].try_into().unwrap());
            let mut map = unsafe { INSTRUCTION_MAP.borrow_mut() };
            if instructions > map.max_instructions {
                let prev = map.max_instructions;
                map.increased = true;
                map.max_instructions = instructions;

                let input_bytes: Vec<u8> = input.clone().into();
                let input_len = input_bytes.len();
                let timestamp = Local::now().format("%Y-%m-%d %H:%M:%S");

                // Save the input and log to the corpus directory
                let corpus_dir = self.corpus_dir();
                let corpus_file = corpus_dir.join(format!("max_instructions_{instructions}"));
                if let Ok(mut f) = File::create(&corpus_file) {
                    let _ = f.write_all(&input_bytes);
                }

                // Hex preview of input bytes (first 64 bytes)
                let hex_preview: String = input_bytes
                    .iter()
                    .take(64)
                    .map(|b| format!("{b:02x}"))
                    .collect();
                let truncated = if input_len > 64 { "..." } else { "" };

                let log_line = format!(
                    "[instructions] NEW MAX | timestamp: {timestamp} | instructions: {instructions} (prev: {prev}) | input_len: {input_len} | hex: {hex_preview}{truncated} | corpus_file: {}",
                    corpus_file.display()
                );
                println!("{log_line}");

                // Append to log file
                let log_path = corpus_dir.join("instruction_log.txt");
                if let Ok(mut f) = fs::OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(&log_path)
                {
                    let _ = writeln!(f, "{log_line}");
                }
            } else {
                map.increased = false;
            }
            map.current_instructions = instructions;

            // Check threshold
            if let Some(threshold) = Self::instruction_config().max_instruction_count
                && instructions > threshold
            {
                return true;
            }
        }
        false
    }

    /// 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
    ///    (and optionally the instruction count).
    /// 3. Sets up `libafl` components:
    ///    - A `HitcountsMapObserver` to monitor the `COVERAGE_MAP`.
    ///    - `AflMapFeedback` for coverage-guided feedback and `CrashFeedback` for finding crashes.
    ///    - Optionally, `InstructionCountObserver` and `InstructionCountFeedback` when
    ///      [`instruction_config`](Self::instruction_config) has `enabled: true`.
    ///    - 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.
    #[allow(static_mut_refs)]
    fn run(&mut self) {
        self.init();

        let inst_config = Self::instruction_config();

        let mut harness = |input: &BytesInput| {
            self.setup();
            let result = self.execute(input.clone());
            self.set_coverage_map();
            if inst_config.enabled && self.set_instruction_count(input) {
                return ExitKind::Crash;
            }
            result
        };

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

        // AflMapFeedback must be created before the observer is moved into a tuple.
        let afl_map_feedback = AflMapFeedback::new(&hitcount_map_observer);

        let candid_enabled = Self::get_candid_args().is_some();

        // The observer/feedback/stage types differ at compile time depending on
        // configuration, so we branch into different macro invocations here.
        // Each combination of (instruction count, candid mutator) needs its own
        // call because `tuple_list!` is a compile-time construct.
        match (inst_config.enabled, candid_enabled) {
            (true, true) => {
                use crate::custom::feedback::instruction_count::InstructionCountFeedback;
                use crate::custom::observer::instruction_count::INSTRUCTION_COUNT_OBSERVER_NAME;
                use crate::libafl::observers::RefCellValueObserver;
                use crate::libafl_bolts::ownedref::OwnedRef;
                use std::ptr::addr_of;

                let instruction_count_observer = unsafe {
                    RefCellValueObserver::new(
                        INSTRUCTION_COUNT_OBSERVER_NAME,
                        OwnedRef::from_ptr(addr_of!(INSTRUCTION_MAP)),
                    )
                };
                let candid_mutator = CandidParserMutator::new(Self::get_candid_args());

                let feedback =
                    feedback_or!(afl_map_feedback.clone(), InstructionCountFeedback::new());
                run_fuzzing_loop!(
                    self,
                    &mut harness,
                    hitcount_map_observer,
                    (instruction_count_observer),
                    (StdPowerMutationalStage::new(candid_mutator)),
                    afl_map_feedback,
                    feedback
                );
            }
            (true, false) => {
                use crate::custom::feedback::instruction_count::InstructionCountFeedback;
                use crate::custom::observer::instruction_count::INSTRUCTION_COUNT_OBSERVER_NAME;
                use crate::libafl::observers::RefCellValueObserver;
                use crate::libafl_bolts::ownedref::OwnedRef;
                use std::ptr::addr_of;

                let instruction_count_observer = unsafe {
                    RefCellValueObserver::new(
                        INSTRUCTION_COUNT_OBSERVER_NAME,
                        OwnedRef::from_ptr(addr_of!(INSTRUCTION_MAP)),
                    )
                };

                let feedback =
                    feedback_or!(afl_map_feedback.clone(), InstructionCountFeedback::new());
                run_fuzzing_loop!(
                    self,
                    &mut harness,
                    hitcount_map_observer,
                    (instruction_count_observer),
                    (),
                    afl_map_feedback,
                    feedback
                );
            }
            (false, true) => {
                let candid_mutator = CandidParserMutator::new(Self::get_candid_args());
                let feedback = afl_map_feedback.clone();
                run_fuzzing_loop!(
                    self,
                    &mut harness,
                    hitcount_map_observer,
                    (),
                    (StdPowerMutationalStage::new(candid_mutator)),
                    afl_map_feedback,
                    feedback
                );
            }
            (false, false) => {
                let feedback = afl_map_feedback.clone();
                run_fuzzing_loop!(
                    self,
                    &mut harness,
                    hitcount_map_observer,
                    (),
                    (),
                    afl_map_feedback,
                    feedback
                );
            }
        }
    }

    /// 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:?}");
    }
}

/// Macro to avoid duplicating the fuzzing loop for different observer/feedback/stage
/// type tuples. The observer tuple, feedback composition, and mutation stages differ
/// depending on configuration (instruction count, Candid mutator), but the rest of
/// the loop (state, executor, corpus loading) is identical.
///
/// `$map_observer` is the owned hitcount map observer. It is borrowed by the scheduler
/// constructors, then moved into the observer tuple alongside any `$extra_observers`.
/// `$afl_map_feedback` must be an already-constructed `AflMapFeedback` (created from
/// the hitcount observer before the observer is moved into the tuple).
/// `$extra_stages` are inserted before the havoc mutator stage (e.g. Candid mutator).
#[macro_export]
macro_rules! run_fuzzing_loop {
    ($self:expr, $harness:expr, $map_observer:expr, ($($extra_observer:expr),*), ($($extra_stage:expr),*), $afl_map_feedback:expr, $feedback:expr) => {{
        let map_observer = $map_observer;
        let afl_map_feedback = $afl_map_feedback;
        let mut feedback = $feedback;
        let calibration_stage = CalibrationStage::new(&afl_map_feedback);

        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);

        let stats_stage = AflStatsStage::builder()
            .map_feedback(&afl_map_feedback)
            .build()
            .unwrap();

        let rng_seed = current_nanos();
        let input_dir = $self.input_dir();
        let crashes_dir = $self.crashes_dir();
        let corpus_dir = $self.corpus_dir();

        // Store session info for diagnostic output on exit.
        let _ = SESSION_INFO.set(SessionInfo {
            name: $self.as_ref().name().to_string(),
            corpus_dir: corpus_dir.clone(),
            input_dir: input_dir.clone(),
            crashes_dir: crashes_dir.clone(),
            rng_seed,
        });

        // Print session info at startup so the user can see artifact paths.
        print_session_info();

        // Install a panic hook that prints session info before the default handler.
        let default_hook = std::panic::take_hook();
        std::panic::set_hook(Box::new(move |info| {
            print_session_info();
            default_hook(info);
        }));

        // Install a Ctrl+C handler that prints session info before exiting.
        let _ = ctrlc::set_handler(move || {
            print_session_info();
            std::process::exit(130);
        });

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

        // AFL++-style weighted scheduler with FAST power schedule, wrapped in a
        // corpus minimizer that favors short + fast inputs covering rare edges.
        let weighted = StdWeightedScheduler::with_schedule(
            &mut state,
            &map_observer,
            Some(PowerSchedule::fast()),
        );
        let scheduler = IndexesLenTimeMinimizerScheduler::new(&map_observer, weighted);

        let mon = SimpleMonitor::new(|s| println!("{s}"));
        let mut mgr = SimpleEventManager::new(mon);
        let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);

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

        // Load initial inputs from the corpus directory, skipping non-input files.
        fn is_corpus_entry(name: &str) -> bool {
            name != "instruction_log.txt" && name != ".gitignore"
        }
        let corpus_entries: Vec<_> = fs::read_dir(&corpus_dir)
            .unwrap()
            .filter_map(|e| e.ok())
            .map(|e| e.path())
            .filter(|p| p.file_name().and_then(|n| n.to_str()).is_some_and(is_corpus_entry))
            .collect();
        if corpus_entries.is_empty() {
            use rand::RngCore;
            let mut rng = rand::rng();
            let len = (rng.next_u32() % 1024 + 1) as usize;
            let mut buf = vec![0u8; len];
            rng.fill_bytes(&mut buf);
            println!("Corpus was empty — using a randomly generated seed ({len} bytes)");
            fuzzer.evaluate_input(&mut state, &mut executor, &mut mgr, &BytesInput::new(buf)).unwrap();
        }
        for p in &corpus_entries {
            let mut f = File::open(p).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();
        }

        // Power-aware mutation stages: mutation count per corpus entry is scaled
        // by its score (bitmap size, exec time, rarity) instead of random 1-128.
        let havoc_mutator = HavocScheduledMutator::new(havoc_mutations());
        let mut stages = tuple_list!(
            calibration_stage,
            $($extra_stage,)*
            StdPowerMutationalStage::new(havoc_mutator),
            stats_stage
        );

        fuzzer
            .fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)
            .expect("Error in the fuzzing loop");
    }};
}
// Required for the macro to be usable within trait methods above.
use run_fuzzing_loop;