From 0805a56ec509f23d3197ba93af35bdd87bd95b35 Mon Sep 17 00:00:00 2001
From: romsters <petriv.roma@gmail.com>
Date: Tue, 10 Dec 2024 12:36:22 +0200
Subject: [PATCH] feat: add anvil_setChainId API

---
 src/fork.rs               | 13 +++++++++++++
 src/namespaces/anvil.rs   |  8 ++++++++
 src/node/anvil.rs         |  9 +++++++++
 src/node/in_memory_ext.rs | 10 ++++++++++
 4 files changed, 40 insertions(+)

diff --git a/src/fork.rs b/src/fork.rs
index fcaf1ce2..db43aaae 100644
--- a/src/fork.rs
+++ b/src/fork.rs
@@ -351,6 +351,13 @@ impl<S> ForkStorage<S> {
         let mut mutator = self.inner.write().unwrap();
         mutator.raw_storage.store_factory_dep(hash, bytecode)
     }
+    pub fn set_chain_id(&mut self, id: L2ChainId) {
+        self.chain_id = id;
+        let mut mutator = self.inner.write().unwrap();
+        if let Some(fork) = &mut mutator.fork {
+            fork.set_chain_id(id)
+        }
+    }
 }
 
 /// Trait that provides necessary data when
@@ -776,6 +783,12 @@ impl ForkDetails {
     pub fn set_rpc_url(&mut self, url: String) {
         self.fork_source = Box::new(HttpForkSource::new(url, self.cache_config.clone()));
     }
+
+    // Sets fork's chain id.
+    pub fn set_chain_id(&mut self, id: L2ChainId) {
+        self.chain_id = id;
+        self.overwrite_chain_id = Some(id);
+    }
 }
 
 /// Serializable representation of [`ForkStorage`]'s state.
diff --git a/src/namespaces/anvil.rs b/src/namespaces/anvil.rs
index 2a76e678..f191c047 100644
--- a/src/namespaces/anvil.rs
+++ b/src/namespaces/anvil.rs
@@ -310,6 +310,14 @@ pub trait AnvilNamespaceT {
     /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation.
     #[rpc(name = "anvil_setStorageAt")]
     fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> RpcResult<bool>;
+
+    /// Sets the chain id.
+    ///
+    /// # Arguments
+    ///
+    /// * `id` - The chain id to be set.
+    #[rpc(name = "anvil_setChainId")]
+    fn set_chain_id(&self, id: u32) -> RpcResult<()>;
 }
 
 #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
diff --git a/src/node/anvil.rs b/src/node/anvil.rs
index 133375c6..6f0e3b4c 100644
--- a/src/node/anvil.rs
+++ b/src/node/anvil.rs
@@ -257,4 +257,13 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> AnvilNames
             })
             .into_boxed_future()
     }
+
+    fn set_chain_id(&self, id: u32) -> RpcResult<()> {
+        self.set_chain_id(id)
+            .map_err(|err| {
+                tracing::error!("failed setting chain id: {:?}", err);
+                into_jsrpc_error(Web3Error::InternalError(err))
+            })
+            .into_boxed_future()
+    }
 }
diff --git a/src/node/in_memory_ext.rs b/src/node/in_memory_ext.rs
index 0d89cb9c..69671bd2 100644
--- a/src/node/in_memory_ext.rs
+++ b/src/node/in_memory_ext.rs
@@ -479,6 +479,16 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
         }
         Ok(())
     }
+
+    pub fn set_chain_id(&self, id: u32) -> Result<()> {
+        let mut inner =  self.inner
+            .write()
+            .map_err(|_| anyhow::anyhow!("Failed to acquire write lock"))?;
+
+        inner.config.update_chain_id(Some(id));
+        inner.fork_storage.set_chain_id(id.into());
+        Ok(())
+    }
 }
 
 #[cfg(test)]