feat(anvil): add support for injecting precompiles (#7589)

alexfertel · mattsse · web-flow · commit 0df7fb19e971 · 2024-04-09T23:40:53.000+02:00
* feat(anvil): add support for injecting precompiles

* test: check precompiles get injected

* feat(docs): add a few doc comments

* feat(docs): document with_extra_precompiles

* ref: localize changes to the anvil crate

* ref: rename with_extra_precompiles -&gt; with_precompile_factory

* lint(fmt): fix formatting

* ref: fix invalid comment

* ref: remove unnecessary generic bound

* ref: revert formatting change

* ref: extract evm creation to a method

* fix: inject precompiles to the executor

* lint(fmt): fix formatting

* chore: add doc

* nit

---------

Co-authored-by: Matthias Seitz &lt;matthias.seitz@outlook.de&gt;
diff --git a/crates/anvil/src/config.rs b/crates/anvil/src/config.rs
@@ -11,9 +11,8 @@ use crate::{
         fees::{INITIAL_BASE_FEE, INITIAL_GAS_PRICE},
         pool::transactions::TransactionOrder,
     },
-    mem,
-    mem::in_memory_db::MemDb,
-    FeeManager, Hardfork,
+    mem::{self, in_memory_db::MemDb},
+    FeeManager, Hardfork, PrecompileFactory,
 };
 use alloy_genesis::Genesis;
 use alloy_network::AnyNetwork;
@@ -176,6 +175,8 @@ pub struct NodeConfig {
     pub slots_in_an_epoch: u64,
     /// The memory limit per EVM execution in bytes.
     pub memory_limit: Option<u64>,
+    /// Factory used by `anvil` to extend the EVM's precompiles.
+    pub precompile_factory: Option<Arc<dyn PrecompileFactory>>,
 }
 
 impl NodeConfig {
@@ -422,6 +423,7 @@ impl Default for NodeConfig {
             enable_optimism: false,
             slots_in_an_epoch: 32,
             memory_limit: None,
+            precompile_factory: None,
         }
     }
 }
@@ -834,6 +836,13 @@ impl NodeConfig {
         self
     }
 
+    /// Injects precompiles to `anvil`'s EVM.
+    #[must_use]
+    pub fn with_precompile_factory(mut self, factory: impl PrecompileFactory + 'static) -> Self {
+        self.precompile_factory = Some(Arc::new(factory));
+        self
+    }
+
     /// Configures everything related to env, backend and database and returns the
     /// [Backend](mem::Backend)
     ///
diff --git a/crates/anvil/src/eth/backend/executor.rs b/crates/anvil/src/eth/backend/executor.rs
@@ -4,7 +4,9 @@ use crate::{
         error::InvalidTransactionError,
         pool::transactions::PoolTransaction,
     },
+    inject_precompiles,
     mem::inspector::Inspector,
+    PrecompileFactory,
 };
 use alloy_consensus::{Header, Receipt, ReceiptWithBloom};
 use alloy_primitives::{Bloom, BloomInput, Log, B256};
@@ -96,6 +98,8 @@ pub struct TransactionExecutor<'a, Db: ?Sized, Validator: TransactionValidator>
     /// Cumulative gas used by all executed transactions
     pub gas_used: u128,
     pub enable_steps_tracing: bool,
