feat: disable prevrandao (#106)

* feat: disable prevrandao

* doc: more of em

* refactor: generalize

* fix: forbid handler

* feat: set precompiles

* Update src/helpers.rs

Co-authored-by: evalir <[email protected]>

---------

Co-authored-by: evalir <[email protected]>

Author: prestwich

src/evm.rs

@@ -1,7 +1,7 @@
 use crate::{
     db::{StateAcc, TryStateAcc},
     driver::DriveBlockResult,
-    helpers::{Ctx, Evm},
+    helpers::{Ctx, Evm, Instruction},
     Block, BlockDriver, BundleDriver, Cfg, ChainDriver, DriveBundleResult, DriveChainResult,
     ErroredState, EvmErrored, EvmExtUnchecked, EvmNeedsBlock, EvmNeedsCfg, EvmNeedsTx, EvmReady,
     EvmTransacted, HasBlock, HasCfg, HasTx, NeedsCfg, NeedsTx, TransactedState, Tx,
@@ -12,13 +12,15 @@ use alloy::{
 };
 use core::{convert::Infallible, fmt};
 use revm::{
+    bytecode::opcode::DIFFICULTY,
     context::{
         result::{EVMError, ExecutionResult, InvalidTransaction, ResultAndState},
         Block as _, BlockEnv, Cfg as _, ContextSetters, ContextTr, Transaction as _, TxEnv,
     },
     database::{states::bundle_state::BundleRetention, BundleState, TryDatabaseCommit},
+    handler::EthPrecompiles,
     inspector::NoOpInspector,
-    interpreter::gas::calculate_initial_tx_gas_for_tx,
+    interpreter::{gas::calculate_initial_tx_gas_for_tx, instructions::block_info},
     primitives::{hardfork::SpecId, TxKind},
     state::{AccountInfo, Bytecode, EvmState},
     Database, DatabaseCommit, DatabaseRef, InspectEvm, Inspector,
@@ -183,6 +185,81 @@ where
             Ok(self)
         }
     }
+
+    /// Overide an opcode with a custom handler. Returns the previous
+    /// instruction handler for the opcode.
+    pub fn override_opcode(&mut self, opcode: u8, handler: Instruction<Db>) -> Instruction<Db> {
+        std::mem::replace(&mut self.inner.instruction.instruction_table[opcode as usize], handler)
+    }
+
+    /// Disable an opcode by replacing it with unknown opcode behavior. This is
+    /// a shortcut for [`Self::override_opcode`] with [`crate::helpers::forbidden`].
+    pub fn disable_opcode(&mut self, opcode: u8) -> Instruction<Db> {
+        self.override_opcode(opcode, crate::helpers::forbidden)
+    }
+
+    /// Run some closure with an opcode override, then restore the previous
+    /// setting.
+    pub fn with_opcode_override<F, NewState>(
+        mut self,
+        opcode: u8,
+        handler: Instruction<Db>,
+        f: F,
+    ) -> Trevm<Db, Insp, NewState>
+    where
+        F: FnOnce(Self) -> Trevm<Db, Insp, NewState>,
+    {
+        let old = self.override_opcode(opcode, handler);
+        self.inner.instruction.insert_instruction(opcode, handler);
+        let mut this = f(self);
+        this.override_opcode(opcode, old);
+        this
+    }
+
+    /// Disable the prevrandao opcode, by replacing it with unknown opcode
+    /// behavior. This is useful for block simulation, where the prevrandao
+    /// opcode may produce incorrect results.
+    pub fn disable_prevrandao(&mut self) -> Instruction<Db> {
+        self.disable_opcode(DIFFICULTY)
+    }
+
+    /// Enable the prevrandao opcode. If the prevrandao opcode was not
+    /// previously disabled or replaced, this will have no effect on behavior.
+    pub fn enable_prevrandao(&mut self) -> Instruction<Db> {
+        self.override_opcode(DIFFICULTY, block_info::difficulty)
+    }
+
+    /// Run some code with the prevrandao opcode disabled, then restore the
+    /// previous setting. This is useful for block simulation, where the
+    /// prevrandao opcode may produce incorrect results.
+    pub fn without_prevrandao<F, NewState>(self, f: F) -> Trevm<Db, Insp, NewState>
+    where
+        F: FnOnce(Self) -> Trevm<Db, Insp, NewState>,
+    {
+        self.with_opcode_override(DIFFICULTY, crate::helpers::forbidden, f)
+    }
+
+    /// Set the precompiles for the EVM. This will replace the current
+    /// precompiles with the provided ones.
+    pub fn override_precompiles(&mut self, precompiles: EthPrecompiles) -> EthPrecompiles {
+        std::mem::replace(&mut self.inner.precompiles, precompiles)
+    }
+
+    /// Run a closure with a different set of precompiles, then restore the
+    /// previous setting.
+    pub fn with_precompiles<F, NewState>(
+        mut self,
+        precompiles: EthPrecompiles,
+        f: F,
+    ) -> Trevm<Db, Insp, NewState>
+    where
+        F: FnOnce(Self) -> Trevm<Db, Insp, NewState>,
+    {
+        let old = self.override_precompiles(precompiles);
+        let mut this = f(self);
+        this.override_precompiles(old);
+        this
+    }
 }
 
 // Fallible DB Reads with &mut self

