feat: anvil set chain id (#486)

* feat: add anvil_setChainId API

* chore: add tests

* fix: merge

* chore: add e2e tests

* chore: fix lint

* fix: update chain id in storage

* chore: add api to the doc

* chore: improve e2e tests

SUPPORTED_APIS.md

@@ -41,6 +41,7 @@ The `status` options are:
 | `ANVIL` | `anvil_setBalance` | `SUPPORTED` | Modifies the balance of an account |
 | `ANVIL` | `anvil_setCode` | `SUPPORTED` | Sets the bytecode of a given account |
 | `ANVIL` | `anvil_setStorageAt` | `SUPPORTED` | Sets the storage value at a given key for a given account |
+| `ANVIL` | `anvil_setChainId` | `SUPPORTED` | Sets the chain id |
 | [`CONFIG`](#config-namespace) | [`config_getShowCalls`](#config_getshowcalls) | `SUPPORTED` | Gets the current value of `show_calls` that's originally set with `--show-calls` option |
 | [`CONFIG`](#config-namespace) | [`config_getShowOutputs`](#config_getshowoutputs) | `SUPPORTED` | Gets the current value of `show_outputs` that's originally set with `--show-outputs` option |
 | [`CONFIG`](#config-namespace) | [`config_getCurrentTimestamp`](#config_getcurrenttimestamp) | `SUPPORTED` | Gets the value of `current_timestamp` for the node |

crates/core/src/fork.rs

@@ -31,9 +31,10 @@ use zksync_types::{
         TransactionDetails, TransactionVariant,
     },
     fee_model::FeeParams,
+    get_system_context_key,
     l2::L2Tx,
     url::SensitiveUrl,
-    ProtocolVersionId, StorageKey,
+    ProtocolVersionId, StorageKey, SYSTEM_CONTEXT_CHAIN_ID_POSITION,
 };
 use zksync_types::{
     Address, L1BatchNumber, L2BlockNumber, L2ChainId, StorageValue, H256, U256, U64,
@@ -348,6 +349,17 @@ 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)
+        }
+        mutator.raw_storage.set_value(
+            get_system_context_key(SYSTEM_CONTEXT_CHAIN_ID_POSITION),
+            H256::from_low_u64_be(id.as_u64()),
+        );
+    }
 }
 
 /// Trait that provides necessary data when
@@ -773,6 +785,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.
@@ -865,7 +883,10 @@ mod tests {
     use anvil_zksync_config::types::{CacheConfig, SystemContractsOptions};
     use zksync_multivm::interface::storage::ReadStorage;
     use zksync_types::{api::TransactionVariant, StorageKey};
-    use zksync_types::{AccountTreeId, L1BatchNumber, H256};
+    use zksync_types::{
+        get_system_context_key, AccountTreeId, L1BatchNumber, L2ChainId, H256,
+        SYSTEM_CONTEXT_CHAIN_ID_POSITION,
+    };
 
     #[test]
     fn test_initial_writes() {
@@ -942,4 +963,51 @@ mod tests {
 
         assert_eq!(actual_result, expected_result);
     }
+
+    #[test]
+    fn test_fork_storage_set_chain_id() {
+        let fork_details = ForkDetails {
+            fork_source: Box::new(testing::ExternalStorage {
+                raw_storage: InMemoryStorage::default(),
+            }),
+            chain_id: TEST_NODE_NETWORK_ID.into(),
+            l1_block: L1BatchNumber(0),
+            l2_block: zksync_types::api::Block::<TransactionVariant>::default(),
+            l2_miniblock: 0,
+            l2_miniblock_hash: H256::zero(),
+            block_timestamp: 0,
+            overwrite_chain_id: None,
+            l1_gas_price: 0,
+            l2_fair_gas_price: 0,
+            fair_pubdata_price: 0,
+            estimate_gas_price_scale_factor: 0.0,
+            estimate_gas_scale_factor: 0.0,
+            fee_params: None,
+            cache_config: CacheConfig::None,
+        };
+        let mut fork_storage: ForkStorage<testing::ExternalStorage> = ForkStorage::new(
+            Some(fork_details),
+            &SystemContractsOptions::default(),
+            false,
+            None,
+        );
+        let new_chain_id = L2ChainId::from(261);
+        fork_storage.set_chain_id(new_chain_id);
+
+        let inner = fork_storage.inner.read().unwrap();
+
+        assert_eq!(new_chain_id, fork_storage.chain_id);
+        assert_eq!(
+            new_chain_id,
+            inner.fork.as_ref().map(|f| f.chain_id).unwrap()
+        );
+        assert_eq!(
+            H256::from_low_u64_be(new_chain_id.as_u64()),
+            *inner
+                .raw_storage
+                .state
+                .get(&get_system_context_key(SYSTEM_CONTEXT_CHAIN_ID_POSITION))
+                .unwrap()
+        );
+    }
 }

crates/core/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)]

