FEAT: Add new known addresses file argument to init (#93)

* feat init

* Expand test

* Fix formatting

---------

Co-authored-by: Lothaire Aubergeon <lothaire.aubergeon@chainsecurity.com>

Author: GiovanniTorrisi-ChainSecurity

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,24 +16,51 @@ use crate::utils::pretty::{AddressType, ResolvedAddress};
 pub struct Registry {
     dvf_storage: PathBuf,
     trusted_signers: Vec<Address>,
+    known_addresses: HashMap<u64, HashMap<Address, ResolvedAddress>>,
 }
 
 impl Registry {
     pub fn from_config(config: &DVFConfig) -> Result<Self, ValidationError> {
+        Self::from_config_with_known(config, None)
+    }
+
+    pub fn from_config_with_known(
+        config: &DVFConfig,
+        known_addresses: Option<HashMap<u64, HashMap<Address, ResolvedAddress>>>,
+    ) -> Result<Self, ValidationError> {
         let dvf_storage = config.dvf_storage.clone();
 
         Ok(Registry {
             dvf_storage,
             trusted_signers: config.trusted_signers.clone(),
+            known_addresses: known_addresses.unwrap_or_default(),
         })
     }
 
     pub fn is_trusted_signer(&self, address: &Address) -> bool {
         self.trusted_signers.contains(address)
     }
 
+    pub fn load_known_addresses_from_file(
+        path: &Path,
+    ) -> Result<HashMap<u64, HashMap<Address, ResolvedAddress>>, ValidationError> {
+        let f = fs::File::open(path)?;
+        serde_json::from_reader::<_, HashMap<u64, HashMap<Address, ResolvedAddress>>>(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<Address, ResolvedAddress> {
         let mut res: HashMap<Address, ResolvedAddress> = 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<Vec<PathBuf>, 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());
+    }
+}

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,

src/dvf.rs

@@ -485,6 +485,11 @@ fn main() {
                     arg!(<OUTPUT>)
                         .help("Path of the generated DVF file")
                         .required(true),
+                )
+                .arg(
+                    arg!(--knownaddressesfile <PATH>)
+                        .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::<PathBuf>("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(&registry));
 
             // Parse optional initblock or take deployment_block_num + 1

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);
 

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;