src/helpers.rs

@@ -1,18 +1,30 @@
 use revm::{
     context::{BlockEnv, CfgEnv, TxEnv},
+    context_interface::context::ContextError,
     handler::{instructions::EthInstructions, EthPrecompiles},
     inspector::NoOpInspector,
-    interpreter::interpreter::EthInterpreter,
-    Context, Journal,
+    interpreter::{interpreter::EthInterpreter, Interpreter, InterpreterTypes},
+    Context, Database, Journal,
 };
 
 /// [`revm::Context`] with default env types and adjustable DB
 pub type Ctx<Db, J = Journal<Db>, C = ()> = Context<BlockEnv, TxEnv, CfgEnv, Db, J, C>;
 
 /// EVM with default env types and adjustable DB.
-pub type Evm<
-    Db,
-    Insp = NoOpInspector,
-    Inst = EthInstructions<EthInterpreter, Ctx<Db>>,
-    Prec = EthPrecompiles,
-> = revm::context::Evm<Ctx<Db>, Insp, Inst, Prec>;
+pub type Evm<Db, Insp = NoOpInspector, Inst = Instructions<Db>, Prec = EthPrecompiles> =
+    revm::context::Evm<Ctx<Db>, Insp, Inst, Prec>;
+
+/// Handler table for EVM opcodes.
+pub type Instructions<Db> = EthInstructions<EthInterpreter, Ctx<Db>>;
+
+/// The handler type for an EVM opcode.
+pub type Instruction<Db> = revm::interpreter::Instruction<EthInterpreter, Ctx<Db>>;
+
+/// An [`Instruction`] that sets a [`ContextError`] in the [`Ctx`] whenever it
+/// is executed.
+pub fn forbidden<Db: Database, Int: InterpreterTypes>(
+    _interpreter: &mut Interpreter<Int>,
+    ctx: &mut Ctx<Db>,
+) {
+    ctx.error = Err(ContextError::Custom("forbidden opcode".to_string()));
+}

src/inspectors/mod.rs

@@ -1,3 +1,6 @@
+mod layer;
+pub use layer::Layered;
+
 mod timeout;
 pub use timeout::TimeLimit;
 
@@ -7,9 +10,6 @@ pub use set::InspectorSet;
 mod spanning;
 pub use spanning::SpanningInspector;
 
-mod layer;
-pub use layer::Layered;
-
 #[cfg(test)]
 mod test {
     use super::*;