From 8fdcdf7851c7e66ec8e435d555ef4a4fd217a3b6 Mon Sep 17 00:00:00 2001 From: Nick <49566599+nick199910@users.noreply.github.com> Date: Mon, 29 Jul 2024 09:51:30 +0800 Subject: [PATCH] Add constructor args (#524) * add constructor-args value * fix constructor empty bug * cargo fmt * update bsc rpc * add real_balance test * fix some value bugs --- src/evm/contract_utils.rs | 68 ++++++++++++++++++-- src/evm/corpus_initializer.rs | 13 +++- src/evm/mod.rs | 14 ++++ tests/evm_real_balance/contract1.abi | 1 + tests/evm_real_balance/contract1.bin | 1 + tests/evm_real_balance/contract2.abi | 1 + tests/evm_real_balance/contract2.bin | 1 + tests/evm_real_balance/contract3.abi | 1 + tests/evm_real_balance/contract3.bin | 1 + tests/evm_real_balance/test_real_balance.sol | 41 ++++++++++++ 10 files changed, 137 insertions(+), 5 deletions(-) create mode 100644 tests/evm_real_balance/contract1.abi create mode 100644 tests/evm_real_balance/contract1.bin create mode 100644 tests/evm_real_balance/contract2.abi create mode 100644 tests/evm_real_balance/contract2.bin create mode 100644 tests/evm_real_balance/contract3.abi create mode 100644 tests/evm_real_balance/contract3.bin create mode 100644 tests/evm_real_balance/test_real_balance.sol diff --git a/src/evm/contract_utils.rs b/src/evm/contract_utils.rs index d9fe3cf4..4fb8ae7c 100644 --- a/src/evm/contract_utils.rs +++ b/src/evm/contract_utils.rs @@ -88,6 +88,7 @@ pub struct ContractInfo { pub code: Vec, pub abi: Vec, pub is_code_deployed: bool, + pub balance: EVMU256, pub constructor_args: Vec, pub deployed_address: EVMAddress, pub build_artifact: Option, @@ -173,6 +174,27 @@ pub fn set_hash(name: &str, out: &mut [u8]) { hasher.result(out) } +fn parse_and_convert(input: &str) -> String { + let parts: Vec<&str> = input.trim().split_whitespace().collect(); + if parts.len() != 2 && (parts[1] != "wei" || parts[1] != "gwei" || parts[1] != "ether") { + return 0.to_string(); + } + + let number = match u128::from_str(parts[0]) { + Ok(num) => num, + Err(_) => return 0.to_string(), + }; + + let value = match parts[1].to_lowercase().as_str() { + "wei" => number, + "gwei" => number.checked_mul(10u128.pow(9)).unwrap_or(0), + "ether" => number.checked_mul(10u128.pow(18)).unwrap_or(0), + _ => 0, + } + .to_string(); + value +} + impl ContractLoader { fn parse_abi(path: &Path) -> Vec { let mut file = File::open(path).unwrap(); @@ -322,9 +344,30 @@ impl ContractLoader { raw_source_maps: HashMap, // contract name -> raw source map ) -> Self { let contract_name = prefix.split('/').last().unwrap().replace('*', ""); + // number ether, number wei, number gwei + let contract_balance = if !constructor_args.is_empty() { + let bal: String = constructor_args.last().unwrap().to_owned(); + if bal.ends_with("wei") || bal.ends_with("gwei") || bal.ends_with("ether") { + EVMU256::from_str(parse_and_convert(bal.as_str()).as_str()).unwrap() + } else { + EVMU256::from(0) + } + } else { + EVMU256::from(0) + }; + + let mut real_constructor_args = constructor_args.to_owned(); + + if !constructor_args.is_empty() && + (constructor_args[constructor_args.len() - 1].ends_with("wei") || + constructor_args[constructor_args.len() - 1].ends_with("gwei") || + constructor_args[constructor_args.len() - 1].ends_with("ether")) + { + real_constructor_args.remove(constructor_args.len() - 1); + } // get constructor args - let constructor_args_in_bytes: Vec = Self::constructor_args_encode(constructor_args); + let constructor_args_in_bytes: Vec = Self::constructor_args_encode(&real_constructor_args); // create dummy contract info let mut contract_result = ContractInfo { @@ -338,6 +381,7 @@ impl ContractLoader { files, source_map_replacements, raw_source_map: raw_source_maps.get(&contract_name).cloned(), + balance: contract_balance, }; let mut abi_result = ABIInfo { source: prefix.to_string(), @@ -380,6 +424,7 @@ impl ContractLoader { contract_result.constructor_args = abi_instance.get().get_bytes(); } // debug!("Constructor args: {:?}", result.constructor_args); + // set constructor args balance contract_result.code.extend(contract_result.constructor_args.clone()); } else { debug!("No constructor in ABI found, skipping"); @@ -499,6 +544,7 @@ impl ContractLoader { } } } + let prefix_loader = Self::from_prefix( (prefix.to_owned() + &String::from('*')).as_str(), state, @@ -584,6 +630,7 @@ impl ContractLoader { files, source_map_replacements, raw_source_map, + balance: EVMU256::from(0), }); abis.push(ABIInfo { source: addr.to_string(), @@ -642,6 +689,7 @@ impl ContractLoader { files: sources, source_map_replacements: Some(more_info.source_map_replacements), raw_source_map: Some(more_info.source_map.clone()), + balance: EVMU256::from(0), }); } @@ -752,6 +800,7 @@ impl ContractLoader { files: artifact.sources.clone(), source_map_replacements: Some(more_info.source_map_replacements.clone()), raw_source_map: Some(more_info.source_map.clone()), + balance: EVMU256::from(0), }); } Self { @@ -889,6 +938,7 @@ impl ContractLoader { files: vec![], source_map_replacements: None, raw_source_map: None, + balance: EVMU256::from(0), }); abis.push(ABIInfo { @@ -950,6 +1000,7 @@ impl ContractLoader { files: artifact.sources.clone(), source_map_replacements: Some(more_info.source_map_replacements.clone()), raw_source_map: Some(more_info.source_map.clone()), + balance: EVMU256::from(0), }); } @@ -1390,13 +1441,22 @@ pub fn to_hex_string(bytes: &[u8]) -> String { mod tests { use super::*; - use crate::{skip_cbor, state::FuzzState}; + use crate::{evm::parse_constructor_args_string, skip_cbor, state::FuzzState}; #[test] fn test_load() { let codes: Vec = vec![]; - let args: HashMap> = HashMap::new(); - let loader = ContractLoader::from_glob("demo/*", &mut FuzzState::new(0), &codes, &args, String::from(""), None); + // let mut args: HashMap> = HashMap::new(); + let input = "contract1:88,97C6D26d7E0D316850A967b46845E15a32666d25,1800;contract2:88,97C6D26d7E0D316850A967b46845E15a32666d25,1800".to_string(); + let args = parse_constructor_args_string(input); + let loader = ContractLoader::from_glob( + "tests/evm/real_balance/*", + &mut FuzzState::new(0), + &codes, + &args, + String::from(""), + None, + ); debug!( "{:?}", loader.contracts.iter().map(|x| x.name.clone()).collect::>() diff --git a/src/evm/corpus_initializer.rs b/src/evm/corpus_initializer.rs index aadcf2b3..9eb9aba8 100644 --- a/src/evm/corpus_initializer.rs +++ b/src/evm/corpus_initializer.rs @@ -200,6 +200,7 @@ where .host .evmstate .set_balance(self.executor.deployer, EVMU256::from(INITIAL_BALANCE)); + // deploy for contract in &mut loader.contracts { info!("Deploying contract: {}", contract.name); @@ -228,7 +229,17 @@ where contract.deployed_address }; contract.deployed_address = deployed_address; - info!("Contract {} deployed to: {deployed_address:?}", contract.name); + + // set all contracts real balance + self.executor + .host + .evmstate + .set_balance(deployed_address, contract.balance); + + info!( + "Contract {} deployed to: {deployed_address:?} -> balacne is {:?}", + contract.name, contract.balance + ); if deployed_address != CHEATCODE_ADDRESS { self.state.add_address(&deployed_address); diff --git a/src/evm/mod.rs b/src/evm/mod.rs index 4f189615..8b5847e0 100644 --- a/src/evm/mod.rs +++ b/src/evm/mod.rs @@ -1015,3 +1015,17 @@ fn test_evm_offchain_setup() { utils::try_write_file(&abis_json, &json_str, true).unwrap(); evm_fuzzer(config, &mut state) } + +#[cfg(test)] +mod test { + use super::parse_constructor_args_string; + + #[test] + fn test_parse_constructor_args_string() { + let input = + "Test1:88,0x97C6D26d7E0D316850A967b46845E15a32666d25;Test2:88,0x97C6D26d7E0D316850A967b46845E15a32666d25" + .to_string(); + let ret = parse_constructor_args_string(input); + // println!("constructor args: {:?}", ret); + } +} diff --git a/tests/evm_real_balance/contract1.abi b/tests/evm_real_balance/contract1.abi new file mode 100644 index 00000000..1fb30f1c --- /dev/null +++ b/tests/evm_real_balance/contract1.abi @@ -0,0 +1 @@ +[{"type":"constructor","inputs":[{"name":"_x1","type":"uint256","internalType":"uint256"},{"name":"_t","type":"address","internalType":"address"}],"stateMutability":"payable"},{"type":"function","name":"t","inputs":[],"outputs":[{"name":"","type":"address","internalType":"address"}],"stateMutability":"view"},{"type":"function","name":"test1","inputs":[],"outputs":[],"stateMutability":"payable"},{"type":"function","name":"x1","inputs":[],"outputs":[{"name":"","type":"uint256","internalType":"uint256"}],"stateMutability":"view"},{"type":"event","name":"AssertionFailed","inputs":[{"name":"message","type":"string","indexed":false,"internalType":"string"}],"anonymous":false}] \ No newline at end of file diff --git a/tests/evm_real_balance/contract1.bin b/tests/evm_real_balance/contract1.bin new file mode 100644 index 00000000..3d5cf3e7 --- /dev/null +++ b/tests/evm_real_balance/contract1.bin @@ -0,0 +1 @@ +60806040525f80556040516101be3803806101be833981016040819052602391604b565b5f91909155600180546001600160a01b0319166001600160a01b039092169190911790556083565b5f8060408385031215605b575f80fd5b825160208401519092506001600160a01b03811681146078575f80fd5b809150509250929050565b61012e806100905f395ff3fe608060405260043610602f575f3560e01c8063343943bd1460335780636b59084d14605857806392d0d153146060575b5f80fd5b348015603d575f80fd5b5060455f5481565b6040519081526020015b60405180910390f35b605e6094565b005b348015606a575f80fd5b50600154607d906001600160a01b031681565b6040516001600160a01b039091168152602001604f565b60015f55476107080360a65760a660a8565b565b60a67fb42604cb105a16c8f6db8a41e6b00c0c1b4826465e8bc504b3eb3e88b3e6a4a060405160ee9060208082526003908201526242756760e81b604082015260600190565b60405180910390a156fea264697066735822122036ed630f0c26eee1d9ac1ee1b34a6d817190be6bcb5d6c9cdab1d10dff0795d264736f6c63430008190033 \ No newline at end of file diff --git a/tests/evm_real_balance/contract2.abi b/tests/evm_real_balance/contract2.abi new file mode 100644 index 00000000..4f94f9ae --- /dev/null +++ b/tests/evm_real_balance/contract2.abi @@ -0,0 +1 @@ +[{"type":"constructor","inputs":[{"name":"_x1","type":"uint256","internalType":"uint256"},{"name":"_t","type":"address","internalType":"address"}],"stateMutability":"payable"},{"type":"function","name":"t","inputs":[],"outputs":[{"name":"","type":"address","internalType":"address"}],"stateMutability":"view"},{"type":"function","name":"test1","inputs":[],"outputs":[],"stateMutability":"nonpayable"},{"type":"function","name":"x1","inputs":[],"outputs":[{"name":"","type":"uint256","internalType":"uint256"}],"stateMutability":"view"}] \ No newline at end of file diff --git a/tests/evm_real_balance/contract2.bin b/tests/evm_real_balance/contract2.bin new file mode 100644 index 00000000..b020fcdb --- /dev/null +++ b/tests/evm_real_balance/contract2.bin @@ -0,0 +1 @@ +60806040525f8055604051610131380380610131833981016040819052602391604b565b5f91909155600180546001600160a01b0319166001600160a01b039092169190911790556083565b5f8060408385031215605b575f80fd5b825160208401519092506001600160a01b03811681146078575f80fd5b809150509250929050565b60a28061008f5f395ff3fe6080604052348015600e575f80fd5b50600436106030575f3560e01c80636b59084d14603457806392d0d15314603e575b5f80fd5b603c60025f55565b005b6001546050906001600160a01b031681565b6040516001600160a01b03909116815260200160405180910390f3fea2646970667358221220defb85baec4ed35ad480839ba541008dc7ca40059a79a7ed2023c71a1226ff2b64736f6c63430008190033 \ No newline at end of file diff --git a/tests/evm_real_balance/contract3.abi b/tests/evm_real_balance/contract3.abi new file mode 100644 index 00000000..17719906 --- /dev/null +++ b/tests/evm_real_balance/contract3.abi @@ -0,0 +1 @@ +[{"type":"function","name":"test1","inputs":[],"outputs":[],"stateMutability":"nonpayable"}] \ No newline at end of file diff --git a/tests/evm_real_balance/contract3.bin b/tests/evm_real_balance/contract3.bin new file mode 100644 index 00000000..a03c47d1 --- /dev/null +++ b/tests/evm_real_balance/contract3.bin @@ -0,0 +1 @@ +60806040525f80553480156011575f80fd5b50606a80601d5f395ff3fe6080604052348015600e575f80fd5b50600436106026575f3560e01c80636b59084d14602a575b5f80fd5b603260025f55565b00fea26469706673582212201f1f74f6a1949919571720fbfc08bc69eb34a2dd22176f10cefc91651203c53a64736f6c63430008190033 \ No newline at end of file diff --git a/tests/evm_real_balance/test_real_balance.sol b/tests/evm_real_balance/test_real_balance.sol new file mode 100644 index 00000000..d01ae47e --- /dev/null +++ b/tests/evm_real_balance/test_real_balance.sol @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.3; +import "solidity_utils/lib.sol"; + +contract Test1 { + uint public x1 = 0; + address public t; + constructor(uint _x1, address _t) payable { + x1 = _x1; + t = _t; + } + + function test1() payable public { + x1 = 1; + if (address(this).balance == 1800) { + bug(); + } + } +} + + +contract Test2 { + uint x2 = 0; + address public t; + constructor(uint _x2, address _t) payable { + x2 = _x2; + t = _t; + } + + function test1() public { + x2 = 2; + } +} + +contract Test3 { + uint x2 = 0; + + function test1() public { + x2 = 2; + } +} \ No newline at end of file