fix: split code from remote account in (#517)

.changeset/twelve-yaks-unite.md

@@ -0,0 +1,5 @@
+---
+"@nomicfoundation/edr": patch
+---
+
+Fixed missing remote contract code when setting storage

crates/edr_evm/src/state/remote/cached.rs

@@ -42,22 +42,9 @@ impl<ChainSpecT: RpcSpec> State for CachedRemoteState<ChainSpecT> {
             return Ok(Some(account.info.clone()));
         }
 
-        if let Some(mut account_info) = self.remote.basic(address)? {
-            // Split code and store separately
-            // Always cache code regardless of the block number for two reasons:
-            // 1. It's an invariant of this trait getting an `AccountInfo` by calling
-            //    `basic`,
-            // one can call `code_by_hash` with `AccountInfo.code_hash` and get the code.
-            // 2. Since the code is identified by its hash, it never goes stale.
-            if let Some(code) = account_info.code.take() {
-                let block_code = self
-                    .code_cache
-                    .entry(self.remote.block_number())
-                    .or_default();
-
-                block_code.entry(account_info.code_hash).or_insert(code);
-            }
-
+        if let Some(account_info) =
+            fetch_remote_account(address, &self.remote, &mut self.code_cache)?
+        {
             if self.remote.is_cacheable()? {
                 block_accounts.insert(address, account_info.clone().into());
             }
@@ -101,10 +88,9 @@ impl<ChainSpecT: RpcSpec> State for CachedRemoteState<ChainSpecT> {
             }
             Entry::Vacant(account_entry) => {
                 // account needs to be loaded for us to access slots.
-                let mut account = self
-                    .remote
-                    .basic(address)?
-                    .map_or_else(EdrAccount::default, EdrAccount::from);
+                let mut account =
+                    fetch_remote_account(address, &self.remote, &mut self.code_cache)?
+                        .map_or_else(EdrAccount::default, EdrAccount::from);
 
                 let value = self.remote.storage(address, index)?;
 
@@ -119,6 +105,30 @@ impl<ChainSpecT: RpcSpec> State for CachedRemoteState<ChainSpecT> {
     }
 }
 
+/// Fetches an account from the remote state. If it exists, code is split off
+/// and stored separately in the provided cache.
+fn fetch_remote_account(
+    address: Address,
+    remote: &RemoteState,
+    code_cache: &mut HashMap<u64, HashMap<B256, Bytecode>>,
+) -> Result<Option<AccountInfo>, StateError> {
+    let account = remote.basic(address)?.map(|mut account_info| {
+        // Always cache code regardless of the block number for two reasons:
+        // 1. It's an invariant of this trait getting an `AccountInfo` by calling
+        //    `basic`,
+        // one can call `code_by_hash` with `AccountInfo.code_hash` and get the code.
+        // 2. Since the code is identified by its hash, it never goes stale.
+        if let Some(code) = account_info.code.take() {
+            let block_code = code_cache.entry(remote.block_number()).or_default();
+
+            block_code.entry(account_info.code_hash).or_insert(code);
+        }
+        account_info
+    });
+
+    Ok(account)
+}
+
 #[cfg(all(test, feature = "test-remote"))]
 mod tests {
     use std::{str::FromStr, sync::Arc};

crates/edr_provider/tests/issues/issue_503.rs

@@ -0,0 +1,47 @@
+use std::str::FromStr as _;
+
+use edr_eth::{Address, SpecId, U256};
+use edr_provider::{
+    hardhat_rpc_types::ForkConfig, test_utils::create_test_config_with_fork, time::CurrentTime,
+    MethodInvocation, NoopLogger, Provider, ProviderRequest,
+};
+use edr_test_utils::env::get_alchemy_url;
+use tokio::runtime;
+
+// https://github.com/NomicFoundation/edr/issues/503
+#[tokio::test(flavor = "multi_thread")]
+async fn issue_503() -> anyhow::Result<()> {
+    let logger = Box::new(NoopLogger);
+    let subscriber = Box::new(|_event| {});
+
+    let mut config = create_test_config_with_fork(Some(ForkConfig {
+        json_rpc_url: get_alchemy_url(),
+        block_number: Some(19_909_475),
+        http_headers: None,
+    }));
+    config.hardfork = SpecId::CANCUN;
+
+    let provider = Provider::new(
+        runtime::Handle::current(),
+        logger,
+        subscriber,
+        config,
+        CurrentTime,
+    )?;
+
+    let address = Address::from_str("0xbe9895146f7af43049ca1c1ae358b0541ea49704")?;
+    let index =
+        U256::from_str("0x4f039c94bc7b6c8e7867b9fbd2890a637837fea1c829f434a649c572b15b2969")?;
+
+    provider.handle_request(ProviderRequest::Single(MethodInvocation::GetStorageAt(
+        address, index, None,
+    )))?;
+
+    provider.handle_request(ProviderRequest::Single(MethodInvocation::SetStorageAt(
+        address,
+        index,
+        U256::from(1u64),
+    )))?;
+
+    Ok(())
+}

crates/edr_provider/tests/issues/mod.rs

@@ -8,3 +8,4 @@ mod issue_356;
 mod issue_361;
 mod issue_384;
 mod issue_407;
+mod issue_503;