example_cheatcode_inspector/
main.rs

1//! An example that shows how to implement a Foundry-style Solidity test cheatcode inspector.
2//!
3//! The code below mimics relevant parts of the implementation of the [`transact`](https://book.getfoundry.sh/cheatcodes/transact)
4//! and [`rollFork(uint256 forkId, bytes32 transaction)`](https://book.getfoundry.sh/cheatcodes/roll-fork#rollfork) cheatcodes.
5//! Both of these cheatcodes initiate transactions from a call step in the cheatcode inspector which is the most
6//! advanced cheatcode use-case.
7#![cfg_attr(not(test), warn(unused_crate_dependencies))]
8
9use revm::{
10    context::{
11        result::InvalidTransaction, BlockEnv, Cfg, CfgEnv, ContextTr, Evm, LocalContext, TxEnv,
12    },
13    context_interface::{
14        journaled_state::{AccountLoad, JournalCheckpoint, TransferError},
15        result::EVMError,
16        Block, JournalTr, Transaction,
17    },
18    database::InMemoryDB,
19    handler::{
20        instructions::{EthInstructions, InstructionProvider},
21        EthPrecompiles, PrecompileProvider,
22    },
23    inspector::{inspectors::TracerEip3155, JournalExt},
24    interpreter::{
25        interpreter::EthInterpreter, CallInputs, CallOutcome, InterpreterResult, SStoreResult,
26        SelfDestructResult, StateLoad,
27    },
28    primitives::{hardfork::SpecId, Address, HashSet, Log, StorageKey, StorageValue, B256, U256},
29    state::{Account, Bytecode, EvmState},
30    Context, Database, DatabaseCommit, InspectEvm, Inspector, Journal, JournalEntry,
31};
32use std::{convert::Infallible, fmt::Debug};
33
34/// Backend for cheatcodes.
35/// The problematic cheatcodes are only supported in fork mode, so we'll omit the non-fork behavior of the Foundry
36/// `Backend`.
37#[derive(Clone, Debug)]
38struct Backend {
39    /// In fork mode, Foundry stores (`JournaledState`, `Database`) pairs for each fork.
40    journaled_state: Journal<InMemoryDB>,
41    /// Counters to be able to assert that we mutated the object that we expected to mutate.
42    method_with_inspector_counter: usize,
43    method_without_inspector_counter: usize,
44}
45
46impl Backend {
47    fn new(spec: SpecId, db: InMemoryDB) -> Self {
48        let mut journaled_state = Journal::new(db);
49        journaled_state.set_spec_id(spec);
50        Self {
51            journaled_state,
52            method_with_inspector_counter: 0,
53            method_without_inspector_counter: 0,
54        }
55    }
56}
57
58impl JournalTr for Backend {
59    type Database = InMemoryDB;
60    type State = EvmState;
61
62    fn new(database: InMemoryDB) -> Self {
63        Self::new(SpecId::default(), database)
64    }
65
66    fn db(&self) -> &Self::Database {
67        self.journaled_state.db()
68    }
69
70    fn db_mut(&mut self) -> &mut Self::Database {
71        self.journaled_state.db_mut()
72    }
73
74    fn sload(
75        &mut self,
76        address: Address,
77        key: StorageKey,
78    ) -> Result<StateLoad<StorageValue>, <Self::Database as Database>::Error> {
79        self.journaled_state.sload(address, key)
80    }
81
82    fn sstore(
83        &mut self,
84        address: Address,
85        key: StorageKey,
86        value: StorageValue,
87    ) -> Result<StateLoad<SStoreResult>, <Self::Database as Database>::Error> {
88        self.journaled_state.sstore(address, key, value)
89    }
90
91    fn tload(&mut self, address: Address, key: StorageKey) -> StorageValue {
92        self.journaled_state.tload(address, key)
93    }
94
95    fn tstore(&mut self, address: Address, key: StorageKey, value: StorageValue) {
96        self.journaled_state.tstore(address, key, value)
97    }
98
99    fn log(&mut self, log: Log) {
100        self.journaled_state.log(log)
101    }
102
103    fn selfdestruct(
104        &mut self,
105        address: Address,
106        target: Address,
107    ) -> Result<StateLoad<SelfDestructResult>, Infallible> {
108        self.journaled_state.selfdestruct(address, target)
109    }
110
111    fn warm_account_and_storage(
112        &mut self,
113        address: Address,
114        storage_keys: impl IntoIterator<Item = StorageKey>,
115    ) -> Result<(), <Self::Database as Database>::Error> {
116        self.journaled_state
117            .warm_account_and_storage(address, storage_keys)
118    }
119
120    fn warm_account(&mut self, address: Address) {
121        self.journaled_state
122            .warm_preloaded_addresses
123            .insert(address);
124    }
125
126    fn warm_coinbase_account(&mut self, address: Address) {
127        self.journaled_state.warm_coinbase_address = Some(address)
128    }
129
130    fn warm_precompiles(&mut self, addresses: HashSet<Address>) {
131        self.journaled_state.warm_precompiles(addresses)
132    }
133
134    fn precompile_addresses(&self) -> &HashSet<Address> {
135        self.journaled_state.precompile_addresses()
136    }
137
138    fn set_spec_id(&mut self, spec_id: SpecId) {
139        self.journaled_state.set_spec_id(spec_id);
140    }
141
142    fn touch_account(&mut self, address: Address) {
143        self.journaled_state.touch_account(address);
144    }
145
146    fn transfer(
147        &mut self,
148        from: Address,
149        to: Address,
150        balance: U256,
151    ) -> Result<Option<TransferError>, Infallible> {
152        self.journaled_state.transfer(from, to, balance)
153    }
154
155    fn load_account(&mut self, address: Address) -> Result<StateLoad<&mut Account>, Infallible> {
156        self.journaled_state.load_account(address)
157    }
158
159    fn load_account_code(
160        &mut self,
161        address: Address,
162    ) -> Result<StateLoad<&mut Account>, Infallible> {
163        self.journaled_state.load_account_code(address)
164    }
165
166    fn load_account_delegated(
167        &mut self,
168        address: Address,
169    ) -> Result<StateLoad<AccountLoad>, Infallible> {
170        self.journaled_state.load_account_delegated(address)
171    }
172
173    fn set_code_with_hash(&mut self, address: Address, code: Bytecode, hash: B256) {
174        self.journaled_state.set_code_with_hash(address, code, hash);
175    }
176
177    fn code(
178        &mut self,
179        address: Address,
180    ) -> Result<StateLoad<revm::primitives::Bytes>, <Self::Database as Database>::Error> {
181        self.journaled_state.code(address)
182    }
183
184    fn code_hash(
185        &mut self,
186        address: Address,
187    ) -> Result<StateLoad<B256>, <Self::Database as Database>::Error> {
188        self.journaled_state.code_hash(address)
189    }
190
191    fn clear(&mut self) {
192        self.journaled_state.clear();
193    }
194
195    fn checkpoint(&mut self) -> JournalCheckpoint {
196        self.journaled_state.checkpoint()
197    }
198
199    fn checkpoint_commit(&mut self) {
200        self.journaled_state.checkpoint_commit()
201    }
202
203    fn checkpoint_revert(&mut self, checkpoint: JournalCheckpoint) {
204        self.journaled_state.checkpoint_revert(checkpoint)
205    }
206
207    fn create_account_checkpoint(
208        &mut self,
209        caller: Address,
210        address: Address,
211        balance: U256,
212        spec_id: SpecId,
213    ) -> Result<JournalCheckpoint, TransferError> {
214        self.journaled_state
215            .create_account_checkpoint(caller, address, balance, spec_id)
216    }
217
218    /// Returns call depth.
219    #[inline]
220    fn depth(&self) -> usize {
221        self.journaled_state.depth()
222    }
223
224    fn finalize(&mut self) -> Self::State {
225        self.journaled_state.finalize()
226    }
227
228    fn caller_accounting_journal_entry(
229        &mut self,
230        address: Address,
231        old_balance: U256,
232        bump_nonce: bool,
233    ) {
234        self.journaled_state
235            .caller_accounting_journal_entry(address, old_balance, bump_nonce)
236    }
237
238    fn balance_incr(
239        &mut self,
240        address: Address,
241        balance: U256,
242    ) -> Result<(), <Self::Database as Database>::Error> {
243        self.journaled_state.balance_incr(address, balance)
244    }
245
246    fn nonce_bump_journal_entry(&mut self, address: Address) {
247        self.journaled_state.nonce_bump_journal_entry(address)
248    }
249
250    fn take_logs(&mut self) -> Vec<Log> {
251        self.journaled_state.take_logs()
252    }
253
254    fn commit_tx(&mut self) {
255        self.journaled_state.commit_tx()
256    }
257
258    fn discard_tx(&mut self) {
259        self.journaled_state.discard_tx()
260    }
261}
262
263impl JournalExt for Backend {
264    fn logs(&self) -> &[Log] {
265        self.journaled_state.logs()
266    }
267
268    fn journal(&self) -> &[JournalEntry] {
269        self.journaled_state.journal()
270    }
271
272    fn evm_state(&self) -> &EvmState {
273        self.journaled_state.evm_state()
274    }
275
276    fn evm_state_mut(&mut self) -> &mut EvmState {
277        self.journaled_state.evm_state_mut()
278    }
279}
280
281/// Used in Foundry to provide extended functionality to cheatcodes.
282/// The methods are called from the `Cheatcodes` inspector.
283trait DatabaseExt: JournalTr {
284    /// Mimics `DatabaseExt::transact`
285    /// See `commit_transaction` for the generics
286    fn method_that_takes_inspector_as_argument<
287        InspectorT,
288        BlockT,
289        TxT,
290        CfgT,
291        InstructionProviderT,
292        PrecompileT,
293    >(
294        &mut self,
295        env: Env<BlockT, TxT, CfgT>,
296        inspector: InspectorT,
297    ) -> anyhow::Result<()>
298    where
299        InspectorT: Inspector<Context<BlockT, TxT, CfgT, InMemoryDB, Backend>, EthInterpreter>,
300        BlockT: Block,
301        TxT: Transaction + Clone,
302        CfgT: Cfg,
303        InstructionProviderT: InstructionProvider<
304                Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
305                InterpreterTypes = EthInterpreter,
306            > + Default,
307        PrecompileT: PrecompileProvider<
308                Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
309                Output = InterpreterResult,
310            > + Default;
311
312    /// Mimics `DatabaseExt::roll_fork_to_transaction`
313    fn method_that_constructs_inspector<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>(
314        &mut self,
315        env: Env<BlockT, TxT, CfgT>,
316    ) -> anyhow::Result<()>
317    where
318        BlockT: Block,
319        TxT: Transaction + Clone,
320        CfgT: Cfg,
321        InstructionProviderT: InstructionProvider<
322                Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
323                InterpreterTypes = EthInterpreter,
324            > + Default,
325        PrecompileT: PrecompileProvider<
326                Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
327                Output = InterpreterResult,
328            > + Default;
329}
330
331impl DatabaseExt for Backend {
332    fn method_that_takes_inspector_as_argument<
333        InspectorT,
334        BlockT,
335        TxT,
336        CfgT,
337        InstructionProviderT,
338        PrecompileT,
339    >(
340        &mut self,
341        env: Env<BlockT, TxT, CfgT>,
342        inspector: InspectorT,
343    ) -> anyhow::Result<()>
344    where
345        InspectorT: Inspector<Context<BlockT, TxT, CfgT, InMemoryDB, Backend>, EthInterpreter>,
346        BlockT: Block,
347        TxT: Transaction + Clone,
348        CfgT: Cfg,
349        InstructionProviderT: InstructionProvider<
350                Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
351                InterpreterTypes = EthInterpreter,
352            > + Default,
353        PrecompileT: PrecompileProvider<
354                Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
355                Output = InterpreterResult,
356            > + Default,
357    {
358        commit_transaction::<InspectorT, BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>(
359            self, env, inspector,
360        )?;
361        self.method_with_inspector_counter += 1;
362        Ok(())
363    }
364
365    fn method_that_constructs_inspector<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>(
366        &mut self,
367        env: Env<BlockT, TxT, CfgT>,
368    ) -> anyhow::Result<()>
369    where
370        BlockT: Block,
371        TxT: Transaction + Clone,
372        CfgT: Cfg,
373        InstructionProviderT: InstructionProvider<
374                Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
375                InterpreterTypes = EthInterpreter,
376            > + Default,
377        PrecompileT: PrecompileProvider<
378                Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
379                Output = InterpreterResult,
380            > + Default,
381    {
382        let inspector = TracerEip3155::new(Box::new(std::io::sink()));
383        commit_transaction::<
384            // Generic interpreter types are not supported yet in the `Evm` implementation
385            TracerEip3155,
386            BlockT,
387            TxT,
388            CfgT,
389            InstructionProviderT,
390            PrecompileT,
391        >(self, env, inspector)?;
392
393        self.method_without_inspector_counter += 1;
394        Ok(())
395    }
396}
397
398/// An REVM inspector that intercepts calls to the cheatcode address and executes them with the help of the
399/// `DatabaseExt` trait.
400#[derive(Clone, Default)]
401struct Cheatcodes<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT> {
402    call_count: usize,
403    phantom: core::marker::PhantomData<(BlockT, TxT, CfgT, InstructionProviderT, PrecompileT)>,
404}
405
406impl<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>
407    Cheatcodes<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>
408where
409    BlockT: Block + Clone,
410    TxT: Transaction + Clone,
411    CfgT: Cfg + Clone,
412    InstructionProviderT: InstructionProvider<
413            Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
414            InterpreterTypes = EthInterpreter,
415        > + Default,
416    PrecompileT: PrecompileProvider<
417            Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
418            Output = InterpreterResult,
419        > + Default,
420{
421    fn apply_cheatcode(
422        &mut self,
423        context: &mut Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
424    ) -> anyhow::Result<()> {
425        // We cannot avoid cloning here, because we need to mutably borrow the context to get the journal.
426        let block = context.block.clone();
427        let tx = context.tx.clone();
428        let cfg = context.cfg.clone();
429
430        // `transact` cheatcode would do this
431        context
432            .journal_mut()
433            .method_that_takes_inspector_as_argument::<_, _, _, _, InstructionProviderT, PrecompileT>(
434                Env {
435                    block: block.clone(),
436                    tx: tx.clone(),
437                    cfg: cfg.clone(),
438                },
439                self,
440            )?;
441
442        // `rollFork(bytes32 transaction)` cheatcode would do this
443        context
444            .journal_mut()
445            .method_that_constructs_inspector::<_, _, _, InstructionProviderT, PrecompileT>(
446                Env { block, tx, cfg },
447            )?;
448        Ok(())
449    }
450}
451
452impl<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>
453    Inspector<Context<BlockT, TxT, CfgT, InMemoryDB, Backend>>
454    for Cheatcodes<BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>
455where
456    BlockT: Block + Clone,
457    TxT: Transaction + Clone,
458    CfgT: Cfg + Clone,
459    InstructionProviderT: InstructionProvider<
460            Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
461            InterpreterTypes = EthInterpreter,
462        > + Default,
463    PrecompileT: PrecompileProvider<
464            Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
465            Output = InterpreterResult,
466        > + Default,
467{
468    /// Note that precompiles are no longer accessible via `EvmContext::precompiles`.
469    fn call(
470        &mut self,
471        context: &mut Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
472        _inputs: &mut CallInputs,
473    ) -> Option<CallOutcome> {
474        self.call_count += 1;
475        // Don't apply cheatcodes recursively.
476        if self.call_count == 1 {
477            // Instead of calling unwrap here, we would want to return an appropriate call outcome based on the result
478            // in a real project.
479            self.apply_cheatcode(context).unwrap();
480        }
481        None
482    }
483}
484
485/// EVM environment
486#[derive(Clone, Debug)]
487struct Env<BlockT, TxT, CfgT> {
488    block: BlockT,
489    tx: TxT,
490    cfg: CfgT,
491}
492
493impl Env<BlockEnv, TxEnv, CfgEnv> {
494    fn mainnet() -> Self {
495        // `CfgEnv` is non-exhaustive, so we need to set the field after construction.
496        let mut cfg = CfgEnv::default();
497        cfg.disable_nonce_check = true;
498
499        Self {
500            block: BlockEnv::default(),
501            tx: TxEnv::default(),
502            cfg,
503        }
504    }
505}
506
507/// Executes a transaction and runs the inspector using the `Backend` as the state.
508/// Mimics `commit_transaction` <https://github.com/foundry-rs/foundry/blob/25cc1ac68b5f6977f23d713c01ec455ad7f03d21/crates/evm/core/src/backend/mod.rs#L1931>
509fn commit_transaction<InspectorT, BlockT, TxT, CfgT, InstructionProviderT, PrecompileT>(
510    backend: &mut Backend,
511    env: Env<BlockT, TxT, CfgT>,
512    inspector: InspectorT,
513) -> Result<(), EVMError<Infallible, InvalidTransaction>>
514where
515    InspectorT: Inspector<Context<BlockT, TxT, CfgT, InMemoryDB, Backend>, EthInterpreter>,
516    BlockT: Block,
517    TxT: Transaction + Clone,
518    CfgT: Cfg,
519    InstructionProviderT: InstructionProvider<
520            Context = Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
521            InterpreterTypes = EthInterpreter,
522        > + Default,
523    PrecompileT: PrecompileProvider<
524            Context<BlockT, TxT, CfgT, InMemoryDB, Backend>,
525            Output = InterpreterResult,
526        > + Default,
527{
528    // Create new journaled state and backend with the same DB and journaled state as the original for the transaction.
529    // This new backend and state will be discarded after the transaction is done and the changes are applied to the
530    // original backend.
531    // Mimics https://github.com/foundry-rs/foundry/blob/25cc1ac68b5f6977f23d713c01ec455ad7f03d21/crates/evm/core/src/backend/mod.rs#L1950-L1953
532    let new_backend = backend.clone();
533    let tx = env.tx.clone();
534
535    let context = Context {
536        tx: env.tx,
537        block: env.block,
538        cfg: env.cfg,
539        journaled_state: new_backend,
540        chain: (),
541        local: LocalContext::default(),
542        error: Ok(()),
543    };
544
545    let mut evm = Evm::new_with_inspector(
546        context,
547        inspector,
548        InstructionProviderT::default(),
549        PrecompileT::default(),
550    );
551
552    let state = evm.inspect_tx(tx)?.state;
553
554    // Persist the changes to the original backend.
555    backend.journaled_state.database.commit(state);
556    update_state(
557        &mut backend.journaled_state.inner.state,
558        &mut backend.journaled_state.database,
559    )?;
560
561    Ok(())
562}
563
564/// Mimics <https://github.com/foundry-rs/foundry/blob/25cc1ac68b5f6977f23d713c01ec455ad7f03d21/crates/evm/core/src/backend/mod.rs#L1968>
565/// Omits persistent accounts (accounts that should be kept persistent when switching forks) for simplicity.
566fn update_state<DB: Database>(state: &mut EvmState, db: &mut DB) -> Result<(), DB::Error> {
567    for (addr, acc) in state.iter_mut() {
568        acc.info = db.basic(*addr)?.unwrap_or_default();
569        for (key, val) in acc.storage.iter_mut() {
570            val.present_value = db.storage(*addr, *key)?;
571        }
572    }
573
574    Ok(())
575}
576
577fn main() -> anyhow::Result<()> {
578    let backend = Backend::new(SpecId::default(), InMemoryDB::default());
579    let mut inspector = Cheatcodes::<
580        BlockEnv,
581        TxEnv,
582        CfgEnv,
583        EthInstructions<EthInterpreter, Context<BlockEnv, TxEnv, CfgEnv, InMemoryDB, Backend>>,
584        EthPrecompiles,
585    >::default();
586    let env = Env::mainnet();
587    let tx = env.tx.clone();
588
589    let context = Context {
590        tx: env.tx,
591        block: env.block,
592        cfg: env.cfg,
593        journaled_state: backend,
594        chain: (),
595        local: LocalContext::default(),
596        error: Ok(()),
597    };
598
599    let mut evm = Evm::new_with_inspector(
600        context,
601        &mut inspector,
602        EthInstructions::default(),
603        EthPrecompiles::default(),
604    );
605    evm.inspect_tx(tx)?;
606
607    // Sanity check
608    assert_eq!(evm.inspector.call_count, 2);
609    assert_eq!(evm.journaled_state.method_with_inspector_counter, 1);
610    assert_eq!(evm.journaled_state.method_without_inspector_counter, 1);
611
612    Ok(())
613}
example_cheatcode_inspector/main.rs

example_cheatcode_inspector/
main.rs
