diff --git a/lib/dvf/registry.rs b/lib/dvf/registry.rs index 5880a73..240832c 100644 --- a/lib/dvf/registry.rs +++ b/lib/dvf/registry.rs @@ -5,6 +5,7 @@ use std::path::Path; use std::path::PathBuf; use alloy::primitives::Address; +use serde_json; use tracing::debug; use crate::dvf::config::DVFConfig; @@ -15,15 +16,24 @@ use crate::utils::pretty::{AddressType, ResolvedAddress}; pub struct Registry { dvf_storage: PathBuf, trusted_signers: Vec
, + known_addresses: HashMap>, } impl Registry { pub fn from_config(config: &DVFConfig) -> Result { + Self::from_config_with_known(config, None) + } + + pub fn from_config_with_known( + config: &DVFConfig, + known_addresses: Option>>, + ) -> Result { let dvf_storage = config.dvf_storage.clone(); Ok(Registry { dvf_storage, trusted_signers: config.trusted_signers.clone(), + known_addresses: known_addresses.unwrap_or_default(), }) } @@ -31,8 +41,26 @@ impl Registry { self.trusted_signers.contains(address) } + pub fn load_known_addresses_from_file( + path: &Path, + ) -> Result>, ValidationError> { + let f = fs::File::open(path)?; + serde_json::from_reader::<_, HashMap>>(f).map_err( + |e| { + ValidationError::Error(format!( + "Could not parse known addresses file {}: {}", + path.display(), + e + )) + }, + ) + } + pub fn collect_name_resolution(&self, chain_id: u64) -> HashMap { let mut res: HashMap = HashMap::new(); + if let Some(extra) = self.known_addresses.get(&chain_id) { + res.extend(extra.clone()); + } self.collect_names_inner(&self.dvf_storage, &mut res, chain_id); res } @@ -175,3 +203,61 @@ fn search_for_id(dir: &Path, id: &String) -> Result, ValidationErro } Ok(results) } + +#[cfg(test)] +mod tests { + use super::Registry; + use crate::dvf::config::DVFConfig; + use alloy::primitives::Address; + use std::fs; + use std::str::FromStr; + use tempfile::tempdir; + + #[test] + fn load_known_addresses_and_collect() { + let dir = tempdir().unwrap(); + let known_path = dir.path().join("known.json"); + + let address = "0x1111111111111111111111111111111111111111"; + let address2 = "0x1111111111111111111111111111111111111110"; + let name = "TestContract"; + let name2 = "TestContract2"; + + let json = format!( + r#" + {{ + "1": {{ + "{address}": {{ + "name": "{name}", + "address_type": "Contract" + }} + }}, + "2": {{ + "{address2}": {{ + "name": "{name2}", + "address_type": "Contract" + }} + }} + }} + "# + ); + fs::write(&known_path, json).unwrap(); + + let known = Registry::load_known_addresses_from_file(&known_path).unwrap(); + + let mut config = DVFConfig::default(); + config.dvf_storage = dir.path().to_path_buf(); + + let registry = Registry::from_config_with_known(&config, Some(known)).unwrap(); + + let resolved = registry.collect_name_resolution(1); + let addr = Address::from_str(address).unwrap(); + assert_eq!(resolved.get(&addr).unwrap().name, name); + + let resolved2 = registry.collect_name_resolution(2); + let addr2 = Address::from_str(address2).unwrap(); + assert_eq!(resolved2.get(&addr2).unwrap().name, name2); + + assert!(registry.collect_name_resolution(137).is_empty()); + } +} diff --git a/lib/utils/pretty.rs b/lib/utils/pretty.rs index ccf5bdc..476a2b2 100644 --- a/lib/utils/pretty.rs +++ b/lib/utils/pretty.rs @@ -17,7 +17,7 @@ use crate::dvf::config::DVFConfig; use crate::dvf::registry::Registry; use crate::state::contract_state::ContractState; -#[derive(Deserialize, Debug)] +#[derive(Deserialize, Debug, Clone)] pub enum AddressType { Token, Contract, @@ -26,7 +26,7 @@ pub enum AddressType { Eoa, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Clone)] pub struct ResolvedAddress { pub address_type: AddressType, pub name: String, diff --git a/src/dvf.rs b/src/dvf.rs index 9422bd3..c179cfb 100644 --- a/src/dvf.rs +++ b/src/dvf.rs @@ -485,6 +485,11 @@ fn main() { arg!() .help("Path of the generated DVF file") .required(true), + ) + .arg( + arg!(--knownaddressesfile ) + .help("Optional JSON file with known addresses to include in name resolution") + .value_parser(is_valid_path), ), ) .subcommand( @@ -865,7 +870,12 @@ fn process(matches: ArgMatches) -> Result<(), ValidationError> { let mut dumped = parse::CompleteDVF::from_cli(sub_m)?; config.set_chain_id(dumped.chain_id)?; - let registry = registry::Registry::from_config(&config)?; + let known_addresses = sub_m + .get_one::("knownaddressesfile") + .map(|path| registry::Registry::load_known_addresses_from_file(path)) + .transpose()?; + + let registry = registry::Registry::from_config_with_known(&config, known_addresses)?; let pretty_printer = PrettyPrinter::new(&config, Some(®istry)); // Parse optional initblock or take deployment_block_num + 1 diff --git a/tests/Contracts/src/CrazyHiddenStruct.sol b/tests/Contracts/src/CrazyHiddenStruct.sol index 8856a53..e14f9bd 100644 --- a/tests/Contracts/src/CrazyHiddenStruct.sol +++ b/tests/Contracts/src/CrazyHiddenStruct.sol @@ -39,8 +39,10 @@ contract CrazyHiddenStruct { bytes32 private constant StorageLocation2 = 0x852cbd6b186221cbf354c68826ab57cef1512cf2f5d959ca4501e155cbea7ae8; bytes32 private constant StorageLocation3 = 0x9482765040f1c978ae595e69b3ad0e4697ca0d1e0581a09be85cfb4a8462e752; bytes32 private constant StorageLocation4 = 0xe82aa111a62567be9a414850f7168d2e6c9f9d61a82b90598df0a59035cd53a6; - bytes32 private constant DirectStorageLocation1 = 0xbfbceebbfa6e5996c6a04ac6db0e347528756a4f073935304cc6139dcc2fb653; - bytes32 private constant DirectStorageLocation2 = 0x42d0407cb447148fd182bf527909ab1ba2fbaefe3f25cbe9851153586910b294; + bytes32 private constant DirectStorageLocation1 = + 0xbfbceebbfa6e5996c6a04ac6db0e347528756a4f073935304cc6139dcc2fb653; + bytes32 private constant DirectStorageLocation2 = + 0x42d0407cb447148fd182bf527909ab1ba2fbaefe3f25cbe9851153586910b294; bytes32 private constant KeccakStorageLocation1 = keccak256("keccak1"); bytes32 private constant KeccakStorageLocation2 = bytes32(uint256(keccak256("keccak2")) - 1); diff --git a/tests/Contracts/src/MyToken.sol b/tests/Contracts/src/MyToken.sol index 0c28cbf..b50741b 100644 --- a/tests/Contracts/src/MyToken.sol +++ b/tests/Contracts/src/MyToken.sol @@ -78,6 +78,7 @@ interface IERC20Upgradeable { */ function transferFrom(address from, address to, uint256 amount) external returns (bool); } + // OpenZeppelin Contracts v4.4.1 (token/ERC20/extensions/IERC20Metadata.sol) /** @@ -101,6 +102,7 @@ interface IERC20MetadataUpgradeable is IERC20Upgradeable { */ function decimals() external view returns (uint8); } + // OpenZeppelin Contracts v4.4.1 (utils/Context.sol) // OpenZeppelin Contracts (last updated v4.9.0) (proxy/utils/Initializable.sol) @@ -422,7 +424,8 @@ abstract contract Initializable { modifier initializer() { bool isTopLevelCall = !_initializing; require( - (isTopLevelCall && _initialized < 1) || (!AddressUpgradeable.isContract(address(this)) && _initialized == 1), + (isTopLevelCall && _initialized < 1) + || (!AddressUpgradeable.isContract(address(this)) && _initialized == 1), "Initializable: contract is already initialized" ); _initialized = 1;