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

example_cheatcode_inspector/
main.rs
