diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 1f3d8d7da6139..8e21191f4e3b8 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -1346,8 +1346,8 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap BTreeMap Option { - // Check if we have available artifacts to match against - let artifacts = ccx.state.config.available_artifacts.as_ref()?; - - // Try to load the account and get its code - let account = ccx.ecx.journaled_state.load_account(address).ok()?; - let code = account.info.code.as_ref()?; - - // Skip if code is empty - if code.is_empty() { - return None; - } - - // Try to find the artifact by deployed code - let code_bytes = code.original_bytes(); - if let Some((artifact_id, _)) = artifacts.find_by_deployed_code_exact(&code_bytes) { - return Some(artifact_id.identifier()); - } +/// EIP-1967 implementation storage slot +const EIP1967_IMPL_SLOT: &str = "360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc"; - // Fallback to fuzzy matching if exact match fails - if let Some((artifact_id, _)) = artifacts.find_by_deployed_code(&code_bytes) { - return Some(artifact_id.identifier()); - } - - None -} +/// EIP-1822 UUPS implementation storage slot: keccak256("PROXIABLE") +const EIP1822_PROXIABLE_SLOT: &str = + "c5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7"; /// Helper function to get the contract data from the deployed code at an address. fn get_contract_data<'a>( @@ -1518,6 +1497,22 @@ fn get_contract_data<'a>( // Try to find the artifact by deployed code let code_bytes = code.original_bytes(); + // First check for proxy patterns + let hex_str = hex::encode(&code_bytes); + let find_by_suffix = + |suffix: &str| artifacts.iter().find(|(a, _)| a.identifier().ends_with(suffix)); + // Simple proxy detection based on storage slot patterns + if hex_str.contains(EIP1967_IMPL_SLOT) + && let Some(result) = find_by_suffix(":TransparentUpgradeableProxy") + { + return Some(result); + } else if hex_str.contains(EIP1822_PROXIABLE_SLOT) + && let Some(result) = find_by_suffix(":UUPSUpgradeable") + { + return Some(result); + } + + // Try exact match if let Some(result) = artifacts.find_by_deployed_code_exact(&code_bytes) { return Some(result); }