From 2f9ac345e4686ec415d7093f6da9148b005231f5 Mon Sep 17 00:00:00 2001 From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com> Date: Fri, 22 Aug 2025 17:31:38 +0530 Subject: [PATCH 1/3] fix(cheatcodes): identify proxies in state diffs --- crates/cheatcodes/src/evm.rs | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 1f3d8d7da6139..b681b94a4363e 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -1471,11 +1471,17 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> 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()?; @@ -1487,11 +1493,27 @@ fn get_contract_name(ccx: &mut CheatsCtxt, address: Address) -> Option { // Try to find the artifact by deployed code let code_bytes = code.original_bytes(); + + // Check if this is a proxy contract by looking for proxy storage slot patterns + let hex_str = hex::encode(&code_bytes); + + let find_by_suffix = |suffix: &str| { + artifacts.iter().find_map(|(a, _)| a.identifier().ends_with(suffix).then(|| a.identifier())) + }; + + // Simple proxy detection based on storage slot patterns + if hex_str.contains(EIP1967_IMPL_SLOT) { + return find_by_suffix(":TransparentUpgradeableProxy"); + } else if hex_str.contains(EIP1822_PROXIABLE_SLOT) { + return find_by_suffix(":UUPSUpgradeable"); + } + + // Try exact match if let Some((artifact_id, _)) = artifacts.find_by_deployed_code_exact(&code_bytes) { return Some(artifact_id.identifier()); } - // Fallback to fuzzy matching if exact match fails + // Try fuzzy matching if let Some((artifact_id, _)) = artifacts.find_by_deployed_code(&code_bytes) { return Some(artifact_id.identifier()); } @@ -1518,6 +1540,7 @@ fn get_contract_data<'a>( // Try to find the artifact by deployed code let code_bytes = code.original_bytes(); + if let Some(result) = artifacts.find_by_deployed_code_exact(&code_bytes) { return Some(result); } From 7389e0c38d10beada5eb9758d4f74c59a37ae7b0 Mon Sep 17 00:00:00 2001 From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com> Date: Fri, 22 Aug 2025 17:36:44 +0530 Subject: [PATCH 2/3] dedup get_contract_name and get_contract_data --- crates/cheatcodes/src/evm.rs | 62 ++++++++++-------------------------- 1 file changed, 17 insertions(+), 45 deletions(-) diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index b681b94a4363e..e05b805b8a7fd 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 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(); - - // Check if this is a proxy contract by looking for proxy storage slot patterns - let hex_str = hex::encode(&code_bytes); - - let find_by_suffix = |suffix: &str| { - artifacts.iter().find_map(|(a, _)| a.identifier().ends_with(suffix).then(|| a.identifier())) - }; - - // Simple proxy detection based on storage slot patterns - if hex_str.contains(EIP1967_IMPL_SLOT) { - return find_by_suffix(":TransparentUpgradeableProxy"); - } else if hex_str.contains(EIP1822_PROXIABLE_SLOT) { - return find_by_suffix(":UUPSUpgradeable"); - } - - // Try exact match - if let Some((artifact_id, _)) = artifacts.find_by_deployed_code_exact(&code_bytes) { - return Some(artifact_id.identifier()); - } - - // Try fuzzy matching - if let Some((artifact_id, _)) = artifacts.find_by_deployed_code(&code_bytes) { - return Some(artifact_id.identifier()); - } - - None -} - /// Helper function to get the contract data from the deployed code at an address. fn get_contract_data<'a>( ccx: &'a mut CheatsCtxt, @@ -1540,7 +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) { + if let Some(result) = find_by_suffix(":TransparentUpgradeableProxy") { + return Some(result); + } + } else if hex_str.contains(EIP1822_PROXIABLE_SLOT) { + if 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); } From 66363b2474d35f5f10a5b4e4bb6294ce7ecdb887 Mon Sep 17 00:00:00 2001 From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com> Date: Fri, 22 Aug 2025 17:44:28 +0530 Subject: [PATCH 3/3] clippy --- crates/cheatcodes/src/evm.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index e05b805b8a7fd..8e21191f4e3b8 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -1502,14 +1502,14 @@ fn get_contract_data<'a>( 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) { - if let Some(result) = find_by_suffix(":TransparentUpgradeableProxy") { - return Some(result); - } - } else if hex_str.contains(EIP1822_PROXIABLE_SLOT) { - if let Some(result) = find_by_suffix(":UUPSUpgradeable") { - return Some(result); - } + 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