From dd055a4615c2e8fc42f63103bb3c3d369349a85d Mon Sep 17 00:00:00 2001 From: Federico Gimenez Date: Mon, 2 Dec 2024 13:11:08 +0100 Subject: [PATCH] feat(executor): call hook with state changes after post block balance increments (#13050) --- crates/ethereum/evm/src/execute.rs | 78 +++++++++++++++-- crates/evm/src/execute.rs | 131 ++++++++++++++++++++++++++++- crates/optimism/evm/src/execute.rs | 10 ++- 3 files changed, 208 insertions(+), 11 deletions(-) diff --git a/crates/ethereum/evm/src/execute.rs b/crates/ethereum/evm/src/execute.rs index 3bfc3cb2e..65fbbdd25 100644 --- a/crates/ethereum/evm/src/execute.rs +++ b/crates/ethereum/evm/src/execute.rs @@ -13,8 +13,9 @@ use reth_consensus::ConsensusError; use reth_ethereum_consensus::validate_block_post_execution; use reth_evm::{ execute::{ - BasicBlockExecutorProvider, BlockExecutionError, BlockExecutionStrategy, - BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput, ProviderError, + balance_increment_state, BasicBlockExecutorProvider, BlockExecutionError, + BlockExecutionStrategy, BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput, + ProviderError, }, state_change::post_block_balance_increments, system_calls::{OnStateHook, SystemCaller}, @@ -263,8 +264,11 @@ where } // increment balances self.state - .increment_balances(balance_increments) + .increment_balances(balance_increments.clone()) .map_err(|_| BlockValidationError::IncrementBalanceFailed)?; + // call state hook with changes due to balance increments. + let balance_state = balance_increment_state(&balance_increments, &mut self.state)?; + self.system_caller.on_state(&balance_state); Ok(requests) } @@ -317,6 +321,7 @@ mod tests { use alloy_eips::{ eip2935::{HISTORY_STORAGE_ADDRESS, HISTORY_STORAGE_CODE}, eip4788::{BEACON_ROOTS_ADDRESS, BEACON_ROOTS_CODE, SYSTEM_ADDRESS}, + eip4895::Withdrawal, eip7002::{WITHDRAWAL_REQUEST_PREDEPLOY_ADDRESS, WITHDRAWAL_REQUEST_PREDEPLOY_CODE}, eip7685::EMPTY_REQUESTS_HASH, }; @@ -333,9 +338,9 @@ mod tests { database::StateProviderDatabase, test_utils::StateProviderTest, TransitionState, }; use reth_testing_utils::generators::{self, sign_tx_with_key_pair}; - use revm_primitives::BLOCKHASH_SERVE_WINDOW; + use revm_primitives::{address, EvmState, BLOCKHASH_SERVE_WINDOW}; use secp256k1::{Keypair, Secp256k1}; - use std::collections::HashMap; + use std::{collections::HashMap, sync::mpsc}; fn create_state_provider_with_beacon_root_contract() -> StateProviderTest { let mut db = StateProviderTest::default(); @@ -1220,4 +1225,67 @@ mod tests { ), } } + + #[test] + fn test_balance_increment_not_duplicated() { + let chain_spec = Arc::new( + ChainSpecBuilder::from(&*MAINNET) + .shanghai_activated() + .with_fork(EthereumHardfork::Prague, ForkCondition::Timestamp(0)) + .build(), + ); + + let withdrawal_recipient = address!("1000000000000000000000000000000000000000"); + + let mut db = StateProviderTest::default(); + let initial_balance = 100; + db.insert_account( + withdrawal_recipient, + Account { balance: U256::from(initial_balance), nonce: 1, bytecode_hash: None }, + None, + HashMap::default(), + ); + + let withdrawal = + Withdrawal { index: 0, validator_index: 0, address: withdrawal_recipient, amount: 1 }; + + let header = Header { timestamp: 1, number: 1, ..Header::default() }; + + let block = BlockWithSenders { + block: Block { + header, + body: BlockBody { + transactions: vec![], + ommers: vec![], + withdrawals: Some(vec![withdrawal].into()), + }, + }, + senders: vec![], + }; + + let provider = executor_provider(chain_spec); + let executor = provider.executor(StateProviderDatabase::new(&db)); + + let (tx, rx) = mpsc::channel(); + let tx_clone = tx.clone(); + + let _output = executor + .execute_with_state_hook((&block, U256::ZERO).into(), move |state: &EvmState| { + if let Some(account) = state.get(&withdrawal_recipient) { + let _ = tx_clone.send(account.info.balance); + } + }) + .expect("Block execution should succeed"); + + drop(tx); + let balance_changes: Vec = rx.try_iter().collect(); + + if let Some(final_balance) = balance_changes.last() { + let expected_final_balance = U256::from(initial_balance) + U256::from(1_000_000_000); // initial + 1 Gwei in Wei + assert_eq!( + *final_balance, expected_final_balance, + "Final balance should match expected value after withdrawal" + ); + } + } } diff --git a/crates/evm/src/execute.rs b/crates/evm/src/execute.rs index 7d477d219..8c3e0108f 100644 --- a/crates/evm/src/execute.rs +++ b/crates/evm/src/execute.rs @@ -12,7 +12,10 @@ pub use reth_storage_errors::provider::ProviderError; use crate::{system_calls::OnStateHook, TxEnvOverrides}; use alloc::{boxed::Box, vec::Vec}; use alloy_eips::eip7685::Requests; -use alloy_primitives::BlockNumber; +use alloy_primitives::{ + map::{DefaultHashBuilder, HashMap}, + Address, BlockNumber, +}; use core::fmt::Display; use reth_consensus::ConsensusError; use reth_primitives::{BlockWithSenders, NodePrimitives, Receipt}; @@ -22,7 +25,7 @@ use revm::{ db::{states::bundle_state::BundleRetention, BundleState}, State, }; -use revm_primitives::{db::Database, U256}; +use revm_primitives::{db::Database, Account, AccountStatus, EvmState, U256}; /// A general purpose executor trait that executes an input (e.g. block) and produces an output /// (e.g. state changes and receipts). @@ -499,6 +502,42 @@ where } } +/// Creates an `EvmState` from a map of balance increments and the current state +/// to load accounts from. No balance increment is done in the function. +/// Zero balance increments are ignored and won't create state entries. +pub fn balance_increment_state( + balance_increments: &HashMap, + state: &mut State, +) -> Result +where + DB: Database, +{ + let mut load_account = |address: &Address| -> Result<(Address, Account), BlockExecutionError> { + let cache_account = state.load_cache_account(*address).map_err(|_| { + BlockExecutionError::msg("could not load account for balance increment") + })?; + + let account = cache_account.account.as_ref().ok_or_else(|| { + BlockExecutionError::msg("could not load account for balance increment") + })?; + + Ok(( + *address, + Account { + info: account.info.clone(), + storage: Default::default(), + status: AccountStatus::Touched, + }, + )) + }; + + balance_increments + .iter() + .filter(|(_, &balance)| balance != 0) + .map(|(addr, _)| load_account(addr)) + .collect::>() +} + #[cfg(test)] mod tests { use super::*; @@ -507,7 +546,7 @@ mod tests { use reth_chainspec::{ChainSpec, MAINNET}; use reth_primitives::EthPrimitives; use revm::db::{CacheDB, EmptyDBTyped}; - use revm_primitives::{bytes, TxEnv}; + use revm_primitives::{address, bytes, AccountInfo, TxEnv, KECCAK_EMPTY}; use std::sync::Arc; #[derive(Clone, Default)] @@ -760,4 +799,90 @@ mod tests { let result = executor.execute(BlockExecutionInput::new(&Default::default(), U256::ZERO)); assert!(result.is_ok()); } + + fn setup_state_with_account( + addr: Address, + balance: u128, + nonce: u64, + ) -> State>> { + let db = CacheDB::>::default(); + let mut state = State::builder().with_database(db).with_bundle_update().build(); + + let account_info = AccountInfo { + balance: U256::from(balance), + nonce, + code_hash: KECCAK_EMPTY, + code: None, + }; + state.insert_account(addr, account_info); + state + } + + #[test] + fn test_balance_increment_state_zero() { + let addr = address!("1000000000000000000000000000000000000000"); + let mut state = setup_state_with_account(addr, 100, 1); + + let mut increments = HashMap::::default(); + increments.insert(addr, 0); + + let result = balance_increment_state(&increments, &mut state).unwrap(); + assert!(result.is_empty(), "Zero increments should be ignored"); + } + + #[test] + fn test_balance_increment_state_empty_increments_map() { + let mut state = State::builder() + .with_database(CacheDB::>::default()) + .with_bundle_update() + .build(); + + let increments = HashMap::::default(); + let result = balance_increment_state(&increments, &mut state).unwrap(); + assert!(result.is_empty(), "Empty increments map should return empty state"); + } + + #[test] + fn test_balance_increment_state_multiple_valid_increments() { + let addr1 = address!("1000000000000000000000000000000000000000"); + let addr2 = address!("2000000000000000000000000000000000000000"); + + let mut state = setup_state_with_account(addr1, 100, 1); + + let account2 = + AccountInfo { balance: U256::from(200), nonce: 1, code_hash: KECCAK_EMPTY, code: None }; + state.insert_account(addr2, account2); + + let mut increments = HashMap::::default(); + increments.insert(addr1, 50); + increments.insert(addr2, 100); + + let result = balance_increment_state(&increments, &mut state).unwrap(); + + assert_eq!(result.len(), 2); + assert_eq!(result.get(&addr1).unwrap().info.balance, U256::from(100)); + assert_eq!(result.get(&addr2).unwrap().info.balance, U256::from(200)); + } + + #[test] + fn test_balance_increment_state_mixed_zero_and_nonzero_increments() { + let addr1 = address!("1000000000000000000000000000000000000000"); + let addr2 = address!("2000000000000000000000000000000000000000"); + + let mut state = setup_state_with_account(addr1, 100, 1); + + let account2 = + AccountInfo { balance: U256::from(200), nonce: 1, code_hash: KECCAK_EMPTY, code: None }; + state.insert_account(addr2, account2); + + let mut increments = HashMap::::default(); + increments.insert(addr1, 0); + increments.insert(addr2, 100); + + let result = balance_increment_state(&increments, &mut state).unwrap(); + + assert_eq!(result.len(), 1, "Only non-zero increments should be included"); + assert!(!result.contains_key(&addr1), "Zero increment account should not be included"); + assert_eq!(result.get(&addr2).unwrap().info.balance, U256::from(200)); + } } diff --git a/crates/optimism/evm/src/execute.rs b/crates/optimism/evm/src/execute.rs index a333978f0..549f52c89 100644 --- a/crates/optimism/evm/src/execute.rs +++ b/crates/optimism/evm/src/execute.rs @@ -10,8 +10,9 @@ use reth_chainspec::EthereumHardforks; use reth_consensus::ConsensusError; use reth_evm::{ execute::{ - BasicBlockExecutorProvider, BlockExecutionError, BlockExecutionStrategy, - BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput, ProviderError, + balance_increment_state, BasicBlockExecutorProvider, BlockExecutionError, + BlockExecutionStrategy, BlockExecutionStrategyFactory, BlockValidationError, ExecuteOutput, + ProviderError, }, state_change::post_block_balance_increments, system_calls::{OnStateHook, SystemCaller}, @@ -260,8 +261,11 @@ where post_block_balance_increments(&self.chain_spec.clone(), block, total_difficulty); // increment balances self.state - .increment_balances(balance_increments) + .increment_balances(balance_increments.clone()) .map_err(|_| BlockValidationError::IncrementBalanceFailed)?; + // call state hook with changes due to balance increments. + let balance_state = balance_increment_state(&balance_increments, &mut self.state)?; + self.system_caller.on_state(&balance_state); Ok(Requests::default()) }