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

foundry_cheatcodes/
script.rs
