Skip to content

Commit

Permalink
Support forcing ABI (#361)
Browse files Browse the repository at this point in the history 
* support forcing abi for contracts

* remove euler

* clippy

* fmt
  • Loading branch information
@shouc
shouc authored Nov 18, 2023
1 parent a7657bd commit ad4e8ec
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 56 deletions.
22 changes: 21 additions & 1 deletion src/evm/contract_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,26 @@ pub struct ContractLoader {
pub setup_data: Option<SetupData>,
}

impl ContractLoader {
pub fn force_abi(&mut self, mapping: HashMap<String, String>) {
for (filename_or_address, abi) in mapping {
let mut addr = None;
if filename_or_address.starts_with("0x") {
addr = EVMAddress::from_str(&filename_or_address).map(Some).unwrap_or(None);
}

for contract in &mut self.contracts {
if contract.name == filename_or_address ||
(addr.is_some() && contract.deployed_address == addr.unwrap())
{
debug!("Forcing ABI for {}", contract.name);
contract.abi = ContractLoader::parse_abi_str(&abi);
}
}
}
}
}

#[derive(Debug, Clone)]
pub struct SetupData {
pub evmstate: EVMState,
Expand Down Expand Up @@ -130,7 +150,7 @@ impl ContractLoader {
return format!("({})", v);
} else if ty.ends_with("[]") {
return format!("{}[]", Self::process_input(ty[..ty.len() - 2].to_string(), input));
} else if ty.ends_with("]") && ty.contains("[") {
} else if ty.ends_with(']') && ty.contains('[') {
let split = ty.rsplit_once('[').unwrap();
let name = split.0.to_string();
let len = split
Expand Down
133 changes: 78 additions & 55 deletions src/evm/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ use config::{Config, FuzzerTypes, StorageFetchingMode};
use contract_utils::ContractLoader;
use ethers::types::Transaction;
use input::{ConciseEVMInput, EVMInput};
use itertools::Itertools;
use num_cpus;
use onchain::{
endpoints::{Chain, OnChainConfig},
Expand Down Expand Up @@ -310,6 +311,12 @@ pub struct EvmArgs {
#[arg(long, default_value = "")]
setup_file: String,

/// Forcing a contract to use the given abi. This is useful when the
/// contract is a complex proxy or decompiler has trouble to detect the abi.
/// Format: address:abi_file,...
#[arg(long, default_value = "")]
force_abi: String,

/// Preset file. If specified, will load the preset file and match past
/// exploit template.
#[cfg(feature = "use_presets")]
Expand Down Expand Up @@ -530,65 +537,81 @@ pub fn evm_main(args: EvmArgs) {
None
};

let config = Config {
fuzzer_type: FuzzerTypes::from_str(args.fuzzer_type.as_str()).expect("unknown fuzzer"),
contract_loader: match target_type {
EVMTargetType::Glob => ContractLoader::from_glob(
args.target.as_str(),
&mut state,
&proxy_deploy_codes,
&constructor_args_map,
),
EVMTargetType::Config => ContractLoader::from_config(
let force_abis = args
.force_abi
.split(',')
.filter(|s| !s.is_empty())
.map(|x| {
let runes = x.split(':').collect_vec();
assert_eq!(runes.len(), 2, "Invalid force abi format");
let abi = std::fs::read_to_string(runes[1]).expect("Failed to read abi file");
(runes[0].to_string(), abi)
})
.collect::<HashMap<_, _>>();

let mut contract_loader = match target_type {
EVMTargetType::Glob => ContractLoader::from_glob(
args.target.as_str(),
&mut state,
&proxy_deploy_codes,
&constructor_args_map,
),
EVMTargetType::Config => ContractLoader::from_config(
&offchain_artifacts.expect("offchain artifacts is required for config target type"),
&offchain_config.expect("offchain config is required for config target type"),
),
EVMTargetType::AnvilFork => {
let addresses: Vec<EVMAddress> = args
.target
.split(',')
.map(|s| EVMAddress::from_str(s).unwrap())
.collect();
ContractLoader::from_fork(
&offchain_artifacts.expect("offchain artifacts is required for config target type"),
&offchain_config.expect("offchain config is required for config target type"),
),
EVMTargetType::AnvilFork => {
let addresses: Vec<EVMAddress> = args
.target
.split(',')
.map(|s| EVMAddress::from_str(s).unwrap())
.collect();
ContractLoader::from_fork(
&offchain_artifacts.expect("offchain artifacts is required for config target type"),
onchain.as_mut().expect("onchain is required to fork anvil"),
HashSet::from_iter(addresses),
)
onchain.as_mut().expect("onchain is required to fork anvil"),
HashSet::from_iter(addresses),
)
}
EVMTargetType::Setup => ContractLoader::from_setup(
&offchain_artifacts.expect("offchain artifacts is required for config target type"),
args.setup_file,
),
EVMTargetType::Address => {
if onchain.is_none() {
panic!("Onchain is required for address target type");
}
EVMTargetType::Setup => ContractLoader::from_setup(
&offchain_artifacts.expect("offchain artifacts is required for config target type"),
args.setup_file,
),
EVMTargetType::Address => {
if onchain.is_none() {
panic!("Onchain is required for address target type");
}
let mut args_target = args.target.clone();

if args.ierc20_oracle || args.flashloan {
const ETH_ADDRESS: &str = "0x7a250d5630b4cf539739df2c5dacb4c659f2488d";
const BSC_ADDRESS: &str = "0x10ed43c718714eb63d5aa57b78b54704e256024e";
if "bsc" == onchain.as_ref().unwrap().chain_name {
if !args_target.contains(BSC_ADDRESS) {
args_target.push(',');
args_target.push_str(BSC_ADDRESS);
}
} else if "eth" == onchain.as_ref().unwrap().chain_name && !args_target.contains(ETH_ADDRESS) {
let mut args_target = args.target.clone();

if args.ierc20_oracle || args.flashloan {
const ETH_ADDRESS: &str = "0x7a250d5630b4cf539739df2c5dacb4c659f2488d";
const BSC_ADDRESS: &str = "0x10ed43c718714eb63d5aa57b78b54704e256024e";
if "bsc" == onchain.as_ref().unwrap().chain_name {
if !args_target.contains(BSC_ADDRESS) {
args_target.push(',');
args_target.push_str(ETH_ADDRESS);
args_target.push_str(BSC_ADDRESS);
}
} else if "eth" == onchain.as_ref().unwrap().chain_name && !args_target.contains(ETH_ADDRESS) {
args_target.push(',');
args_target.push_str(ETH_ADDRESS);
}
let addresses: Vec<EVMAddress> = args_target
.split(',')
.map(|s| EVMAddress::from_str(s).unwrap())
.collect();
ContractLoader::from_address(
onchain.as_mut().unwrap(),
HashSet::from_iter(addresses),
builder.clone(),
)
}
},
let addresses: Vec<EVMAddress> = args_target
.split(',')
.map(|s| EVMAddress::from_str(s).unwrap())
.collect();
ContractLoader::from_address(
onchain.as_mut().unwrap(),
HashSet::from_iter(addresses),
builder.clone(),
)
}
};

contract_loader.force_abi(force_abis);

let config = Config {
fuzzer_type: FuzzerTypes::from_str(args.fuzzer_type.as_str()).expect("unknown fuzzer"),
contract_loader,
only_fuzz: if !args.only_fuzz.is_empty() {
args.only_fuzz
.split(',')
Expand Down Expand Up @@ -668,13 +691,13 @@ pub fn evm_main(args: EvmArgs) {
.iter()
.map(|config| {
json!({
hex::encode(&config.function): format!("{}{}", &config.function_name, &config.abi)
hex::encode(config.function): format!("{}{}", &config.function_name, &config.abi)
})
})
.collect();
abis_map
.entry(hex::encode(contract_info.deployed_address))
.or_insert_with(Vec::new)
.or_default()
.push(abis);
}

Expand Down

0 comments on commit ad4e8ec

Please sign in to comment.