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

example_cheatcode_inspector/
main.rs