+    /// Precompiles to inject to the EVM.
+    pub precompile_factory: Option<Arc<dyn PrecompileFactory>>,
 }
 
 impl<'a, DB: Db + ?Sized, Validator: TransactionValidator> TransactionExecutor<'a, DB, Validator> {
@@ -265,6 +269,9 @@ impl<'a, 'b, DB: Db + ?Sized, Validator: TransactionValidator> Iterator
         let exec_result = {
             let mut evm =
                 foundry_evm::utils::new_evm_with_inspector(&mut *self.db, env, &mut inspector);
+            if let Some(ref factory) = self.precompile_factory {
+                inject_precompiles(&mut evm, factory.precompiles());
+            }
 
             trace!(target: "backend", "[{:?}] executing", transaction.hash());
             // transact and commit the transaction
diff --git a/crates/anvil/src/eth/backend/mem/mod.rs b/crates/anvil/src/eth/backend/mem/mod.rs
@@ -23,6 +23,7 @@ use crate::{
         pool::transactions::PoolTransaction,
         util::get_precompiles_for,
     },
+    inject_precompiles,
     mem::{
         inspector::Inspector,
         storage::{BlockchainStorage, InMemoryBlockStates, MinedBlockOutcome},
@@ -31,7 +32,7 @@ use crate::{
         db::DatabaseRef,
         primitives::{AccountInfo, U256 as rU256},
     },
-    NodeConfig,
+    NodeConfig, PrecompileFactory,
 };
 use alloy_consensus::{Header, Receipt, ReceiptWithBloom};
 use alloy_primitives::{keccak256, Address, Bytes, TxHash, B256, U256, U64};
@@ -78,7 +79,10 @@ use foundry_evm::{
 };
 use futures::channel::mpsc::{unbounded, UnboundedSender};
 use parking_lot::{Mutex, RwLock};
-use revm::primitives::{HashMap, ResultAndState};
+use revm::{
+    db::WrapDatabaseRef,
+    primitives::{HashMap, ResultAndState},
+};
 use std::{
     collections::BTreeMap,
     io::{Read, Write},
@@ -168,6 +172,8 @@ pub struct Backend {
     node_config: Arc<AsyncRwLock<NodeConfig>>,
     /// Slots in an epoch
     slots_in_an_epoch: u64,
+    /// Precompiles to inject to the EVM.
+    precompile_factory: Option<Arc<dyn PrecompileFactory>>,
 }
 
 impl Backend {
@@ -214,7 +220,10 @@ impl Backend {
             Default::default()
         };
 
-        let slots_in_an_epoch = node_config.read().await.slots_in_an_epoch;
+        let (slots_in_an_epoch, precompile_factory) = {
+            let cfg = node_config.read().await;
+            (cfg.slots_in_an_epoch, cfg.precompile_factory.clone())
+        };
 
         let backend = Self {
             db,
@@ -233,6 +242,7 @@ impl Backend {
             transaction_block_keeper,
             node_config,
             slots_in_an_epoch,
+            precompile_factory,
         };
 
         if let Some(interval_block_time) = automine_block_time {
@@ -800,6 +810,24 @@ impl Backend {
         env
     }
 
+    /// Creates an EVM instance with optionally injected precompiles.
+    fn new_evm_with_inspector_ref<DB, I>(
+        &self,
+        db: DB,
+        env: EnvWithHandlerCfg,
+        inspector: I,
+    ) -> revm::Evm<'_, I, WrapDatabaseRef<DB>>
+    where
+        DB: revm::DatabaseRef,
+        I: revm::Inspector<WrapDatabaseRef<DB>>,
+    {
+        let mut evm = new_evm_with_inspector_ref(db, env, inspector);
+        if let Some(ref factory) = self.precompile_factory {
+            inject_precompiles(&mut evm, factory.precompiles());
+        }
+        evm
+    }
+
     /// executes the transactions without writing to the underlying database
     pub async fn inspect_tx(
         &self,
@@ -812,9 +840,8 @@ impl Backend {
         env.tx = tx.pending_transaction.to_revm_tx_env();
         let db = self.db.read().await;
         let mut inspector = Inspector::default();
-
-        let ResultAndState { result, state } =
-            new_evm_with_inspector_ref(&*db, env, &mut inspector).transact()?;
+        let mut evm = self.new_evm_with_inspector_ref(&*db, env, &mut inspector);
+        let ResultAndState { result, state } = evm.transact()?;
         let (exit_reason, gas_used, out, logs) = match result {
             ExecutionResult::Success { reason, gas_used, logs, output, .. } => {
                 (reason.into(), gas_used, Some(output), Some(logs))
@@ -825,6 +852,7 @@ impl Backend {
             ExecutionResult::Halt { reason, gas_used } => (reason.into(), gas_used, None, None),
         };
 
+        drop(evm);
         inspector.print_logs();
 
         Ok((exit_reason, out, gas_used, state, logs.unwrap_or_default()))
@@ -865,6 +893,7 @@ impl Backend {
             parent_hash: storage.best_hash,
             gas_used: 0,
             enable_steps_tracing: self.enable_steps_tracing,
+            precompile_factory: self.precompile_factory.clone(),
         };
 
         // create a new pending block
@@ -924,6 +953,7 @@ impl Backend {
                     parent_hash: best_hash,
                     gas_used: 0,
                     enable_steps_tracing: self.enable_steps_tracing,
+                    precompile_factory: self.precompile_factory.clone(),
                 };
                 let executed_tx = executor.execute();
 
@@ -1123,8 +1153,8 @@ impl Backend {
         let mut inspector = Inspector::default();
 
         let env = self.build_call_env(request, fee_details, block_env);
-        let ResultAndState { result, state } =
-            new_evm_with_inspector_ref(state, env, &mut inspector).transact()?;
+        let mut evm = self.new_evm_with_inspector_ref(state, env, &mut inspector);
+        let ResultAndState { result, state } = evm.transact()?;
         let (exit_reason, gas_used, out) = match result {
             ExecutionResult::Success { reason, gas_used, output, .. } => {
                 (reason.into(), gas_used, Some(output))
@@ -1134,6 +1164,7 @@ impl Backend {
             }
             ExecutionResult::Halt { reason, gas_used } => (reason.into(), gas_used, None),
         };
+        drop(evm);
         inspector.print_logs();
         Ok((exit_reason, out, gas_used as u128, state))
     }
@@ -1150,8 +1181,9 @@ impl Backend {
             let block_number = block.number;
 
             let env = self.build_call_env(request, fee_details, block);
-            let ResultAndState { result, state: _ } =
-                new_evm_with_inspector_ref(state, env, &mut inspector).transact()?;
+            let mut evm = self.new_evm_with_inspector_ref(state, env, &mut inspector);
+            let ResultAndState { result, state: _ } = evm.transact()?;
+
             let (exit_reason, gas_used, out) = match result {
                 ExecutionResult::Success { reason, gas_used, output, .. } => {
                     (reason.into(), gas_used, Some(output))
@@ -1161,6 +1193,8 @@ impl Backend {
                 }
                 ExecutionResult::Halt { reason, gas_used } => (reason.into(), gas_used, None),
             };
+
+            drop(evm);
             let tracer = inspector.tracer.expect("tracer disappeared");
             let return_value = out.as_ref().map(|o| o.data().clone()).unwrap_or_default();
             let res = tracer.into_geth_builder().geth_traces(gas_used, return_value, opts);
@@ -1196,8 +1230,8 @@ impl Backend {
         );
 
         let env = self.build_call_env(request, fee_details, block_env);
-        let ResultAndState { result, state: _ } =
-            new_evm_with_inspector_ref(state, env, &mut inspector).transact()?;
+        let mut evm = self.new_evm_with_inspector_ref(state, env, &mut inspector);
+        let ResultAndState { result, state: _ } = evm.transact()?;
         let (exit_reason, gas_used, out) = match result {
             ExecutionResult::Success { reason, gas_used, output, .. } => {
                 (reason.into(), gas_used, Some(output))
@@ -1207,6 +1241,7 @@ impl Backend {
             }
             ExecutionResult::Halt { reason, gas_used } => (reason.into(), gas_used, None),
         };
+        drop(evm);
         let access_list = inspector.access_list();
         Ok((exit_reason, out, gas_used, access_list))
     }
diff --git a/crates/anvil/src/evm.rs b/crates/anvil/src/evm.rs
@@ -0,0 +1,92 @@
+use std::{fmt::Debug, sync::Arc};
+
+use alloy_primitives::Address;
+use foundry_evm::revm::{self, precompile::Precompile, ContextPrecompile, ContextPrecompiles};
+
+/// Object-safe trait that enables injecting extra precompiles when using
+/// `anvil` as a library.
+pub trait PrecompileFactory: Send + Sync + Unpin + Debug {
+    /// Returns a set of precompiles to extend the EVM with.
+    fn precompiles(&self) -> Vec<(Address, Precompile)>;
+}
+
+/// Appends a handler register to `evm` that injects the given `precompiles`.
+///
+/// This will add an additional handler that extends the default precompiles with the given set of
+/// precompiles.
+pub fn inject_precompiles<DB, I>(
+    evm: &mut revm::Evm<'_, I, DB>,
+    precompiles: Vec<(Address, Precompile)>,
+) where
+    DB: revm::Database,
+{
+    evm.handler.append_handler_register_box(Box::new(move |handler| {
+        let precompiles = precompiles.clone();
+        let loaded_precompiles = handler.pre_execution().load_precompiles();
+        handler.pre_execution.load_precompiles = Arc::new(move || {
+            let mut loaded_precompiles = loaded_precompiles.clone();
+            loaded_precompiles.extend(
+                precompiles
+                    .clone()
+                    .into_iter()
+                    .map(|(addr, p)| (addr, ContextPrecompile::Ordinary(p))),
+            );
+            let mut default_precompiles = ContextPrecompiles::default();
+            default_precompiles.extend(loaded_precompiles);
+            default_precompiles
+        });
+    }));
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{evm::inject_precompiles, PrecompileFactory};
+    use alloy_primitives::Address;
+    use foundry_evm::revm::{
+        self,
+        primitives::{address, Bytes, Precompile, PrecompileResult, SpecId},
+    };
+
+    #[test]
+    fn build_evm_with_extra_precompiles() {
+        const PRECOMPILE_ADDR: Address = address!("0000000000000000000000000000000000000071");
+        fn my_precompile(_bytes: &Bytes, _gas_limit: u64) -> PrecompileResult {
+            Ok((0, Bytes::new()))
+        }
+
+        #[derive(Debug)]
+        struct CustomPrecompileFactory;
+
+        impl PrecompileFactory for CustomPrecompileFactory {
+            fn precompiles(&self) -> Vec<(Address, Precompile)> {
+                vec![(PRECOMPILE_ADDR, Precompile::Standard(my_precompile))]
+            }
+        }
+
+        let db = revm::db::EmptyDB::default();
+        let env = Box::<revm::primitives::Env>::default();
+        let spec = SpecId::LATEST;
+        let handler_cfg = revm::primitives::HandlerCfg::new(spec);
+        let inspector = revm::inspectors::NoOpInspector;
+        let context = revm::Context::new(revm::EvmContext::new_with_env(db, env), inspector);
+        let handler = revm::Handler::new(handler_cfg);
+        let mut evm = revm::Evm::new(context, handler);
+        assert!(!evm
+            .handler
+            .pre_execution()
+            .load_precompiles()
+            .addresses()
+            .any(|&addr| addr == PRECOMPILE_ADDR));
+
+        inject_precompiles(&mut evm, CustomPrecompileFactory.precompiles());
+        assert!(evm
+            .handler
+            .pre_execution()
+            .load_precompiles()
+            .addresses()
+            .any(|&addr| addr == PRECOMPILE_ADDR));
+
+        let result = evm.transact().unwrap();
+        assert!(result.result.is_success());
+    }
+}
diff --git a/crates/anvil/src/lib.rs b/crates/anvil/src/lib.rs
@@ -50,6 +50,9 @@ pub use hardfork::Hardfork;
 
 /// ethereum related implementations
 pub mod eth;
+/// Evm related abstractions
+mod evm;
+pub use evm::{inject_precompiles, PrecompileFactory};
 /// support for polling filters
 pub mod filter;
 /// commandline output
