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

foundry_cheatcodes/
script.rs
