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                let value = signature.clone().unwrap_or_default();
63                match *kind {
64                    SelectorKind::Function(selector) => _ = acc.0.insert(selector, value),
65                    SelectorKind::Error(selector) => _ = acc.1.insert(selector, value),
66                    SelectorKind::Event(selector) => _ = acc.2.insert(selector, value),
67                }
68                acc
69            },
70        );
71        Self { functions, errors, events }
72    }
73}
74
75impl Serialize for SignaturesCache {
76    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
77    where
78        S: serde::Serializer,
79    {
80        SignaturesDiskCache::from(self).serialize(serializer)
81    }
82}
83
84impl SignaturesCache {
85    /// Loads the cache from a file.
86    #[instrument(target = "evm::traces", name = "SignaturesCache::load")]
87    pub fn load(path: &Path) -> Self {
88        trace!(target: "evm::traces", ?path, "reading signature cache");
89        fs::read_json_file(path)
90            .inspect_err(
91                |err| warn!(target: "evm::traces", ?path, ?err, "failed to read cache file"),
92            )
93            .unwrap_or_default()
94    }
95
96    /// Saves the cache to a file.
97    #[instrument(target = "evm::traces", name = "SignaturesCache::save", skip(self))]
98    pub fn save(&self, path: &Path) {
99        if let Some(parent) = path.parent()
100            && let Err(err) = std::fs::create_dir_all(parent)
101        {
102            warn!(target: "evm::traces", ?parent, %err, "failed to create cache");
103        }
104        if let Err(err) = fs::write_json_file(path, self) {
105            warn!(target: "evm::traces", %err, "failed to flush signature cache");
106        } else {
107            trace!(target: "evm::traces", "flushed signature cache")
108        }
109    }
110
111    /// Updates the cache from an ABI.
112    pub fn extend_from_abi(&mut self, abi: &JsonAbi) {
113        self.extend(abi.items().filter_map(|item| match item {
114            alloy_json_abi::AbiItem::Function(f) => {
115                Some((SelectorKind::Function(f.selector()), f.signature()))
116            }
117            alloy_json_abi::AbiItem::Error(e) => {
118                Some((SelectorKind::Error(e.selector()), e.signature()))
119            }
120            alloy_json_abi::AbiItem::Event(e) => {
121                Some((SelectorKind::Event(e.selector()), e.full_signature()))
122            }
123            _ => None,
124        }));
125    }
126
127    /// Inserts a single signature into the cache.
128    pub fn insert(&mut self, key: SelectorKind, value: String) {
129        self.extend(std::iter::once((key, value)));
130    }
131
132    /// Extends the cache with multiple signatures.
133    pub fn extend(&mut self, signatures: impl IntoIterator<Item = (SelectorKind, String)>) {
134        self.signatures
135            .extend(signatures.into_iter().map(|(k, v)| (k, (!v.is_empty()).then_some(v))));
136    }
137
138    /// Gets a signature from the cache.
139    pub fn get(&self, key: &SelectorKind) -> Option<Option<String>> {
140        self.signatures.get(key).cloned()
141    }
142
143    /// Returns true if the cache contains a signature.
144    pub fn contains_key(&self, key: &SelectorKind) -> bool {
145        self.signatures.contains_key(key)
146    }
147}
148
149/// An identifier that tries to identify functions and events using signatures found at
150/// `https://openchain.xyz` or a local cache.
151#[derive(Clone, Debug)]
152pub struct SignaturesIdentifier(Arc<SignaturesIdentifierInner>);
153
154#[derive(Debug)]
155struct SignaturesIdentifierInner {
156    /// Cached selectors for functions, events and custom errors.
157    cache: RwLock<SignaturesCache>,
158    /// Location where to save the signature cache.
159    cache_path: Option<PathBuf>,
160    /// The OpenChain client to fetch signatures from. `None` if disabled on construction.
161    client: Option<OpenChainClient>,
162}
163
164impl SignaturesIdentifier {
165    /// Creates a new `SignaturesIdentifier` with the default cache directory.
166    pub fn new(offline: bool) -> Result<Self> {
167        Self::new_with(Config::foundry_cache_dir().as_deref(), offline)
168    }
169
170    /// Creates a new `SignaturesIdentifier` from the global configuration.
171    pub fn from_config(config: &Config) -> Result<Self> {
172        Self::new(config.offline)
173    }
174
175    /// Creates a new `SignaturesIdentifier`.
176    ///
177    /// - `cache_dir` is the cache directory to store the signatures.
178    /// - `offline` disables the OpenChain client.
179    pub fn new_with(cache_dir: Option<&Path>, offline: bool) -> Result<Self> {
180        let client = if !offline { Some(OpenChainClient::new()?) } else { None };
181        let (cache, cache_path) = if let Some(cache_dir) = cache_dir {
182            let path = cache_dir.join("signatures");
183            let cache = SignaturesCache::load(&path);
184            (cache, Some(path))
185        } else {
186            Default::default()
187        };
188        Ok(Self(Arc::new(SignaturesIdentifierInner {
189            cache: RwLock::new(cache),
190            cache_path,
191            client,
192        })))
193    }
194
195    /// Saves the cache to the file system.
196    pub fn save(&self) {
197        self.0.save();
198    }
199
200    /// Identifies `Function`s.
201    pub async fn identify_functions(
202        &self,
203        identifiers: impl IntoIterator<Item = Selector>,
204    ) -> Vec<Option<Function>> {
205        self.identify_map(identifiers.into_iter().map(SelectorKind::Function), get_func).await
206    }
207
208    /// Identifies a `Function`.
209    pub async fn identify_function(&self, identifier: Selector) -> Option<Function> {
210        self.identify_functions([identifier]).await.pop().unwrap()
211    }
212
213    /// Identifies `Event`s.
214    pub async fn identify_events(
215        &self,
216        identifiers: impl IntoIterator<Item = B256>,
217    ) -> Vec<Option<Event>> {
218        self.identify_map(identifiers.into_iter().map(SelectorKind::Event), get_event).await
219    }
220
221    /// Identifies an `Event`.
222    pub async fn identify_event(&self, identifier: B256) -> Option<Event> {
223        self.identify_events([identifier]).await.pop().unwrap()
224    }
225
226    /// Identifies `Error`s.
227    pub async fn identify_errors(
228        &self,
229        identifiers: impl IntoIterator<Item = Selector>,
230    ) -> Vec<Option<Error>> {
231        self.identify_map(identifiers.into_iter().map(SelectorKind::Error), get_error).await
232    }
233
234    /// Identifies an `Error`.
235    pub async fn identify_error(&self, identifier: Selector) -> Option<Error> {
236        self.identify_errors([identifier]).await.pop().unwrap()
237    }
238
239    /// Identifies a list of selectors.
240    pub async fn identify(&self, selectors: &[SelectorKind]) -> Vec<Option<String>> {
241        if selectors.is_empty() {
242            return vec![];
243        }
244        trace!(target: "evm::traces", ?selectors, "identifying selectors");
245
246        let mut cache_r = self.0.cache.read().await;
247        if let Some(client) = &self.0.client {
248            let query =
249                selectors.iter().copied().filter(|v| !cache_r.contains_key(v)).collect::<Vec<_>>();
250            if !query.is_empty() {
251                drop(cache_r);
252                let mut cache_w = self.0.cache.write().await;
253                if let Ok(res) = client.decode_selectors(&query).await {
254                    for (selector, signatures) in std::iter::zip(query, res) {
255                        cache_w.signatures.insert(selector, signatures.into_iter().next());
256                    }
257                }
258                drop(cache_w);
259                cache_r = self.0.cache.read().await;
260            }
261        }
262        selectors.iter().map(|selector| cache_r.get(selector).unwrap_or_default()).collect()
263    }
264
265    async fn identify_map<T>(
266        &self,
267        selectors: impl IntoIterator<Item = SelectorKind>,
268        get_type: impl Fn(&str) -> Result<T>,
269    ) -> Vec<Option<T>> {
270        let results = self.identify(&Vec::from_iter(selectors)).await;
271        results.into_iter().map(|r| r.and_then(|r| get_type(&r).ok())).collect()
272    }
273}
274
275impl SignaturesIdentifierInner {
276    fn save(&self) {
277        // We only identify new signatures if the client is enabled.
278        if let Some(path) = &self.cache_path
279            && self.client.is_some()
280        {
281            self.cache
282                .try_read()
283                .expect("SignaturesIdentifier cache is locked while attempting to save")
284                .save(path);
285        }
286    }
287}
288
289impl Drop for SignaturesIdentifierInner {
290    fn drop(&mut self) {
291        self.save();
292    }
293}
foundry_evm_traces/identifier/signatures.rs

foundry_evm_traces/identifier/
signatures.rs
