1use 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
15static OVERRIDE_PATTERN: LazyLock<Regex> =
18 LazyLock::new(|| Regex::new(r"^([^:]+):([^:]+):([^:]+)$").unwrap());
19
20#[derive(Args, Clone, Debug, Default)]
22pub struct CallOverrideOpts {
23 #[arg(long = "override-balance", value_name = "ADDRESS:BALANCE", value_delimiter = ',')]
26 pub balance_overrides: Option<Vec<String>>,
27
28 #[arg(long = "override-nonce", value_name = "ADDRESS:NONCE", value_delimiter = ',')]
31 pub nonce_overrides: Option<Vec<String>>,
32
33 #[arg(long = "override-code", value_name = "ADDRESS:CODE", value_delimiter = ',')]
36 pub code_overrides: Option<Vec<String>>,
37
38 #[arg(long = "override-state", value_name = "ADDRESS:SLOT:VALUE", value_delimiter = ',')]
41 pub state_overrides: Option<Vec<String>>,
42
43 #[arg(long = "override-state-diff", value_name = "ADDRESS:SLOT:VALUE", value_delimiter = ',')]
46 pub state_diff_overrides: Option<Vec<String>>,
47
48 #[arg(long = "block.time", value_name = "TIME")]
50 pub block_time: Option<u64>,
51
52 #[arg(long = "block.number", value_name = "NUMBER")]
54 pub block_number: Option<u64>,
55}
56
57impl CallOverrideOpts {
58 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 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 pub fn get_state_overrides(&self) -> Result<Option<StateOverride>> {
86 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 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
157fn 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
164fn 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}