diff --git a/Cargo.lock b/Cargo.lock index 909992af7a..13a1c276ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1455,6 +1455,16 @@ dependencies = [ "tokio", ] +[[package]] +name = "example-cheatcode-inspector" +version = "0.0.0" +dependencies = [ + "anyhow", + "revm", + "revm-database", + "revm-inspector", +] + [[package]] name = "example-contract-deployment" version = "0.0.0" diff --git a/Cargo.toml b/Cargo.toml index eb9f16f333..a3e5a42bdf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ members = [ # examples "examples/block_traces", + "examples/cheatcode_inspector", "examples/contract_deployment", "examples/database_components", "examples/uniswap_get_reserves", diff --git a/examples/cheatcode_inspector/Cargo.toml b/examples/cheatcode_inspector/Cargo.toml new file mode 100644 index 0000000000..e53f49088a --- /dev/null +++ b/examples/cheatcode_inspector/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "example-cheatcode-inspector" +version = "0.0.0" +publish = false +authors.workspace = true +edition.workspace = true +keywords.workspace = true +license.workspace = true +repository.workspace = true +readme.workspace = true + +[package.metadata.docs.rs] +all-features = true +rustdoc-args = ["--cfg", "docsrs"] + +[lints.rust] +unreachable_pub = "warn" +unused_must_use = "deny" +rust_2018_idioms = "deny" + +[lints.rustdoc] +all = "warn" + +[dependencies] +revm = {workspace = true, features = ["std"] } +database = { workspace = true, features = ["std"] } +inspector = { workspace = true, features = ["std", "serde-json"] } + +# misc +anyhow = "1.0.89" diff --git a/examples/cheatcode_inspector/src/main.rs b/examples/cheatcode_inspector/src/main.rs new file mode 100644 index 0000000000..10d9898388 --- /dev/null +++ b/examples/cheatcode_inspector/src/main.rs @@ -0,0 +1,600 @@ +//! An example that shows how to implement a Foundry-style Solidity test cheatcode inspector. +//! +//! The code below mimics relevant parts of the implementation of the [`transact`](https://book.getfoundry.sh/cheatcodes/transact) +//! and [`rollFork(uint256 forkId, bytes32 transaction)`](https://book.getfoundry.sh/cheatcodes/roll-fork#rollfork) cheatcodes. +//! Both of these cheatcodes initiate transactions from a call step in the cheatcode inspector which is the most advanced cheatcode use-case. +#![cfg_attr(not(test), warn(unused_crate_dependencies))] + +use std::{convert::Infallible, fmt::Debug}; + +use database::InMemoryDB; +use inspector::{ + inspector_context::InspectorContext, inspector_handler, inspectors::TracerEip3155, + journal::JournalExt, GetInspector, Inspector, InspectorHandler, +}; +use revm::bytecode::Bytecode; +use revm::interpreter::{interpreter::EthInterpreter, CallInputs, CallOutcome, InterpreterResult}; +use revm::primitives::{Log, U256}; +use revm::{ + context::{BlockEnv, Cfg, CfgEnv, TxEnv}, + context_interface::{ + host::{SStoreResult, SelfDestructResult}, + journaled_state::{AccountLoad, JournalCheckpoint, StateLoad, TransferError}, + result::{EVMError, InvalidTransaction}, + Block, Journal, JournalGetter, Transaction, + }, + handler::EthPrecompileProvider, + handler_interface::PrecompileProvider, + precompile::{Address, HashSet, B256}, + specification::hardfork::SpecId, + state::{Account, EvmState, TransientStorage}, + Context, Database, DatabaseCommit, Evm, JournalEntry, JournaledState, +}; + +/// Backend for cheatcodes. +/// The problematic cheatcodes are only supported in fork mode, so we'll omit the non-fork behavior of the Foundry `Backend`. +#[derive(Clone, Debug)] +struct Backend { + /// In fork mode, Foundry stores (`JournaledState`, `Database`) pairs for each fork. + journaled_state: JournaledState, + /// Counters to be able to assert that we mutated the object that we expected to mutate. + method_with_inspector_counter: usize, + method_without_inspector_counter: usize, +} + +impl Backend { + fn new(spec: SpecId, db: InMemoryDB) -> Self { + Self { + journaled_state: JournaledState::new(spec, db), + method_with_inspector_counter: 0, + method_without_inspector_counter: 0, + } + } +} + +impl Journal for Backend { + type Database = InMemoryDB; + type FinalOutput = (EvmState, Vec); + + fn new(database: InMemoryDB) -> Self { + Self::new(SpecId::LATEST, database) + } + + fn db_ref(&self) -> &Self::Database { + &self.journaled_state.database + } + + fn db(&mut self) -> &mut Self::Database { + &mut self.journaled_state.database + } + + fn sload( + &mut self, + address: Address, + key: U256, + ) -> Result, ::Error> { + self.journaled_state.sload(address, key) + } + + fn sstore( + &mut self, + address: Address, + key: U256, + value: U256, + ) -> Result, ::Error> { + self.journaled_state.sstore(address, key, value) + } + + fn tload(&mut self, address: Address, key: U256) -> U256 { + self.journaled_state.tload(address, key) + } + + fn tstore(&mut self, address: Address, key: U256, value: U256) { + self.journaled_state.tstore(address, key, value) + } + + fn log(&mut self, log: Log) { + self.journaled_state.log(log) + } + + fn selfdestruct( + &mut self, + address: Address, + target: Address, + ) -> Result, Infallible> { + self.journaled_state.selfdestruct(address, target) + } + + fn warm_account_and_storage( + &mut self, + address: Address, + storage_keys: impl IntoIterator, + ) -> Result<(), ::Error> { + self.journaled_state + .initial_account_load(address, storage_keys)?; + Ok(()) + } + + fn warm_account(&mut self, address: Address) { + self.journaled_state + .warm_preloaded_addresses + .insert(address); + } + + fn warm_precompiles(&mut self, addresses: HashSet
) { + self.journaled_state.warm_precompiles(addresses) + } + + fn precompile_addresses(&self) -> &HashSet
{ + self.journaled_state.precompile_addresses() + } + + fn set_spec_id(&mut self, spec_id: SpecId) { + self.journaled_state.spec = spec_id; + } + + fn touch_account(&mut self, address: Address) { + self.journaled_state.touch(&address); + } + + fn transfer( + &mut self, + from: &Address, + to: &Address, + balance: U256, + ) -> Result, Infallible> { + self.journaled_state.transfer(from, to, balance) + } + + fn inc_account_nonce(&mut self, address: Address) -> Result, Infallible> { + Ok(self.journaled_state.inc_nonce(address)) + } + + fn load_account(&mut self, address: Address) -> Result, Infallible> { + self.journaled_state.load_account(address) + } + + fn load_account_code( + &mut self, + address: Address, + ) -> Result, Infallible> { + self.journaled_state.load_code(address) + } + + fn load_account_delegated( + &mut self, + address: Address, + ) -> Result, Infallible> { + self.journaled_state.load_account_delegated(address) + } + + fn set_code_with_hash(&mut self, address: Address, code: Bytecode, hash: B256) { + self.journaled_state.set_code_with_hash(address, code, hash); + } + + fn clear(&mut self) { + // Clears the JournaledState. Preserving only the spec. + self.journaled_state.state.clear(); + self.journaled_state.transient_storage.clear(); + self.journaled_state.logs.clear(); + self.journaled_state.journal = vec![vec![]]; + self.journaled_state.depth = 0; + self.journaled_state.warm_preloaded_addresses.clear(); + } + + fn checkpoint(&mut self) -> JournalCheckpoint { + self.journaled_state.checkpoint() + } + + fn checkpoint_commit(&mut self) { + self.journaled_state.checkpoint_commit() + } + + fn checkpoint_revert(&mut self, checkpoint: JournalCheckpoint) { + self.journaled_state.checkpoint_revert(checkpoint) + } + + fn create_account_checkpoint( + &mut self, + caller: Address, + address: Address, + balance: U256, + spec_id: SpecId, + ) -> Result { + // Ignore error. + self.journaled_state + .create_account_checkpoint(caller, address, balance, spec_id) + } + + /// Returns call depth. + #[inline] + fn depth(&self) -> usize { + self.journaled_state.depth + } + + fn finalize(&mut self) -> Result::Error> { + let JournaledState { + state, + transient_storage, + logs, + depth, + journal, + database: _, + spec: _, + warm_preloaded_addresses: _, + precompiles: _, + } = &mut self.journaled_state; + + *transient_storage = TransientStorage::default(); + *journal = vec![vec![]]; + *depth = 0; + let state = std::mem::take(state); + let logs = std::mem::take(logs); + + Ok((state, logs)) + } +} + +impl JournalExt for Backend { + fn logs(&self) -> &[Log] { + &self.journaled_state.logs + } + + fn last_journal(&self) -> &[JournalEntry] { + self.journaled_state + .journal + .last() + .expect("Journal is never empty") + } + + fn evm_state(&self) -> &EvmState { + &self.journaled_state.state + } + + fn evm_state_mut(&mut self) -> &mut EvmState { + &mut self.journaled_state.state + } +} + +/// Used in Foundry to provide extended functionality to cheatcodes. +/// The methods are called from the `Cheatcodes` inspector. +trait DatabaseExt: Journal { + /// Mimics `DatabaseExt::transact` + /// See `commit_transaction` for the generics + fn method_that_takes_inspector_as_argument( + &mut self, + env: Env, + inspector: InspectorT, + ) -> anyhow::Result<()> + where + InspectorT: Inspector, EthInterpreter> + + GetInspector, EthInterpreter>, + BlockT: Block, + TxT: Transaction, + CfgT: Cfg, + PrecompileT: PrecompileProvider< + Context = InspectorContext< + InspectorT, + InMemoryDB, + Context, + >, + Output = InterpreterResult, + Error = EVMError, + >; + + /// Mimics `DatabaseExt::roll_fork_to_transaction` + fn method_that_constructs_inspector( + &mut self, + env: Env, + ) -> anyhow::Result<()> + where + BlockT: Block, + TxT: Transaction, + CfgT: Cfg; + // Can't declare a method that takes the precompile provider as a generic parameter and constructs a + // new inspector, because the `PrecompileProvider` trait needs to know the inspector type + // due to its context being `InspectorContext` instead of `Context`. + // `DatabaseExt::roll_fork_to_transaction` actually creates a noop inspector, so this not working is not a hard + // blocker for multichain cheatcodes. + /* + PrecompileT: PrecompileProvider< + Context = InspectorContext>, + Output = InterpreterResult, + Error = EVMError, + >; + */ +} + +impl DatabaseExt for Backend { + fn method_that_takes_inspector_as_argument( + &mut self, + env: Env, + inspector: InspectorT, + ) -> anyhow::Result<()> + where + InspectorT: Inspector, EthInterpreter> + + GetInspector, EthInterpreter>, + BlockT: Block, + TxT: Transaction, + CfgT: Cfg, + PrecompileT: PrecompileProvider< + Context = InspectorContext< + InspectorT, + InMemoryDB, + Context, + >, + Output = InterpreterResult, + Error = EVMError, + >, + { + commit_transaction::(self, env, inspector)?; + self.method_with_inspector_counter += 1; + Ok(()) + } + + fn method_that_constructs_inspector( + &mut self, + env: Env, + ) -> anyhow::Result<()> + where + BlockT: Block, + TxT: Transaction, + CfgT: Cfg, + { + let inspector = TracerEip3155::new(Box::new(std::io::sink())); + commit_transaction::< + // Generic interpreter types are not supported yet in the `Evm` implementation + TracerEip3155, EthInterpreter>, + BlockT, + TxT, + CfgT, + // Since we can't have a generic precompiles type param as explained in the trait definition, we're using + // concrete type here. + EthPrecompileProvider< + InspectorContext< + TracerEip3155, EthInterpreter>, + InMemoryDB, + Context, + >, + EVMError, + >, + >(self, env, inspector)?; + + self.method_without_inspector_counter += 1; + Ok(()) + } +} + +/// An REVM inspector that intercepts calls to the cheatcode address and executes them with the help of the +/// `DatabaseExt` trait. +#[derive(Clone, Default)] +struct Cheatcodes { + call_count: usize, + phantom: core::marker::PhantomData<(BlockT, TxT, CfgT)>, +} + +impl Cheatcodes +where + BlockT: Block + Clone, + TxT: Transaction + Clone, + CfgT: Cfg + Clone, +{ + fn apply_cheatcode( + &mut self, + context: &mut Context, + ) -> anyhow::Result<()> { + // We cannot avoid cloning here, because we need to mutably borrow the context to get the journal. + let block = context.block.clone(); + let tx = context.tx.clone(); + let cfg = context.cfg.clone(); + + // `transact` cheatcode would do this + context + .journal() + .method_that_takes_inspector_as_argument::<&mut Self, BlockT, TxT, CfgT, EthPrecompileProvider< + InspectorContext<&mut Self, InMemoryDB, Context>, + EVMError, + >>( + Env { + block: block.clone(), + tx: tx.clone(), + cfg: cfg.clone(), + }, + self, + )?; + + // `rollFork(bytes32 transaction)` cheatcode would do this + context + .journal() + .method_that_constructs_inspector::(Env { block, tx, cfg })?; + + Ok(()) + } +} + +impl Inspector, EthInterpreter> + for Cheatcodes +where + BlockT: Block + Clone, + TxT: Transaction + Clone, + CfgT: Cfg + Clone, +{ + /// Note that precompiles are no longer accessible via `EvmContext::precompiles`. + fn call( + &mut self, + context: &mut Context, + _inputs: &mut CallInputs, + ) -> Option { + self.call_count += 1; + // Don't apply cheatcodes recursively. + if self.call_count == 1 { + // Instead of calling unwrap here, we would want to return an appropriate call outcome based on the result in a real project. + self.apply_cheatcode(context).unwrap(); + } + None + } +} + +/// EVM environment +#[derive(Clone, Debug)] +struct Env { + block: BlockT, + tx: TxT, + cfg: CfgT, +} + +impl Env { + fn mainnet() -> Self { + // `CfgEnv` is non-exhaustive, so we need to set the field after construction. + let mut cfg = CfgEnv::default(); + cfg.disable_nonce_check = true; + + Self { + block: BlockEnv::default(), + tx: TxEnv::default(), + cfg, + } + } +} + +/// Executes a transaction and runs the inspector using the `Backend` as the state. +/// Mimics `commit_transaction` +fn commit_transaction( + backend: &mut Backend, + env: Env, + inspector: InspectorT, +) -> Result<(), EVMError> +where + InspectorT: Inspector< + Context, + // Generic interpreter types are not supported yet in the `Evm` implementation + EthInterpreter, + > + GetInspector, EthInterpreter>, + BlockT: Block, + TxT: Transaction, + CfgT: Cfg, + PrecompileT: PrecompileProvider< + Context = InspectorContext< + InspectorT, + InMemoryDB, + Context, + >, + Output = InterpreterResult, + Error = EVMError, + >, +{ + // Create new journaled state and backend with the same DB and journaled state as the original for the transaction. + // This new backend and state will be discarded after the transaction is done and the changes are applied to the + // original backend. + // Mimics https://github.com/foundry-rs/foundry/blob/25cc1ac68b5f6977f23d713c01ec455ad7f03d21/crates/evm/core/src/backend/mod.rs#L1950-L1953 + let new_backend = backend.clone(); + + let context = Context { + tx: env.tx, + block: env.block, + cfg: env.cfg, + journaled_state: new_backend, + chain: (), + error: Ok(()), + }; + + let inspector_context = InspectorContext::< + InspectorT, + InMemoryDB, + Context, + >::new(context, inspector); + + let mut evm = Evm::new( + inspector_context, + inspector_handler::< + InspectorContext< + InspectorT, + InMemoryDB, + Context, + >, + EVMError, + PrecompileT, + >(), + ); + + let result = evm.transact()?; + + // Persist the changes to the original backend. + backend.journaled_state.database.commit(result.state); + update_state( + &mut backend.journaled_state.state, + &mut backend.journaled_state.database, + )?; + + Ok(()) +} + +/// Mimics +/// Omits persistent accounts (accounts that should be kept persistent when switching forks) for simplicity. +fn update_state(state: &mut EvmState, db: &mut DB) -> Result<(), DB::Error> { + for (addr, acc) in state.iter_mut() { + acc.info = db.basic(*addr)?.unwrap_or_default(); + for (key, val) in acc.storage.iter_mut() { + val.present_value = db.storage(*addr, *key)?; + } + } + + Ok(()) +} + +fn main() -> anyhow::Result<()> { + type InspectorT<'cheatcodes> = &'cheatcodes mut Cheatcodes; + type ErrorT = EVMError; + type InspectorContextT<'cheatcodes> = InspectorContext< + InspectorT<'cheatcodes>, + InMemoryDB, + Context, + >; + type PrecompileT<'cheatcodes> = EthPrecompileProvider, ErrorT>; + + let backend = Backend::new(SpecId::LATEST, InMemoryDB::default()); + let mut inspector = Cheatcodes::::default(); + let env = Env::mainnet(); + + let context = Context { + tx: env.tx, + block: env.block, + cfg: env.cfg, + journaled_state: backend, + chain: (), + error: Ok(()), + }; + let inspector_context = InspectorContext::< + InspectorT<'_>, + InMemoryDB, + Context, + >::new(context, &mut inspector); + let handler = inspector_handler::, ErrorT, PrecompileT<'_>>(); + + let mut evm = Evm::< + ErrorT, + InspectorContextT<'_>, + InspectorHandler, ErrorT, PrecompileT<'_>>, + >::new(inspector_context, handler); + + evm.transact()?; + + // Sanity check + assert_eq!(evm.context.inspector.call_count, 2); + assert_eq!( + evm.context + .inner + .journaled_state + .method_with_inspector_counter, + 1 + ); + assert_eq!( + evm.context + .inner + .journaled_state + .method_without_inspector_counter, + 1 + ); + + Ok(()) +}