Skip to main content

cast/cmd/
call_overrides.rs

1//! Shared `eth_call` state and block override options.
2
3use alloy_contract::{CallDecoder, EthCall};
4use alloy_network::Network;
5use alloy_primitives::{Address, B256, Bytes, U256, map::HashMap};
6use alloy_rpc_types::{
7    BlockOverrides,
8    state::{StateOverride, StateOverridesBuilder},
9};
10use clap::Args;
11use eyre::Result;
12use regex::Regex;
13use std::{str::FromStr, sync::LazyLock};
14
15// Matches override pattern <address>:<slot>:<value>.
16// e.g. 0x123:0x1:0x1234.
17static OVERRIDE_PATTERN: LazyLock<Regex> =
18    LazyLock::new(|| Regex::new(r"^([^:]+):([^:]+):([^:]+)$").unwrap());
19
20/// State and block overrides for an `eth_call`.
21#[derive(Args, Clone, Debug, Default)]
22pub struct CallOverrideOpts {
23    /// Override account balances.
24    /// Format: "address:balance,address:balance".
25    #[arg(long = "override-balance", value_name = "ADDRESS:BALANCE", value_delimiter = ',')]
26    pub balance_overrides: Option<Vec<String>>,
27
28    /// Override account nonces.
29    /// Format: "address:nonce,address:nonce".
30    #[arg(long = "override-nonce", value_name = "ADDRESS:NONCE", value_delimiter = ',')]
31    pub nonce_overrides: Option<Vec<String>>,
32
33    /// Override account code.
34    /// Format: "address:code,address:code".
35    #[arg(long = "override-code", value_name = "ADDRESS:CODE", value_delimiter = ',')]
36    pub code_overrides: Option<Vec<String>>,
37
38    /// Override account state and replace the current state entirely with the new one.
39    /// Format: "address:slot:value,address:slot:value".
40    #[arg(long = "override-state", value_name = "ADDRESS:SLOT:VALUE", value_delimiter = ',')]
41    pub state_overrides: Option<Vec<String>>,
42
43    /// Override specific account storage slots and preserve the rest of the state.
44    /// Format: "address:slot:value,address:slot:value".
45    #[arg(long = "override-state-diff", value_name = "ADDRESS:SLOT:VALUE", value_delimiter = ',')]
46    pub state_diff_overrides: Option<Vec<String>>,
47
48    /// Override the block timestamp.
49    #[arg(long = "block.time", value_name = "TIME")]
50    pub block_time: Option<u64>,
51
52    /// Override the block number.
53    #[arg(long = "block.number", value_name = "NUMBER")]
54    pub block_number: Option<u64>,
55}
56
57impl CallOverrideOpts {
58    /// Returns true when no state or block override was provided.
59    pub const fn is_empty(&self) -> bool {
60        self.balance_overrides.is_none()
61            && self.nonce_overrides.is_none()
62            && self.code_overrides.is_none()
63            && self.state_overrides.is_none()
64            && self.state_diff_overrides.is_none()
65            && self.block_time.is_none()
66            && self.block_number.is_none()
67    }
68
69    /// Applies the configured overrides to an `eth_call`.
70    pub fn apply<'a, D, N>(&self, mut call: EthCall<'a, D, N>) -> Result<EthCall<'a, D, N>>
71    where
72        D: CallDecoder,
73        N: Network,
74    {
75        if let Some(state_overrides) = self.get_state_overrides()? {
76            call = call.overrides(state_overrides);
77        }
78        if let Some(block_overrides) = self.get_block_overrides()? {
79            call = call.with_block_overrides(block_overrides);
80        }
81        Ok(call)
82    }
83
84    /// Parses state overrides from command line arguments.
85    pub fn get_state_overrides(&self) -> Result<Option<StateOverride>> {
86        // Early return if no override set - <https://github.com/foundry-rs/foundry/issues/10705>.
87        if [
88            self.balance_overrides.as_ref(),
89            self.nonce_overrides.as_ref(),
90            self.code_overrides.as_ref(),
91            self.state_overrides.as_ref(),
92            self.state_diff_overrides.as_ref(),
93        ]
94        .iter()
95        .all(Option::is_none)
96        {
97            return Ok(None);
98        }
99
100        let mut state_overrides_builder = StateOverridesBuilder::default();
101
102        for override_str in self.balance_overrides.iter().flatten() {
103            let (addr, balance) = address_value_override(override_str)?;
104            state_overrides_builder =
105                state_overrides_builder.with_balance(addr.parse()?, balance.parse()?);
106        }
107
108        for override_str in self.nonce_overrides.iter().flatten() {
109            let (addr, nonce) = address_value_override(override_str)?;
110            state_overrides_builder =
111                state_overrides_builder.with_nonce(addr.parse()?, nonce.parse()?);
112        }
113
114        for override_str in self.code_overrides.iter().flatten() {
115            let (addr, code_str) = address_value_override(override_str)?;
116            state_overrides_builder =
117                state_overrides_builder.with_code(addr.parse()?, Bytes::from_str(code_str)?);
118        }
119
120        type StateOverrides = HashMap<Address, HashMap<B256, B256>>;
121        let parse_state_overrides = |overrides: &Option<Vec<String>>| -> Result<StateOverrides> {
122            let mut state_overrides = StateOverrides::default();
123
124            overrides.iter().flatten().try_for_each(|s| -> Result<()> {
125                let (addr, slot, value) = address_slot_value_override(s)?;
126                state_overrides.entry(addr).or_default().insert(slot.into(), value.into());
127                Ok(())
128            })?;
129
130            Ok(state_overrides)
131        };
132
133        for (addr, entries) in parse_state_overrides(&self.state_overrides)? {
134            state_overrides_builder = state_overrides_builder.with_state(addr, entries);
135        }
136
137        for (addr, entries) in parse_state_overrides(&self.state_diff_overrides)? {
138            state_overrides_builder = state_overrides_builder.with_state_diff(addr, entries)
139        }
140
141        Ok(Some(state_overrides_builder.build()))
142    }
143
144    /// Parses block overrides from command line arguments.
145    pub fn get_block_overrides(&self) -> Result<Option<BlockOverrides>> {
146        let mut overrides = BlockOverrides::default();
147        if let Some(number) = self.block_number {
148            overrides = overrides.with_number(U256::from(number));
149        }
150        if let Some(time) = self.block_time {
151            overrides = overrides.with_time(time);
152        }
153        if overrides.is_empty() { Ok(None) } else { Ok(Some(overrides)) }
154    }
155}
156
157/// Parses an override string in the format address:value.
158fn address_value_override(address_override: &str) -> Result<(&str, &str)> {
159    address_override.split_once(':').ok_or_else(|| {
160        eyre::eyre!("Invalid override {address_override}. Expected <address>:<value>")
161    })
162}
163
164/// Parses an override string in the format address:slot:value.
165fn address_slot_value_override(address_override: &str) -> Result<(Address, U256, U256)> {
166    let captures = OVERRIDE_PATTERN.captures(address_override).ok_or_else(|| {
167        eyre::eyre!("Invalid override {address_override}. Expected <address>:<slot>:<value>")
168    })?;
169
170    Ok((captures[1].parse()?, captures[2].parse()?, captures[3].parse()?))
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use alloy_primitives::{address, b256, fixed_bytes};
177    use clap::Parser;
178
179    #[derive(Debug, Parser)]
180    struct TestArgs {
181        #[command(flatten)]
182        overrides: CallOverrideOpts,
183    }
184
185    #[test]
186    fn test_get_state_overrides() {
187        let args = TestArgs::parse_from([
188            "foundry-cli",
189            "--override-balance",
190            "0x0000000000000000000000000000000000000001:2",
191            "--override-nonce",
192            "0x0000000000000000000000000000000000000001:3",
193            "--override-code",
194            "0x0000000000000000000000000000000000000001:0x04",
195            "--override-state",
196            "0x0000000000000000000000000000000000000001:5:6",
197            "--override-state-diff",
198            "0x0000000000000000000000000000000000000001:7:8",
199        ]);
200        let overrides = args.overrides.get_state_overrides().unwrap().unwrap();
201        let address = address!("0x0000000000000000000000000000000000000001");
202        let account = overrides.get(&address).unwrap();
203
204        assert_eq!(account.balance, Some(U256::from(2)));
205        assert_eq!(account.nonce, Some(3));
206        assert_eq!(account.code, Some(Bytes::from([0x04])));
207        assert_eq!(
208            account
209                .state
210                .as_ref()
211                .unwrap()
212                .get(&b256!("0x0000000000000000000000000000000000000000000000000000000000000005")),
213            Some(&b256!("0x0000000000000000000000000000000000000000000000000000000000000006"))
214        );
215        assert_eq!(
216            account
217                .state_diff
218                .as_ref()
219                .unwrap()
220                .get(&b256!("0x0000000000000000000000000000000000000000000000000000000000000007")),
221            Some(&b256!("0x0000000000000000000000000000000000000000000000000000000000000008"))
222        );
223    }
224
225    #[test]
226    fn test_get_state_overrides_empty() {
227        let args = TestArgs::parse_from([""]);
228        assert_eq!(args.overrides.get_state_overrides().unwrap(), None);
229    }
230
231    #[test]
232    fn test_get_block_overrides() {
233        let args =
234            TestArgs::parse_from(["foundry-cli", "--block.number", "1", "--block.time", "2"]);
235        let overrides = args.overrides.get_block_overrides().unwrap().unwrap();
236        assert_eq!(overrides.number, Some(U256::from(1)));
237        assert_eq!(overrides.time, Some(2));
238    }
239
240    #[test]
241    fn test_get_block_overrides_empty() {
242        let args = TestArgs::parse_from([""]);
243        assert_eq!(args.overrides.get_block_overrides().unwrap(), None);
244    }
245
246    #[test]
247    fn test_address_value_override_success() {
248        let text = "0x0000000000000000000000000000000000000001:2";
249        let (address, value) = address_value_override(text).unwrap();
250        assert_eq!(address, "0x0000000000000000000000000000000000000001");
251        assert_eq!(value, "2");
252    }
253
254    #[test]
255    fn test_address_value_override_error() {
256        let text = "invalid_value";
257        let error = address_value_override(text).unwrap_err();
258        assert_eq!(error.to_string(), "Invalid override invalid_value. Expected <address>:<value>");
259    }
260
261    #[test]
262    fn test_address_slot_value_override_success() {
263        let text = "0x0000000000000000000000000000000000000001:2:3";
264        let (address, slot, value) = address_slot_value_override(text).unwrap();
265        assert_eq!(*address, fixed_bytes!("0x0000000000000000000000000000000000000001"));
266        assert_eq!(slot, U256::from(2));
267        assert_eq!(value, U256::from(3));
268    }
269
270    #[test]
271    fn test_address_slot_value_override_error() {
272        let text = "invalid_value";
273        let error = address_slot_value_override(text).unwrap_err();
274        assert_eq!(
275            error.to_string(),
276            "Invalid override invalid_value. Expected <address>:<slot>:<value>"
277        );
278    }
279}