feat(executor): add init methods to set TxEnv overrides (#12551)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

crates/ethereum/evm/src/execute.rs

@@ -18,7 +18,7 @@ use reth_evm::{
     },
     state_change::post_block_balance_increments,
     system_calls::{OnStateHook, SystemCaller},
-    ConfigureEvm,
+    ConfigureEvm, TxEnvOverrides,
 };
 use reth_primitives::{BlockWithSenders, Receipt};
 use reth_revm::db::State;
@@ -83,6 +83,8 @@ where
     chain_spec: Arc<ChainSpec>,
     /// How to create an EVM.
     evm_config: EvmConfig,
+    /// Optional overrides for the transactions environment.
+    tx_env_overrides: Option<Box<dyn TxEnvOverrides>>,
     /// Current state for block execution.
     state: State<DB>,
     /// Utility to call system smart contracts.
@@ -96,7 +98,7 @@ where
     /// Creates a new [`EthExecutionStrategy`]
     pub fn new(state: State<DB>, chain_spec: Arc<ChainSpec>, evm_config: EvmConfig) -> Self {
         let system_caller = SystemCaller::new(evm_config.clone(), chain_spec.clone());
-        Self { state, chain_spec, evm_config, system_caller }
+        Self { state, chain_spec, evm_config, system_caller, tx_env_overrides: None }
     }
 }
 
