foundry_cheatcodes/
script.rs

Help
1//! Implementations of [`Scripting`](spec::Group::Scripting) cheatcodes.
2
3use crate::{Cheatcode, CheatsCtxt, Result, Vm::*};
4use alloy_consensus::{SidecarBuilder, SimpleCoder};
5use alloy_primitives::{Address, Uint, B256, U256};
6use alloy_rpc_types::Authorization;
7use alloy_signer::SignerSync;
8use alloy_signer_local::PrivateKeySigner;
9use alloy_sol_types::SolValue;
10use foundry_wallets::{multi_wallet::MultiWallet, WalletSigner};
11use parking_lot::Mutex;
12use revm::{
13    bytecode::Bytecode,
14    context::JournalTr,
15    context_interface::transaction::SignedAuthorization,
16    primitives::{hardfork::SpecId, KECCAK_EMPTY},
17};
18use std::sync::Arc;
19
20impl Cheatcode for broadcast_0Call {
21    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
22        let Self {} = self;
23        broadcast(ccx, None, true)
24    }
25}
26
27impl Cheatcode for broadcast_1Call {
28    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
29        let Self { signer } = self;
30        broadcast(ccx, Some(signer), true)
31    }
32}
33
34impl Cheatcode for broadcast_2Call {
35    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
36        let Self { privateKey } = self;
37        broadcast_key(ccx, privateKey, true)
38    }
39}
40
41impl Cheatcode for attachDelegation_0Call {
42    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
43        let Self { signedDelegation } = self;
44        attach_delegation(ccx, signedDelegation, false)
45    }
46}
47
48impl Cheatcode for attachDelegation_1Call {
49    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
50        let Self { signedDelegation, crossChain } = self;
51        attach_delegation(ccx, signedDelegation, *crossChain)
52    }
53}
54
55impl Cheatcode for signDelegation_0Call {
56    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
57        let Self { implementation, privateKey } = *self;
58        sign_delegation(ccx, privateKey, implementation, None, false, false)
59    }
60}
61
62impl Cheatcode for signDelegation_1Call {
63    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
64        let Self { implementation, privateKey, nonce } = *self;
65        sign_delegation(ccx, privateKey, implementation, Some(nonce), false, false)
66    }
67}
68
69impl Cheatcode for signDelegation_2Call {
70    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
71        let Self { implementation, privateKey, crossChain } = *self;
72        sign_delegation(ccx, privateKey, implementation, None, crossChain, false)
73    }
74}
75
76impl Cheatcode for signAndAttachDelegation_0Call {
77    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
78        let Self { implementation, privateKey } = *self;
79        sign_delegation(ccx, privateKey, implementation, None, false, true)
80    }
81}
82
83impl Cheatcode for signAndAttachDelegation_1Call {
84    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
85        let Self { implementation, privateKey, nonce } = *self;
86        sign_delegation(ccx, privateKey, implementation, Some(nonce), false, true)
87    }
88}
89
90impl Cheatcode for signAndAttachDelegation_2Call {
91    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
92        let Self { implementation, privateKey, crossChain } = *self;
93        sign_delegation(ccx, privateKey, implementation, None, crossChain, true)
94    }
95}
96
97/// Helper function to attach an EIP-7702 delegation.
98fn attach_delegation(
99    ccx: &mut CheatsCtxt,
100    delegation: &SignedDelegation,
101    cross_chain: bool,
102) -> Result {
103    let SignedDelegation { v, r, s, nonce, implementation } = delegation;
104    // Set chain id to 0 if universal deployment is preferred.
105    // See https://github.com/ethereum/EIPs/blob/master/EIPS/eip-7702.md#protection-from-malleability-cross-chain
106    let chain_id = if cross_chain { U256::from(0) } else { U256::from(ccx.ecx.cfg.chain_id) };
107
108    let auth = Authorization { address: *implementation, nonce: *nonce, chain_id };
109    let signed_auth = SignedAuthorization::new_unchecked(
110        auth,
111        *v,
112        U256::from_be_bytes(r.0),
113        U256::from_be_bytes(s.0),
114    );
115    write_delegation(ccx, signed_auth.clone())?;
116    ccx.state.active_delegation = Some(signed_auth);
117    Ok(Default::default())
118}
119
120/// Helper function to sign and attach (if needed) an EIP-7702 delegation.
121/// Uses the provided nonce, otherwise retrieves and increments the nonce of the EOA.
122fn sign_delegation(
123    ccx: &mut CheatsCtxt,
124    private_key: Uint<256, 4>,
125    implementation: Address,
126    nonce: Option<u64>,
127    cross_chain: bool,
128    attach: bool,
129) -> Result<Vec<u8>> {
130    let signer = PrivateKeySigner::from_bytes(&B256::from(private_key))?;
131    let nonce = if let Some(nonce) = nonce {
132        nonce
133    } else {
134        let authority_acc = ccx.ecx.journaled_state.load_account(signer.address())?;
135        // If we don't have a nonce then use next auth account nonce.
136        authority_acc.data.info.nonce + 1
137    };
138    let chain_id = if cross_chain { U256::from(0) } else { U256::from(ccx.ecx.cfg.chain_id) };
139
140    let auth = Authorization { address: implementation, nonce, chain_id };
141    let sig = signer.sign_hash_sync(&auth.signature_hash())?;
142    // Attach delegation.
143    if attach {
144        let signed_auth = SignedAuthorization::new_unchecked(auth, sig.v() as u8, sig.r(), sig.s());
145        write_delegation(ccx, signed_auth.clone())?;
146        ccx.state.active_delegation = Some(signed_auth);
147    }
148    Ok(SignedDelegation {
149        v: sig.v() as u8,
150        r: sig.r().into(),
151        s: sig.s().into(),
152        nonce,
153        implementation,
154    }
155    .abi_encode())
156}
157
158fn write_delegation(ccx: &mut CheatsCtxt, auth: SignedAuthorization) -> Result<()> {
159    let authority = auth.recover_authority().map_err(|e| format!("{e}"))?;
160    let authority_acc = ccx.ecx.journaled_state.load_account(authority)?;
161    if authority_acc.data.info.nonce + 1 != auth.nonce {
162        return Err("invalid nonce".into());
163    }
164
165    if auth.address.is_zero() {
166        // Set empty code if the delegation address of authority is 0x.
167        // See https://github.com/ethereum/EIPs/blob/master/EIPS/eip-7702.md#behavior.
168        ccx.ecx.journaled_state.set_code_with_hash(authority, Bytecode::default(), KECCAK_EMPTY);
169    } else {
170        let bytecode = Bytecode::new_eip7702(*auth.address());
171        ccx.ecx.journaled_state.set_code(authority, bytecode);
172    }
173    Ok(())
174}
175
176impl Cheatcode for attachBlobCall {
177    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
178        let Self { blob } = self;
179        ensure!(
180            ccx.ecx.cfg.spec >= SpecId::CANCUN,
181            "`attachBlob` is not supported before the Cancun hard fork; \
182             see EIP-4844: https://eips.ethereum.org/EIPS/eip-4844"
183        );
184        let sidecar: SidecarBuilder<SimpleCoder> = SidecarBuilder::from_slice(blob);
185        let sidecar = sidecar.build().map_err(|e| format!("{e}"))?;
186        ccx.state.active_blob_sidecar = Some(sidecar);
187        Ok(Default::default())
188    }
189}
190
191impl Cheatcode for startBroadcast_0Call {
192    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
193        let Self {} = self;
194        broadcast(ccx, None, false)
195    }
196}
197
198impl Cheatcode for startBroadcast_1Call {
199    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
200        let Self { signer } = self;
201        broadcast(ccx, Some(signer), false)
202    }
203}
204
205impl Cheatcode for startBroadcast_2Call {
206    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
207        let Self { privateKey } = self;
208        broadcast_key(ccx, privateKey, false)
209    }
210}
211
212impl Cheatcode for stopBroadcastCall {
213    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
214        let Self {} = self;
215        let Some(broadcast) = ccx.state.broadcast.take() else {
216            bail!("no broadcast in progress to stop");
217        };
218        debug!(target: "cheatcodes", ?broadcast, "stopped");
219        Ok(Default::default())
220    }
221}
222
223impl Cheatcode for getWalletsCall {
224    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
225        let wallets = ccx.state.wallets().signers().unwrap_or_default();
226        Ok(wallets.abi_encode())
227    }
228}
229
230#[derive(Clone, Debug, Default)]
231pub struct Broadcast {
232    /// Address of the transaction origin
233    pub new_origin: Address,
234    /// Original caller
235    pub original_caller: Address,
236    /// Original `tx.origin`
237    pub original_origin: Address,
238    /// Depth of the broadcast
239    pub depth: usize,
240    /// Whether the prank stops by itself after the next call
241    pub single_call: bool,
242}
243
244/// Contains context for wallet management.
245#[derive(Debug)]
246pub struct WalletsInner {
247    /// All signers in scope of the script.
248    pub multi_wallet: MultiWallet,
249    /// Optional signer provided as `--sender` flag.
250    pub provided_sender: Option<Address>,
251}
252
253/// Clonable wrapper around [`WalletsInner`].
254#[derive(Debug, Clone)]
255pub struct Wallets {
256    /// Inner data.
257    pub inner: Arc<Mutex<WalletsInner>>,
258}
259
260impl Wallets {
261    #[expect(missing_docs)]
262    pub fn new(multi_wallet: MultiWallet, provided_sender: Option<Address>) -> Self {
263        Self { inner: Arc::new(Mutex::new(WalletsInner { multi_wallet, provided_sender })) }
264    }
265
266    /// Consumes [Wallets] and returns [MultiWallet].
267    ///
268    /// Panics if [Wallets] is still in use.
269    pub fn into_multi_wallet(self) -> MultiWallet {
270        Arc::into_inner(self.inner)
271            .map(|m| m.into_inner().multi_wallet)
272            .unwrap_or_else(|| panic!("not all instances were dropped"))
273    }
274
275    /// Locks inner Mutex and adds a signer to the [MultiWallet].
276    pub fn add_private_key(&self, private_key: &B256) -> Result<()> {
277        self.add_local_signer(PrivateKeySigner::from_bytes(private_key)?);
278        Ok(())
279    }
280
281    /// Locks inner Mutex and adds a signer to the [MultiWallet].
282    pub fn add_local_signer(&self, wallet: PrivateKeySigner) {
283        self.inner.lock().multi_wallet.add_signer(WalletSigner::Local(wallet));
284    }
285
286    /// Locks inner Mutex and returns all signer addresses in the [MultiWallet].
287    pub fn signers(&self) -> Result<Vec<Address>> {
288        Ok(self.inner.lock().multi_wallet.signers()?.keys().cloned().collect())
289    }
290
291    /// Number of signers in the [MultiWallet].
292    pub fn len(&self) -> usize {
293        let mut inner = self.inner.lock();
294        let signers = inner.multi_wallet.signers();
295        if signers.is_err() {
296            return 0;
297        }
298        signers.unwrap().len()
299    }
300
301    /// Whether the [MultiWallet] is empty.
302    pub fn is_empty(&self) -> bool {
303        self.len() == 0
304    }
305}
306
307/// Sets up broadcasting from a script using `new_origin` as the sender.
308fn broadcast(ccx: &mut CheatsCtxt, new_origin: Option<&Address>, single_call: bool) -> Result {
309    let depth = ccx.ecx.journaled_state.depth();
310    ensure!(
311        ccx.state.get_prank(depth).is_none(),
312        "you have an active prank; broadcasting and pranks are not compatible"
313    );
314    ensure!(ccx.state.broadcast.is_none(), "a broadcast is active already");
315
316    let mut new_origin = new_origin.cloned();
317
318    if new_origin.is_none() {
319        let mut wallets = ccx.state.wallets().inner.lock();
320        if let Some(provided_sender) = wallets.provided_sender {
321            new_origin = Some(provided_sender);
322        } else {
323            let signers = wallets.multi_wallet.signers()?;
324            if signers.len() == 1 {
325                let address = signers.keys().next().unwrap();
326                new_origin = Some(*address);
327            }
328        }
329    }
330
331    let broadcast = Broadcast {
332        new_origin: new_origin.unwrap_or(ccx.ecx.tx.caller),
333        original_caller: ccx.caller,
334        original_origin: ccx.ecx.tx.caller,
335        depth,
336        single_call,
337    };
338    debug!(target: "cheatcodes", ?broadcast, "started");
339    ccx.state.broadcast = Some(broadcast);
340    Ok(Default::default())
341}
342
343/// Sets up broadcasting from a script with the sender derived from `private_key`.
344/// Adds this private key to `state`'s `wallets` vector to later be used for signing
345/// if broadcast is successful.
346fn broadcast_key(ccx: &mut CheatsCtxt, private_key: &U256, single_call: bool) -> Result {
347    let wallet = super::crypto::parse_wallet(private_key)?;
348    let new_origin = wallet.address();
349
350    let result = broadcast(ccx, Some(&new_origin), single_call);
351    if result.is_ok() {
352        let wallets = ccx.state.wallets();
353        wallets.add_local_signer(wallet);
354    }
355    result
356}
foundry_cheatcodes/script.rs

foundry_cheatcodes/
script.rs
