feat: added hardhat_setStorageAt (#298)

* feat: added hardhat_setStorageAt

* reformatted

* moved import to test module

SUPPORTED_APIS.md

@@ -108,7 +108,7 @@ The `status` options are:
 | `HARDHAT` | `hardhat_setNextBlockBaseFeePerGas` | `NOT IMPLEMENTED` | Sets the base fee per gas for the next block |
 | `HARDHAT` | `hardhat_setPrevRandao` | `NOT IMPLEMENTED` | Sets the PREVRANDAO value of the next block |
 | [`HARDHAT`](#hardhat-namespace) | [`hardhat_setNonce`](#hardhat_setnonce) | `SUPPORTED` | Sets the nonce of a given account |
-| `HARDHAT` | `hardhat_setStorageAt` | `NOT IMPLEMENTED` | Sets the storage value at a given key for a given account |
+| [`HARDHAT`](#hardhat-namespace) | [`hardhat_setStorageAt`](#hardhat_setstorageat) | `SUPPORTED` | Sets the storage value at a given key for a given account |
 | [`HARDHAT`](#hardhat-namespace) | [`hardhat_stopImpersonatingAccount`](#hardhat_stopimpersonatingaccount) | `SUPPORTED` | Stop impersonating an account after having previously used `hardhat_impersonateAccount` |
 | [`NETWORK`](#network-namespace) | [`net_version`](#net_version) | `SUPPORTED` | Returns the current network id <br />_(default is `260`)_ |
 | [`NETWORK`](#network-namespace) | [`net_peerCount`](#net_peercount) | `SUPPORTED` | Returns the number of peers currently connected to the client <br/>_(hard-coded to `0`)_ |

src/namespaces/hardhat.rs

@@ -86,4 +86,18 @@ pub trait HardhatNamespaceT {
     /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation.
     #[rpc(name = "hardhat_setCode")]
     fn set_code(&self, address: Address, code: Vec<u8>) -> RpcResult<()>;
+
+    /// Directly modifies the storage of a contract at a specified slot.
+    ///
+    /// # Arguments
+    ///
+    /// * `address` - The contract address whose storage is to be modified.
+    /// * `slot` - The storage slot to modify.
+    /// * `value` - The value to be set at the specified slot.
+    ///
+    /// # Returns
+    ///
+    /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation.
+    #[rpc(name = "hardhat_setStorageAt")]
+    fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> RpcResult<bool>;
 }

src/node/hardhat.rs

@@ -64,4 +64,13 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> HardhatNam
             })
             .into_boxed_future()
     }
+
+    fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> RpcResult<bool> {
+        self.set_storage_at(address, slot, value)
+            .map_err(|err| {
+                tracing::error!("failed setting storage: {:?}", err);
+                into_jsrpc_error(Web3Error::InternalError(err))
+            })
+            .into_boxed_future()
+    }
 }

src/node/in_memory_ext.rs

@@ -1,8 +1,9 @@
 use anyhow::anyhow;
-use zksync_basic_types::{Address, U256, U64};
+use zksync_basic_types::{AccountTreeId, Address, U256, U64};
 use zksync_types::{
     get_code_key, get_nonce_key,
     utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
+    StorageKey,
 };
 use zksync_utils::{h256_to_u256, u256_to_h256};
 
@@ -323,6 +324,17 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
                 Ok(())
             })
     }
+
+    pub fn set_storage_at(&self, address: Address, slot: U256, value: U256) -> Result<bool> {
+        self.get_inner()
+            .write()
+            .map_err(|err| anyhow!("failed acquiring lock: {:?}", err))
+            .map(|mut writer| {
+                let key = StorageKey::new(AccountTreeId::new(address), u256_to_h256(slot));
+                writer.fork_storage.set_value(key, u256_to_h256(value));
+                true
+            })
+    }
 }
 
 #[cfg(test)]
@@ -332,6 +344,7 @@ mod tests {
     use crate::{http_fork_source::HttpForkSource, node::InMemoryNode};
     use std::str::FromStr;
     use zksync_basic_types::{Nonce, H256};
+    use zksync_state::ReadStorage;
     use zksync_types::{api::BlockNumber, fee::Fee, l2::L2Tx, PackedEthSignature};
 
     #[tokio::test]
@@ -529,6 +542,36 @@ mod tests {
         assert_eq!(new_code, code_after);
     }
 
+    #[tokio::test]
+    async fn test_set_storage_at() {
+        let node = InMemoryNode::<HttpForkSource>::default();
+        let address = Address::repeat_byte(0x1);
+        let slot = U256::from(37);
+        let value = U256::from(42);
+
+        let key = StorageKey::new(AccountTreeId::new(address), u256_to_h256(slot));
+        let value_before = node
+            .get_inner()
+            .write()
+            .unwrap()
+            .fork_storage
+            .read_value(&key);
+        assert_eq!(H256::default(), value_before);
+
+        let result = node
+            .set_storage_at(address, slot, value)
+            .expect("failed setting value");
+        assert!(result);
+
+        let value_after = node
+            .get_inner()
+            .write()
+            .unwrap()
+            .fork_storage
+            .read_value(&key);
+        assert_eq!(value, h256_to_u256(value_after));
+    }
+
     #[tokio::test]
     async fn test_increase_time_zero_value() {
         let node = InMemoryNode::<HttpForkSource>::default();