crates/core/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()
+    }
 }

crates/core/src/node/in_memory_ext.rs

@@ -481,6 +481,17 @@ 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)]
@@ -495,7 +506,7 @@ mod tests {
     use std::sync::{Arc, RwLock};
     use zksync_multivm::interface::storage::ReadStorage;
     use zksync_types::{api::BlockNumber, fee::Fee, l2::L2Tx, PackedEthSignature};
-    use zksync_types::{Nonce, H256};
+    use zksync_types::{L2ChainId, Nonce, H256};
     use zksync_utils::h256_to_u256;
 
     #[tokio::test]
@@ -1070,4 +1081,19 @@ mod tests {
         let result = node.revert_snapshot(U64::from(100));
         assert!(result.is_err());
     }
+
+    #[tokio::test]
+    async fn test_node_set_chain_id() {
+        let node = InMemoryNode::<HttpForkSource>::default();
+        let new_chain_id = 261;
+
+        let _ = node.set_chain_id(new_chain_id);
+
+        let node_inner = node.inner.read().unwrap();
+        assert_eq!(new_chain_id, node_inner.config.chain_id.unwrap());
+        assert_eq!(
+            L2ChainId::from(new_chain_id),
+            node_inner.fork_storage.chain_id
+        );
+    }
 }

e2e-tests-rust/src/ext.rs

@@ -1,7 +1,7 @@
 //! Module containing extensions of existing alloy abstractions.
 
-use alloy::network::ReceiptResponse;
-use alloy::primitives::{Address, BlockHash};
+use alloy::network::{ReceiptResponse,TxSigner};
+use alloy::primitives::{Address,BlockHash,PrimitiveSignature};
 use alloy::providers::WalletProvider;
 use alloy::signers::local::PrivateKeySigner;
 use alloy_zksync::network::Zksync;
@@ -78,6 +78,11 @@ pub trait ZksyncWalletProviderExt: WalletProvider<Zksync, Wallet = ZksyncWallet>
         self.wallet_mut().register_signer(signer);
         address
     }
+
+    /// Registers provider signer.
+    fn register_signer<T: TxSigner<PrimitiveSignature> + Send + Sync + 'static>(&mut self, signer: T) {
+        self.wallet_mut().register_signer(signer);
+    }
 }
 
 impl<T: WalletProvider<Zksync, Wallet = ZksyncWallet>> ZksyncWalletProviderExt for T {}

e2e-tests-rust/src/provider/testing.rs

@@ -383,6 +383,24 @@ where
         self
     }
 
+    /// Builder-pattern method for setting the chain id.
+    pub fn with_chain_id(mut self, id: u64) -> Self {
+        self.inner = self.inner.with_chain_id(id);
+        self
+    }
+
+    /// Builder-pattern method for setting the recipient.
+    pub fn with_to(mut self, to: Address) -> Self {
+        self.inner = self.inner.with_to(to);
+        self
+    }
+
+    /// Builder-pattern method for setting the value.
+    pub fn with_value(mut self, value: U256) -> Self {
+        self.inner = self.inner.with_value(value);
+        self
+    }
+
     /// Submits transaction to the node.
     ///
     /// This does not wait for the transaction to be confirmed, but returns a [`PendingTransactionFinalizable`]

e2e-tests-rust/tests/lib.rs

@@ -1,8 +1,13 @@
 use alloy::network::ReceiptResponse;
 use alloy::providers::ext::AnvilApi;
+use alloy::providers::Provider;
 use anvil_zksync_e2e_tests::{
     init_testing_provider, AnvilZKsyncApi, ReceiptExt, ZksyncWalletProviderExt, DEFAULT_TX_VALUE,
 };
+use alloy::{
+    primitives::U256,
+    signers::local::PrivateKeySigner,
+};
 use std::convert::identity;
 use std::time::Duration;
 
@@ -333,3 +338,28 @@ async fn cant_load_into_existing_state() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn set_chain_id() -> anyhow::Result<()> {
+    let mut provider = init_testing_provider(identity).await?;
+
+    let random_signer = PrivateKeySigner::random();
+    let random_signer_address = random_signer.address();
+
+    // Send transaction before changing chain id
+    provider.tx().with_to(random_signer_address).with_value(U256::from(1e18 as u64)).finalize().await?;
+
+    // Change chain id
+    let new_chain_id = 123;
+    provider.anvil_set_chain_id(new_chain_id).await?;
+
+    // Verify new chain id
+    assert_eq!(new_chain_id, provider.get_chain_id().await?);
+
+    // Verify transactions can be executed after chain id change
+    // Registering and using new signer to get new chain id applied
+    provider.register_signer(random_signer);
+    provider.tx().with_from(random_signer_address).with_chain_id(new_chain_id).finalize().await?;
+
+    Ok(())
+}