foundry_evm_fuzz/invariant/
call_override.rs

Help
1use super::{BasicTxDetails, CallDetails};
2use alloy_primitives::Address;
3use parking_lot::{Mutex, RwLock};
4use proptest::{
5    option::weighted,
6    strategy::{SBoxedStrategy, Strategy, ValueTree},
7    test_runner::TestRunner,
8};
9use std::sync::Arc;
10
11/// Given a TestRunner and a strategy, it generates calls. Used inside the Fuzzer inspector to
12/// override external calls to test for potential reentrancy vulnerabilities..
13#[derive(Clone, Debug)]
14pub struct RandomCallGenerator {
15    /// Address of the test contract.
16    pub test_address: Address,
17    /// Runner that will generate the call from the strategy.
18    pub runner: Arc<Mutex<TestRunner>>,
19    /// Strategy to be used to generate calls from `target_reference`.
20    pub strategy: SBoxedStrategy<Option<CallDetails>>,
21    /// Reference to which contract we want a fuzzed calldata from.
22    pub target_reference: Arc<RwLock<Address>>,
23    /// Flag to know if a call has been overridden. Don't allow nesting for now.
24    pub used: bool,
25    /// If set to `true`, consumes the next call from `last_sequence`, otherwise queries it from
26    /// the strategy.
27    pub replay: bool,
28    /// Saves the sequence of generated calls that can be replayed later on.
29    pub last_sequence: Arc<RwLock<Vec<Option<BasicTxDetails>>>>,
30}
31
32impl RandomCallGenerator {
33    pub fn new(
34        test_address: Address,
35        runner: TestRunner,
36        strategy: impl Strategy<Value = CallDetails> + Send + Sync + 'static,
37        target_reference: Arc<RwLock<Address>>,
38    ) -> Self {
39        Self {
40            test_address,
41            runner: Arc::new(Mutex::new(runner)),
42            strategy: weighted(0.9, strategy).sboxed(),
43            target_reference,
44            last_sequence: Arc::default(),
45            replay: false,
46            used: false,
47        }
48    }
49
50    /// All `self.next()` calls will now pop `self.last_sequence`. Used to replay an invariant
51    /// failure.
52    pub fn set_replay(&mut self, status: bool) {
53        self.replay = status;
54        if status {
55            // So it can later be popped.
56            self.last_sequence.write().reverse();
57        }
58    }
59
60    /// Gets the next call. Random if replay is not set. Otherwise, it pops from `last_sequence`.
61    pub fn next(
62        &mut self,
63        original_caller: Address,
64        original_target: Address,
65    ) -> Option<BasicTxDetails> {
66        if self.replay {
67            self.last_sequence.write().pop().expect(
68                "to have same size as the number of (unsafe) external calls of the sequence.",
69            )
70        } else {
71            // TODO: Do we want it to be 80% chance only too ?
72            let sender = original_target;
73
74            // Set which contract we mostly (80% chance) want to generate calldata from.
75            *self.target_reference.write() = original_caller;
76
77            // `original_caller` has a 80% chance of being the `new_target`.
78            let choice = self
79                .strategy
80                .new_tree(&mut self.runner.lock())
81                .unwrap()
82                .current()
83                .map(|call_details| BasicTxDetails { sender, call_details });
84
85            self.last_sequence.write().push(choice.clone());
86            choice
87        }
88    }
89}
foundry_evm_fuzz/invariant/call_override.rs

foundry_evm_fuzz/invariant/
call_override.rs
