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, B256, U256, Uint};
6use alloy_rpc_types::Authorization;
7use alloy_signer::SignerSync;
8use alloy_signer_local::PrivateKeySigner;
9use alloy_sol_types::SolValue;
10use foundry_wallets::{WalletSigner, multi_wallet::MultiWallet};
11use parking_lot::Mutex;
12use revm::{
13    bytecode::Bytecode,
14    context::JournalTr,
15    context_interface::transaction::SignedAuthorization,
16    primitives::{KECCAK_EMPTY, hardfork::SpecId},
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.add_delegation(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        // Calculate next nonce considering existing active delegations
136        next_delegation_nonce(
137            &ccx.state.active_delegations,
138            signer.address(),
139            &ccx.state.broadcast,
140            authority_acc.data.info.nonce,
141        )
142    };
143    let chain_id = if cross_chain { U256::from(0) } else { U256::from(ccx.ecx.cfg.chain_id) };
144
145    let auth = Authorization { address: implementation, nonce, chain_id };
146    let sig = signer.sign_hash_sync(&auth.signature_hash())?;
147    // Attach delegation.
148    if attach {
149        let signed_auth = SignedAuthorization::new_unchecked(auth, sig.v() as u8, sig.r(), sig.s());
150        write_delegation(ccx, signed_auth.clone())?;
151        ccx.state.add_delegation(signed_auth);
152    }
153    Ok(SignedDelegation {
154        v: sig.v() as u8,
155        r: sig.r().into(),
156        s: sig.s().into(),
157        nonce,
158        implementation,
159    }
160    .abi_encode())
161}
162
163/// Returns the next valid nonce for a delegation, considering existing active delegations.
164fn next_delegation_nonce(
165    active_delegations: &[SignedAuthorization],
166    authority: Address,
167    broadcast: &Option<Broadcast>,
168    account_nonce: u64,
169) -> u64 {
170    match active_delegations
171        .iter()
172        .rfind(|auth| auth.recover_authority().is_ok_and(|recovered| recovered == authority))
173    {
174        Some(auth) => {
175            // Increment nonce of last recorded delegation.
176            auth.nonce + 1
177        }
178        None => {
179            // First time a delegation is added for this authority.
180            if let Some(broadcast) = broadcast {
181                // Increment nonce if authority is the sender of transaction.
182                if broadcast.new_origin == authority {
183                    return account_nonce + 1;
184                }
185            }
186            // Return current nonce if authority is not the sender of transaction.
187            account_nonce
188        }
189    }
190}
191
192fn write_delegation(ccx: &mut CheatsCtxt, auth: SignedAuthorization) -> Result<()> {
193    let authority = auth.recover_authority().map_err(|e| format!("{e}"))?;
194    let authority_acc = ccx.ecx.journaled_state.load_account(authority)?;
195
196    let expected_nonce = next_delegation_nonce(
197        &ccx.state.active_delegations,
198        authority,
199        &ccx.state.broadcast,
200        authority_acc.data.info.nonce,
201    );
202
203    if expected_nonce != auth.nonce {
204        return Err(format!(
205            "invalid nonce for {authority:?}: expected {expected_nonce}, got {}",
206            auth.nonce
207        )
208        .into());
209    }
210
211    if auth.address.is_zero() {
212        // Set empty code if the delegation address of authority is 0x.
213        // See https://github.com/ethereum/EIPs/blob/master/EIPS/eip-7702.md#behavior.
214        ccx.ecx.journaled_state.set_code_with_hash(authority, Bytecode::default(), KECCAK_EMPTY);
215    } else {
216        let bytecode = Bytecode::new_eip7702(*auth.address());
217        ccx.ecx.journaled_state.set_code(authority, bytecode);
218    }
219    Ok(())
220}
221
222impl Cheatcode for attachBlobCall {
223    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
224        let Self { blob } = self;
225        ensure!(
226            ccx.ecx.cfg.spec >= SpecId::CANCUN,
227            "`attachBlob` is not supported before the Cancun hard fork; \
228             see EIP-4844: https://eips.ethereum.org/EIPS/eip-4844"
229        );
230        let sidecar: SidecarBuilder<SimpleCoder> = SidecarBuilder::from_slice(blob);
231        let sidecar = sidecar.build().map_err(|e| format!("{e}"))?;
232        ccx.state.active_blob_sidecar = Some(sidecar);
233        Ok(Default::default())
234    }
235}
236
237impl Cheatcode for startBroadcast_0Call {
238    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
239        let Self {} = self;
240        broadcast(ccx, None, false)
241    }
242}
243
244impl Cheatcode for startBroadcast_1Call {
245    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
246        let Self { signer } = self;
247        broadcast(ccx, Some(signer), false)
248    }
249}
250
251impl Cheatcode for startBroadcast_2Call {
252    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
253        let Self { privateKey } = self;
254        broadcast_key(ccx, privateKey, false)
255    }
256}
257
258impl Cheatcode for stopBroadcastCall {
259    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
260        let Self {} = self;
261        let Some(broadcast) = ccx.state.broadcast.take() else {
262            bail!("no broadcast in progress to stop");
263        };
264        debug!(target: "cheatcodes", ?broadcast, "stopped");
265        Ok(Default::default())
266    }
267}
268
269impl Cheatcode for getWalletsCall {
270    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
271        let wallets = ccx.state.wallets().signers().unwrap_or_default();
272        Ok(wallets.abi_encode())
273    }
274}
275
276#[derive(Clone, Debug, Default)]
277pub struct Broadcast {
278    /// Address of the transaction origin
279    pub new_origin: Address,
280    /// Original caller
281    pub original_caller: Address,
282    /// Original `tx.origin`
283    pub original_origin: Address,
284    /// Depth of the broadcast
285    pub depth: usize,
286    /// Whether the prank stops by itself after the next call
287    pub single_call: bool,
288}
289
290/// Contains context for wallet management.
291#[derive(Debug)]
292pub struct WalletsInner {
293    /// All signers in scope of the script.
294    pub multi_wallet: MultiWallet,
295    /// Optional signer provided as `--sender` flag.
296    pub provided_sender: Option<Address>,
297}
298
299/// Cloneable wrapper around [`WalletsInner`].
300#[derive(Debug, Clone)]
301pub struct Wallets {
302    /// Inner data.
303    pub inner: Arc<Mutex<WalletsInner>>,
304}
305
306impl Wallets {
307    #[expect(missing_docs)]
308    pub fn new(multi_wallet: MultiWallet, provided_sender: Option<Address>) -> Self {
309        Self { inner: Arc::new(Mutex::new(WalletsInner { multi_wallet, provided_sender })) }
310    }
311
312    /// Consumes [Wallets] and returns [MultiWallet].
313    ///
314    /// Panics if [Wallets] is still in use.
315    pub fn into_multi_wallet(self) -> MultiWallet {
316        Arc::into_inner(self.inner)
317            .map(|m| m.into_inner().multi_wallet)
318            .unwrap_or_else(|| panic!("not all instances were dropped"))
319    }
320
321    /// Locks inner Mutex and adds a signer to the [MultiWallet].
322    pub fn add_private_key(&self, private_key: &B256) -> Result<()> {
323        self.add_local_signer(PrivateKeySigner::from_bytes(private_key)?);
324        Ok(())
325    }
326
327    /// Locks inner Mutex and adds a signer to the [MultiWallet].
328    pub fn add_local_signer(&self, wallet: PrivateKeySigner) {
329        self.inner.lock().multi_wallet.add_signer(WalletSigner::Local(wallet));
330    }
331
332    /// Locks inner Mutex and returns all signer addresses in the [MultiWallet].
333    pub fn signers(&self) -> Result<Vec<Address>> {
334        Ok(self.inner.lock().multi_wallet.signers()?.keys().copied().collect())
335    }
336
337    /// Number of signers in the [MultiWallet].
338    pub fn len(&self) -> usize {
339        let mut inner = self.inner.lock();
340        let signers = inner.multi_wallet.signers();
341        if signers.is_err() {
342            return 0;
343        }
344        signers.unwrap().len()
345    }
346
347    /// Whether the [MultiWallet] is empty.
348    pub fn is_empty(&self) -> bool {
349        self.len() == 0
350    }
351}
352
353/// Sets up broadcasting from a script using `new_origin` as the sender.
354fn broadcast(ccx: &mut CheatsCtxt, new_origin: Option<&Address>, single_call: bool) -> Result {
355    let depth = ccx.ecx.journaled_state.depth();
356    ensure!(
357        ccx.state.get_prank(depth).is_none(),
358        "you have an active prank; broadcasting and pranks are not compatible"
359    );
360    ensure!(ccx.state.broadcast.is_none(), "a broadcast is active already");
361
362    let mut new_origin = new_origin.copied();
363
364    if new_origin.is_none() {
365        let mut wallets = ccx.state.wallets().inner.lock();
366        if let Some(provided_sender) = wallets.provided_sender {
367            new_origin = Some(provided_sender);
368        } else {
369            let signers = wallets.multi_wallet.signers()?;
370            if signers.len() == 1 {
371                let address = signers.keys().next().unwrap();
372                new_origin = Some(*address);
373            }
374        }
375    }
376
377    let broadcast = Broadcast {
378        new_origin: new_origin.unwrap_or(ccx.ecx.tx.caller),
379        original_caller: ccx.caller,
380        original_origin: ccx.ecx.tx.caller,
381        depth,
382        single_call,
383    };
384    debug!(target: "cheatcodes", ?broadcast, "started");
385    ccx.state.broadcast = Some(broadcast);
386    Ok(Default::default())
387}
388
389/// Sets up broadcasting from a script with the sender derived from `private_key`.
390/// Adds this private key to `state`'s `wallets` vector to later be used for signing
391/// if broadcast is successful.
392fn broadcast_key(ccx: &mut CheatsCtxt, private_key: &U256, single_call: bool) -> Result {
393    let wallet = super::crypto::parse_wallet(private_key)?;
394    let new_origin = wallet.address();
395
396    let result = broadcast(ccx, Some(&new_origin), single_call);
397    if result.is_ok() {
398        let wallets = ccx.state.wallets();
399        wallets.add_local_signer(wallet);
400    }
401    result
402}
foundry_cheatcodes/script.rs

foundry_cheatcodes/
script.rs
