diff --git a/fuzzamoto-ir/src/compiler.rs b/fuzzamoto-ir/src/compiler.rs index ef673713..c199a37c 100644 --- a/fuzzamoto-ir/src/compiler.rs +++ b/fuzzamoto-ir/src/compiler.rs @@ -66,6 +66,8 @@ pub enum CompiledAction { /// Set mock time for all nodes in the test SetTime(u64), Probe, + /// Take an incremental snapshot + IncrementalSnapshot, } #[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] @@ -396,6 +398,12 @@ impl Compiler { self.handle_addr_operations(instruction)?; } + Operation::IncrementalSnapshot => { + self.output + .actions + .push(CompiledAction::IncrementalSnapshot); + } + Operation::BeginWitnessStack | Operation::AddWitness | Operation::EndWitnessStack => { diff --git a/fuzzamoto-ir/src/instruction.rs b/fuzzamoto-ir/src/instruction.rs index 6c975a56..dd8a7d2d 100644 --- a/fuzzamoto-ir/src/instruction.rs +++ b/fuzzamoto-ir/src/instruction.rs @@ -162,6 +162,7 @@ impl Instruction { | Operation::TakeTxo => true, Operation::Nop { .. } + | Operation::IncrementalSnapshot | Operation::BeginBuildTx | Operation::EndBuildTx | Operation::BeginBuildTxInputs diff --git a/fuzzamoto-ir/src/mutators/combine.rs b/fuzzamoto-ir/src/mutators/combine.rs index 5a842c67..f1335aa0 100644 --- a/fuzzamoto-ir/src/mutators/combine.rs +++ b/fuzzamoto-ir/src/mutators/combine.rs @@ -27,10 +27,20 @@ impl Splicer for CombineMutator { program: &mut Program, splice_with: &Program, rng: &mut R, + ) -> MutatorResult { + self.splice_from(program, splice_with, rng, 0) + } + + fn splice_from( + &mut self, + program: &mut Program, + splice_with: &Program, + rng: &mut R, + min_index: usize, ) -> MutatorResult { let combine_index = program - .get_random_instruction_index(rng, &InstructionContext::Global) - .expect("Global instruction index should always exist"); + .get_random_instruction_index_from(rng, &InstructionContext::Global, min_index) + .ok_or(MutatorError::NoMutationsAvailable)?; let mut builder = ProgramBuilder::new(program.context.clone()); diff --git a/fuzzamoto-ir/src/mutators/input.rs b/fuzzamoto-ir/src/mutators/input.rs index fd2f0618..2c9b0f5a 100644 --- a/fuzzamoto-ir/src/mutators/input.rs +++ b/fuzzamoto-ir/src/mutators/input.rs @@ -11,16 +11,26 @@ pub struct InputMutator; impl Mutator for InputMutator { fn mutate( + &mut self, + program: &mut Program, + rng: &mut R, + meta: Option<&PerTestcaseMetadata>, + ) -> MutatorResult { + self.mutate_from(program, rng, meta, 0) + } + + fn mutate_from( &mut self, program: &mut Program, rng: &mut R, _meta: Option<&PerTestcaseMetadata>, + min_index: usize, ) -> MutatorResult { let Some(candidate_instruction) = program .instructions .iter() .enumerate() - .filter(|(_, instruction)| instruction.is_input_mutable()) + .filter(|(i, instruction)| *i >= min_index && instruction.is_input_mutable()) .choose(rng) else { return Err(MutatorError::NoMutationsAvailable); diff --git a/fuzzamoto-ir/src/mutators/mod.rs b/fuzzamoto-ir/src/mutators/mod.rs index a3ea7c27..514effca 100644 --- a/fuzzamoto-ir/src/mutators/mod.rs +++ b/fuzzamoto-ir/src/mutators/mod.rs @@ -25,6 +25,18 @@ pub trait Mutator { rng: &mut R, meta: Option<&PerTestcaseMetadata>, ) -> MutatorResult; + + /// Mutate the program, only considering instructions with index >= `min_index`. + fn mutate_from( + &mut self, + program: &mut Program, + rng: &mut R, + meta: Option<&PerTestcaseMetadata>, + _min_index: usize, + ) -> MutatorResult { + self.mutate(program, rng, meta) + } + fn name(&self) -> &'static str; } @@ -36,4 +48,15 @@ pub trait Splicer: Mutator { splice_with: &Program, rng: &mut R, ) -> MutatorResult; + + /// Splice two programs together, only considering instructions with index >= `min_index`. + fn splice_from( + &mut self, + program: &mut Program, + splice_with: &Program, + rng: &mut R, + _min_index: usize, + ) -> MutatorResult { + self.splice(program, splice_with, rng) + } } diff --git a/fuzzamoto-ir/src/mutators/operation.rs b/fuzzamoto-ir/src/mutators/operation.rs index cbb87d56..871f6bec 100644 --- a/fuzzamoto-ir/src/mutators/operation.rs +++ b/fuzzamoto-ir/src/mutators/operation.rs @@ -30,20 +30,30 @@ pub struct OperationMutator { } impl Mutator for OperationMutator { + fn mutate( + &mut self, + program: &mut Program, + rng: &mut R, + meta: Option<&PerTestcaseMetadata>, + ) -> MutatorResult { + self.mutate_from(program, rng, meta, 0) + } + #[expect(clippy::cast_sign_loss)] #[expect(clippy::cast_possible_truncation)] #[expect(clippy::cast_precision_loss)] - fn mutate( + fn mutate_from( &mut self, program: &mut Program, rng: &mut R, _meta: Option<&PerTestcaseMetadata>, + min_index: usize, ) -> MutatorResult { let Some(candidate_instruction) = program .instructions .iter_mut() .enumerate() - .filter(|(_, instr)| instr.is_operation_mutable()) + .filter(|(i, instr)| *i >= min_index && instr.is_operation_mutable()) .choose(rng) else { return Err(super::MutatorError::NoMutationsAvailable); diff --git a/fuzzamoto-ir/src/operation.rs b/fuzzamoto-ir/src/operation.rs index 8f3b4e0a..f3c58470 100644 --- a/fuzzamoto-ir/src/operation.rs +++ b/fuzzamoto-ir/src/operation.rs @@ -205,6 +205,9 @@ pub enum Operation { /// None = key-path only spend; Some = script-path with one spendable leaf script_leaf: Option, }, + + /// Snapshot creation + IncrementalSnapshot, // TODO: SendGetBlockTxn // TODO: SendGetBlocks // TODO: SendGetHeaders @@ -437,6 +440,8 @@ impl fmt::Display for Operation { } write!(f, ")") } + + Operation::IncrementalSnapshot => write!(f, "IncrementalSnapshot"), } } } @@ -584,7 +589,8 @@ impl Operation { | Operation::Probe | Operation::TaprootScriptsUseAnnex | Operation::TaprootTxoUseAnnex - | Operation::BuildTaprootTree { .. } => false, + | Operation::BuildTaprootTree { .. } + | Operation::IncrementalSnapshot => false, } } @@ -743,7 +749,8 @@ impl Operation { | Operation::BuildCoinbaseTxInput | Operation::AddCoinbaseTxOutput | Operation::SendBlockTxn - | Operation::Probe => false, + | Operation::Probe + | Operation::IncrementalSnapshot => false, } } @@ -916,6 +923,8 @@ impl Operation { Operation::SendCompactBlock => vec![], Operation::SendBlockTxn => vec![], Operation::Probe => vec![], + + Operation::IncrementalSnapshot => vec![], } } @@ -1095,7 +1104,8 @@ impl Operation { | Operation::BeginBlockTransactions | Operation::BeginWitnessStack | Operation::BuildPayToAnchor - | Operation::Probe => vec![], + | Operation::Probe + | Operation::IncrementalSnapshot => vec![], } } @@ -1216,7 +1226,8 @@ impl Operation { | Operation::EndBuildBlockTxn | Operation::AddTxToBlockTxn | Operation::SendBlockTxn - | Operation::Probe => vec![], + | Operation::Probe + | Operation::IncrementalSnapshot => vec![], } } } diff --git a/fuzzamoto-libafl/src/feedbacks/mod.rs b/fuzzamoto-libafl/src/feedbacks/mod.rs index 628d3262..fafe4273 100644 --- a/fuzzamoto-libafl/src/feedbacks/mod.rs +++ b/fuzzamoto-libafl/src/feedbacks/mod.rs @@ -82,7 +82,10 @@ where if *self.enabled.borrow() && matches!(exit_kind, ExitKind::Timeout) { let timeouts = state.metadata_or_insert_with(TimeoutsToVerify::new); log::info!("Timeout detected, adding to verification queue!"); - timeouts.push(input.clone()); + // If we're using incremental snapshots, clear frozen_prefix_len. + let mut timeout_input = input.clone(); + timeout_input.frozen_prefix_len = None; + timeouts.push(timeout_input); return Ok(false); } diff --git a/fuzzamoto-libafl/src/input.rs b/fuzzamoto-libafl/src/input.rs index 11e86a27..f9673d83 100644 --- a/fuzzamoto-libafl/src/input.rs +++ b/fuzzamoto-libafl/src/input.rs @@ -1,6 +1,6 @@ use std::{fs::File, hash::Hash, io::Read, path::PathBuf}; -use fuzzamoto_ir::Program; +use fuzzamoto_ir::{Instruction, Operation, Program}; use libafl::inputs::{HasTargetBytes, Input}; use libafl_bolts::{HasLen, ownedref::OwnedSlice}; @@ -8,13 +8,20 @@ use libafl_bolts::{HasLen, ownedref::OwnedSlice}; #[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Hash)] pub struct IrInput { ir: Program, + /// Program index where the incremental snapshot is taken. The snapshot operation + /// is injected here during execution. + #[serde(skip)] + pub frozen_prefix_len: Option, } impl Input for IrInput {} impl IrInput { pub fn new(ir: Program) -> Self { - Self { ir } + Self { + ir, + frozen_prefix_len: None, + } } pub fn ir(&self) -> &Program { @@ -31,7 +38,29 @@ impl IrInput { file.read_to_end(&mut bytes).unwrap(); let program = postcard::from_bytes(&bytes).unwrap(); - Self { ir: program } + Self { + ir: program, + frozen_prefix_len: None, + } + } + + fn insert_snapshot(&self) -> Program { + if let Some(prefix_len) = self.frozen_prefix_len { + let mut instructions = self.ir.instructions.clone(); + + // Insert snapshot opcode at the frozen prefix position + let snapshot_instr = Instruction { + inputs: vec![], + operation: Operation::IncrementalSnapshot, + }; + + let insert_pos = prefix_len.min(instructions.len()); + instructions.insert(insert_pos, snapshot_instr); + + Program::unchecked_new(self.ir.context.clone(), instructions) + } else { + self.ir.clone() + } } } @@ -43,12 +72,14 @@ impl HasLen for IrInput { impl HasTargetBytes for IrInput { fn target_bytes(&self) -> OwnedSlice<'_, u8> { + let program = self.insert_snapshot(); + #[cfg(not(feature = "compile_in_vm"))] { let mut compiler = fuzzamoto_ir::compiler::Compiler::new(); let compiled_input = compiler - .compile(self.ir()) + .compile(&program) .expect("Compilation should never fail"); let mut bytes = @@ -64,7 +95,7 @@ impl HasTargetBytes for IrInput { #[cfg(feature = "compile_in_vm")] { let mut bytes = - postcard::to_allocvec(self.ir()).expect("serialization should never fail"); + postcard::to_allocvec(&program).expect("serialization should never fail"); log::trace!("Input size: {}", bytes.len()); if bytes.len() > 1 * 1024 * 1024 { bytes = Vec::new(); diff --git a/fuzzamoto-libafl/src/instance.rs b/fuzzamoto-libafl/src/instance.rs index b2c9ad26..3f92f4cc 100644 --- a/fuzzamoto-libafl/src/instance.rs +++ b/fuzzamoto-libafl/src/instance.rs @@ -51,7 +51,10 @@ use crate::{ mutators::{IrGenerator, IrMutator, IrSpliceMutator, LibAflByteMutator}, options::FuzzerOptions, schedulers::SupportedSchedulers, - stages::{IrMinimizerStage, ProbingStage, StabilityCheckStage, VerifyTimeoutsStage}, + stages::{ + IncrementalSnapshotStage, IrMinimizerStage, ProbingStage, SnapshotPlacementPolicy, + StabilityCheckStage, VerifyTimeoutsStage, + }, }; #[cfg(feature = "bench")] @@ -442,6 +445,16 @@ where let probing = ProbingStage::new(&stdout_observer_handle); let stability = StabilityCheckStage::new(&map_observer_handle, &map_feedback_name, 8); + + let mutation_stage = TuneableMutationalStage::new(&mut state, mutator); + + let incremental_snapshot_stage = IncrementalSnapshotStage::new( + self.options.incremental_snapshots, + mutation_stage, + SnapshotPlacementPolicy::Balanced, + 50, + ); + let mut stages = tuple_list!( ClosureStage::new(|_a: &mut _, _b: &mut _, _c: &mut _, _d: &mut _| { // Always try minimizing at least for one pass @@ -482,7 +495,7 @@ where tuple_list!( stability, probing, - TuneableMutationalStage::new(&mut state, mutator), + incremental_snapshot_stage, timeout_verify_stage, bench_stats_stage, ) diff --git a/fuzzamoto-libafl/src/mutators.rs b/fuzzamoto-libafl/src/mutators.rs index 05660c00..c32a98ee 100644 --- a/fuzzamoto-libafl/src/mutators.rs +++ b/fuzzamoto-libafl/src/mutators.rs @@ -74,11 +74,15 @@ where None }; + let min_index = input.frozen_prefix_len.unwrap_or(0); + Ok( - match self - .mutator - .mutate(input.ir_mut(), &mut self.rng, tc_data.as_deref()) - { + match self.mutator.mutate_from( + input.ir_mut(), + &mut self.rng, + tc_data.as_deref(), + min_index, + ) { Ok(()) => MutationResult::Mutated, _ => MutationResult::Skipped, }, @@ -148,10 +152,12 @@ where let other = other_testcase.load_input(state.corpus())?; + let min_index = input.frozen_prefix_len.unwrap_or(0); + let mut input_clone = input.clone(); if self .mutator - .splice(input_clone.ir_mut(), other.ir(), &mut self.rng) + .splice_from(input_clone.ir_mut(), other.ir(), &mut self.rng, min_index) .is_err() { return Ok(MutationResult::Skipped); @@ -231,6 +237,12 @@ where return Ok(MutationResult::Skipped); }; + let min_index = input.frozen_prefix_len.unwrap_or(0); + + if index < min_index { + return Ok(MutationResult::Skipped); + } + let mut builder = fuzzamoto_ir::ProgramBuilder::new(input.ir().context.clone()); builder diff --git a/fuzzamoto-libafl/src/options.rs b/fuzzamoto-libafl/src/options.rs index a8098bae..897c9199 100644 --- a/fuzzamoto-libafl/src/options.rs +++ b/fuzzamoto-libafl/src/options.rs @@ -100,7 +100,7 @@ pub struct FuzzerOptions { long, help = "Number of corpus entries cached in memory", env = "FUZZAMOTO_CORPUS_CACHE", - default_value_t = 100 + default_value_t = 5000 )] pub corpus_cache: usize, @@ -165,6 +165,9 @@ pub struct FuzzerOptions { help = "Profile that defines which generators are enabled" )] pub profile: Profile, + + #[arg(long, help = "Enable incremental snapshots", default_value_t = false)] + pub incremental_snapshots: bool, } fn unix_time() -> u64 { diff --git a/fuzzamoto-libafl/src/stages/incremental_snapshot_stage.rs b/fuzzamoto-libafl/src/stages/incremental_snapshot_stage.rs new file mode 100644 index 00000000..30ca0870 --- /dev/null +++ b/fuzzamoto-libafl/src/stages/incremental_snapshot_stage.rs @@ -0,0 +1,243 @@ +use std::marker::PhantomData; +use std::num::NonZeroUsize; + +use libafl::{ + Error, HasMetadata, + corpus::HasCurrentCorpusId, + executors::Executor, + stages::{Restartable, Stage, mutational::MutatedTransform}, + state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand}, +}; +use libafl_bolts::rands::Rand; +use libafl_nyx::executor::NyxExecutor; + +use fuzzamoto_ir::Program; + +use crate::input::IrInput; + +#[derive(Debug, Clone, Copy)] +pub enum SnapshotPlacementPolicy { + Balanced, +} + +pub struct IncrementalSnapshotStage { + enabled: bool, + inner_stage: IS, + policy: SnapshotPlacementPolicy, + max_reuse_count: usize, + phantom: PhantomData<(S, OT)>, +} + +impl IncrementalSnapshotStage { + pub fn new( + enabled: bool, + inner_stage: IS, + policy: SnapshotPlacementPolicy, + max_reuse_count: usize, + ) -> Self { + Self { + enabled, + inner_stage, + policy, + max_reuse_count, + phantom: PhantomData, + } + } + + /// Choose where to take the snapshot based on the placement policy + fn choose_position(&self, rand: &mut impl Rand, program_len: usize) -> Option { + if program_len == 0 { + return None; + } + + match self.policy { + SnapshotPlacementPolicy::Balanced => { + if program_len == 1 { + Some(0) + } else if rand.coinflip(0.5_f64) { + // First half + let half = (program_len / 2).max(1); + let nz_half = NonZeroUsize::new(half).expect("half should be non-zero"); + Some(rand.below(nz_half)) + } else { + // Second half + let half = program_len / 2; + let range = program_len - half; + let nz_range = NonZeroUsize::new(range).expect("range should be non-zero"); + Some(half + rand.below(nz_range)) + } + } + } + } +} + +impl Stage, EM, S, Z> for IncrementalSnapshotStage +where + IS: Stage, EM, S, Z> + Restartable, + S: HasCorpus + + HasRand + + HasMetadata + + HasCurrentTestcase + + HasCurrentCorpusId + + HasExecutions, + OT: libafl::observers::ObserversTuple, +{ + fn perform( + &mut self, + fuzzer: &mut Z, + executor: &mut NyxExecutor, + state: &mut S, + manager: &mut EM, + ) -> Result<(), Error> { + if !self.enabled { + if self.inner_stage.should_restart(state)? { + self.inner_stage.perform(fuzzer, executor, state, manager)?; + } + let _ = self.inner_stage.clear_progress(state); + + return Ok(()); + } + + // No incremental snapshot should exist at this point. + assert!(!executor.helper.nyx_process.aux_tmp_snapshot_created()); + + // Load input in case of eviction + { + let mut testcase = state.current_testcase_mut()?; + let _ = IrInput::try_transform_from(&mut testcase, state)?; + } + + let program_len = { + let testcase = state.current_testcase()?; + let input = testcase.input().as_ref().unwrap(); + input.ir().instructions.len() + }; + + if program_len == 0 || state.rand_mut().coinflip(0.04) { + // Skip creating an incremental snapshot if we're using the empty program or + // randomly decide to use the root. + if self.inner_stage.should_restart(state)? { + self.inner_stage.perform(fuzzer, executor, state, manager)?; + } + let _ = self.inner_stage.clear_progress(state); + + return Ok(()); + } + + let chosen_pos = self.choose_position(state.rand_mut(), program_len); + + let new_prefix_len = { + let testcase = state.current_testcase()?; + let input = testcase.input().as_ref().unwrap(); + chosen_pos.and_then(|pos| find_valid_snapshot_position(input.ir(), pos)) + }; + + if let Some(prefix_len) = new_prefix_len { + executor + .helper + .nyx_process + .option_set_delete_incremental_snapshot(false); + executor.helper.nyx_process.option_apply(); + + // Set frozen_prefix_len on the input so inner_stage is aware of it + { + let mut testcase = state.current_testcase_mut()?; + let input = testcase.input_mut().as_mut().unwrap(); + input.frozen_prefix_len = Some(prefix_len); + } + + log::info!("Created incremental snapshot at position {prefix_len}"); + + for reuse_count in 1..=self.max_reuse_count { + if reuse_count == self.max_reuse_count { + // Discard the incremental snapshot at the end of the last iteration. + // The inner mutational stage may not call run_target if the mutation + // is skipped, so call it explicitly. Without this, the assert at the + // top of this function is hit. + executor + .helper + .nyx_process + .option_set_delete_incremental_snapshot(true); + executor.helper.nyx_process.option_apply(); + + let testcase = state.current_testcase()?; + let input = testcase.input().as_ref().unwrap().clone(); + drop(testcase); + executor.run_target(fuzzer, state, manager, &input)?; + } else { + if self.inner_stage.should_restart(state)? { + self.inner_stage.perform(fuzzer, executor, state, manager)?; + } + let _ = self.inner_stage.clear_progress(state); + } + } + + // Reset frozen_prefix_len + { + let mut testcase = state.current_testcase_mut()?; + let input = testcase.input_mut().as_mut().unwrap(); + input.frozen_prefix_len = None; + } + } else { + log::info!("No valid position to create incremental snapshot",); + + return Ok(()); + } + + Ok(()) + } +} + +impl Restartable for IncrementalSnapshotStage +where + S: HasMetadata, + IS: Restartable, +{ + fn should_restart(&mut self, _state: &mut S) -> Result { + Ok(true) + } + + fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> { + Ok(()) + } +} + +/// Find a valid position for the snapshot that's not inside a block. +#[expect(clippy::cast_possible_wrap)] +fn find_valid_snapshot_position(program: &Program, target_pos: usize) -> Option { + let instructions = &program.instructions; + + if instructions.is_empty() { + return None; + } + + let target_pos = target_pos.min(instructions.len()); + + let mut block_depth: usize = 0; + let mut valid_positions = Vec::new(); + + for (i, instr) in instructions.iter().enumerate() { + if block_depth == 0 { + valid_positions.push(i); + } + + if instr.operation.is_block_begin() { + block_depth += 1; + } + if instr.operation.is_block_end() { + block_depth = block_depth.saturating_sub(1); + } + } + + if block_depth == 0 { + valid_positions.push(instructions.len()); + } + + if valid_positions.is_empty() { + return None; + } + + valid_positions + .into_iter() + .min_by_key(|&pos| (pos as isize - target_pos as isize).unsigned_abs()) +} diff --git a/fuzzamoto-libafl/src/stages/mod.rs b/fuzzamoto-libafl/src/stages/mod.rs index 9ce7d51e..af9f5cda 100644 --- a/fuzzamoto-libafl/src/stages/mod.rs +++ b/fuzzamoto-libafl/src/stages/mod.rs @@ -13,6 +13,9 @@ pub mod verify_timeouts; pub use verify_timeouts::*; +pub mod incremental_snapshot_stage; +pub use incremental_snapshot_stage::*; + use std::{borrow::Borrow, cell::RefCell, marker::PhantomData}; use fuzzamoto_ir::Minimizer; diff --git a/fuzzamoto-nyx-sys/src/lib.rs b/fuzzamoto-nyx-sys/src/lib.rs index b7b6127f..9af30162 100644 --- a/fuzzamoto-nyx-sys/src/lib.rs +++ b/fuzzamoto-nyx-sys/src/lib.rs @@ -12,6 +12,10 @@ unsafe extern "C" { len: usize, ); pub fn nyx_get_fuzz_input(data: *const c_uchar, max_size: usize) -> usize; + pub fn nyx_create_incremental_and_next( + data: *mut c_uchar, + resume_position: *mut usize, + ) -> usize; pub fn nyx_skip(); pub fn nyx_release(); pub fn nyx_fail(message: *const c_char); diff --git a/fuzzamoto-nyx-sys/src/nyx-agent.c b/fuzzamoto-nyx-sys/src/nyx-agent.c index 8ab0a36e..78c73474 100644 --- a/fuzzamoto-nyx-sys/src/nyx-agent.c +++ b/fuzzamoto-nyx-sys/src/nyx-agent.c @@ -29,6 +29,8 @@ __attribute__((weak)) extern uint32_t __afl_map_size; static uint8_t *trace_buffer = NULL; static size_t trace_buffer_size = 0; +static kAFL_payload *payload_buffer = NULL; + /** Initiliaze the nyx agent and return the maximum size for generated fuzz * inputs. * @@ -170,8 +172,8 @@ void nyx_println(const char *message, size_t message_len) { * * Note: This will take the snapshot on the first call. */ size_t nyx_get_fuzz_input(const uint8_t *data, size_t max_size) { - kAFL_payload *payload_buffer = mmap(NULL, max_size, PROT_READ | PROT_WRITE, - MAP_SHARED | MAP_ANONYMOUS, -1, 0); + payload_buffer = mmap(NULL, max_size, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_ANONYMOUS, -1, 0); mlock(payload_buffer, max_size); memset(payload_buffer, 0, max_size); @@ -202,6 +204,28 @@ size_t nyx_get_fuzz_input(const uint8_t *data, size_t max_size) { return payload_buffer->size; } +/** Create an incremental snapshot, fetch the payload, and save the snapshot + * prefix for later restores. + */ +size_t nyx_create_incremental_and_next(const uint8_t *data, size_t *prefix_index) { + static size_t saved_prefix_index = 0; + + // Save the prefix index before taking the snapshot + saved_prefix_index = *prefix_index; + + // Create the incremental snapshot which also captures saved_prefix_index + kAFL_hypercall(HYPERCALL_KAFL_CREATE_TMP_SNAPSHOT, 0); + kAFL_hypercall(HYPERCALL_KAFL_NEXT_PAYLOAD, 0); + kAFL_hypercall(HYPERCALL_KAFL_ACQUIRE, 0); + + memcpy((void *)data, payload_buffer->data, payload_buffer->size); + + // Update the prefix for the caller + *prefix_index = saved_prefix_index; + + return payload_buffer->size; +} + /** Resets the coverage bitmap and then resets the vm to the snapshot state. */ void nyx_skip() { // TODO: this is racy, we should stop the target from writing to the trace diff --git a/fuzzamoto-scenarios/bin/compact_blocks.rs b/fuzzamoto-scenarios/bin/compact_blocks.rs index 4febd3c9..d0223c11 100644 --- a/fuzzamoto-scenarios/bin/compact_blocks.rs +++ b/fuzzamoto-scenarios/bin/compact_blocks.rs @@ -1,6 +1,7 @@ use fuzzamoto::{ connections::Transport, fuzzamoto_main, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult, generic::GenericScenario}, targets::{BitcoinCoreTarget, Target}, test_utils, @@ -213,7 +214,7 @@ impl> Scenario<'_, TestCase> for CompactBlocksScena }) } - fn run(&mut self, testcase: TestCase) -> ScenarioResult { + fn run(&mut self, testcase: TestCase, _runner: &dyn Runner) -> ScenarioResult { let mut prevs: Vec<(u32, BlockHash, bitcoin::OutPoint)> = self .inner .block_tree diff --git a/fuzzamoto-scenarios/bin/http_server.rs b/fuzzamoto-scenarios/bin/http_server.rs index b447f5ca..c57a1c4d 100644 --- a/fuzzamoto-scenarios/bin/http_server.rs +++ b/fuzzamoto-scenarios/bin/http_server.rs @@ -1,5 +1,6 @@ use fuzzamoto::{ fuzzamoto_main, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult}, targets::{BitcoinCoreTarget, TargetNode}, }; @@ -53,7 +54,7 @@ impl<'a> Scenario<'a, TestCase<'a>> for HttpServerScenario { }) } - fn run(&mut self, input: TestCase) -> ScenarioResult { + fn run(&mut self, input: TestCase, _runner: &dyn Runner) -> ScenarioResult { // Network actions are slow; limit them const MAX_ACTIONS: usize = 128; if input.actions.len() > MAX_ACTIONS { diff --git a/fuzzamoto-scenarios/bin/import_mempool.rs b/fuzzamoto-scenarios/bin/import_mempool.rs index 52fbb3c9..132a481b 100644 --- a/fuzzamoto-scenarios/bin/import_mempool.rs +++ b/fuzzamoto-scenarios/bin/import_mempool.rs @@ -1,6 +1,7 @@ use fuzzamoto::{ connections::Transport, fuzzamoto_main, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult, generic::GenericScenario}, targets::{BitcoinCoreTarget, TargetNode}, }; @@ -48,7 +49,7 @@ where }) } - fn run(&mut self, input: MempoolDotDatBytes) -> ScenarioResult { + fn run(&mut self, input: MempoolDotDatBytes, _runner: &dyn Runner) -> ScenarioResult { if let Ok(mut mempool_file) = std::fs::File::create(&self.mempool_path) { let _ = mempool_file.write_all(input.0); let _ = mempool_file.flush(); diff --git a/fuzzamoto-scenarios/bin/ir.rs b/fuzzamoto-scenarios/bin/ir.rs index 85689536..b8820205 100644 --- a/fuzzamoto-scenarios/bin/ir.rs +++ b/fuzzamoto-scenarios/bin/ir.rs @@ -11,6 +11,7 @@ use fuzzamoto::{ connections::Transport, fuzzamoto_main, oracles::{CrashOracle, Oracle, OracleResult}, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult, generic::GenericScenario}, targets::{ BitcoinCoreTarget, ConnectableTarget, GenerateToAddress, HasBlockChainInterface, Target, @@ -273,10 +274,15 @@ where Ok(()) } - fn process_actions(&mut self, mut program: CompiledProgram) { + fn process_actions( + &mut self, + program: CompiledProgram, + start_index: usize, + runner: &dyn Runner, + ) -> Option<(Vec, usize)> { let message_filter = |(s, _): &(String, Vec)| ["getblocktxn"].contains(&s.as_str()); let mut non_probe_action_count = 0; - for action in program.actions.drain(..) { + for (i, action) in program.actions.into_iter().enumerate().skip(start_index) { match action { CompiledAction::Connect(_node, connection_type) => { let conn_type = match connection_type.as_str() { @@ -335,7 +341,7 @@ where } CompiledAction::SendRawMessage(from, command, message) => { if self.inner.connections.is_empty() { - return; + return None; } let num_connections = self.inner.connections.len(); @@ -376,8 +382,17 @@ where self.futurest = std::cmp::max(self.futurest, time); } + CompiledAction::IncrementalSnapshot => { + // If we're creating a new incremental snapshot, we want to save the index to skip + // ahead to. + let prefix_index = i + 1; + let (new_payload, action_pos) = + runner.create_incremental_and_next(prefix_index); + return Some((new_payload, action_pos)); + } } } + None } fn print_received(&mut self) { @@ -512,9 +527,30 @@ where }) } - fn run(&mut self, testcase: TestCase) -> ScenarioResult { + fn run(&mut self, testcase: TestCase, runner: &dyn Runner) -> ScenarioResult { let metadata = testcase.program.metadata.clone(); - self.process_actions(testcase.program); + let mut program = testcase.program; + let mut start_index = 0; + + while let Some((new_payload, action_pos)) = + self.process_actions(program, start_index, runner) + { + let new_testcase = match TestCase::decode(&new_payload) { + Ok(tc) => tc, + Err(e) => { + log::warn!( + "Failed to decode new payload after creating incremental snapshot: {e}", + ); + runner.skip(); + return ScenarioResult::Skip; + } + }; + + // Resume just after the snapshot prefix. + program = new_testcase.program; + start_index = action_pos; + } + self.ping_connections(); if self.recording_received_messages diff --git a/fuzzamoto-scenarios/bin/rpc_generic.rs b/fuzzamoto-scenarios/bin/rpc_generic.rs index e3234c1f..8e9ae2f6 100644 --- a/fuzzamoto-scenarios/bin/rpc_generic.rs +++ b/fuzzamoto-scenarios/bin/rpc_generic.rs @@ -1,5 +1,6 @@ use fuzzamoto::{ fuzzamoto_main, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult}, targets::{BitcoinCoreTarget, TargetNode}, }; @@ -209,7 +210,7 @@ impl Scenario<'_, TestCase> for RpcScenario { }) } - fn run(&mut self, input: TestCase) -> ScenarioResult { + fn run(&mut self, input: TestCase, _runner: &dyn Runner) -> ScenarioResult { if self.available_rpcs.is_empty() { return ScenarioResult::Fail("File with the RPC commands is empty".to_string()); } diff --git a/fuzzamoto-scenarios/bin/wallet_migration.rs b/fuzzamoto-scenarios/bin/wallet_migration.rs index 2b283349..7270598b 100644 --- a/fuzzamoto-scenarios/bin/wallet_migration.rs +++ b/fuzzamoto-scenarios/bin/wallet_migration.rs @@ -1,6 +1,7 @@ use fuzzamoto::{ connections::Transport, fuzzamoto_main, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult, generic::GenericScenario}, targets::{BitcoinCoreTarget, TargetNode}, }; @@ -60,7 +61,7 @@ where }) } - fn run(&mut self, input: WalletDotDatBytes) -> ScenarioResult { + fn run(&mut self, input: WalletDotDatBytes, _runner: &dyn Runner) -> ScenarioResult { let _ = std::fs::create_dir_all(self.wallet_path.parent().unwrap()); if let Ok(mut wallet_file) = std::fs::File::create(&self.wallet_path) { diff --git a/fuzzamoto/src/runners.rs b/fuzzamoto/src/runners.rs index 1e08fcbd..a74fa00e 100644 --- a/fuzzamoto/src/runners.rs +++ b/fuzzamoto/src/runners.rs @@ -5,9 +5,13 @@ use fuzzamoto_nyx_sys::*; /// libafl-qemu, local system, etc.) pub trait Runner { // Initialize the runner - fn new() -> Self; + fn new() -> Self + where + Self: Sized; // Get the next fuzz input fn get_fuzz_input(&self) -> Vec; + // Create an incremental snapshot, return the payload, and save the snapshot prefix + fn create_incremental_and_next(&self, prefix_index: usize) -> (Vec, usize); // Fail the last test case fn fail(&self, message: &str); // Skip the last test case @@ -37,6 +41,10 @@ impl Runner for LocalRunner { } } + fn create_incremental_and_next(&self, _prefix_index: usize) -> (Vec, usize) { + (vec![], 0) + } + fn fail(&self, message: &str) { log::error!("{message}"); } @@ -66,6 +74,14 @@ impl Runner for NyxRunner { data } + fn create_incremental_and_next(&self, prefix_index: usize) -> (Vec, usize) { + let mut data = vec![0u8; self.max_input_size]; + let mut pos = prefix_index; + let len = unsafe { nyx_create_incremental_and_next(data.as_mut_ptr(), &mut pos) }; + data.truncate(len); + (data, pos) + } + fn fail(&self, message: &str) { let c_message = std::ffi::CString::new(message).unwrap_or_default(); unsafe { @@ -111,6 +127,10 @@ impl Runner for StdRunner { self.runner.get_fuzz_input() } + fn create_incremental_and_next(&self, prefix_index: usize) -> (Vec, usize) { + self.runner.create_incremental_and_next(prefix_index) + } + fn fail(&self, message: &str) { self.runner.fail(message); } diff --git a/fuzzamoto/src/scenarios/generic.rs b/fuzzamoto/src/scenarios/generic.rs index 94a86ef7..80a8c5f8 100644 --- a/fuzzamoto/src/scenarios/generic.rs +++ b/fuzzamoto/src/scenarios/generic.rs @@ -1,6 +1,7 @@ use crate::{ connections::{Connection, ConnectionType, HandshakeOpts, Transport}, dictionaries::{Dictionary, FileDictionary}, + runners::Runner, scenarios::{Scenario, ScenarioInput, ScenarioResult}, targets::Target, test_utils, @@ -225,7 +226,7 @@ impl> Scenario<'_, TestCase> for GenericScenario ScenarioResult { + fn run(&mut self, testcase: TestCase, _runner: &dyn Runner) -> ScenarioResult { for action in testcase.actions { match action { Action::Connect { connection_type: _ } => { diff --git a/fuzzamoto/src/scenarios/mod.rs b/fuzzamoto/src/scenarios/mod.rs index 158422a1..e8e71cf3 100644 --- a/fuzzamoto/src/scenarios/mod.rs +++ b/fuzzamoto/src/scenarios/mod.rs @@ -1,3 +1,5 @@ +use crate::runners::Runner; + pub mod generic; /// `ScenarioInput` is a trait for scenario input types @@ -24,7 +26,7 @@ where /// Create a new instance of the scenario, preparing the initial state of the test fn new(args: &[String]) -> Result; // Run the test - fn run(&mut self, testcase: I) -> ScenarioResult; + fn run(&mut self, testcase: I, runner: &dyn Runner) -> ScenarioResult; } #[macro_export] @@ -68,7 +70,7 @@ macro_rules! fuzzamoto_main { return ExitCode::SUCCESS; }; - match scenario.run(testcase) { + match scenario.run(testcase, &runner) { ScenarioResult::Ok => {} ScenarioResult::Skip => { // TODO drop(target);