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
10pub struct LocalTraceIdentifier<'a> {
12 known_contracts: &'a ContractsByArtifact,
14 ordered_ids: Vec<(&'a ArtifactId, usize)>,
16 contracts_bytecode: Option<&'a HashMap<Address, Bytes>>,
18}
19
20impl<'a> LocalTraceIdentifier<'a> {
21 #[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 #[inline]
40 pub fn contracts(&self) -> &'a ContractsByArtifact {
41 self.known_contracts
42 }
43
44 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 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 if let Some(constructor) = contract.abi.constructor() {
69 let constructor_args = ¤t_bytecode[bytecode.len()..];
70 if constructor.abi_decode_input(constructor_args).is_ok() {
71 current_bytecode = ¤t_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 let max_len = (len * 11) / 10;
93
94 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 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 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 if min_score < 0.85 { min_score_id } else { None }
130 }
131
132 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 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 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}