fix(traces): identify artifacts using both deployed and creation code (foundry-rs#9050)

* Identify by creation code

* Compute score for both creation and runtime code

* Fallback to deployed bytecode only if min creation bytecode score is under threshold

* reuse check closure, add basic test

Author: grandizzy

crates/evm/traces/src/decoder/mod.rs

@@ -262,17 +262,18 @@ impl CallTraceDecoder {
     pub fn trace_addresses<'a>(
         &'a self,
         arena: &'a CallTraceArena,
-    ) -> impl Iterator<Item = (&'a Address, Option<&'a [u8]>)> + Clone + 'a {
+    ) -> impl Iterator<Item = (&'a Address, Option<&'a [u8]>, Option<&'a [u8]>)> + Clone + 'a {
         arena
             .nodes()
             .iter()
             .map(|node| {
                 (
                     &node.trace.address,
                     node.trace.kind.is_any_create().then_some(&node.trace.output[..]),
+                    node.trace.kind.is_any_create().then_some(&node.trace.data[..]),
                 )
             })
-            .filter(|&(address, _)| {
+            .filter(|&(address, _, _)| {
                 !self.labels.contains_key(address) || !self.contracts.contains_key(address)
             })
     }

crates/evm/traces/src/identifier/etherscan.rs

@@ -97,7 +97,7 @@ impl EtherscanIdentifier {
 impl TraceIdentifier for EtherscanIdentifier {
     fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity<'_>>
     where
-        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>,
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>, Option<&'a [u8]>)>,
     {
         trace!(target: "evm::traces", "identify {:?} addresses", addresses.size_hint().1);
 
@@ -114,7 +114,7 @@ impl TraceIdentifier for EtherscanIdentifier {
             Arc::clone(&self.invalid_api_key),
         );
 
-        for (addr, _) in addresses {
+        for (addr, _, _) in addresses {
             if let Some(metadata) = self.contracts.get(addr) {
                 let label = metadata.contract_name.clone();
                 let abi = metadata.abi().ok().map(Cow::Owned);

crates/evm/traces/src/identifier/local.rs

@@ -32,23 +32,34 @@ impl<'a> LocalTraceIdentifier<'a> {
         self.known_contracts
     }
 
-    /// Tries to the bytecode most similar to the given one.
-    pub fn identify_code(&self, code: &[u8]) -> Option<(&'a ArtifactId, &'a JsonAbi)> {
-        let len = code.len();
+    /// Identifies the artifact based on score computed for both creation and deployed bytecodes.
+    pub fn identify_code(
+        &self,
+        runtime_code: &[u8],
+        creation_code: &[u8],
+    ) -> Option<(&'a ArtifactId, &'a JsonAbi)> {
+        let len = runtime_code.len();
 
         let mut min_score = f64::MAX;
         let mut min_score_id = None;
 
-        let mut check = |id| {
+        let mut check = |id, is_creation, min_score: &mut f64| {
             let contract = self.known_contracts.get(id)?;
-            if let Some(deployed_bytecode) = contract.deployed_bytecode() {
-                let score = bytecode_diff_score(deployed_bytecode, code);
+            // Select bytecodes to compare based on `is_creation` flag.
+            let (contract_bytecode, current_bytecode) = if is_creation {
+                (contract.bytecode(), creation_code)
+            } else {
+                (contract.deployed_bytecode(), runtime_code)
+            };
+
+            if let Some(bytecode) = contract_bytecode {
+                let score = bytecode_diff_score(bytecode, current_bytecode);
                 if score == 0.0 {
                     trace!(target: "evm::traces", "found exact match");
                     return Some((id, &contract.abi));
                 }
-                if score < min_score {
-                    min_score = score;
+                if score < *min_score {
+                    *min_score = score;
                     min_score_id = Some((id, &contract.abi));
                 }
             }
@@ -65,7 +76,7 @@ impl<'a> LocalTraceIdentifier<'a> {
             if len > max_len {
                 break;
             }
-            if let found @ Some(_) = check(id) {
+            if let found @ Some(_) = check(id, true, &mut min_score) {
                 return found;
             }
         }
@@ -75,11 +86,20 @@ impl<'a> LocalTraceIdentifier<'a> {
         let idx = self.find_index(min_len);
         for i in idx..same_length_idx {
             let (id, _) = self.ordered_ids[i];
-            if let found @ Some(_) = check(id) {
+            if let found @ Some(_) = check(id, true, &mut min_score) {
                 return found;
             }
         }
 
+        // Fallback to comparing deployed code if min score greater than threshold.
+        if min_score >= 0.85 {
+            for (artifact, _) in &self.ordered_ids {
+                if let found @ Some(_) = check(artifact, false, &mut min_score) {
+                    return found;
+                }
+            }
+        }
+
         trace!(target: "evm::traces", %min_score, "no exact match found");
 
         // Note: the diff score can be inaccurate for small contracts so we're using a relatively
@@ -109,16 +129,16 @@ impl<'a> LocalTraceIdentifier<'a> {
 impl TraceIdentifier for LocalTraceIdentifier<'_> {
     fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity<'_>>
     where
-        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>,
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>, Option<&'a [u8]>)>,
     {
         trace!(target: "evm::traces", "identify {:?} addresses", addresses.size_hint().1);
 
         addresses
-            .filter_map(|(address, code)| {
+            .filter_map(|(address, runtime_code, creation_code)| {
                 let _span = trace_span!(target: "evm::traces", "identify", %address).entered();
 
                 trace!(target: "evm::traces", "identifying");
-                let (id, abi) = self.identify_code(code?)?;
+                let (id, abi) = self.identify_code(runtime_code?, creation_code?)?;
                 trace!(target: "evm::traces", id=%id.identifier(), "identified");
 
                 Some(AddressIdentity {

crates/evm/traces/src/identifier/mod.rs

@@ -35,7 +35,7 @@ pub trait TraceIdentifier {
     /// Attempts to identify an address in one or more call traces.
     fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity<'_>>
     where
-        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)> + Clone;
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>, Option<&'a [u8]>)> + Clone;
 }
 
 /// A collection of trace identifiers.
@@ -55,7 +55,7 @@ impl Default for TraceIdentifiers<'_> {
 impl TraceIdentifier for TraceIdentifiers<'_> {
     fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity<'_>>
     where
-        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)> + Clone,
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>, Option<&'a [u8]>)> + Clone,
     {
         let mut identities = Vec::new();
         if let Some(local) = &mut self.local {

crates/forge/tests/cli/test_cmd.rs

@@ -2274,3 +2274,102 @@ Logs:
 ...
 "#]]);
 });
+
+// <https://github.com/foundry-rs/foundry/issues/8995>
+forgetest_init!(metadata_bytecode_traces, |prj, cmd| {
+    prj.add_source(
+        "ParentProxy.sol",
+        r#"
+import {Counter} from "./Counter.sol";
+
+abstract contract ParentProxy {
+    Counter impl;
+    bytes data;
+
+    constructor(Counter _implementation, bytes memory _data) {
+        impl = _implementation;
+        data = _data;
+    }
+}
+   "#,
+    )
+    .unwrap();
+    prj.add_source(
+        "Proxy.sol",
+        r#"
+import {ParentProxy} from "./ParentProxy.sol";
+import {Counter} from "./Counter.sol";
+
+contract Proxy is ParentProxy {
+    constructor(Counter _implementation, bytes memory _data)
+        ParentProxy(_implementation, _data)
+    {}
+}
+   "#,
+    )
+    .unwrap();
+
+    prj.add_test(
+        "MetadataTraceTest.t.sol",
+        r#"
+import {Counter} from "src/Counter.sol";
+import {Proxy} from "src/Proxy.sol";
+
+import {Test} from "forge-std/Test.sol";
+
+contract MetadataTraceTest is Test {
+    function test_proxy_trace() public {
+        Counter counter = new Counter();
+        new Proxy(counter, "");
+    }
+}
+   "#,
+    )
+    .unwrap();
+
+    cmd.args(["test", "--mt", "test_proxy_trace", "-vvvv"]).assert_success().stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+
+Ran 1 test for test/MetadataTraceTest.t.sol:MetadataTraceTest
+[PASS] test_proxy_trace() ([GAS])
+Traces:
+  [152142] MetadataTraceTest::test_proxy_trace()
+    ├─ [49499] → new Counter@0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f
+    │   └─ ← [Return] 247 bytes of code
+    ├─ [37978] → new Proxy@0x2e234DAe75C793f67A35089C9d99245E1C58470b
+    │   └─ ← [Return] 63 bytes of code
+    └─ ← [Stop] 
+
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+
+Ran 1 test suite [ELAPSED]: 1 tests passed, 0 failed, 0 skipped (1 total tests)
+
+"#]]);
+
+    // Check consistent traces for running with no metadata.
+    cmd.forge_fuse()
+        .args(["test", "--mt", "test_proxy_trace", "-vvvv", "--no-metadata"])
+        .assert_success()
+        .stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+
+Ran 1 test for test/MetadataTraceTest.t.sol:MetadataTraceTest
+[PASS] test_proxy_trace() ([GAS])
+Traces:
+  [130521] MetadataTraceTest::test_proxy_trace()
+    ├─ [38693] → new Counter@0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f
+    │   └─ ← [Return] 193 bytes of code
+    ├─ [27175] → new Proxy@0x2e234DAe75C793f67A35089C9d99245E1C58470b
+    │   └─ ← [Return] 9 bytes of code
+    └─ ← [Stop] 
+
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+
+Ran 1 test suite [ELAPSED]: 1 tests passed, 0 failed, 0 skipped (1 total tests)
+
+"#]]);
+});