foundry_evm/executors/invariant/
mod.rs

1use crate::{
2    executors::{Executor, RawCallResult},
3    inspectors::Fuzzer,
4};
5use alloy_primitives::{
6    Address, Bytes, FixedBytes, Selector, U256,
7    map::{AddressMap, HashMap},
8};
9use alloy_sol_types::{SolCall, sol};
10use eyre::{ContextCompat, Result, eyre};
11use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
12use foundry_config::InvariantConfig;
13use foundry_evm_core::{
14    constants::{
15        CALLER, CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME,
16    },
17    precompiles::PRECOMPILES,
18};
19use foundry_evm_fuzz::{
20    BasicTxDetails, FuzzCase, FuzzFixtures, FuzzedCases,
21    invariant::{
22        ArtifactFilters, FuzzRunIdentifiedContracts, InvariantContract, RandomCallGenerator,
23        SenderFilters, TargetedContract, TargetedContracts,
24    },
25    strategies::{EvmFuzzState, invariant_strat, override_call_strat},
26};
27use foundry_evm_traces::{CallTraceArena, SparsedTraceArena};
28use indicatif::ProgressBar;
29use parking_lot::RwLock;
30use proptest::{strategy::Strategy, test_runner::TestRunner};
31use result::{assert_after_invariant, assert_invariants, can_continue};
32use revm::state::Account;
33use std::{
34    collections::{HashMap as Map, btree_map::Entry},
35    sync::Arc,
36    time::{Instant, SystemTime, UNIX_EPOCH},
37};
38
39mod error;
40pub use error::{InvariantFailures, InvariantFuzzError};
41use foundry_evm_coverage::HitMaps;
42
43mod replay;
44pub use replay::{replay_error, replay_run};
45
46mod result;
47use foundry_common::{TestFunctionExt, sh_println};
48pub use result::InvariantFuzzTestResult;
49use serde::{Deserialize, Serialize};
50use serde_json::json;
51
52mod shrink;
53use crate::executors::{
54    DURATION_BETWEEN_METRICS_REPORT, EarlyExit, EvmError, FuzzTestTimer, corpus::CorpusManager,
55};
56pub use shrink::check_sequence;
57
58sol! {
59    interface IInvariantTest {
60        #[derive(Default)]
61        struct FuzzSelector {
62            address addr;
63            bytes4[] selectors;
64        }
65
66        #[derive(Default)]
67        struct FuzzArtifactSelector {
68            string artifact;
69            bytes4[] selectors;
70        }
71
72        #[derive(Default)]
73        struct FuzzInterface {
74            address addr;
75            string[] artifacts;
76        }
77
78        function afterInvariant() external;
79
80        #[derive(Default)]
81        function excludeArtifacts() public view returns (string[] memory excludedArtifacts);
82
83        #[derive(Default)]
84        function excludeContracts() public view returns (address[] memory excludedContracts);
85
86        #[derive(Default)]
87        function excludeSelectors() public view returns (FuzzSelector[] memory excludedSelectors);
88
89        #[derive(Default)]
90        function excludeSenders() public view returns (address[] memory excludedSenders);
91
92        #[derive(Default)]
93        function targetArtifacts() public view returns (string[] memory targetedArtifacts);
94
95        #[derive(Default)]
96        function targetArtifactSelectors() public view returns (FuzzArtifactSelector[] memory targetedArtifactSelectors);
97
98        #[derive(Default)]
99        function targetContracts() public view returns (address[] memory targetedContracts);
100
101        #[derive(Default)]
102        function targetSelectors() public view returns (FuzzSelector[] memory targetedSelectors);
103
104        #[derive(Default)]
105        function targetSenders() public view returns (address[] memory targetedSenders);
106
107        #[derive(Default)]
108        function targetInterfaces() public view returns (FuzzInterface[] memory targetedInterfaces);
109    }
110}
111
112/// Contains invariant metrics for a single fuzzed selector.
113#[derive(Default, Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
114pub struct InvariantMetrics {
115    // Count of fuzzed selector calls.
116    pub calls: usize,
117    // Count of fuzzed selector reverts.
118    pub reverts: usize,
119    // Count of fuzzed selector discards (through assume cheatcodes).
120    pub discards: usize,
121}
122
123/// Contains data collected during invariant test runs.
124struct InvariantTestData {
125    // Consumed gas and calldata of every successful fuzz call.
126    fuzz_cases: Vec<FuzzedCases>,
127    // Data related to reverts or failed assertions of the test.
128    failures: InvariantFailures,
129    // Calldata in the last invariant run.
130    last_run_inputs: Vec<BasicTxDetails>,
131    // Additional traces for gas report.
132    gas_report_traces: Vec<Vec<CallTraceArena>>,
133    // Last call results of the invariant test.
134    last_call_results: Option<RawCallResult>,
135    // Line coverage information collected from all fuzzed calls.
136    line_coverage: Option<HitMaps>,
137    // Metrics for each fuzzed selector.
138    metrics: Map<String, InvariantMetrics>,
139
140    // Proptest runner to query for random values.
141    // The strategy only comes with the first `input`. We fill the rest of the `inputs`
142    // until the desired `depth` so we can use the evolving fuzz dictionary
143    // during the run.
144    branch_runner: TestRunner,
145}
146
147/// Contains invariant test data.
148struct InvariantTest {
149    // Fuzz state of invariant test.
150    fuzz_state: EvmFuzzState,
151    // Contracts fuzzed by the invariant test.
152    targeted_contracts: FuzzRunIdentifiedContracts,
153    // Data collected during invariant runs.
154    test_data: InvariantTestData,
155}
156
157impl InvariantTest {
158    /// Instantiates an invariant test.
159    fn new(
160        fuzz_state: EvmFuzzState,
161        targeted_contracts: FuzzRunIdentifiedContracts,
162        failures: InvariantFailures,
163        last_call_results: Option<RawCallResult>,
164        branch_runner: TestRunner,
165    ) -> Self {
166        let mut fuzz_cases = vec![];
167        if last_call_results.is_none() {
168            fuzz_cases.push(FuzzedCases::new(vec![]));
169        }
170        let test_data = InvariantTestData {
171            fuzz_cases,
172            failures,
173            last_run_inputs: vec![],
174            gas_report_traces: vec![],
175            last_call_results,
176            line_coverage: None,
177            metrics: Map::default(),
178            branch_runner,
179        };
180        Self { fuzz_state, targeted_contracts, test_data }
181    }
182
183    /// Returns number of invariant test reverts.
184    fn reverts(&self) -> usize {
185        self.test_data.failures.reverts
186    }
187
188    /// Whether invariant test has errors or not.
189    fn has_errors(&self) -> bool {
190        self.test_data.failures.error.is_some()
191    }
192
193    /// Set invariant test error.
194    fn set_error(&mut self, error: InvariantFuzzError) {
195        self.test_data.failures.error = Some(error);
196    }
197
198    /// Set last invariant test call results.
199    fn set_last_call_results(&mut self, call_result: Option<RawCallResult>) {
200        self.test_data.last_call_results = call_result;
201    }
202
203    /// Set last invariant run call sequence.
204    fn set_last_run_inputs(&mut self, inputs: &Vec<BasicTxDetails>) {
205        self.test_data.last_run_inputs.clone_from(inputs);
206    }
207
208    /// Merge current collected line coverage with the new coverage from last fuzzed call.
209    fn merge_line_coverage(&mut self, new_coverage: Option<HitMaps>) {
210        HitMaps::merge_opt(&mut self.test_data.line_coverage, new_coverage);
211    }
212
213    /// Update metrics for a fuzzed selector, extracted from tx details.
214    /// Always increments number of calls; discarded runs (through assume cheatcodes) are tracked
215    /// separated from reverts.
216    fn record_metrics(&mut self, tx_details: &BasicTxDetails, reverted: bool, discarded: bool) {
217        if let Some(metric_key) =
218            self.targeted_contracts.targets.lock().fuzzed_metric_key(tx_details)
219        {
220            let test_metrics = &mut self.test_data.metrics;
221            let invariant_metrics = test_metrics.entry(metric_key).or_default();
222            invariant_metrics.calls += 1;
223            if discarded {
224                invariant_metrics.discards += 1;
225            } else if reverted {
226                invariant_metrics.reverts += 1;
227            }
228        }
229    }
230
231    /// End invariant test run by collecting results, cleaning collected artifacts and reverting
232    /// created fuzz state.
233    fn end_run(&mut self, run: InvariantTestRun, gas_samples: usize) {
234        // We clear all the targeted contracts created during this run.
235        self.targeted_contracts.clear_created_contracts(run.created_contracts);
236
237        if self.test_data.gas_report_traces.len() < gas_samples {
238            self.test_data
239                .gas_report_traces
240                .push(run.run_traces.into_iter().map(|arena| arena.arena).collect());
241        }
242        self.test_data.fuzz_cases.push(FuzzedCases::new(run.fuzz_runs));
243
244        // Revert state to not persist values between runs.
245        self.fuzz_state.revert();
246    }
247}
248
249/// Contains data for an invariant test run.
250struct InvariantTestRun {
251    // Invariant run call sequence.
252    inputs: Vec<BasicTxDetails>,
253    // Current invariant run executor.
254    executor: Executor,
255    // Invariant run stat reports (eg. gas usage).
256    fuzz_runs: Vec<FuzzCase>,
257    // Contracts created during current invariant run.
258    created_contracts: Vec<Address>,
259    // Traces of each call of the invariant run call sequence.
260    run_traces: Vec<SparsedTraceArena>,
261    // Current depth of invariant run.
262    depth: u32,
263    // Current assume rejects of the invariant run.
264    rejects: u32,
265    // Whether new coverage was discovered during this run.
266    new_coverage: bool,
267}
268
269impl InvariantTestRun {
270    /// Instantiates an invariant test run.
271    fn new(first_input: BasicTxDetails, executor: Executor, depth: usize) -> Self {
272        Self {
273            inputs: vec![first_input],
274            executor,
275            fuzz_runs: Vec::with_capacity(depth),
276            created_contracts: vec![],
277            run_traces: vec![],
278            depth: 0,
279            rejects: 0,
280            new_coverage: false,
281        }
282    }
283}
284
285/// Wrapper around any [`Executor`] implementer which provides fuzzing support using [`proptest`].
286///
287/// After instantiation, calling `invariant_fuzz` will proceed to hammer the deployed smart
288/// contracts with inputs, until it finds a counterexample sequence. The provided [`TestRunner`]
289/// contains all the configuration which can be overridden via [environment
290/// variables](proptest::test_runner::Config)
291pub struct InvariantExecutor<'a> {
292    pub executor: Executor,
293    /// Proptest runner.
294    runner: TestRunner,
295    /// The invariant configuration
296    config: InvariantConfig,
297    /// Contracts deployed with `setUp()`
298    setup_contracts: &'a ContractsByAddress,
299    /// Contracts that are part of the project but have not been deployed yet. We need the bytecode
300    /// to identify them from the stateset changes.
301    project_contracts: &'a ContractsByArtifact,
302    /// Filters contracts to be fuzzed through their artifact identifiers.
303    artifact_filters: ArtifactFilters,
304}
305
306impl<'a> InvariantExecutor<'a> {
307    /// Instantiates a fuzzed executor EVM given a testrunner
308    pub fn new(
309        executor: Executor,
310        runner: TestRunner,
311        config: InvariantConfig,
312        setup_contracts: &'a ContractsByAddress,
313        project_contracts: &'a ContractsByArtifact,
314    ) -> Self {
315        Self {
316            executor,
317            runner,
318            config,
319            setup_contracts,
320            project_contracts,
321            artifact_filters: ArtifactFilters::default(),
322        }
323    }
324
325    pub fn config(self) -> InvariantConfig {
326        self.config
327    }
328
329    /// Fuzzes any deployed contract and checks any broken invariant at `invariant_address`.
330    pub fn invariant_fuzz(
331        &mut self,
332        invariant_contract: InvariantContract<'_>,
333        fuzz_fixtures: &FuzzFixtures,
334        fuzz_state: EvmFuzzState,
335        progress: Option<&ProgressBar>,
336        early_exit: &EarlyExit,
337    ) -> Result<InvariantFuzzTestResult> {
338        // Throw an error to abort test run if the invariant function accepts input params
339        if !invariant_contract.invariant_function.inputs.is_empty() {
340            return Err(eyre!("Invariant test function should have no inputs"));
341        }
342
343        let (mut invariant_test, mut corpus_manager) =
344            self.prepare_test(&invariant_contract, fuzz_fixtures, fuzz_state)?;
345
346        // Start timer for this invariant test.
347        let mut runs = 0;
348        let timer = FuzzTestTimer::new(self.config.timeout);
349        let mut last_metrics_report = Instant::now();
350        let continue_campaign = |runs: u32| {
351            if early_exit.should_stop() {
352                return false;
353            }
354
355            if timer.is_enabled() { !timer.is_timed_out() } else { runs < self.config.runs }
356        };
357
358        // Invariant runs with edge coverage if corpus dir is set or showing edge coverage.
359        let edge_coverage_enabled = self.config.corpus.collect_edge_coverage();
360
361        'stop: while continue_campaign(runs) {
362            let initial_seq = corpus_manager.new_inputs(
363                &mut invariant_test.test_data.branch_runner,
364                &invariant_test.fuzz_state,
365                &invariant_test.targeted_contracts,
366            )?;
367
368            // Create current invariant run data.
369            let mut current_run = InvariantTestRun::new(
370                initial_seq[0].clone(),
371                // Before each run, we must reset the backend state.
372                self.executor.clone(),
373                self.config.depth as usize,
374            );
375
376            // We stop the run immediately if we have reverted, and `fail_on_revert` is set.
377            if self.config.fail_on_revert && invariant_test.reverts() > 0 {
378                return Err(eyre!("call reverted"));
379            }
380
381            while current_run.depth < self.config.depth {
382                // Check if the timeout has been reached.
383                if timer.is_timed_out() {
384                    // Since we never record a revert here the test is still considered
385                    // successful even though it timed out. We *want*
386                    // this behavior for now, so that's ok, but
387                    // future developers should be aware of this.
388                    break 'stop;
389                }
390
391                let tx = current_run
392                    .inputs
393                    .last()
394                    .ok_or_else(|| eyre!("no input generated to call fuzzed target."))?;
395
396                // Execute call from the randomly generated sequence without committing state.
397                // State is committed only if call is not a magic assume.
398                let mut call_result = execute_tx(&mut current_run.executor, tx)?;
399                let discarded = call_result.result.as_ref() == MAGIC_ASSUME;
400                if self.config.show_metrics {
401                    invariant_test.record_metrics(tx, call_result.reverted, discarded);
402                }
403
404                // Collect line coverage from last fuzzed call.
405                invariant_test.merge_line_coverage(call_result.line_coverage.clone());
406                // Collect edge coverage and set the flag in the current run.
407                if corpus_manager.merge_edge_coverage(&mut call_result) {
408                    current_run.new_coverage = true;
409                }
410
411                if discarded {
412                    current_run.inputs.pop();
413                    current_run.rejects += 1;
414                    if current_run.rejects > self.config.max_assume_rejects {
415                        invariant_test.set_error(InvariantFuzzError::MaxAssumeRejects(
416                            self.config.max_assume_rejects,
417                        ));
418                        break 'stop;
419                    }
420                } else {
421                    // Commit executed call result.
422                    current_run.executor.commit(&mut call_result);
423
424                    // Collect data for fuzzing from the state changeset.
425                    // This step updates the state dictionary and therefore invalidates the
426                    // ValueTree in use by the current run. This manifestsitself in proptest
427                    // observing a different input case than what it was called with, and creates
428                    // inconsistencies whenever proptest tries to use the input case after test
429                    // execution.
430                    // See <https://github.com/foundry-rs/foundry/issues/9764>.
431                    let mut state_changeset = call_result.state_changeset.clone();
432                    if !call_result.reverted {
433                        collect_data(
434                            &invariant_test,
435                            &mut state_changeset,
436                            tx,
437                            &call_result,
438                            self.config.depth,
439                        );
440                    }
441
442                    // Collect created contracts and add to fuzz targets only if targeted contracts
443                    // are updatable.
444                    if let Err(error) =
445                        &invariant_test.targeted_contracts.collect_created_contracts(
446                            &state_changeset,
447                            self.project_contracts,
448                            self.setup_contracts,
449                            &self.artifact_filters,
450                            &mut current_run.created_contracts,
451                        )
452                    {
453                        warn!(target: "forge::test", "{error}");
454                    }
455                    current_run.fuzz_runs.push(FuzzCase {
456                        calldata: tx.call_details.calldata.clone(),
457                        gas: call_result.gas_used,
458                        stipend: call_result.stipend,
459                    });
460
461                    // Determine if test can continue or should exit.
462                    let result = can_continue(
463                        &invariant_contract,
464                        &mut invariant_test,
465                        &mut current_run,
466                        &self.config,
467                        call_result,
468                        &state_changeset,
469                    )
470                    .map_err(|e| eyre!(e.to_string()))?;
471                    if !result.can_continue || current_run.depth == self.config.depth - 1 {
472                        invariant_test.set_last_run_inputs(&current_run.inputs);
473                    }
474                    // If test cannot continue then stop current run and exit test suite.
475                    if !result.can_continue {
476                        break 'stop;
477                    }
478
479                    invariant_test.set_last_call_results(result.call_result);
480                    current_run.depth += 1;
481                }
482
483                current_run.inputs.push(corpus_manager.generate_next_input(
484                    &mut invariant_test.test_data.branch_runner,
485                    &initial_seq,
486                    discarded,
487                    current_run.depth as usize,
488                )?);
489            }
490
491            // Extend corpus with current run data.
492            corpus_manager.process_inputs(&current_run.inputs, current_run.new_coverage);
493
494            // Call `afterInvariant` only if it is declared and test didn't fail already.
495            if invariant_contract.call_after_invariant && !invariant_test.has_errors() {
496                assert_after_invariant(
497                    &invariant_contract,
498                    &mut invariant_test,
499                    &current_run,
500                    &self.config,
501                )
502                .map_err(|_| eyre!("Failed to call afterInvariant"))?;
503            }
504
505            // End current invariant test run.
506            invariant_test.end_run(current_run, self.config.gas_report_samples as usize);
507            if let Some(progress) = progress {
508                // If running with progress then increment completed runs.
509                progress.inc(1);
510                // Display metrics in progress bar.
511                if edge_coverage_enabled {
512                    progress.set_message(format!("{}", &corpus_manager.metrics));
513                }
514            } else if edge_coverage_enabled
515                && last_metrics_report.elapsed() > DURATION_BETWEEN_METRICS_REPORT
516            {
517                // Display metrics inline if corpus dir set.
518                let metrics = json!({
519                    "timestamp": SystemTime::now()
520                        .duration_since(UNIX_EPOCH)?
521                        .as_secs(),
522                    "invariant": invariant_contract.invariant_function.name,
523                    "metrics": &corpus_manager.metrics,
524                });
525                let _ = sh_println!("{}", serde_json::to_string(&metrics)?);
526                last_metrics_report = Instant::now();
527            }
528
529            runs += 1;
530        }
531
532        trace!(?fuzz_fixtures);
533        invariant_test.fuzz_state.log_stats();
534
535        let result = invariant_test.test_data;
536        Ok(InvariantFuzzTestResult {
537            error: result.failures.error,
538            cases: result.fuzz_cases,
539            reverts: result.failures.reverts,
540            last_run_inputs: result.last_run_inputs,
541            gas_report_traces: result.gas_report_traces,
542            line_coverage: result.line_coverage,
543            metrics: result.metrics,
544            failed_corpus_replays: corpus_manager.failed_replays(),
545        })
546    }
547
548    /// Prepares certain structures to execute the invariant tests:
549    /// * Invariant Fuzz Test.
550    /// * Invariant Corpus Manager.
551    fn prepare_test(
552        &mut self,
553        invariant_contract: &InvariantContract<'_>,
554        fuzz_fixtures: &FuzzFixtures,
555        fuzz_state: EvmFuzzState,
556    ) -> Result<(InvariantTest, CorpusManager)> {
557        // Finds out the chosen deployed contracts and/or senders.
558        self.select_contract_artifacts(invariant_contract.address)?;
559        let (targeted_senders, targeted_contracts) =
560            self.select_contracts_and_senders(invariant_contract.address)?;
561
562        // Creates the invariant strategy.
563        let strategy = invariant_strat(
564            fuzz_state.clone(),
565            targeted_senders,
566            targeted_contracts.clone(),
567            self.config.clone(),
568            fuzz_fixtures.clone(),
569        )
570        .no_shrink();
571
572        // Allows `override_call_strat` to use the address given by the Fuzzer inspector during
573        // EVM execution.
574        let mut call_generator = None;
575        if self.config.call_override {
576            let target_contract_ref = Arc::new(RwLock::new(Address::ZERO));
577
578            call_generator = Some(RandomCallGenerator::new(
579                invariant_contract.address,
580                self.runner.clone(),
581                override_call_strat(
582                    fuzz_state.clone(),
583                    targeted_contracts.clone(),
584                    target_contract_ref.clone(),
585                    fuzz_fixtures.clone(),
586                ),
587                target_contract_ref,
588            ));
589        }
590
591        // If any of the targeted contracts have the storage layout enabled then we can sample
592        // mapping values. To accomplish, we need to record the mapping storage slots and keys.
593        let fuzz_state =
594            if targeted_contracts.targets.lock().iter().any(|(_, t)| t.storage_layout.is_some()) {
595                fuzz_state.with_mapping_slots(AddressMap::default())
596            } else {
597                fuzz_state
598            };
599
600        self.executor.inspector_mut().set_fuzzer(Fuzzer {
601            call_generator,
602            fuzz_state: fuzz_state.clone(),
603            collect: true,
604        });
605
606        // Let's make sure the invariant is sound before actually starting the run:
607        // We'll assert the invariant in its initial state, and if it fails, we'll
608        // already know if we can early exit the invariant run.
609        // This does not count as a fuzz run. It will just register the revert.
610        let mut failures = InvariantFailures::new();
611        let last_call_results = assert_invariants(
612            invariant_contract,
613            &self.config,
614            &targeted_contracts,
615            &self.executor,
616            &[],
617            &mut failures,
618        )?;
619        if let Some(error) = failures.error {
620            return Err(eyre!(error.revert_reason().unwrap_or_default()));
621        }
622
623        let corpus_manager = CorpusManager::new(
624            self.config.corpus.clone(),
625            strategy.boxed(),
626            &self.executor,
627            None,
628            Some(&targeted_contracts),
629        )?;
630        let invariant_test = InvariantTest::new(
631            fuzz_state,
632            targeted_contracts,
633            failures,
634            last_call_results,
635            self.runner.clone(),
636        );
637
638        Ok((invariant_test, corpus_manager))
639    }
640
641    /// Fills the `InvariantExecutor` with the artifact identifier filters (in `path:name` string
642    /// format). They will be used to filter contracts after the `setUp`, and more importantly,
643    /// during the runs.
644    ///
645    /// Also excludes any contract without any mutable functions.
646    ///
647    /// Priority:
648    ///
649    /// targetArtifactSelectors > excludeArtifacts > targetArtifacts
650    pub fn select_contract_artifacts(&mut self, invariant_address: Address) -> Result<()> {
651        let targeted_artifact_selectors = self
652            .executor
653            .call_sol_default(invariant_address, &IInvariantTest::targetArtifactSelectorsCall {});
654
655        // Insert them into the executor `targeted_abi`.
656        for IInvariantTest::FuzzArtifactSelector { artifact, selectors } in
657            targeted_artifact_selectors
658        {
659            let identifier = self.validate_selected_contract(artifact, &selectors)?;
660            self.artifact_filters.targeted.entry(identifier).or_default().extend(selectors);
661        }
662
663        let targeted_artifacts = self
664            .executor
665            .call_sol_default(invariant_address, &IInvariantTest::targetArtifactsCall {});
666        let excluded_artifacts = self
667            .executor
668            .call_sol_default(invariant_address, &IInvariantTest::excludeArtifactsCall {});
669
670        // Insert `excludeArtifacts` into the executor `excluded_abi`.
671        for contract in excluded_artifacts {
672            let identifier = self.validate_selected_contract(contract, &[])?;
673
674            if !self.artifact_filters.excluded.contains(&identifier) {
675                self.artifact_filters.excluded.push(identifier);
676            }
677        }
678
679        // Exclude any artifact without mutable functions.
680        for (artifact, contract) in self.project_contracts.iter() {
681            if contract
682                .abi
683                .functions()
684                .filter(|func| {
685                    !matches!(
686                        func.state_mutability,
687                        alloy_json_abi::StateMutability::Pure
688                            | alloy_json_abi::StateMutability::View
689                    )
690                })
691                .count()
692                == 0
693                && !self.artifact_filters.excluded.contains(&artifact.identifier())
694            {
695                self.artifact_filters.excluded.push(artifact.identifier());
696            }
697        }
698
699        // Insert `targetArtifacts` into the executor `targeted_abi`, if they have not been seen
700        // before.
701        for contract in targeted_artifacts {
702            let identifier = self.validate_selected_contract(contract, &[])?;
703
704            if !self.artifact_filters.targeted.contains_key(&identifier)
705                && !self.artifact_filters.excluded.contains(&identifier)
706            {
707                self.artifact_filters.targeted.insert(identifier, vec![]);
708            }
709        }
710        Ok(())
711    }
712
713    /// Makes sure that the contract exists in the project. If so, it returns its artifact
714    /// identifier.
715    fn validate_selected_contract(
716        &mut self,
717        contract: String,
718        selectors: &[FixedBytes<4>],
719    ) -> Result<String> {
720        if let Some((artifact, contract_data)) =
721            self.project_contracts.find_by_name_or_identifier(&contract)?
722        {
723            // Check that the selectors really exist for this contract.
724            for selector in selectors {
725                contract_data
726                    .abi
727                    .functions()
728                    .find(|func| func.selector().as_slice() == selector.as_slice())
729                    .wrap_err(format!("{contract} does not have the selector {selector:?}"))?;
730            }
731
732            return Ok(artifact.identifier());
733        }
734        eyre::bail!(
735            "{contract} not found in the project. Allowed format: `contract_name` or `contract_path:contract_name`."
736        );
737    }
738
739    /// Selects senders and contracts based on the contract methods `targetSenders() -> address[]`,
740    /// `targetContracts() -> address[]` and `excludeContracts() -> address[]`.
741    pub fn select_contracts_and_senders(
742        &self,
743        to: Address,
744    ) -> Result<(SenderFilters, FuzzRunIdentifiedContracts)> {
745        let targeted_senders =
746            self.executor.call_sol_default(to, &IInvariantTest::targetSendersCall {});
747        let mut excluded_senders =
748            self.executor.call_sol_default(to, &IInvariantTest::excludeSendersCall {});
749        // Extend with default excluded addresses - https://github.com/foundry-rs/foundry/issues/4163
750        excluded_senders.extend([
751            CHEATCODE_ADDRESS,
752            HARDHAT_CONSOLE_ADDRESS,
753            DEFAULT_CREATE2_DEPLOYER,
754        ]);
755        // Extend with precompiles - https://github.com/foundry-rs/foundry/issues/4287
756        excluded_senders.extend(PRECOMPILES);
757        let sender_filters = SenderFilters::new(targeted_senders, excluded_senders);
758
759        let selected = self.executor.call_sol_default(to, &IInvariantTest::targetContractsCall {});
760        let excluded = self.executor.call_sol_default(to, &IInvariantTest::excludeContractsCall {});
761
762        let contracts = self
763            .setup_contracts
764            .iter()
765            .filter(|&(addr, (identifier, _))| {
766                // Include to address if explicitly set as target.
767                if *addr == to && selected.contains(&to) {
768                    return true;
769                }
770
771                *addr != to
772                    && *addr != CHEATCODE_ADDRESS
773                    && *addr != HARDHAT_CONSOLE_ADDRESS
774                    && (selected.is_empty() || selected.contains(addr))
775                    && (excluded.is_empty() || !excluded.contains(addr))
776                    && self.artifact_filters.matches(identifier)
777            })
778            .map(|(addr, (identifier, abi))| {
779                (
780                    *addr,
781                    TargetedContract::new(identifier.clone(), abi.clone())
782                        .with_project_contracts(self.project_contracts),
783                )
784            })
785            .collect();
786        let mut contracts = TargetedContracts { inner: contracts };
787
788        self.target_interfaces(to, &mut contracts)?;
789
790        self.select_selectors(to, &mut contracts)?;
791
792        // There should be at least one contract identified as target for fuzz runs.
793        if contracts.is_empty() {
794            eyre::bail!("No contracts to fuzz.");
795        }
796
797        Ok((sender_filters, FuzzRunIdentifiedContracts::new(contracts, selected.is_empty())))
798    }
799
800    /// Extends the contracts and selectors to fuzz with the addresses and ABIs specified in
801    /// `targetInterfaces() -> (address, string[])[]`. Enables targeting of addresses that are
802    /// not deployed during `setUp` such as when fuzzing in a forked environment. Also enables
803    /// targeting of delegate proxies and contracts deployed with `create` or `create2`.
804    pub fn target_interfaces(
805        &self,
806        invariant_address: Address,
807        targeted_contracts: &mut TargetedContracts,
808    ) -> Result<()> {
809        let interfaces = self
810            .executor
811            .call_sol_default(invariant_address, &IInvariantTest::targetInterfacesCall {});
812
813        // Since `targetInterfaces` returns a tuple array there is no guarantee
814        // that the addresses are unique this map is used to merge functions of
815        // the specified interfaces for the same address. For example:
816        // `[(addr1, ["IERC20", "IOwnable"])]` and `[(addr1, ["IERC20"]), (addr1, ("IOwnable"))]`
817        // should be equivalent.
818        let mut combined = TargetedContracts::new();
819
820        // Loop through each address and its associated artifact identifiers.
821        // We're borrowing here to avoid taking full ownership.
822        for IInvariantTest::FuzzInterface { addr, artifacts } in &interfaces {
823            // Identifiers are specified as an array, so we loop through them.
824            for identifier in artifacts {
825                // Try to find the contract by name or identifier in the project's contracts.
826                if let Some((_, contract_data)) =
827                    self.project_contracts.iter().find(|(artifact, _)| {
828                        &artifact.name == identifier || &artifact.identifier() == identifier
829                    })
830                {
831                    let abi = &contract_data.abi;
832                    combined
833                        // Check if there's an entry for the given key in the 'combined' map.
834                        .entry(*addr)
835                        // If the entry exists, extends its ABI with the function list.
836                        .and_modify(|entry| {
837                            // Extend the ABI's function list with the new functions.
838                            entry.abi.functions.extend(abi.functions.clone());
839                        })
840                        // Otherwise insert it into the map.
841                        .or_insert_with(|| {
842                            let mut contract =
843                                TargetedContract::new(identifier.to_string(), abi.clone());
844                            contract.storage_layout =
845                                contract_data.storage_layout.as_ref().map(Arc::clone);
846                            contract
847                        });
848                }
849            }
850        }
851
852        targeted_contracts.extend(combined.inner);
853
854        Ok(())
855    }
856
857    /// Selects the functions to fuzz based on the contract method `targetSelectors()` and
858    /// `targetArtifactSelectors()`.
859    pub fn select_selectors(
860        &self,
861        address: Address,
862        targeted_contracts: &mut TargetedContracts,
863    ) -> Result<()> {
864        for (address, (identifier, _)) in self.setup_contracts {
865            if let Some(selectors) = self.artifact_filters.targeted.get(identifier) {
866                self.add_address_with_functions(*address, selectors, false, targeted_contracts)?;
867            }
868        }
869
870        let mut target_test_selectors = vec![];
871        let mut excluded_test_selectors = vec![];
872
873        // Collect contract functions marked as target for fuzzing campaign.
874        let selectors =
875            self.executor.call_sol_default(address, &IInvariantTest::targetSelectorsCall {});
876        for IInvariantTest::FuzzSelector { addr, selectors } in selectors {
877            if addr == address {
878                target_test_selectors = selectors.clone();
879            }
880            self.add_address_with_functions(addr, &selectors, false, targeted_contracts)?;
881        }
882
883        // Collect contract functions excluded from fuzzing campaign.
884        let excluded_selectors =
885            self.executor.call_sol_default(address, &IInvariantTest::excludeSelectorsCall {});
886        for IInvariantTest::FuzzSelector { addr, selectors } in excluded_selectors {
887            if addr == address {
888                // If fuzz selector address is the test contract, then record selectors to be
889                // later excluded if needed.
890                excluded_test_selectors = selectors.clone();
891            }
892            self.add_address_with_functions(addr, &selectors, true, targeted_contracts)?;
893        }
894
895        if target_test_selectors.is_empty()
896            && let Some(target) = targeted_contracts.get(&address)
897        {
898            // If test contract is marked as a target and no target selector explicitly set, then
899            // include only state-changing functions that are not reserved and selectors that are
900            // not explicitly excluded.
901            let selectors: Vec<_> = target
902                .abi
903                .functions()
904                .filter_map(|func| {
905                    if matches!(
906                        func.state_mutability,
907                        alloy_json_abi::StateMutability::Pure
908                            | alloy_json_abi::StateMutability::View
909                    ) || func.is_reserved()
910                        || excluded_test_selectors.contains(&func.selector())
911                    {
912                        None
913                    } else {
914                        Some(func.selector())
915                    }
916                })
917                .collect();
918            self.add_address_with_functions(address, &selectors, false, targeted_contracts)?;
919        }
920
921        Ok(())
922    }
923
924    /// Adds the address and fuzzed or excluded functions to `TargetedContracts`.
925    fn add_address_with_functions(
926        &self,
927        address: Address,
928        selectors: &[Selector],
929        should_exclude: bool,
930        targeted_contracts: &mut TargetedContracts,
931    ) -> eyre::Result<()> {
932        // Do not add address in target contracts if no function selected.
933        if selectors.is_empty() {
934            return Ok(());
935        }
936
937        let contract = match targeted_contracts.entry(address) {
938            Entry::Occupied(entry) => entry.into_mut(),
939            Entry::Vacant(entry) => {
940                let (identifier, abi) = self.setup_contracts.get(&address).ok_or_else(|| {
941                    eyre::eyre!(
942                        "[{}] address does not have an associated contract: {}",
943                        if should_exclude { "excludeSelectors" } else { "targetSelectors" },
944                        address
945                    )
946                })?;
947                entry.insert(
948                    TargetedContract::new(identifier.clone(), abi.clone())
949                        .with_project_contracts(self.project_contracts),
950                )
951            }
952        };
953        contract.add_selectors(selectors.iter().copied(), should_exclude)?;
954        Ok(())
955    }
956}
957
958/// Collects data from call for fuzzing. However, it first verifies that the sender is not an EOA
959/// before inserting it into the dictionary. Otherwise, we flood the dictionary with
960/// randomly generated addresses.
961fn collect_data(
962    invariant_test: &InvariantTest,
963    state_changeset: &mut HashMap<Address, Account>,
964    tx: &BasicTxDetails,
965    call_result: &RawCallResult,
966    run_depth: u32,
967) {
968    // Verify it has no code.
969    let mut has_code = false;
970    if let Some(Some(code)) =
971        state_changeset.get(&tx.sender).map(|account| account.info.code.as_ref())
972    {
973        has_code = !code.is_empty();
974    }
975
976    // We keep the nonce changes to apply later.
977    let mut sender_changeset = None;
978    if !has_code {
979        sender_changeset = state_changeset.remove(&tx.sender);
980    }
981
982    // Collect values from fuzzed call result and add them to fuzz dictionary.
983    invariant_test.fuzz_state.collect_values_from_call(
984        &invariant_test.targeted_contracts,
985        tx,
986        &call_result.result,
987        &call_result.logs,
988        &*state_changeset,
989        run_depth,
990    );
991
992    // Re-add changes
993    if let Some(changed) = sender_changeset {
994        state_changeset.insert(tx.sender, changed);
995    }
996}
997
998/// Calls the `afterInvariant()` function on a contract.
999/// Returns call result and if call succeeded.
1000/// The state after the call is not persisted.
1001pub(crate) fn call_after_invariant_function(
1002    executor: &Executor,
1003    to: Address,
1004) -> Result<(RawCallResult, bool), EvmError> {
1005    let calldata = Bytes::from_static(&IInvariantTest::afterInvariantCall::SELECTOR);
1006    let mut call_result = executor.call_raw(CALLER, to, calldata, U256::ZERO)?;
1007    let success = executor.is_raw_call_mut_success(to, &mut call_result, false);
1008    Ok((call_result, success))
1009}
1010
1011/// Calls the invariant function and returns call result and if succeeded.
1012pub(crate) fn call_invariant_function(
1013    executor: &Executor,
1014    address: Address,
1015    calldata: Bytes,
1016) -> Result<(RawCallResult, bool)> {
1017    let mut call_result = executor.call_raw(CALLER, address, calldata, U256::ZERO)?;
1018    let success = executor.is_raw_call_mut_success(address, &mut call_result, false);
1019    Ok((call_result, success))
1020}
1021
1022/// Calls the invariant selector and returns call result and if succeeded.
1023/// Updates the block number and block timestamp if configured.
1024pub(crate) fn execute_tx(executor: &mut Executor, tx: &BasicTxDetails) -> Result<RawCallResult> {
1025    // Apply pre-call block adjustments.
1026    if let Some(warp) = tx.warp {
1027        executor.env_mut().evm_env.block_env.timestamp += warp;
1028    }
1029    if let Some(roll) = tx.roll {
1030        executor.env_mut().evm_env.block_env.number += roll;
1031    }
1032
1033    // Perform the raw call.
1034    let mut call_result = executor
1035        .call_raw(tx.sender, tx.call_details.target, tx.call_details.calldata.clone(), U256::ZERO)
1036        .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")))?;
1037
1038    // Propagate block adjustments to call result which will be committed.
1039    if let Some(warp) = tx.warp {
1040        call_result.env.evm_env.block_env.timestamp += warp;
1041    }
1042    if let Some(roll) = tx.roll {
1043        call_result.env.evm_env.block_env.number += roll;
1044    }
1045    Ok(call_result)
1046}