foundry_evm_traces/identifier/
signatures.rs

1use alloy_json_abi::{Error, Event, Function, JsonAbi};
2use alloy_primitives::{B256, Selector, map::HashMap};
3use eyre::Result;
4use foundry_common::{
5    abi::{get_error, get_event, get_func},
6    fs,
7    selectors::{OpenChainClient, SelectorKind},
8};
9use foundry_config::Config;
10use serde::{Deserialize, Serialize};
11use std::{
12    collections::BTreeMap,
13    path::{Path, PathBuf},
14    sync::Arc,
15};
16use tokio::sync::RwLock;
17
18/// Cache for function, event and error signatures. Used by [`SignaturesIdentifier`].
19#[derive(Debug, Default, Deserialize)]
20#[serde(try_from = "SignaturesDiskCache")]
21pub struct SignaturesCache {
22    signatures: HashMap<SelectorKind, Option<String>>,
23}
24
25/// Disk representation of the signatures cache.
26#[derive(Serialize, Deserialize)]
27struct SignaturesDiskCache {
28    functions: BTreeMap<Selector, String>,
29    errors: BTreeMap<Selector, String>,
30    events: BTreeMap<B256, String>,
31}
32
33impl From<SignaturesDiskCache> for SignaturesCache {
34    fn from(value: SignaturesDiskCache) -> Self {
35        let functions = value
36            .functions
37            .into_iter()
38            .map(|(selector, signature)| (SelectorKind::Function(selector), signature));
39        let errors = value
40            .errors
41            .into_iter()
42            .map(|(selector, signature)| (SelectorKind::Error(selector), signature));
43        let events = value
44            .events
45            .into_iter()
46            .map(|(selector, signature)| (SelectorKind::Event(selector), signature));
47        Self {
48            signatures: functions
49                .chain(errors)
50                .chain(events)
51                .map(|(sel, sig)| (sel, (!sig.is_empty()).then_some(sig)))
52                .collect(),
53        }
54    }
55}
56
57impl From<&SignaturesCache> for SignaturesDiskCache {
58    fn from(value: &SignaturesCache) -> Self {
59        let (functions, errors, events) = value.signatures.iter().fold(
60            (BTreeMap::new(), BTreeMap::new(), BTreeMap::new()),
61            |mut acc, (kind, signature)| {
62                // Only persist resolved signatures. Unknown selectors (None) are kept
63                // in-memory for session dedup but not written to disk, so they can be
64                // re-queried in future sessions once the signature database is updated.
65                if let Some(value) = signature.clone() {
66                    match *kind {
67                        SelectorKind::Function(selector) => _ = acc.0.insert(selector, value),
68                        SelectorKind::Error(selector) => _ = acc.1.insert(selector, value),
69                        SelectorKind::Event(selector) => _ = acc.2.insert(selector, value),
70                    }
71                }
72                acc
73            },
74        );
75        Self { functions, errors, events }
76    }
77}
78
79impl Serialize for SignaturesCache {
80    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
81    where
82        S: serde::Serializer,
83    {
84        SignaturesDiskCache::from(self).serialize(serializer)
85    }
86}
87
88impl SignaturesCache {
89    /// Loads the cache from a file.
90    #[instrument(target = "evm::traces", name = "SignaturesCache::load")]
91    pub fn load(path: &Path) -> Self {
92        trace!(target: "evm::traces", ?path, "reading signature cache");
93        fs::read_json_file(path)
94            .inspect_err(
95                |err| warn!(target: "evm::traces", ?path, ?err, "failed to read cache file"),
96            )
97            .unwrap_or_default()
98    }
99
100    /// Saves the cache to a file.
101    #[instrument(target = "evm::traces", name = "SignaturesCache::save", skip(self))]
102    pub fn save(&self, path: &Path) {
103        if let Some(parent) = path.parent()
104            && let Err(err) = std::fs::create_dir_all(parent)
105        {
106            warn!(target: "evm::traces", ?parent, %err, "failed to create cache");
107        }
108        if let Err(err) = fs::write_json_file(path, self) {
109            warn!(target: "evm::traces", %err, "failed to flush signature cache");
110        } else {
111            trace!(target: "evm::traces", "flushed signature cache")
112        }
113    }
114
115    /// Updates the cache from an ABI.
116    pub fn extend_from_abi(&mut self, abi: &JsonAbi) {
117        self.extend(abi.items().filter_map(|item| match item {
118            alloy_json_abi::AbiItem::Function(f) => {
119                Some((SelectorKind::Function(f.selector()), f.signature()))
120            }
121            alloy_json_abi::AbiItem::Error(e) => {
122                Some((SelectorKind::Error(e.selector()), e.signature()))
123            }
124            alloy_json_abi::AbiItem::Event(e) => {
125                Some((SelectorKind::Event(e.selector()), e.full_signature()))
126            }
127            _ => None,
128        }));
129    }
130
131    /// Inserts a single signature into the cache.
132    pub fn insert(&mut self, key: SelectorKind, value: String) {
133        self.extend(std::iter::once((key, value)));
134    }
135
136    /// Extends the cache with multiple signatures.
137    pub fn extend(&mut self, signatures: impl IntoIterator<Item = (SelectorKind, String)>) {
138        self.signatures
139            .extend(signatures.into_iter().map(|(k, v)| (k, (!v.is_empty()).then_some(v))));
140    }
141
142    /// Gets a signature from the cache.
143    pub fn get(&self, key: &SelectorKind) -> Option<Option<String>> {
144        self.signatures.get(key).cloned()
145    }
146
147    /// Returns true if the cache contains a signature.
148    pub fn contains_key(&self, key: &SelectorKind) -> bool {
149        self.signatures.contains_key(key)
150    }
151}
152
153/// An identifier that tries to identify functions and events using signatures found at
154/// `https://openchain.xyz` or a local cache.
155#[derive(Clone, Debug)]
156pub struct SignaturesIdentifier(Arc<SignaturesIdentifierInner>);
157
158#[derive(Debug)]
159struct SignaturesIdentifierInner {
160    /// Cached selectors for functions, events and custom errors.
161    cache: RwLock<SignaturesCache>,
162    /// Location where to save the signature cache.
163    cache_path: Option<PathBuf>,
164    /// The OpenChain client to fetch signatures from. `None` if disabled on construction.
165    client: Option<OpenChainClient>,
166}
167
168impl SignaturesIdentifier {
169    /// Creates a new `SignaturesIdentifier` with the default cache directory.
170    pub fn new(offline: bool) -> Result<Self> {
171        Self::new_with(Config::foundry_cache_dir().as_deref(), offline)
172    }
173
174    /// Creates a new `SignaturesIdentifier` from the global configuration.
175    pub fn from_config(config: &Config) -> Result<Self> {
176        Self::new(config.offline)
177    }
178
179    /// Creates a new `SignaturesIdentifier`.
180    ///
181    /// - `cache_dir` is the cache directory to store the signatures.
182    /// - `offline` disables the OpenChain client.
183    pub fn new_with(cache_dir: Option<&Path>, offline: bool) -> Result<Self> {
184        let client = if offline { None } else { Some(OpenChainClient::new()?) };
185        let (cache, cache_path) = if let Some(cache_dir) = cache_dir {
186            let path = cache_dir.join("signatures");
187            let cache = SignaturesCache::load(&path);
188            (cache, Some(path))
189        } else {
190            Default::default()
191        };
192        Ok(Self(Arc::new(SignaturesIdentifierInner {
193            cache: RwLock::new(cache),
194            cache_path,
195            client,
196        })))
197    }
198
199    /// Saves the cache to the file system.
200    pub fn save(&self) {
201        self.0.save();
202    }
203
204    /// Identifies `Function`s.
205    pub async fn identify_functions(
206        &self,
207        identifiers: impl IntoIterator<Item = Selector>,
208    ) -> Vec<Option<Function>> {
209        self.identify_map(identifiers.into_iter().map(SelectorKind::Function), get_func).await
210    }
211
212    /// Identifies a `Function`.
213    pub async fn identify_function(&self, identifier: Selector) -> Option<Function> {
214        self.identify_functions([identifier]).await.pop().unwrap()
215    }
216
217    /// Identifies `Event`s.
218    pub async fn identify_events(
219        &self,
220        identifiers: impl IntoIterator<Item = B256>,
221    ) -> Vec<Option<Event>> {
222        self.identify_map(identifiers.into_iter().map(SelectorKind::Event), get_event).await
223    }
224
225    /// Identifies an `Event`.
226    pub async fn identify_event(&self, identifier: B256) -> Option<Event> {
227        self.identify_events([identifier]).await.pop().unwrap()
228    }
229
230    /// Identifies `Error`s.
231    pub async fn identify_errors(
232        &self,
233        identifiers: impl IntoIterator<Item = Selector>,
234    ) -> Vec<Option<Error>> {
235        self.identify_map(identifiers.into_iter().map(SelectorKind::Error), get_error).await
236    }
237
238    /// Identifies an `Error`.
239    pub async fn identify_error(&self, identifier: Selector) -> Option<Error> {
240        self.identify_errors([identifier]).await.pop().unwrap()
241    }
242
243    /// Identifies a list of selectors.
244    pub async fn identify(&self, selectors: &[SelectorKind]) -> Vec<Option<String>> {
245        if selectors.is_empty() {
246            return vec![];
247        }
248        trace!(target: "evm::traces", ?selectors, "identifying selectors");
249
250        let mut cache_r = self.0.cache.read().await;
251        if let Some(client) = &self.0.client {
252            let query =
253                selectors.iter().copied().filter(|v| !cache_r.contains_key(v)).collect::<Vec<_>>();
254            if !query.is_empty() {
255                drop(cache_r);
256                let mut cache_w = self.0.cache.write().await;
257                if let Ok(res) = client.decode_selectors(&query).await {
258                    for (selector, signatures) in std::iter::zip(query, res) {
259                        cache_w.signatures.insert(selector, signatures.into_iter().next());
260                    }
261                }
262                drop(cache_w);
263                cache_r = self.0.cache.read().await;
264            }
265        }
266        selectors.iter().map(|selector| cache_r.get(selector).unwrap_or_default()).collect()
267    }
268
269    async fn identify_map<T>(
270        &self,
271        selectors: impl IntoIterator<Item = SelectorKind>,
272        get_type: impl Fn(&str) -> Result<T>,
273    ) -> Vec<Option<T>> {
274        let results = self.identify(&Vec::from_iter(selectors)).await;
275        results.into_iter().map(|r| r.and_then(|r| get_type(&r).ok())).collect()
276    }
277}
278
279impl SignaturesIdentifierInner {
280    fn save(&self) {
281        // We only identify new signatures if the client is enabled.
282        if let Some(path) = &self.cache_path
283            && self.client.is_some()
284        {
285            self.cache
286                .try_read()
287                .expect("SignaturesIdentifier cache is locked while attempting to save")
288                .save(path);
289        }
290    }
291}
292
293impl Drop for SignaturesIdentifierInner {
294    fn drop(&mut self) {
295        self.save();
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn unknown_signatures_not_persisted_to_disk() {
305        let known_selector = SelectorKind::Function(Selector::from([0xaa, 0xbb, 0xcc, 0xdd]));
306        let unknown_selector = SelectorKind::Error(Selector::from([0x11, 0x22, 0x33, 0x44]));
307
308        let mut cache = SignaturesCache::default();
309        cache.signatures.insert(known_selector, Some("transfer(address,uint256)".into()));
310        cache.signatures.insert(unknown_selector, None);
311
312        // Verify both are in memory.
313        assert!(cache.contains_key(&known_selector));
314        assert!(cache.contains_key(&unknown_selector));
315
316        // Round-trip through the disk format.
317        let disk: SignaturesDiskCache = (&cache).into();
318        let reloaded = SignaturesCache::from(disk);
319
320        // Known signature survives the round-trip.
321        assert_eq!(reloaded.get(&known_selector), Some(Some("transfer(address,uint256)".into())));
322        // Unknown signature is gone — it will be re-queried next session.
323        assert_eq!(reloaded.get(&unknown_selector), None);
324        assert!(!reloaded.contains_key(&unknown_selector));
325    }
326}
foundry_evm_traces/identifier/signatures.rs

foundry_evm_traces/identifier/
signatures.rs
