1use super::{IdentifiedAddress, TraceIdentifier};
2use crate::debug::ContractSources;
3use alloy_json_abi::JsonAbi;
4use alloy_primitives::{
5 Address,
6 map::{AddressSet, Entry, HashMap, HashSet},
7};
8use eyre::WrapErr;
9use foundry_block_explorers::{contract::Metadata, errors::EtherscanError};
10use foundry_common::compile::etherscan_project;
11use foundry_config::{Chain, Config};
12use futures::{
13 future::join_all,
14 stream::{FuturesUnordered, Stream, StreamExt},
15 task::{Context, Poll},
16};
17use revm_inspectors::tracing::types::CallTraceNode;
18use serde::Deserialize;
19use std::{
20 borrow::Cow,
21 pin::Pin,
22 sync::{
23 Arc,
24 atomic::{AtomicBool, Ordering},
25 },
26};
27use tokio::time::{Duration, Interval};
28
29pub struct ExternalIdentifier {
31 fetchers: Vec<Arc<dyn ExternalFetcherT>>,
32 contracts: HashMap<Address, (FetcherKind, Option<Metadata>)>,
34 remaining_budget: Duration,
36}
37
38impl ExternalIdentifier {
39 pub fn new(config: &Config, mut chain: Option<Chain>) -> eyre::Result<Option<Self>> {
41 let timeout = config.tracing.external_identification_timeout;
42 if config.offline || timeout == 0 {
43 return Ok(None);
44 }
45
46 let no_proxy = config.eth_rpc_no_proxy;
47 let config = match config.get_etherscan_config_with_chain(chain) {
48 Ok(Some(config)) => {
49 chain = config.chain;
50 Some(config)
51 }
52 Ok(None) => {
53 warn!(target: "evm::traces::external", "etherscan config not found");
54 None
55 }
56 Err(err) => {
57 warn!(target: "evm::traces::external", ?err, "failed to get etherscan config");
58 None
59 }
60 };
61
62 let mut fetchers = Vec::<Arc<dyn ExternalFetcherT>>::new();
63 if let Some(chain) = chain {
64 debug!(target: "evm::traces::external", ?chain, "using sourcify identifier");
65 fetchers.push(Arc::new(SourcifyFetcher::new(chain)));
66 }
67 if let Some(config) = config {
68 debug!(target: "evm::traces::external", chain=?config.chain, url=?config.api_url, "using etherscan identifier");
69 match config.into_client_with_no_proxy(no_proxy) {
70 Ok(client) => {
71 fetchers.push(Arc::new(EtherscanFetcher::new(client)));
72 }
73 Err(err) => {
74 warn!(target: "evm::traces::external", ?err, "failed to create etherscan client");
75 }
76 }
77 }
78 if fetchers.is_empty() {
79 debug!(target: "evm::traces::external", "no fetchers enabled");
80 return Ok(None);
81 }
82
83 Ok(Some(Self {
84 fetchers,
85 contracts: Default::default(),
86 remaining_budget: Duration::from_secs(timeout),
87 }))
88 }
89
90 pub async fn get_compiled_contracts(&self) -> eyre::Result<ContractSources> {
93 let contracts_info: Vec<_> = self
95 .contracts
96 .iter()
97 .filter_map(|(addr, (_, metadata))| {
99 if let Some(metadata) = metadata.as_ref()
100 && !metadata.is_vyper()
101 {
102 Some((*addr, metadata))
103 } else {
104 None
105 }
106 })
107 .collect();
108
109 let outputs_fut = contracts_info
110 .iter()
111 .map(|(addr, metadata)| async move {
112 sh_println!("Compiling: {} {addr}", metadata.contract_name)?;
113 let root = tempfile::tempdir()?;
114 let root_path = root.path();
115 let project = etherscan_project(metadata, root_path)?;
116 let output = project.compile()?;
117 if output.has_compiler_errors() {
118 eyre::bail!("{output}");
119 }
120
121 Ok((project, output, root))
122 })
123 .collect::<Vec<_>>();
124
125 let outputs = join_all(outputs_fut).await;
127
128 let mut sources: ContractSources = Default::default();
129
130 for (idx, res) in outputs.into_iter().enumerate() {
132 let (addr, metadata) = &contracts_info[idx];
133 let name = &metadata.contract_name;
134 let (project, output, _) =
135 res.wrap_err_with(|| format!("Failed to compile contract {name} at {addr}"))?;
136 sources
137 .insert(&output, project.root(), None)
138 .wrap_err_with(|| format!("Failed to insert contract {name} at {addr}"))?;
139 }
140
141 Ok(sources)
142 }
143
144 fn identify_from_metadata(
145 &self,
146 address: Address,
147 metadata: &Metadata,
148 ) -> IdentifiedAddress<'static> {
149 let label = metadata.contract_name.clone();
150 let abi = metadata.abi().ok().map(Cow::Owned);
151 IdentifiedAddress {
152 address,
153 label: Some(label.clone()),
154 contract: Some(label),
155 abi,
156 constructor_args_offset: None,
157 artifact_id: None,
158 }
159 }
160
161 fn cache_fetched(&mut self, address: Address, value: (FetcherKind, Option<Metadata>)) {
162 match self.contracts.entry(address) {
163 Entry::Occupied(mut occupied_entry) => {
164 let old = occupied_entry.get();
165 let should_replace = match (&old.1, &value.1) {
170 (None, Some(_)) => true,
171 (Some(_), None) => false,
172 _ => {
173 matches!(value.0, FetcherKind::Etherscan)
174 && !matches!(old.0, FetcherKind::Etherscan)
175 }
176 };
177 if should_replace {
178 occupied_entry.insert(value);
179 }
180 }
181 Entry::Vacant(vacant_entry) => {
182 vacant_entry.insert(value);
183 }
184 }
185 }
186
187 async fn fetch_addresses_async(&mut self, addresses: &[Address]) {
188 if addresses.is_empty() || self.remaining_budget.is_zero() {
189 return;
190 }
191
192 let fetchers = self
193 .fetchers
194 .clone()
195 .into_iter()
196 .map(|fetcher| ExternalFetcher::new(fetcher, addresses));
197 let started = tokio::time::Instant::now();
198 let timed_out = tokio::time::timeout(self.remaining_budget, async {
199 let mut fetched = futures::stream::select_all(fetchers);
200 while let Some((address, value)) = fetched.next().await {
201 self.cache_fetched(address, value);
202 }
203 })
204 .await
205 .is_err();
206 self.remaining_budget = self.remaining_budget.saturating_sub(started.elapsed());
207 if timed_out {
208 self.remaining_budget = Duration::ZERO;
209 warn!(target: "evm::traces::external", "external identification timed out; disabling it for the remainder of this session");
210 }
211 }
212
213 pub async fn get_abis(
215 &mut self,
216 addresses: &[Address],
217 ) -> Vec<(Address, eyre::Result<(Vec<JsonAbi>, bool)>)> {
218 const MAX_PROXY_DEPTH: usize = 16;
219
220 struct Chain {
221 current: Option<Address>,
222 visited: HashSet<Address>,
223 abis: Vec<JsonAbi>,
224 complete: bool,
225 }
226
227 let mut chains = addresses
228 .iter()
229 .map(|&address| Chain {
230 current: Some(address),
231 visited: HashSet::default(),
232 abis: Vec::new(),
233 complete: true,
234 })
235 .collect::<Vec<_>>();
236
237 for _ in 0..MAX_PROXY_DEPTH {
238 let to_fetch = chains
239 .iter()
240 .filter_map(|chain| chain.current)
241 .filter(|address| !self.contracts.contains_key(address))
242 .collect::<HashSet<_>>()
243 .into_iter()
244 .collect::<Vec<_>>();
245 self.fetch_addresses_async(&to_fetch).await;
246
247 let mut has_next = false;
248 for chain in &mut chains {
249 let Some(current) = chain.current else { continue };
250 if !chain.visited.insert(current) {
251 chain.current = None;
252 chain.complete = false;
253 continue;
254 }
255 let Some((_, Some(metadata))) = self.contracts.get(¤t) else {
256 chain.current = None;
257 chain.complete = false;
258 continue;
259 };
260 if let Ok(abi) = metadata.abi() {
261 chain.abis.push(abi);
262 } else {
263 chain.complete = false;
264 }
265 chain.current = (metadata.proxy != 0).then_some(metadata.implementation).flatten();
266 if metadata.proxy != 0 && chain.current.is_none() {
267 chain.complete = false;
268 }
269 has_next |= chain.current.is_some();
270 }
271 if !has_next {
272 break;
273 }
274 }
275
276 chains
277 .into_iter()
278 .zip(addresses.iter().copied())
279 .map(|(mut chain, address)| {
280 chain.complete &= chain.current.is_none();
281 let result = if chain.abis.is_empty() {
282 Err(eyre::eyre!("external ABI lookup failed"))
283 } else {
284 Ok((chain.abis.into_iter().rev().collect(), chain.complete))
285 };
286 (address, result)
287 })
288 .collect()
289 }
290}
291
292impl TraceIdentifier for ExternalIdentifier {
293 fn identify_addresses(&mut self, nodes: &[&CallTraceNode]) -> Vec<IdentifiedAddress<'_>> {
294 if nodes.is_empty() {
295 return Vec::new();
296 }
297
298 trace!(target: "evm::traces::external", "identify {} addresses", nodes.len());
299
300 let mut identities = Vec::new();
301 let mut to_fetch = AddressSet::default();
302
303 for &node in nodes {
305 let address = node.trace.address;
306 if let Some((_, metadata)) = self.contracts.get(&address) {
307 if let Some(metadata) = metadata {
308 identities.push(self.identify_from_metadata(address, metadata));
309 } else {
310 }
312 } else {
313 to_fetch.insert(address);
314 }
315 }
316
317 if to_fetch.is_empty() {
318 return identities;
319 }
320 if self.remaining_budget.is_zero() {
321 return identities;
322 }
323 trace!(target: "evm::traces::external", "fetching {} addresses", to_fetch.len());
324
325 let to_fetch = to_fetch.into_iter().collect::<Vec<_>>();
326 foundry_common::block_on(self.fetch_addresses_async(&to_fetch));
327
328 for address in to_fetch {
329 if let Some((_, Some(metadata))) = self.contracts.get(&address) {
330 identities.push(self.identify_from_metadata(address, metadata));
331 }
332 }
333 trace!(target: "evm::traces::external", "identified {} addresses", identities.len());
334 identities
335 }
336}
337
338type FetchFuture =
339 Pin<Box<dyn Future<Output = (Address, Result<Option<Metadata>, EtherscanError>)>>>;
340
341const MAX_CLOUDFLARE_RETRIES: u32 = 5;
344
345fn backoff_interval(period: Duration) -> Interval {
346 tokio::time::interval_at(tokio::time::Instant::now() + period, period)
347}
348
349struct ExternalFetcher {
353 fetcher: Arc<dyn ExternalFetcherT>,
355 timeout: Duration,
357 backoff: Option<Interval>,
359 concurrency: usize,
361 queue: Vec<Address>,
363 in_progress: FuturesUnordered<FetchFuture>,
365 attempts: HashMap<Address, u32>,
367}
368
369impl ExternalFetcher {
370 fn new(fetcher: Arc<dyn ExternalFetcherT>, to_fetch: &[Address]) -> Self {
371 Self {
372 timeout: fetcher.timeout(),
373 backoff: None,
374 concurrency: fetcher.concurrency(),
375 fetcher,
376 queue: to_fetch.to_vec(),
377 in_progress: FuturesUnordered::new(),
378 attempts: HashMap::default(),
379 }
380 }
381
382 fn queue_next_reqs(&mut self) {
383 while self.in_progress.len() < self.concurrency {
384 let Some(addr) = self.queue.pop() else { break };
385 let fetcher = Arc::clone(&self.fetcher);
386 self.in_progress.push(Box::pin(async move {
387 trace!(target: "evm::traces::external", ?addr, "fetching info");
388 let res = fetcher.fetch(addr).await;
389 (addr, res)
390 }));
391 }
392 }
393}
394
395impl Stream for ExternalFetcher {
396 type Item = (Address, (FetcherKind, Option<Metadata>));
397
398 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
399 let pin = self.get_mut();
400
401 let _guard =
402 info_span!("evm::traces::external", kind=?pin.fetcher.kind(), "ExternalFetcher")
403 .entered();
404
405 if pin.fetcher.invalid_api_key().load(Ordering::Relaxed) {
406 return Poll::Ready(None);
407 }
408
409 loop {
410 if let Some(mut backoff) = pin.backoff.take()
411 && backoff.poll_tick(cx).is_pending()
412 {
413 pin.backoff = Some(backoff);
414 return Poll::Pending;
415 }
416
417 pin.queue_next_reqs();
418
419 let mut made_progress_this_iter = false;
420 match pin.in_progress.poll_next_unpin(cx) {
421 Poll::Pending => {}
422 Poll::Ready(None) => return Poll::Ready(None),
423 Poll::Ready(Some((addr, res))) => {
424 made_progress_this_iter = true;
425 match res {
426 Ok(metadata) => {
427 return Poll::Ready(Some((addr, (pin.fetcher.kind(), metadata))));
428 }
429 Err(EtherscanError::ContractCodeNotVerified(_)) => {
430 return Poll::Ready(Some((addr, (pin.fetcher.kind(), None))));
431 }
432 Err(EtherscanError::RateLimitExceeded) => {
433 warn!(target: "evm::traces::external", "rate limit exceeded on attempt");
434 pin.backoff = Some(backoff_interval(pin.timeout));
435 pin.queue.push(addr);
436 }
437 Err(EtherscanError::InvalidApiKey) => {
438 warn!(target: "evm::traces::external", "invalid api key");
439 pin.fetcher.invalid_api_key().store(true, Ordering::Relaxed);
441 return Poll::Ready(None);
442 }
443 Err(EtherscanError::BlockedByCloudflare) => {
444 let attempts = {
450 let entry = pin.attempts.entry(addr).or_default();
451 *entry += 1;
452 *entry
453 };
454 if attempts <= MAX_CLOUDFLARE_RETRIES {
455 warn!(target: "evm::traces::external", attempts, "blocked by cloudflare, backing off");
456 pin.backoff = Some(backoff_interval(pin.timeout));
457 pin.queue.push(addr);
458 } else {
459 warn!(target: "evm::traces::external", "blocked by cloudflare, giving up on address");
460 return Poll::Ready(Some((addr, (pin.fetcher.kind(), None))));
461 }
462 }
463 Err(err) => {
464 warn!(target: "evm::traces::external", ?err, "could not get info");
465 return Poll::Ready(Some((addr, (pin.fetcher.kind(), None))));
467 }
468 }
469 }
470 }
471
472 if !made_progress_this_iter {
473 return Poll::Pending;
474 }
475 }
476 }
477}
478
479#[derive(Debug, Clone, Copy, PartialEq, Eq)]
480enum FetcherKind {
481 Etherscan,
482 Sourcify,
483}
484
485#[async_trait::async_trait]
486trait ExternalFetcherT: Send + Sync {
487 fn kind(&self) -> FetcherKind;
488 fn timeout(&self) -> Duration;
489 fn concurrency(&self) -> usize;
490 fn invalid_api_key(&self) -> &AtomicBool;
491 async fn fetch(&self, address: Address) -> Result<Option<Metadata>, EtherscanError>;
492}
493
494struct EtherscanFetcher {
495 client: foundry_block_explorers::Client,
496 invalid_api_key: AtomicBool,
497}
498
499impl EtherscanFetcher {
500 const fn new(client: foundry_block_explorers::Client) -> Self {
501 Self { client, invalid_api_key: AtomicBool::new(false) }
502 }
503}
504
505#[async_trait::async_trait]
506impl ExternalFetcherT for EtherscanFetcher {
507 fn kind(&self) -> FetcherKind {
508 FetcherKind::Etherscan
509 }
510
511 fn timeout(&self) -> Duration {
512 Duration::from_secs(1)
513 }
514
515 fn concurrency(&self) -> usize {
516 5
517 }
518
519 fn invalid_api_key(&self) -> &AtomicBool {
520 &self.invalid_api_key
521 }
522
523 async fn fetch(&self, address: Address) -> Result<Option<Metadata>, EtherscanError> {
524 self.client.contract_source_code(address).await.map(|mut metadata| metadata.items.pop())
525 }
526}
527
528struct SourcifyFetcher {
529 client: reqwest::Client,
530 url: String,
531 invalid_api_key: AtomicBool,
532}
533
534impl SourcifyFetcher {
535 fn new(chain: Chain) -> Self {
536 Self {
537 client: reqwest::Client::new(),
538 url: format!("https://sourcify.dev/server/v2/contract/{}", chain.id()),
539 invalid_api_key: AtomicBool::new(false),
540 }
541 }
542}
543
544#[async_trait::async_trait]
545impl ExternalFetcherT for SourcifyFetcher {
546 fn kind(&self) -> FetcherKind {
547 FetcherKind::Sourcify
548 }
549
550 fn timeout(&self) -> Duration {
551 Duration::from_secs(1)
552 }
553
554 fn concurrency(&self) -> usize {
555 5
556 }
557
558 fn invalid_api_key(&self) -> &AtomicBool {
559 &self.invalid_api_key
560 }
561
562 async fn fetch(&self, address: Address) -> Result<Option<Metadata>, EtherscanError> {
563 let url = format!("{url}/{address}?fields=abi,compilation", url = self.url);
564 let response = self
565 .client
566 .get(url)
567 .send()
568 .await
569 .map_err(|e| EtherscanError::Unknown(e.to_string()))?;
570 let code = response.status();
571 match code.as_u16() {
572 404 => return Err(EtherscanError::ContractCodeNotVerified(address)),
574 429 => return Err(EtherscanError::RateLimitExceeded),
576 _ => {}
577 }
578 let response: SourcifyResponse =
579 response.json().await.map_err(|e| EtherscanError::Unknown(e.to_string()))?;
580 trace!(target: "evm::traces::external", "Sourcify response for {address}: {response:#?}");
581 match response {
582 SourcifyResponse::Success(metadata) => Ok(Some(metadata.into())),
583 SourcifyResponse::Error(error) => Err(EtherscanError::Unknown(format!("{error:#?}"))),
584 }
585 }
586}
587
588#[derive(Debug, Clone, Deserialize)]
590#[serde(untagged)]
591enum SourcifyResponse {
592 Success(SourcifyMetadata),
593 Error(SourcifyError),
594}
595
596#[derive(Debug, Clone, Deserialize)]
597#[serde(rename_all = "camelCase")]
598#[expect(dead_code)] struct SourcifyError {
600 custom_code: String,
601 message: String,
602 error_id: String,
603}
604
605#[derive(Debug, Clone, Deserialize)]
606#[serde(rename_all = "camelCase")]
607struct SourcifyMetadata {
608 #[serde(default)]
609 abi: Option<Box<serde_json::value::RawValue>>,
610 #[serde(default)]
611 compilation: Option<Compilation>,
612}
613
614#[derive(Debug, Clone, Deserialize)]
615#[serde(rename_all = "camelCase")]
616struct Compilation {
617 #[serde(default)]
618 compiler_version: String,
619 #[serde(default)]
620 name: String,
621}
622
623impl From<SourcifyMetadata> for Metadata {
624 fn from(metadata: SourcifyMetadata) -> Self {
625 let SourcifyMetadata { abi, compilation } = metadata;
626 let (contract_name, compiler_version) = compilation
627 .map(|c| (c.name, c.compiler_version))
628 .unwrap_or_else(|| (String::new(), String::new()));
629 Self {
631 source_code: foundry_block_explorers::contract::SourceCodeMetadata::Sources(
632 Default::default(),
633 ),
634 abi: Box::<str>::from(abi.unwrap_or_default()).into(),
635 contract_name,
636 compiler_version,
637 optimization_used: 0,
638 runs: 0,
639 constructor_arguments: Default::default(),
640 evm_version: String::new(),
641 library: String::new(),
642 license_type: String::new(),
643 proxy: 0,
644 implementation: None,
645 swarm_source: String::new(),
646 }
647 }
648}
649
650#[cfg(test)]
651mod tests {
652 use super::*;
653 use std::{
654 collections::HashSet as StdHashSet,
655 future::pending,
656 sync::{
657 Mutex,
658 atomic::{AtomicUsize, Ordering as AtomicOrdering},
659 },
660 };
661
662 struct TestFetcher {
663 kind: FetcherKind,
664 delay: Option<Duration>,
665 contract_name: Option<&'static str>,
666 calls: Arc<AtomicUsize>,
667 invalid: AtomicBool,
668 }
669
670 #[async_trait::async_trait]
671 impl ExternalFetcherT for TestFetcher {
672 fn kind(&self) -> FetcherKind {
673 self.kind
674 }
675
676 fn timeout(&self) -> Duration {
677 Duration::from_millis(1)
678 }
679
680 fn concurrency(&self) -> usize {
681 1
682 }
683
684 fn invalid_api_key(&self) -> &AtomicBool {
685 &self.invalid
686 }
687
688 async fn fetch(&self, _address: Address) -> Result<Option<Metadata>, EtherscanError> {
689 self.calls.fetch_add(1, AtomicOrdering::Relaxed);
690 let Some(delay) = self.delay else { return pending().await };
691 if !delay.is_zero() {
692 tokio::time::sleep(delay).await;
693 }
694 Ok(self.contract_name.map(metadata))
695 }
696 }
697
698 struct RateLimitedFetcher {
699 calls: Arc<AtomicUsize>,
700 invalid: AtomicBool,
701 }
702
703 #[async_trait::async_trait]
704 impl ExternalFetcherT for RateLimitedFetcher {
705 fn kind(&self) -> FetcherKind {
706 FetcherKind::Sourcify
707 }
708
709 fn timeout(&self) -> Duration {
710 Duration::from_millis(5)
711 }
712
713 fn concurrency(&self) -> usize {
714 1
715 }
716
717 fn invalid_api_key(&self) -> &AtomicBool {
718 &self.invalid
719 }
720
721 async fn fetch(&self, _address: Address) -> Result<Option<Metadata>, EtherscanError> {
722 self.calls.fetch_add(1, AtomicOrdering::Relaxed);
723 Err(EtherscanError::RateLimitExceeded)
724 }
725 }
726
727 fn metadata(contract_name: &str) -> Metadata {
728 SourcifyMetadata {
729 abi: None,
730 compilation: Some(Compilation {
731 compiler_version: String::new(),
732 name: contract_name.to_string(),
733 }),
734 }
735 .into()
736 }
737
738 fn test_identifier(
739 fetchers: Vec<Arc<dyn ExternalFetcherT>>,
740 remaining_budget: Duration,
741 ) -> ExternalIdentifier {
742 ExternalIdentifier { fetchers, contracts: Default::default(), remaining_budget }
743 }
744
745 #[test]
746 fn zero_timeout_disables_external_identification() {
747 let mut config = Config::default();
748 config.tracing.external_identification_timeout = 0;
749
750 assert!(ExternalIdentifier::new(&config, Some(Chain::mainnet())).unwrap().is_none());
751 }
752
753 struct FlakyCloudflareFetcher {
756 seen: Mutex<StdHashSet<Address>>,
757 invalid: AtomicBool,
758 }
759
760 #[async_trait::async_trait]
761 impl ExternalFetcherT for FlakyCloudflareFetcher {
762 fn kind(&self) -> FetcherKind {
763 FetcherKind::Etherscan
764 }
765 fn timeout(&self) -> Duration {
766 Duration::from_millis(1)
767 }
768 fn concurrency(&self) -> usize {
769 1
770 }
771 fn invalid_api_key(&self) -> &AtomicBool {
772 &self.invalid
773 }
774 async fn fetch(&self, address: Address) -> Result<Option<Metadata>, EtherscanError> {
775 let first_time = self.seen.lock().unwrap().insert(address);
776 if first_time { Err(EtherscanError::BlockedByCloudflare) } else { Ok(None) }
777 }
778 }
779
780 #[tokio::test]
784 async fn cloudflare_block_retries_instead_of_abandoning_queue() {
785 let addrs: Vec<Address> = (1u8..=4).map(Address::with_last_byte).collect();
786 let fetcher: Arc<dyn ExternalFetcherT> = Arc::new(FlakyCloudflareFetcher {
787 seen: Mutex::new(StdHashSet::new()),
788 invalid: AtomicBool::new(false),
789 });
790
791 let collected: Vec<_> = ExternalFetcher::new(fetcher, &addrs).collect().await;
792
793 let got: StdHashSet<Address> = collected.into_iter().map(|(addr, _)| addr).collect();
794 let want: StdHashSet<Address> = addrs.into_iter().collect();
795 assert_eq!(got, want, "every address must be yielded despite a transient cloudflare block");
796 }
797
798 #[tokio::test(start_paused = true)]
799 async fn timeout_keeps_partial_results_and_opens_circuit() {
800 let successful_calls = Arc::new(AtomicUsize::new(0));
801 let stalled_calls = Arc::new(AtomicUsize::new(0));
802 let fetchers: Vec<Arc<dyn ExternalFetcherT>> = vec![
803 Arc::new(TestFetcher {
804 kind: FetcherKind::Sourcify,
805 delay: Some(Duration::ZERO),
806 contract_name: Some("PartialResult"),
807 calls: Arc::clone(&successful_calls),
808 invalid: AtomicBool::new(false),
809 }),
810 Arc::new(TestFetcher {
811 kind: FetcherKind::Etherscan,
812 delay: None,
813 contract_name: None,
814 calls: Arc::clone(&stalled_calls),
815 invalid: AtomicBool::new(false),
816 }),
817 ];
818 let mut identifier = test_identifier(fetchers, Duration::from_millis(20));
819 let address = Address::with_last_byte(1);
820
821 identifier.fetch_addresses_async(&[address]).await;
822
823 assert!(identifier.remaining_budget.is_zero());
824 assert_eq!(
825 identifier.contracts[&address].1.as_ref().unwrap().contract_name,
826 "PartialResult"
827 );
828 assert_eq!(successful_calls.load(AtomicOrdering::Relaxed), 1);
829 assert_eq!(stalled_calls.load(AtomicOrdering::Relaxed), 1);
830
831 identifier.fetch_addresses_async(&[Address::with_last_byte(2)]).await;
832 assert_eq!(successful_calls.load(AtomicOrdering::Relaxed), 1);
833 assert_eq!(stalled_calls.load(AtomicOrdering::Relaxed), 1);
834 }
835
836 #[tokio::test(flavor = "multi_thread")]
837 async fn timeout_returns_partial_identity() {
838 let fetchers: Vec<Arc<dyn ExternalFetcherT>> = vec![
839 Arc::new(TestFetcher {
840 kind: FetcherKind::Sourcify,
841 delay: Some(Duration::ZERO),
842 contract_name: Some("PartialResult"),
843 calls: Arc::new(AtomicUsize::new(0)),
844 invalid: AtomicBool::new(false),
845 }),
846 Arc::new(TestFetcher {
847 kind: FetcherKind::Etherscan,
848 delay: None,
849 contract_name: None,
850 calls: Arc::new(AtomicUsize::new(0)),
851 invalid: AtomicBool::new(false),
852 }),
853 ];
854 let mut identifier = test_identifier(fetchers, Duration::from_millis(20));
855 let mut node = CallTraceNode::default();
856 node.trace.address = Address::with_last_byte(1);
857
858 let identities = identifier.identify_addresses(&[&node]);
859
860 assert_eq!(identities.len(), 1);
861 assert_eq!(identities[0].label.as_deref(), Some("PartialResult"));
862 }
863
864 #[tokio::test(start_paused = true)]
865 async fn timeout_budget_is_cumulative_across_fetches() {
866 let calls = Arc::new(AtomicUsize::new(0));
867 let fetcher: Arc<dyn ExternalFetcherT> = Arc::new(TestFetcher {
868 kind: FetcherKind::Sourcify,
869 delay: Some(Duration::from_millis(20)),
870 contract_name: Some("FirstResult"),
871 calls: Arc::clone(&calls),
872 invalid: AtomicBool::new(false),
873 });
874 let mut identifier = test_identifier(vec![fetcher], Duration::from_millis(30));
875 let first = Address::with_last_byte(1);
876 let second = Address::with_last_byte(2);
877
878 identifier.fetch_addresses_async(&[first]).await;
879 assert!(identifier.contracts[&first].1.is_some());
880 assert!(identifier.remaining_budget < Duration::from_millis(15));
881
882 identifier.fetch_addresses_async(&[second]).await;
883 assert!(identifier.remaining_budget.is_zero());
884 assert!(!identifier.contracts.contains_key(&second));
885 assert_eq!(calls.load(AtomicOrdering::Relaxed), 2);
886 }
887
888 #[tokio::test(start_paused = true)]
889 async fn rate_limit_retries_cannot_escape_timeout_budget() {
890 let calls = Arc::new(AtomicUsize::new(0));
891 let fetcher: Arc<dyn ExternalFetcherT> = Arc::new(RateLimitedFetcher {
892 calls: Arc::clone(&calls),
893 invalid: AtomicBool::new(false),
894 });
895 let mut identifier = test_identifier(vec![fetcher], Duration::from_millis(20));
896
897 identifier.fetch_addresses_async(&[Address::with_last_byte(1)]).await;
898
899 assert!(identifier.remaining_budget.is_zero());
900 assert!(calls.load(AtomicOrdering::Relaxed) > 1);
901 }
902
903 #[test]
904 fn etherscan_metadata_takes_precedence() {
905 let address = Address::with_last_byte(1);
906 let mut identifier = test_identifier(Vec::new(), Duration::ZERO);
907
908 identifier
909 .cache_fetched(address, (FetcherKind::Sourcify, Some(metadata("SourcifyResult"))));
910 identifier.cache_fetched(address, (FetcherKind::Etherscan, None));
911 assert_eq!(
912 identifier.contracts[&address].1.as_ref().unwrap().contract_name,
913 "SourcifyResult"
914 );
915
916 identifier
917 .cache_fetched(address, (FetcherKind::Etherscan, Some(metadata("EtherscanResult"))));
918 assert_eq!(
919 identifier.contracts[&address].1.as_ref().unwrap().contract_name,
920 "EtherscanResult"
921 );
922 }
923
924 #[tokio::test]
925 async fn proxy_metadata_preserves_address_identity_and_all_abis() {
926 let proxy = Address::with_last_byte(1);
927 let implementation_address = Address::with_last_byte(2);
928 let mut proxy_metadata = metadata("Proxy");
929 proxy_metadata.abi =
930 r#"[{"anonymous":false,"inputs":[],"name":"ProxyEvent","type":"event"}]"#.to_string();
931 proxy_metadata.proxy = 1;
932 proxy_metadata.implementation = Some(implementation_address);
933 let mut implementation = metadata("Implementation");
934 implementation.abi =
935 r#"[{"anonymous":false,"inputs":[],"name":"ImplementationEvent","type":"event"}]"#
936 .to_string();
937 let mut identifier = test_identifier(Vec::new(), Duration::from_secs(1));
938 let identity = identifier.identify_from_metadata(proxy, &proxy_metadata);
939 assert_eq!(identity.contract.as_deref(), Some("Proxy"));
940 identifier.cache_fetched(proxy, (FetcherKind::Etherscan, Some(proxy_metadata)));
941 identifier
942 .cache_fetched(implementation_address, (FetcherKind::Etherscan, Some(implementation)));
943
944 let mut results = identifier.get_abis(&[proxy]).await;
945 let (result_address, result) = results.pop().unwrap();
946 let (abis, complete) = result.unwrap();
947 let event_names =
948 abis.into_iter().map(|abi| abi.events.into_keys().next().unwrap()).collect::<Vec<_>>();
949
950 assert_eq!(result_address, proxy);
951 assert!(complete);
952 assert_eq!(event_names, ["ImplementationEvent", "ProxyEvent"]);
953
954 identifier.contracts.remove(&implementation_address);
955 let (_, result) = identifier.get_abis(&[proxy]).await.pop().unwrap();
956 let (abis, complete) = result.unwrap();
957 assert_eq!(abis.len(), 1);
958 assert!(!complete);
959 }
960}