From cdf6fd6893f955355f4ebdd9e92f4919ddead21b Mon Sep 17 00:00:00 2001 From: tokatoka Date: Fri, 2 Jan 2026 12:07:13 +0100 Subject: [PATCH] Add PredicateTxGenerator --- fuzzamoto-ir/src/compiler.rs | 21 ++- fuzzamoto-ir/src/generators/block.rs | 4 +- fuzzamoto-ir/src/generators/block_txn.rs | 2 +- fuzzamoto-ir/src/generators/mod.rs | 2 +- fuzzamoto-ir/src/generators/tx.rs | 159 +++++++++++++++++++++-- fuzzamoto-ir/src/lib.rs | 32 ++++- fuzzamoto-ir/src/metadata.rs | 19 ++- fuzzamoto-libafl/src/instance.rs | 12 +- fuzzamoto-libafl/src/mutators.rs | 4 +- fuzzamoto-libafl/src/stages/probe.rs | 9 ++ fuzzamoto-scenarios/bin/ir.rs | 32 ++++- 11 files changed, 271 insertions(+), 25 deletions(-) diff --git a/fuzzamoto-ir/src/compiler.rs b/fuzzamoto-ir/src/compiler.rs index 553d3407..1c6afbd4 100644 --- a/fuzzamoto-ir/src/compiler.rs +++ b/fuzzamoto-ir/src/compiler.rs @@ -71,6 +71,8 @@ pub type ConnectionId = usize; pub struct CompiledMetadata { // Map from blockhash to (block variable index, list of transaction variable indices) block_tx_var_map: HashMap)>, + // Map from txid to tx variable index + txo_var_map: HashMap, // Map from connection ids to connection variable indices. connection_map: HashMap, // List of instruction indices that correspond to actions in the compiled program (does not include probe operation) @@ -86,6 +88,7 @@ impl CompiledMetadata { Self { block_tx_var_map: HashMap::new(), connection_map: HashMap::new(), + txo_var_map: HashMap::new(), action_indices: Vec::new(), variable_indices: Vec::new(), instructions: 0, @@ -102,9 +105,8 @@ impl CompiledMetadata { .map(|(header_var, block_var, tx_vars)| (*header_var, *block_var, tx_vars.as_slice())) } - // Get the list of instruction indices that correspond to actions in the compiled program - pub fn instruction_indices(&self) -> &[InstructionIndex] { - &self.action_indices + pub fn txo_variables(&self, txid: Txid) -> Option<&VariableIndex> { + self.txo_var_map.get(&txid) } // Get the list of instruction indices that correspond to variables in the compiled program @@ -112,6 +114,11 @@ impl CompiledMetadata { &self.variable_indices } + // Get the list of instruction indices that correspond to actions in the compiled program + pub fn instruction_indices(&self) -> &[InstructionIndex] { + &self.action_indices + } + pub fn connection_map(&self) -> &HashMap { &self.connection_map } @@ -1211,12 +1218,20 @@ impl Compiler { } Operation::TakeTxo | Operation::TakeCoinbaseTxo => { let tx_var = self.get_input_mut::(&instruction.inputs, 0)?; + let txid = tx_var.id; let num_txos = tx_var.txos.len(); let mut txo = Txo::new(); if num_txos != 0 { txo = tx_var.txos[tx_var.output_selector % num_txos].clone(); tx_var.output_selector += 1; } + let txo_index = self.variables.len(); + + self.output + .metadata + .txo_var_map + .entry(txid) + .or_insert(txo_index); self.append_variable(txo); } diff --git a/fuzzamoto-ir/src/generators/block.rs b/fuzzamoto-ir/src/generators/block.rs index 1f3a1448..db3167d0 100644 --- a/fuzzamoto-ir/src/generators/block.rs +++ b/fuzzamoto-ir/src/generators/block.rs @@ -229,7 +229,7 @@ impl Generator for TipBlockGenerator { &self, program: &crate::Program, rng: &mut R, - meta: Option<&PerTestcaseMetadata>, + meta: Option<&mut PerTestcaseMetadata>, ) -> Option { if let Some(meta) = meta.as_ref() && let Some(nth) = meta.recent_blocks.iter().max() @@ -291,7 +291,7 @@ impl Generator for ReorgBlockGenerator { &self, program: &crate::Program, rng: &mut R, - meta: Option<&PerTestcaseMetadata>, + meta: Option<&mut PerTestcaseMetadata>, ) -> Option { if let Some(meta) = meta.as_ref() && let Some(max) = meta.recent_blocks.iter().max_by_key(|i| i.defining_block.1) diff --git a/fuzzamoto-ir/src/generators/block_txn.rs b/fuzzamoto-ir/src/generators/block_txn.rs index 6a23e480..a9e3a15f 100644 --- a/fuzzamoto-ir/src/generators/block_txn.rs +++ b/fuzzamoto-ir/src/generators/block_txn.rs @@ -123,7 +123,7 @@ impl Generator for BlockTxnGenerator { &self, program: &crate::Program, rng: &mut R, - meta: Option<&PerTestcaseMetadata>, + meta: Option<&mut PerTestcaseMetadata>, ) -> Option { if let Some(meta) = meta && !meta.block_txn_request().is_empty() diff --git a/fuzzamoto-ir/src/generators/mod.rs b/fuzzamoto-ir/src/generators/mod.rs index 0934e491..67a61abf 100644 --- a/fuzzamoto-ir/src/generators/mod.rs +++ b/fuzzamoto-ir/src/generators/mod.rs @@ -66,7 +66,7 @@ pub trait Generator { &self, program: &Program, rng: &mut R, - _meta: Option<&PerTestcaseMetadata>, + _meta: Option<&mut PerTestcaseMetadata>, ) -> Option { program.get_random_instruction_index(rng, self.requested_context()) } diff --git a/fuzzamoto-ir/src/generators/tx.rs b/fuzzamoto-ir/src/generators/tx.rs index 7f71d11e..866c42a7 100644 --- a/fuzzamoto-ir/src/generators/tx.rs +++ b/fuzzamoto-ir/src/generators/tx.rs @@ -1,5 +1,7 @@ +use std::marker::PhantomData; + use crate::{ - IndexedVariable, Operation, PerTestcaseMetadata, TaprootLeafSpec, + IndexedVariable, MempoolTxo, Operation, PerTestcaseMetadata, TaprootLeafSpec, generators::{Generator, ProgramBuilder}, }; use bitcoin::{ @@ -133,12 +135,23 @@ fn build_outputs( Ok(()) } -fn build_tx( +fn build_tx_from_txos( builder: &mut ProgramBuilder, rng: &mut R, funding_txos: &[IndexedVariable], tx_version: u32, output_amounts: &[(u64, OutputType)], +) -> Result<(IndexedVariable, Vec), GeneratorError> { + let txos: Vec = funding_txos.iter().map(|txo| txo.index).collect(); + build_tx(builder, rng, &txos, tx_version, output_amounts) +} + +fn build_tx( + builder: &mut ProgramBuilder, + rng: &mut R, + funding_txos: &[usize], + tx_version: u32, + output_amounts: &[(u64, OutputType)], ) -> Result<(IndexedVariable, Vec), GeneratorError> { let tx_version_var = builder.force_append_expect_output(vec![], Operation::LoadTxVersion(tx_version)); @@ -154,7 +167,7 @@ fn build_tx( let sequence_var = builder.force_append_expect_output(vec![], Operation::LoadSequence(0xffffffff)); builder.force_append( - vec![mut_inputs_var.index, funding_txo.index, sequence_var.index], + vec![mut_inputs_var.index, *funding_txo, sequence_var.index], Operation::AddTxInput, ); } @@ -198,6 +211,135 @@ fn build_tx( Ok((const_tx_var, outputs)) } +/// `PredicateTxGenerator` generates transactions that spends the transactions in mempool depending on the given predicate. +pub struct PredicateTxGenerator { + predicate: F, + phantom: PhantomData, +} + +impl Generator for PredicateTxGenerator +where + F: Fn(&MempoolTxo) -> bool, +{ + fn generate( + &self, + builder: &mut ProgramBuilder, + rng: &mut R, + meta: Option<&PerTestcaseMetadata>, + ) -> GeneratorResult { + if let Some(meta) = meta + && !meta.txo_metadata().txo_entry.is_empty() + { + let chosen = meta + .txo_metadata + .choice + .expect("We should've chosen a txo to spend by now"); + let txo = meta + .txo_metadata + .txo_entry + .get(chosen) + .expect("Getting the chosen spent txo should always succeed") + .definition + .0; + + let tx_version = *[1, 2, 3].choose(rng).unwrap(); // 3 = TRUC-violation + let amount = ( + rng.gen_range(5000..100_000_000), + get_random_output_type(rng), + ); + let (const_tx_var, _) = build_tx(builder, rng, &[txo], tx_version, &[amount])?; + + let conn_var = builder.get_or_create_random_connection(rng); + + let mut_inventory_var = + builder.force_append_expect_output(vec![], Operation::BeginBuildInventory); + builder.force_append( + vec![mut_inventory_var.index, const_tx_var.index], + Operation::AddWtxidInv, + ); + let const_inventory_var = builder.force_append_expect_output( + vec![mut_inventory_var.index], + Operation::EndBuildInventory, + ); + + builder.force_append( + vec![conn_var.index, const_inventory_var.index], + Operation::SendInv, + ); + builder.force_append(vec![conn_var.index, const_tx_var.index], Operation::SendTx); + + Ok(()) + } else { + Ok(()) + } + } + + fn choose_index( + &self, + program: &crate::Program, + rng: &mut R, + meta: Option<&mut PerTestcaseMetadata>, + ) -> Option { + let meta = meta?; + let txo_meta = meta.txo_metadata(); + let filtered = txo_meta + .txo_entry + .iter() + .enumerate() + .filter(|(_, x)| (self.predicate)(*x)) + .collect::>(); + + let chosen = filtered.choose(rng); + if let Some((idx, txo)) = chosen { + let (_, inst) = txo.definition; + meta.txo_metadata_mut().choice = Some(*idx); + program.get_random_instruction_index_from( + rng, + >::requested_context(self), + inst + 1, + ) + } else { + None + } + } + + fn name(&self) -> &'static str { + "PredicateTxGenerator" + } +} + +impl PredicateTxGenerator { + pub fn new(predicate: F) -> Self { + Self { + predicate, + phantom: PhantomData, + } + } +} + +impl PredicateTxGenerator bool> { + pub fn double_spend() -> Self { + Self { + predicate: |x: &MempoolTxo| !x.spentby.is_empty(), + phantom: PhantomData, + } + } + + pub fn chain_spend() -> Self { + Self { + predicate: |x: &MempoolTxo| x.depends.is_empty(), + phantom: PhantomData, + } + } + + pub fn any() -> Self { + Self { + predicate: |_: &MempoolTxo| true, + phantom: PhantomData, + } + } +} + /// `SingleTxGenerator` generates instructions for a single new transaction into a program pub struct SingleTxGenerator; @@ -225,7 +367,8 @@ impl Generator for SingleTxGenerator { } amounts }; - let (const_tx_var, _) = build_tx(builder, rng, &funding_txos, tx_version, &output_amounts)?; + let (const_tx_var, _) = + build_tx_from_txos(builder, rng, &funding_txos, tx_version, &output_amounts)?; if rng.gen_bool(0.5) { let conn_var = builder.get_or_create_random_connection(rng); @@ -277,7 +420,7 @@ impl Generator for OneParentOneChildGenerator { return Err(GeneratorError::MissingVariables); }; - let (parent_tx_var, parent_output_vars) = build_tx( + let (parent_tx_var, parent_output_vars) = build_tx_from_txos( builder, rng, &funding_txos, @@ -287,7 +430,7 @@ impl Generator for OneParentOneChildGenerator { (10000, OutputType::PayToAnchor), ], )?; - let (child_tx_var, _) = build_tx( + let (child_tx_var, _) = build_tx_from_txos( builder, rng, &[parent_output_vars.last().unwrap().clone()], @@ -354,7 +497,7 @@ impl Generator for LongChainGenerator { // transaction spends the output of the previous transaction let mut tx_vars = Vec::new(); for i in 0..25 { - let (tx_var, outputs) = build_tx( + let (tx_var, outputs) = build_tx_from_txos( builder, rng, &funding_txos, @@ -423,7 +566,7 @@ impl Generator for LargeTxGenerator { let conn_var = builder.get_or_create_random_connection(rng); for utxo in funding_txos { - let (tx_var, _) = build_tx( + let (tx_var, _) = build_tx_from_txos( builder, rng, &[utxo.clone()], diff --git a/fuzzamoto-ir/src/lib.rs b/fuzzamoto-ir/src/lib.rs index 078de32a..64f4a23b 100644 --- a/fuzzamoto-ir/src/lib.rs +++ b/fuzzamoto-ir/src/lib.rs @@ -20,11 +20,11 @@ pub use minimizers::*; pub use mutators::*; pub use operation::*; +use bitcoin::Txid; pub use fuzzamoto::taproot::*; use rand::{RngCore, seq::IteratorRandom}; -pub use variable::*; - use std::{collections::HashMap, fmt, hash::Hash}; +pub use variable::*; /// Program represent a sequence of operations to perform on target nodes. #[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Hash)] @@ -334,11 +334,39 @@ pub struct GetBlockTxn { pub tx_indices_variables: Vec, } +/// The metadata holds the txid of transactions in mempool which is already spent by another tx in the mempool +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct MempoolTxo { + pub txid: Txid, + pub definition: (usize, usize), + pub spentby: Vec, + pub depends: Vec, +} + +/// The metadata holds the `MempoolTxo` list and the next txo used to mutate. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct TxoMetadata { + pub txo_entry: Vec, + pub choice: Option, +} + +impl Default for TxoMetadata { + fn default() -> Self { + Self { + txo_entry: Vec::new(), + choice: None, + } + } +} + #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub enum ProbeResult { GetBlockTxn { get_block_txn: GetBlockTxn, }, + Mempool { + txo_entry: Vec, + }, Failure { /// The command that failed to be decoded command: String, diff --git a/fuzzamoto-ir/src/metadata.rs b/fuzzamoto-ir/src/metadata.rs index d5cc15ec..63626cfd 100644 --- a/fuzzamoto-ir/src/metadata.rs +++ b/fuzzamoto-ir/src/metadata.rs @@ -1,12 +1,13 @@ use serde::{Deserialize, Serialize}; -use crate::{GetBlockTxn, RecentBlock}; +use crate::{GetBlockTxn, MempoolTxo, RecentBlock, TxoMetadata}; /// The runtime data observed during the course of harness execution #[derive(Default, Debug, Clone, Serialize, Deserialize)] pub struct PerTestcaseMetadata { pub block_txn_request: Vec, pub recent_blocks: Vec, + pub txo_metadata: TxoMetadata, } impl PerTestcaseMetadata { @@ -14,6 +15,7 @@ impl PerTestcaseMetadata { Self { block_txn_request: Vec::new(), recent_blocks: Vec::new(), + txo_metadata: TxoMetadata::default(), } } @@ -25,6 +27,14 @@ impl PerTestcaseMetadata { &self.recent_blocks } + pub fn txo_metadata(&self) -> &TxoMetadata { + &self.txo_metadata + } + + pub fn txo_metadata_mut(&mut self) -> &mut TxoMetadata { + &mut self.txo_metadata + } + pub fn add_block_tx_request(&mut self, req: GetBlockTxn) { self.block_txn_request.push(req); } @@ -33,4 +43,11 @@ impl PerTestcaseMetadata { self.recent_blocks = blocks; self.recent_blocks.sort(); } + pub fn add_txo_entry(&mut self, txo_entry: Vec) { + let txo_metadata = TxoMetadata { + txo_entry, + choice: None, + }; + self.txo_metadata = txo_metadata; + } } diff --git a/fuzzamoto-libafl/src/instance.rs b/fuzzamoto-libafl/src/instance.rs index 6c3abf95..6e00d5ff 100644 --- a/fuzzamoto-libafl/src/instance.rs +++ b/fuzzamoto-libafl/src/instance.rs @@ -5,10 +5,10 @@ use fuzzamoto_ir::{ BlockGenerator, BlockTxnGenerator, BloomFilterAddGenerator, BloomFilterClearGenerator, BloomFilterLoadGenerator, CombineMutator, CompactBlockGenerator, CompactFilterQueryGenerator, GetAddrGenerator, GetDataGenerator, HeaderGenerator, InputMutator, InventoryGenerator, - LargeTxGenerator, LongChainGenerator, OneParentOneChildGenerator, OperationMutator, Program, - ReorgBlockGenerator, SendBlockGenerator, SendMessageGenerator, SingleTxGenerator, - TipBlockGenerator, TxoGenerator, WitnessGenerator, cutting::CuttingMinimizer, - instr_block::InstrBlockMinimizer, nopping::NoppingMinimizer, + LargeTxGenerator, LongChainGenerator, OneParentOneChildGenerator, OperationMutator, + PredicateTxGenerator, Program, ReorgBlockGenerator, SendBlockGenerator, SendMessageGenerator, + SingleTxGenerator, TipBlockGenerator, TxoGenerator, WitnessGenerator, + cutting::CuttingMinimizer, instr_block::InstrBlockMinimizer, nopping::NoppingMinimizer, }; use libafl::{ @@ -283,6 +283,10 @@ where rng.clone() ) ), + ( + 100.0, + IrGenerator::new(PredicateTxGenerator::any(), rng.clone()) + ), ( 100.0, IrGenerator::new( diff --git a/fuzzamoto-libafl/src/mutators.rs b/fuzzamoto-libafl/src/mutators.rs index 68bc35cc..8b97ff40 100644 --- a/fuzzamoto-libafl/src/mutators.rs +++ b/fuzzamoto-libafl/src/mutators.rs @@ -216,7 +216,7 @@ where let is_first = rt_data.mutation_idx() == 0; rt_data.increment_idx(); - let tc_data = if is_first + let mut tc_data = if is_first && let Some(id) = current_id && let Some(meta) = rt_data.metadata_mut(id) { @@ -227,7 +227,7 @@ where let Some(index) = self.generator - .choose_index(input.ir(), &mut self.rng, tc_data.as_deref()) + .choose_index(input.ir(), &mut self.rng, tc_data.as_deref_mut()) else { return Ok(MutationResult::Skipped); }; diff --git a/fuzzamoto-libafl/src/stages/probe.rs b/fuzzamoto-libafl/src/stages/probe.rs index 3532dcd7..e35630c4 100644 --- a/fuzzamoto-libafl/src/stages/probe.rs +++ b/fuzzamoto-libafl/src/stages/probe.rs @@ -78,6 +78,15 @@ where txvec.add_block_tx_request(get_block_txn.clone()); } } + ProbeResult::Mempool { txo_entry } => { + let current = *state.corpus().current(); + if let Some(cur) = current + && let Ok(meta) = state.metadata_mut::() + { + let txvec = meta.metadatas.entry(cur).or_default(); + txvec.add_txo_entry(txo_entry.clone()); + } + } ProbeResult::Failure { command, reason } => { log::info!( "Command {:?} couln't be parsed; reason: {:?}", diff --git a/fuzzamoto-scenarios/bin/ir.rs b/fuzzamoto-scenarios/bin/ir.rs index 5f5047f3..51e89efd 100644 --- a/fuzzamoto-scenarios/bin/ir.rs +++ b/fuzzamoto-scenarios/bin/ir.rs @@ -23,7 +23,7 @@ use fuzzamoto::oracles::{NetSplitContext, NetSplitOracle}; use fuzzamoto::oracles::{ConsensusContext, ConsensusOracle}; use fuzzamoto_ir::{ - ProbeResult, ProbeResults, Program, ProgramContext, RecentBlock, + MempoolTxo, ProbeResult, ProbeResults, Program, ProgramContext, RecentBlock, compiler::{CompiledAction, CompiledMetadata, CompiledProgram, Compiler}, }; @@ -321,6 +321,34 @@ where } } + fn print_mempool(&mut self, meta: &CompiledMetadata) { + let query = self.inner.target.get_mempool_entries(); + if let Ok(mempool) = query { + let mut ret = Vec::new(); + for tx in mempool { + let txid = tx.txid(); + let definition = if let Some(txo) = meta.txo_variables(*txid) + && let Some(inst) = meta.variable_indices().get(*txo) + { + (*txo, *inst) + } else { + continue; + }; + + let txo_entry = MempoolTxo { + txid: *txid, + definition: definition, + spentby: tx.spentby().to_vec(), + depends: tx.depends().to_vec(), + }; + + ret.push(txo_entry); + } + self.probe_results + .push(ProbeResult::Mempool { txo_entry: ret }); + } + } + fn evaluate_oracles(&mut self) -> ScenarioResult { let crash_oracle = CrashOracle::::default(); if let OracleResult::Fail(e) = crash_oracle.evaluate(&self.inner.target) { @@ -437,7 +465,9 @@ where if let Some(ret) = probe_recent_block_hashes(&self.inner.target, &metadata) { self.probe_results.push(ret); } + self.print_mempool(&metadata); } + self.print_received(); self.evaluate_oracles() }