@@ -130,6 +132,10 @@ where
 {
     type Error = BlockExecutionError;
 
+    fn init(&mut self, tx_env_overrides: Box<dyn TxEnvOverrides>) {
+        self.tx_env_overrides = Some(tx_env_overrides);
+    }
+
     fn apply_pre_execution_changes(
         &mut self,
         block: &BlockWithSenders,
@@ -172,6 +178,10 @@ where
 
             self.evm_config.fill_tx_env(evm.tx_mut(), transaction, *sender);
 
+            if let Some(tx_env_overrides) = &mut self.tx_env_overrides {
+                tx_env_overrides.apply(evm.tx_mut());
+            }
+
             // Execute transaction.
             let result_and_state = evm.transact().map_err(move |err| {
                 let new_err = err.map_db_err(|e| e.into());

crates/evm/src/either.rs

@@ -6,6 +6,7 @@ use crate::{
     execute::{BatchExecutor, BlockExecutorProvider, Executor},
     system_calls::OnStateHook,
 };
+use alloc::boxed::Box;
 use alloy_primitives::BlockNumber;
 use reth_execution_errors::BlockExecutionError;
 use reth_execution_types::{BlockExecutionInput, BlockExecutionOutput, ExecutionOutcome};
@@ -70,6 +71,13 @@ where
     type Output = BlockExecutionOutput<Receipt>;
     type Error = BlockExecutionError;
 
+    fn init(&mut self, tx_env_overrides: Box<dyn crate::TxEnvOverrides>) {
+        match self {
+            Self::Left(a) => a.init(tx_env_overrides),
+            Self::Right(b) => b.init(tx_env_overrides),
+        }
+    }
+
     fn execute(self, input: Self::Input<'_>) -> Result<Self::Output, Self::Error> {
         match self {
             Self::Left(a) => a.execute(input),

crates/evm/src/execute.rs

@@ -6,9 +6,8 @@ pub use reth_execution_errors::{
 };
 pub use reth_execution_types::{BlockExecutionInput, BlockExecutionOutput, ExecutionOutcome};
 pub use reth_storage_errors::provider::ProviderError;
-use revm::db::states::bundle_state::BundleRetention;
 
-use crate::system_calls::OnStateHook;
+use crate::{system_calls::OnStateHook, TxEnvOverrides};
 use alloc::{boxed::Box, vec::Vec};
 use alloy_eips::eip7685::Requests;
 use alloy_primitives::BlockNumber;
@@ -17,7 +16,10 @@ use reth_consensus::ConsensusError;
 use reth_primitives::{BlockWithSenders, Receipt};
 use reth_prune_types::PruneModes;
 use reth_revm::batch::BlockBatchRecord;
-use revm::{db::BundleState, State};
+use revm::{
+    db::{states::bundle_state::BundleRetention, BundleState},
+    State,
+};
 use revm_primitives::{db::Database, U256};
 
 /// A general purpose executor trait that executes an input (e.g. block) and produces an output
@@ -32,6 +34,9 @@ pub trait Executor<DB> {
     /// The error type returned by the executor.
     type Error;
 
+    /// Initialize the executor with the given transaction environment overrides.
+    fn init(&mut self, _tx_env_overrides: Box<dyn TxEnvOverrides>) {}
+
     /// Consumes the type and executes the block.
     ///
     /// # Note
@@ -184,6 +189,9 @@ where
     /// The error type returned by this strategy's methods.
     type Error: From<ProviderError> + core::error::Error;
 
+    /// Initialize the strategy with the given transaction environment overrides.
+    fn init(&mut self, _tx_env_overrides: Box<dyn TxEnvOverrides>) {}
+
     /// Applies any necessary changes before executing the block's transactions.
     fn apply_pre_execution_changes(
         &mut self,
@@ -329,6 +337,10 @@ where
     type Output = BlockExecutionOutput<Receipt>;
     type Error = S::Error;
 
+    fn init(&mut self, env_overrides: Box<dyn TxEnvOverrides>) {
+        self.strategy.init(env_overrides);
+    }
+
     fn execute(mut self, input: Self::Input<'_>) -> Result<Self::Output, Self::Error> {
         let BlockExecutionInput { block, total_difficulty } = input;
 
@@ -480,7 +492,7 @@ mod tests {
     use alloy_primitives::U256;
     use reth_chainspec::{ChainSpec, MAINNET};
     use revm::db::{CacheDB, EmptyDBTyped};
-    use revm_primitives::bytes;
+    use revm_primitives::{bytes, TxEnv};
     use std::sync::Arc;
 
     #[derive(Clone, Default)]
@@ -703,4 +715,28 @@ mod tests {
         assert_eq!(block_execution_output.requests, expected_apply_post_execution_changes_result);
         assert_eq!(block_execution_output.state, expected_finish_result);
     }
+
+    #[test]
+    fn test_tx_env_overrider() {
+        let strategy_factory = TestExecutorStrategyFactory {
+            execute_transactions_result: ExecuteOutput {
+                receipts: vec![Receipt::default()],
+                gas_used: 10,
+            },
+            apply_post_execution_changes_result: Requests::new(vec![bytes!("deadbeef")]),
+            finish_result: BundleState::default(),
+        };
+        let provider = BasicBlockExecutorProvider::new(strategy_factory);
+        let db = CacheDB::<EmptyDBTyped<ProviderError>>::default();
+
+        // if we want to apply tx env overrides the executor must be mut.
+        let mut executor = provider.executor(db);
+        // execute consumes the executor, so we can only call it once.
+        // let result = executor.execute(BlockExecutionInput::new(&Default::default(), U256::ZERO));
+        executor.init(Box::new(|tx_env: &mut TxEnv| {
+            tx_env.nonce.take();
+        }));
+        let result = executor.execute(BlockExecutionInput::new(&Default::default(), U256::ZERO));
+        assert!(result.is_ok());
+    }
 }

crates/evm/src/lib.rs

@@ -212,3 +212,18 @@ pub struct NextBlockEnvAttributes {
     /// The randomness value for the next block.
     pub prev_randao: B256,
 }
+
+/// Function hook that allows to modify a transaction environment.
+pub trait TxEnvOverrides {
+    /// Apply the overrides by modifying the given `TxEnv`.
+    fn apply(&mut self, env: &mut TxEnv);
+}
+
+impl<F> TxEnvOverrides for F
+where
+    F: FnMut(&mut TxEnv),
+{
+    fn apply(&mut self, env: &mut TxEnv) {
+        self(env)
+    }
+}

crates/optimism/evm/src/execute.rs

@@ -15,7 +15,7 @@ use reth_evm::{
     },
     state_change::post_block_balance_increments,
     system_calls::{OnStateHook, SystemCaller},
-    ConfigureEvm,
+    ConfigureEvm, TxEnvOverrides,
 };
 use reth_optimism_chainspec::OpChainSpec;
 use reth_optimism_consensus::validate_block_post_execution;
@@ -78,6 +78,8 @@ where
     chain_spec: Arc<OpChainSpec>,
     /// How to create an EVM.
     evm_config: EvmConfig,
+    /// Optional overrides for the transactions environment.
+    tx_env_overrides: Option<Box<dyn TxEnvOverrides>>,
     /// Current state for block execution.
     state: State<DB>,
     /// Utility to call system smart contracts.
@@ -91,7 +93,7 @@ where
     /// Creates a new [`OpExecutionStrategy`]
     pub fn new(state: State<DB>, chain_spec: Arc<OpChainSpec>, evm_config: EvmConfig) -> Self {
         let system_caller = SystemCaller::new(evm_config.clone(), chain_spec.clone());
-        Self { state, chain_spec, evm_config, system_caller }
+        Self { state, chain_spec, evm_config, system_caller, tx_env_overrides: None }
     }
 }
 
@@ -119,6 +121,10 @@ where
 {
     type Error = BlockExecutionError;
 
+    fn init(&mut self, tx_env_overrides: Box<dyn TxEnvOverrides>) {
+        self.tx_env_overrides = Some(tx_env_overrides);
+    }
+
     fn apply_pre_execution_changes(
         &mut self,
         block: &BlockWithSenders,
@@ -197,6 +203,10 @@ where
 
             self.evm_config.fill_tx_env(evm.tx_mut(), transaction, *sender);
 
+            if let Some(tx_env_overrides) = &mut self.tx_env_overrides {
+                tx_env_overrides.apply(evm.tx_mut());
+            }
+
             // Execute transaction.
             let result_and_state = evm.transact().map_err(move |err| {
                 let new_err = err.map_db_err(|e| e.into());