foundry_evm_traces/identifier/
signatures.rs

1use alloy_json_abi::{Error, Event, Function, JsonAbi};
2use alloy_primitives::{
3    B256, Selector,
4    map::{HashMap, HashSet},
5};
6use eyre::Result;
7use foundry_common::{
8    abi::{get_error, get_event, get_func},
9    fs,
10    selectors::{OpenChainClient, SelectorKind},
11};
12use foundry_config::Config;
13use serde::{Deserialize, Serialize};
14use std::{
15    collections::BTreeMap,
16    path::{Path, PathBuf},
17    sync::Arc,
18};
19use tokio::sync::RwLock;
20
21/// Cache for function, event and error signatures. Used by [`SignaturesIdentifier`].
22#[derive(Debug, Default, Deserialize)]
23#[serde(try_from = "SignaturesDiskCache")]
24pub struct SignaturesCache {
25    signatures: HashMap<SelectorKind, Option<String>>,
26}
27
28/// Disk representation of the signatures cache.
29#[derive(Serialize, Deserialize)]
30struct SignaturesDiskCache {
31    functions: BTreeMap<Selector, String>,
32    errors: BTreeMap<Selector, String>,
33    events: BTreeMap<B256, String>,
34}
35
36impl From<SignaturesDiskCache> for SignaturesCache {
37    fn from(value: SignaturesDiskCache) -> Self {
38        let functions = value
39            .functions
40            .into_iter()
41            .map(|(selector, signature)| (SelectorKind::Function(selector), signature));
42        let errors = value
43            .errors
44            .into_iter()
45            .map(|(selector, signature)| (SelectorKind::Error(selector), signature));
46        let events = value
47            .events
48            .into_iter()
49            .map(|(selector, signature)| (SelectorKind::Event(selector), signature));
50        Self {
51            signatures: functions
52                .chain(errors)
53                .chain(events)
54                .map(|(sel, sig)| (sel, (!sig.is_empty()).then_some(sig)))
55                .collect(),
56        }
57    }
58}
59
60impl From<&SignaturesCache> for SignaturesDiskCache {
61    fn from(value: &SignaturesCache) -> Self {
62        let (functions, errors, events) = value.signatures.iter().fold(
63            (BTreeMap::new(), BTreeMap::new(), BTreeMap::new()),
64            |mut acc, (kind, signature)| {
65                // Only persist resolved signatures. Unknown selectors (None) are kept
66                // in-memory for session dedup but not written to disk, so they can be
67                // re-queried in future sessions once the signature database is updated.
68                if let Some(value) = signature.clone() {
69                    match *kind {
70                        SelectorKind::Function(selector) => _ = acc.0.insert(selector, value),
71                        SelectorKind::Error(selector) => _ = acc.1.insert(selector, value),
72                        SelectorKind::Event(selector) => _ = acc.2.insert(selector, value),
73                    }
74                }
75                acc
76            },
77        );
78        Self { functions, errors, events }
79    }
80}
81
82impl Serialize for SignaturesCache {
83    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
84    where
85        S: serde::Serializer,
86    {
87        SignaturesDiskCache::from(self).serialize(serializer)
88    }
89}
90
91impl SignaturesCache {
92    /// Loads the cache from a file.
93    #[instrument(target = "evm::traces", name = "SignaturesCache::load")]
94    pub fn load(path: &Path) -> Self {
95        trace!(target: "evm::traces", ?path, "reading signature cache");
96        fs::read_json_file(path)
97            .inspect_err(
98                |err| warn!(target: "evm::traces", ?path, ?err, "failed to read cache file"),
99            )
100            .unwrap_or_default()
101    }
102
103    /// Saves the cache to a file.
104    #[instrument(target = "evm::traces", name = "SignaturesCache::save", skip(self))]
105    pub fn save(&self, path: &Path) {
106        if let Some(parent) = path.parent()
107            && let Err(err) = std::fs::create_dir_all(parent)
108        {
109            warn!(target: "evm::traces", ?parent, %err, "failed to create cache");
110        }
111        if let Err(err) = fs::write_json_file(path, self) {
112            warn!(target: "evm::traces", %err, "failed to flush signature cache");
113        } else {
114            trace!(target: "evm::traces", "flushed signature cache")
115        }
116    }
117
118    /// Updates the cache from an ABI.
119    pub fn extend_from_abi(&mut self, abi: &JsonAbi) {
120        self.extend(Self::signatures_from_abi(abi));
121    }
122
123    /// Updates the cache from ABIs without overwriting previous entries from the same batch.
124    fn extend_from_abis_without_collisions<'a>(
125        &mut self,
126        abis: impl IntoIterator<Item = &'a JsonAbi>,
127    ) {
128        let mut seeded: HashSet<SelectorKind> = HashSet::default();
129        for abi in abis {
130            for (selector, signature) in Self::signatures_from_abi(abi) {
131                if seeded.insert(selector) {
132                    self.insert(selector, signature);
133                } else {
134                    trace!(target: "evm::traces", ?selector, %signature, "skipping duplicate ABI signature");
135                }
136            }
137        }
138    }
139
140    fn signatures_from_abi(abi: &JsonAbi) -> impl Iterator<Item = (SelectorKind, String)> + '_ {
141        abi.items().filter_map(|item| match item {
142            alloy_json_abi::AbiItem::Function(f) => {
143                Some((SelectorKind::Function(f.selector()), f.signature()))
144            }
145            alloy_json_abi::AbiItem::Error(e) => {
146                Some((SelectorKind::Error(e.selector()), e.signature()))
147            }
148            alloy_json_abi::AbiItem::Event(e) => {
149                Some((SelectorKind::Event(e.selector()), e.full_signature()))
150            }
151            _ => None,
152        })
153    }
154
155    /// Inserts a single signature into the cache.
156    pub fn insert(&mut self, key: SelectorKind, value: String) {
157        self.extend(std::iter::once((key, value)));
158    }
159
160    /// Extends the cache with multiple signatures.
161    pub fn extend(&mut self, signatures: impl IntoIterator<Item = (SelectorKind, String)>) {
162        self.signatures
163            .extend(signatures.into_iter().map(|(k, v)| (k, (!v.is_empty()).then_some(v))));
164    }
165
166    /// Gets a signature from the cache.
167    pub fn get(&self, key: &SelectorKind) -> Option<Option<String>> {
168        self.signatures.get(key).cloned()
169    }
170
171    /// Returns true if the cache contains a signature.
172    pub fn contains_key(&self, key: &SelectorKind) -> bool {
173        self.signatures.contains_key(key)
174    }
175}
176
177/// An identifier that tries to identify functions and events using signatures found at
178/// `https://openchain.xyz` or a local cache.
179#[derive(Clone, Debug)]
180pub struct SignaturesIdentifier(Arc<SignaturesIdentifierInner>);
181
182#[derive(Debug)]
183struct SignaturesIdentifierInner {
184    /// Cached selectors for functions, events and custom errors.
185    cache: RwLock<SignaturesCache>,
186    /// ABI events keyed by topic0 and indexed topic count.
187    local_events: HashMap<(B256, usize), Vec<Event>>,
188    /// Location where to save the signature cache.
189    cache_path: Option<PathBuf>,
190    /// The OpenChain client to fetch signatures from. `None` if disabled on construction.
191    client: Option<OpenChainClient>,
192}
193
194impl SignaturesIdentifier {
195    /// Creates a new `SignaturesIdentifier` with the default cache directory.
196    pub fn new(offline: bool) -> Result<Self> {
197        Self::new_with(Config::foundry_cache_dir().as_deref(), offline)
198    }
199
200    /// Creates a new `SignaturesIdentifier` from the global configuration.
201    pub fn from_config(config: &Config) -> Result<Self> {
202        Self::new(config.offline)
203    }
204
205    /// Creates an offline `SignaturesIdentifier` with the default cache directory and local ABIs.
206    pub fn new_offline_with_abis<'a>(abis: impl IntoIterator<Item = &'a JsonAbi>) -> Result<Self> {
207        Ok(Self::new_offline_with_abis_from_cache(Config::foundry_cache_dir().as_deref(), abis))
208    }
209
210    /// Creates a new `SignaturesIdentifier`.
211    ///
212    /// - `cache_dir` is the cache directory to store the signatures.
213    /// - `offline` disables the OpenChain client.
214    pub fn new_with(cache_dir: Option<&Path>, offline: bool) -> Result<Self> {
215        let client = if offline { None } else { Some(OpenChainClient::new()?) };
216        Ok(Self::from_cache(Self::load_cache(cache_dir), client))
217    }
218
219    fn new_offline_with_abis_from_cache<'a>(
220        cache_dir: Option<&Path>,
221        abis: impl IntoIterator<Item = &'a JsonAbi>,
222    ) -> Self {
223        let abis = abis.into_iter().collect::<Vec<_>>();
224        let (mut cache, cache_path) = Self::load_cache(cache_dir);
225        cache.extend_from_abis_without_collisions(abis.iter().copied());
226        let local_events = Self::local_events_from_abis(abis);
227        Self::from_cache_and_events((cache, cache_path), None, local_events)
228    }
229
230    fn load_cache(cache_dir: Option<&Path>) -> (SignaturesCache, Option<PathBuf>) {
231        if let Some(cache_dir) = cache_dir {
232            let path = cache_dir.join("signatures");
233            let cache = SignaturesCache::load(&path);
234            (cache, Some(path))
235        } else {
236            Default::default()
237        }
238    }
239
240    fn from_cache(
241        (cache, cache_path): (SignaturesCache, Option<PathBuf>),
242        client: Option<OpenChainClient>,
243    ) -> Self {
244        Self::from_cache_and_events((cache, cache_path), client, Default::default())
245    }
246
247    fn from_cache_and_events(
248        (cache, cache_path): (SignaturesCache, Option<PathBuf>),
249        client: Option<OpenChainClient>,
250        local_events: HashMap<(B256, usize), Vec<Event>>,
251    ) -> Self {
252        Self(Arc::new(SignaturesIdentifierInner {
253            cache: RwLock::new(cache),
254            local_events,
255            cache_path,
256            client,
257        }))
258    }
259
260    fn local_events_from_abis<'a>(
261        abis: impl IntoIterator<Item = &'a JsonAbi>,
262    ) -> HashMap<(B256, usize), Vec<Event>> {
263        let mut local_events: HashMap<(B256, usize), Vec<Event>> = HashMap::default();
264        for abi in abis {
265            for event in abi.events() {
266                local_events
267                    .entry((
268                        event.selector(),
269                        event.inputs.iter().filter(|input| input.indexed).count(),
270                    ))
271                    .or_default()
272                    .push(event.clone());
273            }
274        }
275        local_events
276    }
277
278    /// Saves the cache to the file system.
279    pub fn save(&self) {
280        self.0.save();
281    }
282
283    /// Identifies `Function`s.
284    pub async fn identify_functions(
285        &self,
286        identifiers: impl IntoIterator<Item = Selector>,
287    ) -> Vec<Option<Function>> {
288        self.identify_map(identifiers.into_iter().map(SelectorKind::Function), get_func).await
289    }
290
291    /// Identifies a `Function`.
292    pub async fn identify_function(&self, identifier: Selector) -> Option<Function> {
293        self.identify_functions([identifier]).await.pop().unwrap()
294    }
295
296    /// Identifies `Event`s.
297    pub async fn identify_events(
298        &self,
299        identifiers: impl IntoIterator<Item = B256>,
300    ) -> Vec<Option<Event>> {
301        self.identify_map(identifiers.into_iter().map(SelectorKind::Event), get_event).await
302    }
303
304    /// Identifies an `Event`.
305    pub async fn identify_event(&self, identifier: B256) -> Option<Event> {
306        self.identify_events([identifier]).await.pop().unwrap()
307    }
308
309    /// Identifies an `Event`, preferring local ABI events with the matching indexed topic count.
310    pub async fn identify_event_with_indexed_count(
311        &self,
312        identifier: B256,
313        indexed_count: usize,
314    ) -> Option<Event> {
315        if let Some(events) = self.0.local_events.get(&(identifier, indexed_count))
316            && let Some(event) = events.first()
317        {
318            return Some(event.clone());
319        }
320        self.identify_event(identifier).await
321    }
322
323    /// Identifies `Error`s.
324    pub async fn identify_errors(
325        &self,
326        identifiers: impl IntoIterator<Item = Selector>,
327    ) -> Vec<Option<Error>> {
328        self.identify_map(identifiers.into_iter().map(SelectorKind::Error), get_error).await
329    }
330
331    /// Identifies an `Error`.
332    pub async fn identify_error(&self, identifier: Selector) -> Option<Error> {
333        self.identify_errors([identifier]).await.pop().unwrap()
334    }
335
336    /// Identifies a list of selectors.
337    pub async fn identify(&self, selectors: &[SelectorKind]) -> Vec<Option<String>> {
338        if selectors.is_empty() {
339            return vec![];
340        }
341        trace!(target: "evm::traces", ?selectors, "identifying selectors");
342
343        let mut cache_r = self.0.cache.read().await;
344        if let Some(client) = &self.0.client {
345            let query =
346                selectors.iter().copied().filter(|v| !cache_r.contains_key(v)).collect::<Vec<_>>();
347            if !query.is_empty() {
348                drop(cache_r);
349                let mut cache_w = self.0.cache.write().await;
350                if let Ok(res) = client.decode_selectors(&query).await {
351                    for (selector, signatures) in std::iter::zip(query, res) {
352                        cache_w.signatures.insert(selector, signatures.into_iter().next());
353                    }
354                }
355                drop(cache_w);
356                cache_r = self.0.cache.read().await;
357            }
358        }
359        selectors.iter().map(|selector| cache_r.get(selector).unwrap_or_default()).collect()
360    }
361
362    async fn identify_map<T>(
363        &self,
364        selectors: impl IntoIterator<Item = SelectorKind>,
365        get_type: impl Fn(&str) -> Result<T>,
366    ) -> Vec<Option<T>> {
367        let results = self.identify(&Vec::from_iter(selectors)).await;
368        results.into_iter().map(|r| r.and_then(|r| get_type(&r).ok())).collect()
369    }
370}
371
372impl SignaturesIdentifierInner {
373    fn save(&self) {
374        // We only identify new signatures if the client is enabled.
375        if let Some(path) = &self.cache_path
376            && self.client.is_some()
377        {
378            self.cache
379                .try_read()
380                .expect("SignaturesIdentifier cache is locked while attempting to save")
381                .save(path);
382        }
383    }
384}
385
386impl Drop for SignaturesIdentifierInner {
387    fn drop(&mut self) {
388        self.save();
389    }
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    #[test]
397    fn unknown_signatures_not_persisted_to_disk() {
398        let known_selector = SelectorKind::Function(Selector::from([0xaa, 0xbb, 0xcc, 0xdd]));
399        let unknown_selector = SelectorKind::Error(Selector::from([0x11, 0x22, 0x33, 0x44]));
400
401        let mut cache = SignaturesCache::default();
402        cache.signatures.insert(known_selector, Some("transfer(address,uint256)".into()));
403        cache.signatures.insert(unknown_selector, None);
404
405        // Verify both are in memory.
406        assert!(cache.contains_key(&known_selector));
407        assert!(cache.contains_key(&unknown_selector));
408
409        // Round-trip through the disk format.
410        let disk: SignaturesDiskCache = (&cache).into();
411        let reloaded = SignaturesCache::from(disk);
412
413        // Known signature survives the round-trip.
414        assert_eq!(reloaded.get(&known_selector), Some(Some("transfer(address,uint256)".into())));
415        // Unknown signature is gone — it will be re-queried next session.
416        assert_eq!(reloaded.get(&unknown_selector), None);
417        assert!(!reloaded.contains_key(&unknown_selector));
418    }
419
420    #[tokio::test]
421    async fn abi_seeded_signatures_are_not_persisted_to_disk() {
422        let temp = tempfile::tempdir().unwrap();
423        let event = Event::parse("event CodexEphemeral(uint256 indexed value)").unwrap();
424        let mut abi = JsonAbi::default();
425        abi.events.insert(event.name.clone(), vec![event.clone()]);
426
427        {
428            let identifier =
429                SignaturesIdentifier::new_offline_with_abis_from_cache(Some(temp.path()), [&abi]);
430            let decoded = identifier.identify_event(event.selector()).await;
431            assert_eq!(decoded.as_ref().map(Event::full_signature), Some(event.full_signature()));
432            identifier.save();
433        }
434
435        let reloaded = SignaturesCache::load(&temp.path().join("signatures"));
436        assert!(!reloaded.contains_key(&SelectorKind::Event(event.selector())));
437    }
438
439    #[test]
440    fn abi_seeded_collisions_keep_first_signature() {
441        let first = Event::parse("event CodexCollision(uint256 indexed value)").unwrap();
442        let second = Event::parse("event CodexCollision(uint256 value)").unwrap();
443
444        let mut first_abi = JsonAbi::default();
445        first_abi.events.insert(first.name.clone(), vec![first.clone()]);
446        let mut second_abi = JsonAbi::default();
447        second_abi.events.insert(second.name.clone(), vec![second]);
448
449        let mut cache = SignaturesCache::default();
450        cache.extend_from_abis_without_collisions([&first_abi, &second_abi]);
451
452        assert_eq!(
453            cache.get(&SelectorKind::Event(first.selector())),
454            Some(Some(first.full_signature()))
455        );
456    }
457
458    #[tokio::test]
459    async fn abi_seeded_events_prefer_matching_indexed_count() {
460        let one_topic =
461            Event::parse("event CodexIndexedCount(uint256 indexed marker, uint256 value)").unwrap();
462        let two_topics =
463            Event::parse("event CodexIndexedCount(uint256 indexed marker, uint256 indexed value)")
464                .unwrap();
465
466        let mut one_topic_abi = JsonAbi::default();
467        one_topic_abi.events.insert(one_topic.name.clone(), vec![one_topic.clone()]);
468        let mut two_topics_abi = JsonAbi::default();
469        two_topics_abi.events.insert(two_topics.name.clone(), vec![two_topics.clone()]);
470
471        let identifier = SignaturesIdentifier::new_offline_with_abis_from_cache(
472            None,
473            [&two_topics_abi, &one_topic_abi],
474        );
475
476        let decoded_one_topic =
477            identifier.identify_event_with_indexed_count(one_topic.selector(), 1).await.unwrap();
478        let decoded_two_topics =
479            identifier.identify_event_with_indexed_count(two_topics.selector(), 2).await.unwrap();
480
481        assert_eq!(decoded_one_topic.full_signature(), one_topic.full_signature());
482        assert_eq!(decoded_two_topics.full_signature(), two_topics.full_signature());
483    }
484}
foundry_evm_traces/identifier/signatures.rs

foundry_evm_traces/identifier/
signatures.rs
