foundry_evm_traces/identifier/
local.rs

1use super::{IdentifiedAddress, TraceIdentifier};
2use alloy_dyn_abi::JsonAbiExt;
3use alloy_json_abi::JsonAbi;
4use alloy_primitives::{Address, Bytes, map::HashMap};
5use foundry_common::contracts::{ContractsByArtifact, bytecode_diff_score};
6use foundry_compilers::ArtifactId;
7use revm_inspectors::tracing::types::CallTraceNode;
8use std::borrow::Cow;
9
10/// A trace identifier that tries to identify addresses using local contracts.
11pub struct LocalTraceIdentifier<'a> {
12    /// Known contracts to search through.
13    known_contracts: &'a ContractsByArtifact,
14    /// Vector of pairs of artifact ID and the runtime code length of the given artifact.
15    ordered_ids: Vec<(&'a ArtifactId, usize)>,
16    /// The contracts bytecode.
17    contracts_bytecode: Option<&'a HashMap<Address, Bytes>>,
18}
19
20impl<'a> LocalTraceIdentifier<'a> {
21    /// Creates a new local trace identifier.
22    #[inline]
23    pub fn new(known_contracts: &'a ContractsByArtifact) -> Self {
24        let mut ordered_ids = known_contracts
25            .iter()
26            .filter_map(|(id, contract)| Some((id, contract.deployed_bytecode()?)))
27            .map(|(id, bytecode)| (id, bytecode.len()))
28            .collect::<Vec<_>>();
29        ordered_ids.sort_by_key(|(_, len)| *len);
30        Self { known_contracts, ordered_ids, contracts_bytecode: None }
31    }
32
33    pub fn with_bytecodes(mut self, contracts_bytecode: &'a HashMap<Address, Bytes>) -> Self {
34        self.contracts_bytecode = Some(contracts_bytecode);
35        self
36    }
37
38    /// Returns the known contracts.
39    #[inline]
40    pub fn contracts(&self) -> &'a ContractsByArtifact {
41        self.known_contracts
42    }
43
44    /// Identifies the artifact based on score computed for both creation and deployed bytecodes.
45    pub fn identify_code(
46        &self,
47        runtime_code: &[u8],
48        creation_code: &[u8],
49    ) -> Option<(&'a ArtifactId, &'a JsonAbi)> {
50        let len = runtime_code.len();
51
52        let mut min_score = f64::MAX;
53        let mut min_score_id = None;
54
55        let mut check = |id, is_creation, min_score: &mut f64| {
56            let contract = self.known_contracts.get(id)?;
57            // Select bytecodes to compare based on `is_creation` flag.
58            let (contract_bytecode, current_bytecode) = if is_creation {
59                (contract.bytecode_without_placeholders(), creation_code)
60            } else {
61                (contract.deployed_bytecode_without_placeholders(), runtime_code)
62            };
63
64            if let Some(bytecode) = contract_bytecode {
65                let mut current_bytecode = current_bytecode;
66                if is_creation && current_bytecode.len() > bytecode.len() {
67                    // Try to decode ctor args with contract abi.
68                    if let Some(constructor) = contract.abi.constructor() {
69                        let constructor_args = &current_bytecode[bytecode.len()..];
70                        if constructor.abi_decode_input(constructor_args).is_ok() {
71                            // If we can decode args with current abi then remove args from
72                            // code to compare.
73                            current_bytecode = &current_bytecode[..bytecode.len()]
74                        }
75                    }
76                }
77
78                let score = bytecode_diff_score(&bytecode, current_bytecode);
79                if score == 0.0 {
80                    trace!(target: "evm::traces::local", "found exact match");
81                    return Some((id, &contract.abi));
82                }
83                if score < *min_score {
84                    *min_score = score;
85                    min_score_id = Some((id, &contract.abi));
86                }
87            }
88            None
89        };
90
91        // Check `[len * 0.9, ..., len * 1.1]`.
92        let max_len = (len * 11) / 10;
93
94        // Start at artifacts with the same code length: `len..len*1.1`.
95        let same_length_idx = self.find_index(len);
96        for idx in same_length_idx..self.ordered_ids.len() {
97            let (id, len) = self.ordered_ids[idx];
98            if len > max_len {
99                break;
100            }
101            if let found @ Some(_) = check(id, true, &mut min_score) {
102                return found;
103            }
104        }
105
106        // Iterate over the remaining artifacts with less code length: `len*0.9..len`.
107        let min_len = (len * 9) / 10;
108        let idx = self.find_index(min_len);
109        for i in idx..same_length_idx {
110            let (id, _) = self.ordered_ids[i];
111            if let found @ Some(_) = check(id, true, &mut min_score) {
112                return found;
113            }
114        }
115
116        // Fallback to comparing deployed code if min score greater than threshold.
117        if min_score >= 0.85 {
118            for (artifact, _) in &self.ordered_ids {
119                if let found @ Some(_) = check(artifact, false, &mut min_score) {
120                    return found;
121                }
122            }
123        }
124
125        trace!(target: "evm::traces::local", %min_score, "no exact match found");
126
127        // Note: the diff score can be inaccurate for small contracts so we're using a relatively
128        // high threshold here to avoid filtering out too many contracts.
129        if min_score < 0.85 { min_score_id } else { None }
130    }
131
132    /// Returns the index of the artifact with the given code length, or the index of the first
133    /// artifact with a greater code length if the exact code length is not found.
134    fn find_index(&self, len: usize) -> usize {
135        let (Ok(mut idx) | Err(mut idx)) =
136            self.ordered_ids.binary_search_by_key(&len, |(_, probe)| *probe);
137
138        // In case of multiple artifacts with the same code length, we need to find the first one.
139        while idx > 0 && self.ordered_ids[idx - 1].1 == len {
140            idx -= 1;
141        }
142
143        idx
144    }
145}
146
147impl TraceIdentifier for LocalTraceIdentifier<'_> {
148    fn identify_addresses(&mut self, nodes: &[&CallTraceNode]) -> Vec<IdentifiedAddress<'_>> {
149        if nodes.is_empty() {
150            return Vec::new();
151        }
152
153        trace!(target: "evm::traces::local", "identify {} addresses", nodes.len());
154
155        nodes
156            .iter()
157            .map(|&node| {
158                (
159                    node.trace.address,
160                    node.trace.kind.is_any_create().then_some(&node.trace.output[..]),
161                    node.trace.kind.is_any_create().then_some(&node.trace.data[..]),
162                )
163            })
164            .filter_map(|(address, runtime_code, creation_code)| {
165                let _span =
166                    trace_span!(target: "evm::traces::local", "identify", %address).entered();
167
168                // In order to identify the addresses, we need at least the runtime code. It can be
169                // obtained from the trace itself (if it's a CREATE* call), or from the fetched
170                // bytecodes.
171                let (runtime_code, creation_code) = match (runtime_code, creation_code) {
172                    (Some(runtime_code), Some(creation_code)) => (runtime_code, creation_code),
173                    (Some(runtime_code), _) => (runtime_code, &[] as &[u8]),
174                    _ => {
175                        let code = self.contracts_bytecode?.get(&address)?;
176                        (code.as_ref(), &[] as &[u8])
177                    }
178                };
179                let (id, abi) = self.identify_code(runtime_code, creation_code)?;
180                trace!(target: "evm::traces::local", id=%id.identifier(), "identified");
181
182                Some(IdentifiedAddress {
183                    address,
184                    contract: Some(id.identifier()),
185                    label: Some(id.name.clone()),
186                    abi: Some(Cow::Borrowed(abi)),
187                    artifact_id: Some(id.clone()),
188                })
189            })
190            .collect()
191    }
192}