hardhat_api: Implement hardhat_setBalance API endpoint (#91)

* hardhat_api: Implement hardhat_setBalance API endpoint

* docs: move comments to trait declaration

* fix: use storage_key_for_eth_balance

* feat: add println when manually setting balance

* docs: mark hardhat_setBalance as supported and add docs/example

* doc: add hardhat_setBalance to test_endpoints.http

* fix link in SUPPORTED_APIS.md

SUPPORTED_APIS.md

@@ -95,7 +95,7 @@ The `status` options are:
 | `HARDHAT` | `hardhat_metadata` | `NOT IMPLEMENTED` | Returns the metadata of the current network |
 | `HARDHAT` | `hardhat_mine` | `NOT IMPLEMENTED`<br />[GitHub Issue #75](https://github.com/matter-labs/era-test-node/issues/75) | Mine any number of blocks at once, in constant time |
 | `HARDHAT` | `hardhat_reset` | `NOT IMPLEMENTED` | Resets the state of the network |
-| `HARDHAT` | `hardhat_setBalance` | `NOT IMPLEMENTED`<br />[GitHub Issue #76](https://github.com/matter-labs/era-test-node/issues/76) | Modifies the balance of an account |
+| [`HARDHAT`](#hardhat-namespace) | [`hardhat_setBalance`](#hardhat_setbalance) | `SUPPORTED` | Modifies the balance of an account |
 | `HARDHAT` | `hardhat_setCode` | `NOT IMPLEMENTED` | Sets the bytecode of a given account |
 | `HARDHAT` | `hardhat_setCoinbase` | `NOT IMPLEMENTED` | Sets the coinbase address |
 | `HARDHAT` | `hardhat_setLoggingEnabled` | `NOT IMPLEMENTED` | Enables or disables logging in Hardhat Network |
@@ -683,6 +683,40 @@ curl --request POST \
 }'
 ```
 
+## `HARDHAT NAMESPACE`
+
+### `hardhat_setBalance`
+
+[source](src/hardhat.rs)
+
+Sets the balance of the given address to the given balance.
+
+#### Arguments
+
++ `address: Address` - The `Address` whose balance will be edited
++ `balance: U256` - The balance to set for the given address, in wei
+
+#### Status
+
+`SUPPORTED`
+
+#### Example
+
+```bash
+curl --request POST \
+  --url http://localhost:8011/ \
+  --header 'content-type: application/json' \
+  --data '{
+    "jsonrpc": "2.0",
+      "id": "1",
+      "method": "hardhat_setBalance",
+      "params": [
+        "0x36615Cf349d7F6344891B1e7CA7C72883F5dc049",
+        "0x1337"
+      ]
+  }'
+```
+
 ## `ZKS NAMESPACE`
 
 ### `zks_estimateFee`

src/hardhat.rs

@@ -0,0 +1,97 @@
+use std::sync::{Arc, RwLock};
+
+use crate::{fork::ForkSource, node::InMemoryNodeInner};
+use jsonrpc_core::{BoxFuture, Result};
+use jsonrpc_derive::rpc;
+use zksync_basic_types::{Address, U256};
+use zksync_core::api_server::web3::backend_jsonrpc::error::into_jsrpc_error;
+use zksync_types::utils::storage_key_for_eth_balance;
+use zksync_utils::u256_to_h256;
+use zksync_web3_decl::error::Web3Error;
+
+/// Implementation of HardhatNamespaceImpl
+pub struct HardhatNamespaceImpl<S> {
+    node: Arc<RwLock<InMemoryNodeInner<S>>>,
+}
+
+impl<S> HardhatNamespaceImpl<S> {
+    /// Creates a new `Hardhat` instance with the given `node`.
+    pub fn new(node: Arc<RwLock<InMemoryNodeInner<S>>>) -> Self {
+        Self { node }
+    }
+}
+
+#[rpc]
+pub trait HardhatNamespaceT {
+    /// Sets the balance of the given address to the given balance.
+    ///
+    /// # Arguments
+    ///
+    /// * `address` - The `Address` whose balance will be edited
+    /// * `balance` - The new balance to set for the given address, in wei
+    ///
+    /// # Returns
+    ///
+    /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation.
+    #[rpc(name = "hardhat_setBalance")]
+    fn set_balance(&self, address: Address, balance: U256) -> BoxFuture<Result<bool>>;
+}
+
+impl<S: Send + Sync + 'static + ForkSource + std::fmt::Debug> HardhatNamespaceT
+    for HardhatNamespaceImpl<S>
+{
+    fn set_balance(
+        &self,
+        address: Address,
+        balance: U256,
+    ) -> jsonrpc_core::BoxFuture<jsonrpc_core::Result<bool>> {
+        let inner = Arc::clone(&self.node);
+
+        Box::pin(async move {
+            match inner.write() {
+                Ok(mut inner_guard) => {
+                    let balance_key = storage_key_for_eth_balance(&address);
+                    inner_guard
+                        .fork_storage
+                        .set_value(balance_key, u256_to_h256(balance));
+                    println!(
+                        "👷 Balance for address {:?} has been manually set to {} Wei",
+                        address, balance
+                    );
+                    Ok(true)
+                }
+                Err(_) => {
+                    let web3_error = Web3Error::InternalError;
+                    Err(into_jsrpc_error(web3_error))
+                }
+            }
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{http_fork_source::HttpForkSource, node::InMemoryNode};
+    use std::str::FromStr;
+    use zksync_core::api_server::web3::backend_jsonrpc::namespaces::eth::EthNamespaceT;
+
+    #[tokio::test]
+    async fn test_set_balance() {
+        let address = Address::from_str("0x36615Cf349d7F6344891B1e7CA7C72883F5dc049").unwrap();
+        let node = InMemoryNode::<HttpForkSource>::default();
+        let hardhat = HardhatNamespaceImpl::new(node.get_inner());
+
+        let balance_before = node.get_balance(address, None).await.unwrap();
+
+        let result = hardhat
+            .set_balance(address, U256::from(1337))
+            .await
+            .unwrap();
+        assert!(result);
+
+        let balance_after = node.get_balance(address, None).await.unwrap();
+        assert_eq!(balance_after, U256::from(1337));
+        assert_ne!(balance_before, balance_after);
+    }
+}

src/main.rs

@@ -41,7 +41,7 @@
 //! ## Contributions
 //!
 //! Contributions to improve `era-test-node` are welcome. Please refer to the contribution guidelines for more details.
-
+use crate::hardhat::{HardhatNamespaceImpl, HardhatNamespaceT};
 use crate::node::{ShowGasDetails, ShowStorageLogs, ShowVMDetails};
 use clap::{Parser, Subcommand};
 use configuration_api::ConfigurationApiNamespaceT;
@@ -55,6 +55,7 @@ mod console_log;
 mod deps;
 mod fork;
 mod formatter;
+mod hardhat;
 mod http_fork_source;
 mod node;
 mod resolver;
@@ -139,6 +140,7 @@ async fn build_json_http<
     net: NetNamespace,
     config_api: ConfigurationApiNamespace<S>,
     zks: ZkMockNamespaceImpl<S>,
+    hardhat: HardhatNamespaceImpl<S>,
 ) -> tokio::task::JoinHandle<()> {
     let (sender, recv) = oneshot::channel::<()>();
 
@@ -148,7 +150,7 @@ async fn build_json_http<
         io.extend_with(net.to_delegate());
         io.extend_with(config_api.to_delegate());
         io.extend_with(zks.to_delegate());
-
+        io.extend_with(hardhat.to_delegate());
         io
     };
 
@@ -315,13 +317,15 @@ async fn main() -> anyhow::Result<()> {
     let net = NetNamespace::new(L2ChainId(TEST_NODE_NETWORK_ID));
     let config_api = ConfigurationApiNamespace::new(node.get_inner());
     let zks = ZkMockNamespaceImpl::new(node.get_inner());
+    let hardhat = HardhatNamespaceImpl::new(node.get_inner());
 
     let threads = build_json_http(
         SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), opt.port),
         node,
         net,
         config_api,
         zks,
+        hardhat,
     )
     .await;
 

test_endpoints.http

@@ -265,6 +265,20 @@ content-type: application/json
 POST http://localhost:8011
 content-type: application/json
 
+{
+    "jsonrpc": "2.0",
+    "id": "2",
+    "method": "hardhat_setBalance",
+    "params": [
+        "0x36615Cf349d7F6344891B1e7CA7C72883F5dc049",
+        "0x1337"
+    ]
+}
+
+###
+POST http://localhost:8011
+content-type: application/json
+
 {
     "jsonrpc": "2.0",
     "id": "1",