foundry_cheatcodes/
script.rs

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

foundry_cheatcodes/
script.rs
