Skip to content

Commit

Permalink
Refactor parameters that adjust fuzzer behavior (#477)
Browse files Browse the repository at this point in the history
* refactor: move parameter constants to const.rs

* fix: below is exclusive. change to 256 to fill all u8 range

* fix: add in RANDOM_ADDRESS_CHOICE const

* chore: fmt
  • Loading branch information
plotchy authored May 18, 2024
1 parent 4d06f28 commit 9840724
Show file tree
Hide file tree
Showing 10 changed files with 187 additions and 70 deletions.
76 changes: 76 additions & 0 deletions src/const.rs
Original file line number Diff line number Diff line change
@@ -1 +1,77 @@
pub const DEBUG_PRINT_PERCENT: usize = 8000;

pub const INFANT_STATE_INITIAL_VOTES: usize = 3; // used in fuzzer.rs when infant is added to states
pub const CORPUS_INITIAL_VOTES: usize = 3; // used in fuzzer.rs when infant is added to states
/// used in libafl scheduled mutator (this is at max 1 << MAX_STACK_POW ==
/// 2^MAX_STACK_POW)
pub const MAX_STACK_POW: usize = 7;
pub const MAX_HAVOC_ATTEMPTS: usize = 10;
pub const KNOWN_STATE_MAX_SIZE: usize = 1000;
pub const KNOWN_STATE_SKIP_SIZE: usize = 500;

// src/fuzzer.rs
/// The maximum number of inputs (or VMState) to keep in the corpus before
/// pruning
pub const DROP_THRESHOLD: usize = 500;
/// The number of inputs (or VMState) to prune each time the corpus is pruned
pub const PRUNE_AMT: usize = 250;
/// If inputs (or VMState) has not been visited this many times, it will be
/// ignored during pruning
pub const VISIT_IGNORE_THRESHOLD: usize = 2;

// src/state.rs
/// Amount of accounts and contracts that can be caller during fuzzing.
/// We will generate random addresses for these accounts and contracts.
pub const ACCOUNT_AMT: u8 = 2;
/// Amount of accounts and contracts that can be caller during fuzzing.
/// We will generate random addresses for these accounts and contracts.
pub const CONTRACT_AMT: u8 = 2;
/// Maximum size of the input data
pub const MAX_INPUT_SIZE: usize = 20;

// src/abi.rs
/// Sample will be used to generate a random value with max value
pub const SAMPLE_MAX: u64 = 100;
/// Maximum value of the mutate choice. Related to [SAMPLE_MAX]
pub const MUTATE_CHOICE_MAX: u64 = 80;
/// Maximum value of the expand choice. Related to [SAMPLE_MAX]
pub const EXPAND_CHOICE_MAX: u64 = 90;
/// Maximum value of the random address choice. Related to [SAMPLE_MAX]
pub const RANDOM_ADDRESS_CHOICE: u64 = 90;

// src/evm/corpus_initializer.rs
/// If there are more than 1/UNKNOWN_SIGS_DIVISOR unknown sigs, we will
/// decompile with heimdall
pub const UNKNOWN_SIGS_DIVISOR: usize = 30;

// src/evm/mutator.rs
/// Sample will be used to generate a random value with max value
pub const MUTATOR_SAMPLE_MAX: u64 = 100;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const EXPLOIT_PRESET_CHOICE: u64 = 20;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const ABI_MUTATE_CHOICE: u64 = 96;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const HAVOC_CHOICE: u64 = 60;
/// Maximum number of iterations to try to find a valid havoc mutation
pub const HAVOC_MAX_ITERS: u64 = 10;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const MUTATE_CALLER_CHOICE: u64 = 20;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const TURN_TO_STEP_CHOICE: u64 = 60;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const RANDOMNESS_CHOICE: u64 = 33;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const LIQUIDATE_CHOICE: u64 = 5;
/// Related to [MUTATOR_SAMPLE_MAX]
pub const LIQ_PERCENT_CHOICE: u64 = 80;
pub const LIQ_PERCENT: u64 = 10;
/// Related to [MUTATOR_SAMPLE_MAX] and [LIQUIDATE_CHOICE]
pub const RANDOMNESS_CHOICE_2: u64 = 6;
/// Maximum number of retries to try to find a valid mutation
pub const MUTATION_RETRIES: usize = 20;

// src/evm/scheduler.rs
pub const POWER_MULTIPLIER: f64 = 32.0;
pub const MAX_POWER: f64 = 3200.0;
pub const MIN_POWER: f64 = 32.0;
15 changes: 8 additions & 7 deletions src/evm/abi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ use crate::{
generic_vm::vm_state::VMStateT,
input::ConciseSerde,
mutation_utils::{byte_mutator, byte_mutator_with_expansion},
r#const::{EXPAND_CHOICE_MAX, MUTATE_CHOICE_MAX, RANDOM_ADDRESS_CHOICE, SAMPLE_MAX},
state::{HasCaller, HasItyState},
};

Expand Down Expand Up @@ -344,7 +345,7 @@ where
match state.rand_mut().below(100) % 4 {
// dynamic
0 => BoxedABI::new(Box::new(ADynamic {
data: vec![state.rand_mut().below(255) as u8; vec_size],
data: vec![state.rand_mut().below(256) as u8; vec_size],
multiplier: 32,
})),
// tuple
Expand Down Expand Up @@ -412,7 +413,7 @@ impl BoxedABI {
return MutationResult::Skipped;
}
if a256.is_address {
if state.rand_mut().below(100) < 90 {
if state.rand_mut().below(SAMPLE_MAX) < RANDOM_ADDRESS_CHOICE {
a256.data = state.get_rand_address().0.to_vec();
} else {
a256.data = [0; 20].to_vec();
Expand All @@ -438,13 +439,13 @@ impl BoxedABI {
return MutationResult::Skipped;
}
if aarray.dynamic_size {
match state.rand_mut().below(100) {
0..=80 => {
match state.rand_mut().below(SAMPLE_MAX) {
0..=MUTATE_CHOICE_MAX => {
let index: usize = state.rand_mut().next() as usize % data_len;
let result = aarray.data[index].mutate_with_vm_slots(state, vm_slots);
return result;
}
81..=90 => {
MUTATE_CHOICE_MAX..=EXPAND_CHOICE_MAX => {
// increase size
if state.max_size() <= aarray.data.len() {
return MutationResult::Skipped;
Expand All @@ -453,7 +454,7 @@ impl BoxedABI {
aarray.data.push(aarray.data[0].clone());
}
}
91..=100 => {
EXPAND_CHOICE_MAX..=SAMPLE_MAX => {
// decrease size
if aarray.data.is_empty() {
return MutationResult::Skipped;
Expand All @@ -478,7 +479,7 @@ impl BoxedABI {
a_unknown.concrete = BoxedABI::new(Box::new(AEmpty {}));
return MutationResult::Skipped;
}
if (state.rand_mut().below(100)) < 80 {
if (state.rand_mut().below(SAMPLE_MAX)) < MUTATE_CHOICE_MAX {
a_unknown.concrete.mutate_with_vm_slots(state, vm_slots)
} else {
a_unknown.concrete = sample_abi(state, a_unknown.size);
Expand Down
3 changes: 2 additions & 1 deletion src/evm/corpus_initializer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ use crate::{
fuzzer::REPLAY,
generic_vm::vm_executor::GenericVM,
input::ConciseSerde,
r#const::UNKNOWN_SIGS_DIVISOR,
state::HasCaller,
state_input::StagedVMState,
};
Expand Down Expand Up @@ -310,7 +311,7 @@ where
}
}

if unknown_sigs >= sigs.len() / 30 {
if unknown_sigs >= sigs.len() / UNKNOWN_SIGS_DIVISOR {
info!("Too many unknown function signature for {:?}, we are going to decompile this contract using Heimdall", contract.name);
let abis = fetch_abi_heimdall(contract_code)
.iter()
Expand Down
79 changes: 56 additions & 23 deletions src/evm/mutator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,21 @@ use crate::{
},
generic_vm::vm_state::VMStateT,
input::{ConciseSerde, VMInputT},
r#const::{
ABI_MUTATE_CHOICE,
EXPLOIT_PRESET_CHOICE,
HAVOC_CHOICE,
HAVOC_MAX_ITERS,
LIQUIDATE_CHOICE,
LIQ_PERCENT,
LIQ_PERCENT_CHOICE,
MUTATE_CALLER_CHOICE,
MUTATION_RETRIES,
MUTATOR_SAMPLE_MAX,
RANDOMNESS_CHOICE,
RANDOMNESS_CHOICE_2,
TURN_TO_STEP_CHOICE,
},
state::{HasCaller, HasItyState, HasPresets, InfantStateState},
};

Expand Down Expand Up @@ -149,15 +164,21 @@ where
}
Constraint::Contract(target) => {
let rand_int = state.rand_mut().next();
let always_none = state.rand_mut().next() % 30 == 0;
let always_none = state.rand_mut().below(MUTATOR_SAMPLE_MAX);
let abis = state
.metadata_map()
.get::<ABIAddressToInstanceMap>()
.expect("ABIAddressToInstanceMap not found");
let abi = match abis.map.get(&target) {
Some(abi) => {
if !abi.is_empty() && !always_none {
Some((*abi)[rand_int as usize % abi.len()].clone())
if !abi.is_empty() {
match always_none {
0..=ABI_MUTATE_CHOICE => {
// we return a random abi
Some((*abi)[rand_int as usize % abi.len()].clone())
}
_ => None,
}
} else {
None
}
Expand Down Expand Up @@ -216,7 +237,7 @@ where
}

// use exploit template
if state.has_preset() && state.rand_mut().below(100) < 20 {
if state.has_preset() && state.rand_mut().below(MUTATOR_SAMPLE_MAX) < EXPLOIT_PRESET_CHOICE {
// if flashloan_v2, we don't mutate if it's a borrow
if input.get_input_type() != Borrow {
match state.get_next_call() {
Expand All @@ -237,19 +258,21 @@ where
// abi.b.get_size()).unwrap_or(0) / 32 + 1; if amount_of_args > 6 {
// amount_of_args = 6;
// }
let should_havoc = state.rand_mut().below(100) < 60; // (amount_of_args * 10) as u64;
let should_havoc = state.rand_mut().below(MUTATOR_SAMPLE_MAX) < HAVOC_CHOICE;

// determine how many times we should mutate the input
let havoc_times = if should_havoc {
state.rand_mut().below(10) + 1
state.rand_mut().below(HAVOC_MAX_ITERS) + 1 // (amount_of_args *
// HAVOC_MAX_ITERS) as
// u64;
} else {
1
};

let mut mutated = false;

{
if !input.is_step() && state.rand_mut().below(100) < 20_u64 {
if !input.is_step() && state.rand_mut().below(MUTATOR_SAMPLE_MAX) < MUTATE_CALLER_CHOICE {
let old_idx = input.get_state_idx();
let (idx, new_state) = state.get_infant_state(&mut self.infant_scheduler).unwrap();
if idx != old_idx {
Expand All @@ -266,7 +289,7 @@ where

if input.get_staged_state().state.has_post_execution() &&
!input.is_step() &&
state.rand_mut().below(100) < 60_u64
state.rand_mut().below(MUTATOR_SAMPLE_MAX) < TURN_TO_STEP_CHOICE
{
macro_rules! turn_to_step {
() => {
Expand All @@ -292,12 +315,18 @@ where
// if the input is a step input (resume execution from a control leak)
// we should not mutate the VM state, but only mutate the bytes
if input.is_step() {
let res = match state.rand_mut().below(100) {
0..=5 => {
let res = match state.rand_mut().below(MUTATOR_SAMPLE_MAX) {
0..=LIQUIDATE_CHOICE => {
// only when there are more than one liquidation path, we attempt to liquidate
if unsafe { CAN_LIQUIDATE } {
let prev_percent = input.get_liquidation_percent();
input.set_liquidation_percent(if state.rand_mut().below(100) < 80 { 10 } else { 0 } as u8);
input.set_liquidation_percent(if state.rand_mut().below(MUTATOR_SAMPLE_MAX) <
LIQ_PERCENT_CHOICE
{
LIQ_PERCENT
} else {
0
} as u8);
if prev_percent != input.get_liquidation_percent() {
MutationResult::Mutated
} else {
Expand All @@ -316,9 +345,9 @@ where
// if the input is to borrow token, we should mutate the randomness
// (use to select the paths to buy token), VM state, and bytes
if input.get_input_type() == Borrow {
let rand_u8 = state.rand_mut().below(255) as u8;
return match state.rand_mut().below(3) {
0 => {
let rand_u8 = state.rand_mut().below(256) as u8;
return match state.rand_mut().below(MUTATOR_SAMPLE_MAX) {
0..=RANDOMNESS_CHOICE => {
// mutate the randomness
input.set_randomness(vec![rand_u8; 1]);
MutationResult::Mutated
Expand All @@ -330,18 +359,22 @@ where

// mutate the bytes or VM state or liquidation percent (percentage of token to
// liquidate) by default
match state.rand_mut().below(100) {
6..=10 => {
match state.rand_mut().below(MUTATOR_SAMPLE_MAX) {
0..=LIQUIDATE_CHOICE => {
let prev_percent = input.get_liquidation_percent();
input.set_liquidation_percent(if state.rand_mut().below(100) < 80 { 10 } else { 0 } as u8);
input.set_liquidation_percent(if state.rand_mut().below(MUTATOR_SAMPLE_MAX) < LIQ_PERCENT_CHOICE {
LIQ_PERCENT
} else {
0
} as u8);
if prev_percent != input.get_liquidation_percent() {
MutationResult::Mutated
} else {
MutationResult::Skipped
}
}
11 => {
let rand_u8 = state.rand_mut().below(255) as u8;
LIQUIDATE_CHOICE..=RANDOMNESS_CHOICE_2 => {
let rand_u8 = state.rand_mut().below(256) as u8;
input.set_randomness(vec![rand_u8; 1]);
MutationResult::Mutated
}
Expand All @@ -356,10 +389,10 @@ where
};
let mut tries = 0;

// try to mutate the input for [`havoc_times`] times with 20 retries if
// the input is not mutated
while res != MutationResult::Mutated && tries < 20 {
for _ in 0..havoc_times {
// try to mutate the input for [`havoc_times`] times with MUTATION_RETRIES
// retries if the input is not mutated
while res != MutationResult::Mutated && tries < MUTATION_RETRIES {
for i in 0..havoc_times {
if mutator() == MutationResult::Mutated {
res = MutationResult::Mutated;
}
Expand Down
14 changes: 8 additions & 6 deletions src/evm/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ use crate::{
},
input::VMInputT,
power_sched::{PowerMutationalStageWithId, TestcaseScoreWithId},
r#const::{MAX_POWER, MIN_POWER, POWER_MULTIPLIER},
};

/// The status of the branch, whether it is covered on true, false or both
Expand Down Expand Up @@ -394,14 +395,15 @@ where
meta.testcase_to_uncovered_branches.get(&idx).unwrap_or(&0).to_owned() + 1
};

let mut power = uncov_branch as f64 * 32.0;

if power >= 3200.0 {
power = 3200.0;
let mut power = uncov_branch as f64 * POWER_MULTIPLIER;
// we score based on how a test case uncovered branches. 100 is cap, 1 is always
// min
if power >= MAX_POWER {
power = MAX_POWER;
}

if power <= 32.0 {
power = 32.0;
if power <= MIN_POWER {
power = MIN_POWER;
}

Ok(power)
Expand Down
18 changes: 11 additions & 7 deletions src/feedback.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,11 @@ use crate::{
generic_vm::vm_state::VMStateT,
input::{ConciseSerde, VMInputT},
oracle::{BugMetadata, Oracle, OracleCtx, Producer},
r#const::{INFANT_STATE_INITIAL_VOTES, KNOWN_STATE_MAX_SIZE, KNOWN_STATE_SKIP_SIZE},
scheduler::HasVote,
state::{HasExecutionResult, HasInfantStateState, InfantStateState},
};

const KNOWN_STATE_MAX_SIZE: usize = 1000;
const KNOWN_STATE_SKIP_SIZE: usize = 500;

/// OracleFeedback is a wrapper around a set of oracles and producers.
/// It executes the producers and then oracles after each successful execution.
/// If any of the oracle returns true, then it returns true and report a
Expand Down Expand Up @@ -541,16 +539,22 @@ where
// if the current distance is smaller than the min_map, vote for the state
if cmp_interesting {
debug!("Voted for {} because of CMP", input.get_state_idx());
self.scheduler
.vote(state.get_infant_state_state(), input.get_state_idx(), 3);
self.scheduler.vote(
state.get_infant_state_state(),
input.get_state_idx(),
INFANT_STATE_INITIAL_VOTES,
);
}

// if coverage has increased, vote for the state
if cov_interesting {
debug!("Voted for {} because of COV", input.get_state_idx());

self.scheduler
.vote(state.get_infant_state_state(), input.get_state_idx(), 3);
self.scheduler.vote(
state.get_infant_state_state(),
input.get_state_idx(),
INFANT_STATE_INITIAL_VOTES,
);
}

if self.vm.deref().borrow_mut().state_changed() ||
Expand Down
Loading

0 comments on commit 9840724

Please sign in to comment.