From 7c773a1d3afe8811e82e7fba33490dec160c5cdd Mon Sep 17 00:00:00 2001 From: Dan Cline <6798349+Rjected@users.noreply.github.com> Date: Tue, 18 Jun 2024 07:27:06 -0400 Subject: [PATCH] feat: add stateful precompile example (#8911) Co-authored-by: Matthias Seitz --- Cargo.lock | 17 ++ Cargo.toml | 1 + examples/stateful-precompile/Cargo.toml | 20 ++ examples/stateful-precompile/src/main.rs | 237 +++++++++++++++++++++++ 4 files changed, 275 insertions(+) create mode 100644 examples/stateful-precompile/Cargo.toml create mode 100644 examples/stateful-precompile/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 21c5c73e1..879f8663e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9373,6 +9373,23 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "stateful-precompile" +version = "0.0.0" +dependencies = [ + "eyre", + "parking_lot 0.12.3", + "reth", + "reth-chainspec", + "reth-node-api", + "reth-node-core", + "reth-node-ethereum", + "reth-primitives", + "reth-tracing", + "schnellru", + "tokio", +] + [[package]] name = "static_assertions" version = "1.1.0" diff --git a/Cargo.toml b/Cargo.toml index 7a8183eb8..148eba483 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -109,6 +109,7 @@ members = [ "examples/custom-dev-node/", "examples/custom-engine-types/", "examples/custom-evm/", + "examples/stateful-precompile/", "examples/custom-inspector/", "examples/custom-node-components/", "examples/custom-payload-builder/", diff --git a/examples/stateful-precompile/Cargo.toml b/examples/stateful-precompile/Cargo.toml new file mode 100644 index 000000000..2a248fbef --- /dev/null +++ b/examples/stateful-precompile/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "stateful-precompile" +version = "0.0.0" +publish = false +edition.workspace = true +license.workspace = true + +[dependencies] +reth.workspace = true +reth-chainspec.workspace = true +reth-node-api.workspace = true +reth-node-core.workspace = true +reth-primitives.workspace = true +reth-node-ethereum.workspace = true +reth-tracing.workspace = true + +eyre.workspace = true +parking_lot.workspace = true +schnellru.workspace = true +tokio.workspace = true diff --git a/examples/stateful-precompile/src/main.rs b/examples/stateful-precompile/src/main.rs new file mode 100644 index 000000000..0cd495e85 --- /dev/null +++ b/examples/stateful-precompile/src/main.rs @@ -0,0 +1,237 @@ +//! This example shows how to implement a node with a custom EVM that uses a stateful precompile + +#![cfg_attr(not(test), warn(unused_crate_dependencies))] + +use parking_lot::RwLock; +use reth::{ + builder::{components::ExecutorBuilder, BuilderContext, NodeBuilder}, + primitives::{ + revm_primitives::{CfgEnvWithHandlerCfg, Env, PrecompileResult, TxEnv}, + Address, Bytes, U256, + }, + revm::{ + handler::register::EvmHandler, + inspector_handle_register, + precompile::{Precompile, PrecompileSpecId}, + ContextPrecompile, ContextPrecompiles, Database, Evm, EvmBuilder, GetInspector, + }, + tasks::TaskManager, +}; +use reth_chainspec::{Chain, ChainSpec}; +use reth_node_api::{ConfigureEvm, ConfigureEvmEnv, FullNodeTypes}; +use reth_node_core::{args::RpcServerArgs, node_config::NodeConfig}; +use reth_node_ethereum::{EthEvmConfig, EthExecutorProvider, EthereumNode}; +use reth_primitives::{ + revm_primitives::{SpecId, StatefulPrecompileMut}, + Genesis, Header, TransactionSigned, +}; +use reth_tracing::{RethTracer, Tracer}; +use schnellru::{ByLength, LruMap}; +use std::{collections::HashMap, sync::Arc}; + +/// A cache for precompile inputs / outputs. +/// +/// This assumes that the precompile is a standard precompile, as in `StandardPrecompileFn`, meaning +/// its inputs are only `(Bytes, u64)`. +/// +/// NOTE: This does not work with "context stateful precompiles", ie `ContextStatefulPrecompile` or +/// `ContextStatefulPrecompileMut`. They are explicitly banned. +#[derive(Debug, Default)] +pub struct PrecompileCache { + /// Caches for each precompile input / output. + #[allow(clippy::type_complexity)] + cache: HashMap<(Address, SpecId), Arc>>>, +} + +/// Custom EVM configuration +#[derive(Debug, Clone, Default)] +#[non_exhaustive] +pub struct MyEvmConfig { + precompile_cache: Arc>, +} + +impl MyEvmConfig { + /// Sets the precompiles to the EVM handler + /// + /// This will be invoked when the EVM is created via [ConfigureEvm::evm] or + /// [ConfigureEvm::evm_with_inspector] + /// + /// This will use the default mainnet precompiles and wrap them with a cache. + pub fn set_precompiles( + handler: &mut EvmHandler, + cache: Arc>, + ) where + DB: Database, + { + // first we need the evm spec id, which determines the precompiles + let spec_id = handler.cfg.spec_id; + + let mut loaded_precompiles: ContextPrecompiles = + ContextPrecompiles::new(PrecompileSpecId::from_spec_id(spec_id)); + for (address, precompile) in loaded_precompiles.to_mut().iter_mut() { + // get or insert the cache for this address / spec + let mut cache = cache.write(); + let cache = cache + .cache + .entry((*address, spec_id)) + .or_insert(Arc::new(RwLock::new(LruMap::new(ByLength::new(1024))))); + + *precompile = Self::wrap_precompile(precompile.clone(), cache.clone()); + } + + // install the precompiles + handler.pre_execution.load_precompiles = Arc::new(move || loaded_precompiles.clone()); + } + + /// Given a [`ContextPrecompile`] and cache for a specific precompile, create a new precompile + /// that wraps the precompile with the cache. + fn wrap_precompile( + precompile: ContextPrecompile, + cache: Arc>>, + ) -> ContextPrecompile + where + DB: Database, + { + let ContextPrecompile::Ordinary(precompile) = precompile else { + // context stateful precompiles are not supported, due to lifetime issues or skill + // issues + panic!("precompile is not ordinary"); + }; + + let wrapped = WrappedPrecompile { precompile, cache: cache.clone() }; + + ContextPrecompile::Ordinary(Precompile::StatefulMut(Box::new(wrapped))) + } +} + +/// A custom precompile that contains the cache and precompile it wraps. +#[derive(Clone)] +pub struct WrappedPrecompile { + /// The precompile to wrap. + precompile: Precompile, + /// The cache to use. + cache: Arc>>, +} + +impl StatefulPrecompileMut for WrappedPrecompile { + fn call_mut(&mut self, bytes: &Bytes, gas_price: u64, _env: &Env) -> PrecompileResult { + let mut cache = self.cache.write(); + let key = (bytes.clone(), gas_price); + + // get the result if it exists + if let Some(result) = cache.get(&key) { + return result.clone() + } + + // call the precompile if cache miss + let output = self.precompile.call(bytes, gas_price, _env); + cache.insert(key, output.clone()); + + output + } +} + +impl ConfigureEvmEnv for MyEvmConfig { + fn fill_tx_env(tx_env: &mut TxEnv, transaction: &TransactionSigned, sender: Address) { + EthEvmConfig::fill_tx_env(tx_env, transaction, sender) + } + + fn fill_cfg_env( + cfg_env: &mut CfgEnvWithHandlerCfg, + chain_spec: &ChainSpec, + header: &Header, + total_difficulty: U256, + ) { + EthEvmConfig::fill_cfg_env(cfg_env, chain_spec, header, total_difficulty) + } +} + +impl ConfigureEvm for MyEvmConfig { + type DefaultExternalContext<'a> = (); + + fn evm<'a, DB: Database + 'a>(&self, db: DB) -> Evm<'a, Self::DefaultExternalContext<'a>, DB> { + let new_cache = self.precompile_cache.clone(); + EvmBuilder::default() + .with_db(db) + // add additional precompiles + .append_handler_register_box(Box::new(move |handler| { + MyEvmConfig::set_precompiles(handler, new_cache.clone()) + })) + .build() + } + + fn evm_with_inspector<'a, DB, I>(&self, db: DB, inspector: I) -> Evm<'a, I, DB> + where + DB: Database + 'a, + I: GetInspector, + { + let new_cache = self.precompile_cache.clone(); + EvmBuilder::default() + .with_db(db) + .with_external_context(inspector) + // add additional precompiles + .append_handler_register_box(Box::new(move |handler| { + MyEvmConfig::set_precompiles(handler, new_cache.clone()) + })) + .append_handler_register(inspector_handle_register) + .build() + } +} + +/// Builds a regular ethereum block executor that uses the custom EVM. +#[derive(Debug, Default, Clone)] +#[non_exhaustive] +pub struct MyExecutorBuilder { + /// The precompile cache to use for all executors. + precompile_cache: Arc>, +} + +impl ExecutorBuilder for MyExecutorBuilder +where + Node: FullNodeTypes, +{ + type EVM = MyEvmConfig; + type Executor = EthExecutorProvider; + + async fn build_evm( + self, + ctx: &BuilderContext, + ) -> eyre::Result<(Self::EVM, Self::Executor)> { + let evm_config = MyEvmConfig { precompile_cache: self.precompile_cache.clone() }; + Ok((evm_config.clone(), EthExecutorProvider::new(ctx.chain_spec(), evm_config))) + } +} + +#[tokio::main] +async fn main() -> eyre::Result<()> { + let _guard = RethTracer::new().init()?; + + let tasks = TaskManager::current(); + + // create a custom chain spec + let spec = ChainSpec::builder() + .chain(Chain::mainnet()) + .genesis(Genesis::default()) + .london_activated() + .paris_activated() + .shanghai_activated() + .cancun_activated() + .build(); + + let node_config = + NodeConfig::test().with_rpc(RpcServerArgs::default().with_http()).with_chain(spec); + + let handle = NodeBuilder::new(node_config) + .testing_node(tasks.executor()) + // configure the node with regular ethereum types + .with_types::() + // use default ethereum components but with our executor + .with_components(EthereumNode::components().executor(MyExecutorBuilder::default())) + .launch() + .await + .unwrap(); + + println!("Node started"); + + handle.node_exit_future.await +}