foundry_evm_traces/identifier/
signatures.rs

Help
1use alloy_json_abi::{Error, Event, Function, JsonAbi};
2use alloy_primitives::{map::HashMap, Selector, B256};
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")]
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", skip(self))]
98    pub fn save(&self, path: &Path) {
99        if let Some(parent) = path.parent() {
100            if let Err(err) = std::fs::create_dir_all(parent) {
101                warn!(target: "evm::traces", ?parent, %err, "failed to create cache");
102            }
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 {
153    /// Cached selectors for functions, events and custom errors.
154    cache: Arc<RwLock<SignaturesCache>>,
155    /// Location where to save the signature cache.
156    cache_path: Option<PathBuf>,
157    /// The OpenChain client to fetch signatures from. `None` if disabled on construction.
158    client: Option<OpenChainClient>,
159}
160
161impl SignaturesIdentifier {
162    /// Creates a new `SignaturesIdentifier` with the default cache directory.
163    pub fn new(offline: bool) -> Result<Self> {
164        Self::new_with(Config::foundry_cache_dir().as_deref(), offline)
165    }
166
167    /// Creates a new `SignaturesIdentifier` from the global configuration.
168    pub fn from_config(config: &Config) -> Result<Self> {
169        Self::new(config.offline)
170    }
171
172    /// Creates a new `SignaturesIdentifier`.
173    ///
174    /// - `cache_dir` is the cache directory to store the signatures.
175    /// - `offline` disables the OpenChain client.
176    pub fn new_with(cache_dir: Option<&Path>, offline: bool) -> Result<Self> {
177        let client = if !offline { Some(OpenChainClient::new()?) } else { None };
178        let (cache, cache_path) = if let Some(cache_dir) = cache_dir {
179            let path = cache_dir.join("signatures");
180            let cache = SignaturesCache::load(&path);
181            (cache, Some(path))
182        } else {
183            Default::default()
184        };
185        Ok(Self { cache: Arc::new(RwLock::new(cache)), cache_path, client })
186    }
187
188    /// Saves the cache to the file system.
189    pub fn save(&self) {
190        if let Some(path) = &self.cache_path {
191            foundry_compilers::utils::RuntimeOrHandle::new().block_on(self.cache.read()).save(path);
192        }
193    }
194
195    /// Identifies `Function`s.
196    pub async fn identify_functions(
197        &self,
198        identifiers: impl IntoIterator<Item = Selector>,
199    ) -> Vec<Option<Function>> {
200        self.identify_map(identifiers.into_iter().map(SelectorKind::Function), get_func).await
201    }
202
203    /// Identifies a `Function`.
204    pub async fn identify_function(&self, identifier: Selector) -> Option<Function> {
205        self.identify_functions([identifier]).await.pop().unwrap()
206    }
207
208    /// Identifies `Event`s.
209    pub async fn identify_events(
210        &self,
211        identifiers: impl IntoIterator<Item = B256>,
212    ) -> Vec<Option<Event>> {
213        self.identify_map(identifiers.into_iter().map(SelectorKind::Event), get_event).await
214    }
215
216    /// Identifies an `Event`.
217    pub async fn identify_event(&self, identifier: B256) -> Option<Event> {
218        self.identify_events([identifier]).await.pop().unwrap()
219    }
220
221    /// Identifies `Error`s.
222    pub async fn identify_errors(
223        &self,
224        identifiers: impl IntoIterator<Item = Selector>,
225    ) -> Vec<Option<Error>> {
226        self.identify_map(identifiers.into_iter().map(SelectorKind::Error), get_error).await
227    }
228
229    /// Identifies an `Error`.
230    pub async fn identify_error(&self, identifier: Selector) -> Option<Error> {
231        self.identify_errors([identifier]).await.pop().unwrap()
232    }
233
234    /// Identifies a list of selectors.
235    pub async fn identify(&self, selectors: &[SelectorKind]) -> Vec<Option<String>> {
236        if selectors.is_empty() {
237            return vec![];
238        }
239        trace!(target: "evm::traces", ?selectors, "identifying selectors");
240
241        let mut cache_r = self.cache.read().await;
242        if let Some(client) = &self.client {
243            let query =
244                selectors.iter().copied().filter(|v| !cache_r.contains_key(v)).collect::<Vec<_>>();
245            if !query.is_empty() {
246                drop(cache_r);
247                let mut cache_w = self.cache.write().await;
248                if let Ok(res) = client.decode_selectors(&query).await {
249                    for (selector, signatures) in std::iter::zip(query, res) {
250                        cache_w.signatures.insert(selector, signatures.into_iter().next());
251                    }
252                }
253                drop(cache_w);
254                cache_r = self.cache.read().await;
255            }
256        }
257        selectors.iter().map(|selector| cache_r.get(selector).unwrap_or_default()).collect()
258    }
259
260    async fn identify_map<T>(
261        &self,
262        selectors: impl IntoIterator<Item = SelectorKind>,
263        get_type: impl Fn(&str) -> Result<T>,
264    ) -> Vec<Option<T>> {
265        let results = self.identify(&Vec::from_iter(selectors)).await;
266        results.into_iter().map(|r| r.and_then(|r| get_type(&r).ok())).collect()
267    }
268}
269
270impl Drop for SignaturesIdentifier {
271    fn drop(&mut self) {
272        self.save();
273    }
274}
foundry_evm_traces/identifier/signatures.rs

foundry_evm_traces/identifier/
signatures.rs
