fix: fetch account code if not empty (#5229)

crates/revm/revm-inspectors/src/tracing/builder/geth.rs

@@ -2,6 +2,7 @@
 
 use crate::tracing::{
     types::{CallTraceNode, CallTraceStepStackItem},
+    utils::load_account_code,
     TracingInspectorConfig,
 };
 use reth_primitives::{Address, Bytes, B256, U256};
@@ -9,10 +10,7 @@ use reth_rpc_types::trace::geth::{
     AccountChangeKind, AccountState, CallConfig, CallFrame, DefaultFrame, DiffMode,
     GethDefaultTracingOptions, PreStateConfig, PreStateFrame, PreStateMode, StructLog,
 };
-use revm::{
-    db::DatabaseRef,
-    primitives::{AccountInfo, ResultAndState, KECCAK_EMPTY},
-};
+use revm::{db::DatabaseRef, primitives::ResultAndState};
 use std::collections::{btree_map::Entry, BTreeMap, HashMap, VecDeque};
 
 /// A type for creating geth style traces
@@ -185,25 +183,6 @@ impl GethTraceBuilder {
         prestate_config: PreStateConfig,
         db: DB,
     ) -> Result<PreStateFrame, DB::Error> {
-        // loads the code from the account or the database
-        // Geth always includes the contract code in the prestate. However,
-        // the code hash will be KECCAK_EMPTY if the account is an EOA. Therefore
-        // we need to filter it out.
-        let load_account_code = |db_acc: &AccountInfo| {
-            db_acc
-                .code
-                .as_ref()
-                .map(|code| code.original_bytes())
-                .or_else(|| {
-                    if db_acc.code_hash == KECCAK_EMPTY {
-                        None
-                    } else {
-                        db.code_by_hash_ref(db_acc.code_hash).ok().map(|code| code.original_bytes())
-                    }
-                })
-                .map(Into::into)
-        };
-
         let account_diffs = state.into_iter().map(|(addr, acc)| (*addr, acc));
 
         if prestate_config.is_default_mode() {
@@ -215,7 +194,7 @@ impl GethTraceBuilder {
                 let acc_state = match prestate.0.entry(addr) {
                     Entry::Vacant(entry) => {
                         let db_acc = db.basic_ref(addr)?.unwrap_or_default();
-                        let code = load_account_code(&db_acc);
+                        let code = load_account_code(&db, &db_acc);
                         let acc_state =
                             AccountState::from_account_info(db_acc.nonce, db_acc.balance, code);
                         entry.insert(acc_state)
@@ -240,7 +219,7 @@ impl GethTraceBuilder {
                 let acc_state = match prestate.0.entry(addr) {
                     Entry::Vacant(entry) => {
                         let db_acc = db.basic_ref(addr)?.unwrap_or_default();
-                        let code = load_account_code(&db_acc);
+                        let code = load_account_code(&db, &db_acc);
                         let acc_state =
                             AccountState::from_account_info(db_acc.nonce, db_acc.balance, code);
                         entry.insert(acc_state)
@@ -272,7 +251,7 @@ impl GethTraceBuilder {
             for (addr, changed_acc) in account_diffs {
                 let db_acc = db.basic_ref(addr)?.unwrap_or_default();
 
-                let pre_code = load_account_code(&db_acc);
+                let pre_code = load_account_code(&db, &db_acc);
 
                 let mut pre_state =
                     AccountState::from_account_info(db_acc.nonce, db_acc.balance, pre_code);

crates/revm/revm-inspectors/src/tracing/builder/parity.rs

@@ -1,6 +1,7 @@
 use super::walker::CallTraceNodeWalkerBF;
 use crate::tracing::{
     types::{CallTraceNode, CallTraceStep},
+    utils::load_account_code,
     TracingInspectorConfig,
 };
 use reth_primitives::{Address, U64};
@@ -583,11 +584,10 @@ where
         if changed_acc.is_created() || changed_acc.is_loaded_as_not_existing() {
             entry.balance = Delta::Added(changed_acc.info.balance);
             entry.nonce = Delta::Added(U64::from(changed_acc.info.nonce));
-            if changed_acc.info.code_hash == KECCAK_EMPTY {
-                // this is an additional check to ensure new accounts always get the empty code
-                // marked as added
-                entry.code = Delta::Added(Default::default());
-            }
+
+            // accounts without code are marked as added
+            let account_code = load_account_code(&db, &changed_acc.info).unwrap_or_default();
+            entry.code = Delta::Added(account_code);
 
             // new storage values
             for (key, slot) in changed_acc.storage.iter() {

crates/revm/revm-inspectors/src/tracing/utils.rs

@@ -1,6 +1,6 @@
 //! Util functions for revm related ops
 
-use reth_primitives::{hex, Address, B256};
+use reth_primitives::{hex, revm_primitives::db::DatabaseRef, Address, Bytes, B256, KECCAK_EMPTY};
 use revm::{
     interpreter::CreateInputs,
     primitives::{CreateScheme, SpecId},
@@ -35,3 +35,25 @@ pub(crate) fn get_create_address(call: &CreateInputs, nonce: u64) -> Address {
         }
     }
 }
+
+/// Loads the code for the given account from the account itself or the database
+///
+/// Returns None if the code hash is the KECCAK_EMPTY hash
+#[inline]
+pub(crate) fn load_account_code<DB: DatabaseRef>(
+    db: DB,
+    db_acc: &revm::primitives::AccountInfo,
+) -> Option<Bytes> {
+    db_acc
+        .code
+        .as_ref()
+        .map(|code| code.original_bytes())
+        .or_else(|| {
+            if db_acc.code_hash == KECCAK_EMPTY {
+                None
+            } else {
+                db.code_by_hash_ref(db_acc.code_hash).ok().map(|code| code.original_bytes())
+            }
+        })
+        .map(Into::into)
+}