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, Uint};
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::{Cfg, ContextTr, JournalTr, Transaction},
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::from(0) } else { U256::from(ccx.ecx.cfg().chain_id()) };
108
109    let auth = Authorization { address: *implementation, nonce: *nonce, chain_id };
110    let signed_auth = SignedAuthorization::new_unchecked(
111        auth,
112        *v,
113        U256::from_be_bytes(r.0),
114        U256::from_be_bytes(s.0),
115    );
116    write_delegation(ccx, signed_auth.clone())?;
117    ccx.state.add_delegation(signed_auth);
118    Ok(Default::default())
119}
120
121/// Helper function to sign and attach (if needed) an EIP-7702 delegation.
122/// Uses the provided nonce, otherwise retrieves and increments the nonce of the EOA.
123fn sign_delegation<FEN: FoundryEvmNetwork>(
124    ccx: &mut CheatsCtxt<'_, '_, FEN>,
125    private_key: Uint<256, 4>,
126    implementation: Address,
127    nonce: Option<u64>,
128    cross_chain: bool,
129    attach: bool,
130) -> Result<Vec<u8>> {
131    let signer = PrivateKeySigner::from_bytes(&B256::from(private_key))?;
132    let nonce = if let Some(nonce) = nonce {
133        nonce
134    } else {
135        let account_nonce = {
136            let authority_acc = ccx.ecx.journal_mut().load_account(signer.address())?;
137            authority_acc.info.nonce
138        };
139        // Calculate next nonce considering existing active delegations
140        next_delegation_nonce(
141            &ccx.state.active_delegations,
142            signer.address(),
143            &ccx.state.broadcast,
144            account_nonce,
145        )
146    };
147    let chain_id = if cross_chain { U256::from(0) } else { U256::from(ccx.ecx.cfg().chain_id()) };
148
149    let auth = Authorization { address: implementation, nonce, chain_id };
150    let sig = signer.sign_hash_sync(&auth.signature_hash())?;
151    // Attach delegation.
152    if attach {
153        let signed_auth = SignedAuthorization::new_unchecked(auth, sig.v() as u8, sig.r(), sig.s());
154        write_delegation(ccx, signed_auth.clone())?;
155        ccx.state.add_delegation(signed_auth);
156    }
157    Ok(SignedDelegation {
158        v: sig.v() as u8,
159        r: sig.r().into(),
160        s: sig.s().into(),
161        nonce,
162        implementation,
163    }
164    .abi_encode())
165}
166
167/// Returns the next valid nonce for a delegation, considering existing active delegations.
168fn next_delegation_nonce(
169    active_delegations: &[SignedAuthorization],
170    authority: Address,
171    broadcast: &Option<Broadcast>,
172    account_nonce: u64,
173) -> u64 {
174    match active_delegations
175        .iter()
176        .rfind(|auth| auth.recover_authority().is_ok_and(|recovered| recovered == authority))
177    {
178        Some(auth) => {
179            // Increment nonce of last recorded delegation.
180            auth.nonce + 1
181        }
182        None => {
183            // First time a delegation is added for this authority.
184            if let Some(broadcast) = broadcast {
185                // Increment nonce if authority is the sender of transaction.
186                if broadcast.new_origin == authority {
187                    return account_nonce + 1;
188                }
189            }
190            // Return current nonce if authority is not the sender of transaction.
191            account_nonce
192        }
193    }
194}
195
196fn write_delegation<FEN: FoundryEvmNetwork>(
197    ccx: &mut CheatsCtxt<'_, '_, FEN>,
198    auth: SignedAuthorization,
199) -> Result<()> {
200    let authority = auth.recover_authority().map_err(|e| format!("{e}"))?;
201    let account_nonce = {
202        let authority_acc = ccx.ecx.journal_mut().load_account(authority)?;
203        authority_acc.info.nonce
204    };
205
206    let expected_nonce = next_delegation_nonce(
207        &ccx.state.active_delegations,
208        authority,
209        &ccx.state.broadcast,
210        account_nonce,
211    );
212
213    if expected_nonce != auth.nonce {
214        return Err(format!(
215            "invalid nonce for {authority:?}: expected {expected_nonce}, got {}",
216            auth.nonce
217        )
218        .into());
219    }
220
221    if auth.address.is_zero() {
222        // Set empty code if the delegation address of authority is 0x.
223        // See https://github.com/ethereum/EIPs/blob/master/EIPS/eip-7702.md#behavior.
224        ccx.ecx.journal_mut().set_code_with_hash(authority, Bytecode::default(), KECCAK_EMPTY);
225    } else {
226        let bytecode = Bytecode::new_eip7702(*auth.address());
227        ccx.ecx.journal_mut().set_code(authority, bytecode);
228    }
229    Ok(())
230}
231
232impl Cheatcode for attachBlobCall {
233    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
234        let Self { blob } = self;
235        ensure!(
236            (*ccx.ecx.cfg().spec()).into() >= SpecId::CANCUN,
237            "`attachBlob` is not supported before the Cancun hard fork; \
238             see EIP-4844: https://eips.ethereum.org/EIPS/eip-4844"
239        );
240        let sidecar: SidecarBuilder<SimpleCoder> = SidecarBuilder::from_slice(blob);
241        let sidecar_variant = if (*ccx.ecx.cfg().spec()).into() < SpecId::OSAKA {
242            sidecar.build_4844().map_err(|e| format!("{e}"))?.into()
243        } else {
244            sidecar.build_7594().map_err(|e| format!("{e}"))?.into()
245        };
246        ccx.state.active_blob_sidecar = Some(sidecar_variant);
247        Ok(Default::default())
248    }
249}
250
251impl Cheatcode for startBroadcast_0Call {
252    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
253        let Self {} = self;
254        broadcast(ccx, None, false)
255    }
256}
257
258impl Cheatcode for startBroadcast_1Call {
259    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
260        let Self { signer } = self;
261        broadcast(ccx, Some(signer), false)
262    }
263}
264
265impl Cheatcode for startBroadcast_2Call {
266    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
267        let Self { privateKey } = self;
268        broadcast_key(ccx, privateKey, false)
269    }
270}
271
272impl Cheatcode for stopBroadcastCall {
273    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
274        let Self {} = self;
275        let Some(broadcast) = ccx.state.broadcast.take() else {
276            bail!("no broadcast in progress to stop");
277        };
278        debug!(target: "cheatcodes", ?broadcast, "stopped");
279        Ok(Default::default())
280    }
281}
282
283impl Cheatcode for getWalletsCall {
284    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
285        let wallets = ccx.state.wallets().signers().unwrap_or_default();
286        Ok(wallets.abi_encode())
287    }
288}
289
290#[derive(Clone, Debug, Default)]
291pub struct Broadcast {
292    /// Address of the transaction origin
293    pub new_origin: Address,
294    /// Original caller
295    pub original_caller: Address,
296    /// Original `tx.origin`
297    pub original_origin: Address,
298    /// Depth of the broadcast
299    pub depth: usize,
300    /// Whether the prank stops by itself after the next call
301    pub single_call: bool,
302    /// Whether `vm.deployCode` cheatcode is used to deploy from code.
303    pub deploy_from_code: bool,
304}
305
306/// Contains context for wallet management.
307#[derive(Debug)]
308pub struct WalletsInner {
309    /// All signers in scope of the script.
310    pub multi_wallet: MultiWallet,
311    /// Optional signer provided as `--sender` flag.
312    pub provided_sender: Option<Address>,
313}
314
315/// Cloneable wrapper around [`WalletsInner`].
316#[derive(Debug, Clone)]
317pub struct Wallets {
318    /// Inner data.
319    pub inner: Arc<Mutex<WalletsInner>>,
320}
321
322impl Wallets {
323    #[expect(missing_docs)]
324    pub fn new(multi_wallet: MultiWallet, provided_sender: Option<Address>) -> Self {
325        Self { inner: Arc::new(Mutex::new(WalletsInner { multi_wallet, provided_sender })) }
326    }
327
328    /// Consumes [Wallets] and returns [MultiWallet].
329    ///
330    /// Panics if [Wallets] is still in use.
331    pub fn into_multi_wallet(self) -> MultiWallet {
332        Arc::into_inner(self.inner)
333            .map(|m| m.into_inner().multi_wallet)
334            .unwrap_or_else(|| panic!("not all instances were dropped"))
335    }
336
337    /// Locks inner Mutex and adds a signer to the [MultiWallet].
338    pub fn add_local_signer(&self, wallet: PrivateKeySigner) {
339        self.inner.lock().multi_wallet.add_signer(WalletSigner::Local(wallet));
340    }
341
342    /// Locks inner Mutex and returns all signer addresses in the [MultiWallet].
343    pub fn signers(&self) -> Result<Vec<Address>> {
344        Ok(self.inner.lock().multi_wallet.signers()?.keys().copied().collect())
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.ecx.journal().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.ecx.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.ecx.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
