feat: add hardhat_setCode

SUPPORTED_APIS.md

@@ -97,7 +97,7 @@ The `status` options are:
 | [`HARDHAT`](#hardhat-namespace) | [`hardhat_mine`](#hardhat_mine) | Mine any number of blocks at once, in constant time |
 | `HARDHAT` | `hardhat_reset` | `NOT IMPLEMENTED` | Resets the state of the network |
 | [`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-namespace) | [`hardhat_setCode`](#hardhat_setcode) | `SUPPORTED` | 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 |
 | `HARDHAT` | `hardhat_setMinGasPrice` | `NOT IMPLEMENTED` | Sets the minimum gas price |
@@ -1489,6 +1489,38 @@ curl --request POST \
 }'
 ```
 
+### `hardhat_setCode`
+
+[source](src/hardhat.rs)
+
+Sets the code for a given address.
+
+#### Arguments
+
++ `address: Address` - The `Address` whose code will be updated
++ `code: Bytes` - The code to set to
+
+#### Status
+
+`SUPPORTED`
+
+#### Example
+
+```bash
+curl --request POST \
+  --url http://localhost:8011/ \
+  --header 'content-type: application/json' \
+  --data '{
+    "jsonrpc": "2.0",
+      "id": "1",
+      "method": "hardhat_setCode",
+      "params": [
+        "0x36615Cf349d7F6344891B1e7CA7C72883F5dc049",
+        [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
+      ]
+  }'
+```
+
 ## `EVM NAMESPACE`
 
 ### `evm_mine`

e2e-tests/contracts/Return5.sol

@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity >=0.8.2 <0.9.0;
+
+contract Return5 {
+  function value() public pure returns (uint8) {
+    return 5;
+  }
+}

e2e-tests/test/hardhat-apis.test.ts

@@ -1,8 +1,12 @@
 import { expect } from "chai";
 import { Wallet } from "zksync-web3";
-import { getTestProvider } from "../helpers/utils";
+import { deployContract, getTestProvider } from "../helpers/utils";
 import { RichAccounts } from "../helpers/constants";
 import { ethers } from "hardhat";
+import { Deployer } from "@matterlabs/hardhat-zksync-deploy";
+import * as hre from "hardhat";
+import { keccak256 } from "ethers/lib/utils";
+import { BigNumber } from "ethers";
 
 const provider = getTestProvider();
 
@@ -84,3 +88,62 @@ xdescribe("hardhat_impersonateAccount & hardhat_stopImpersonatingAccount", funct
     expect(await provider.getBalance(RichAccounts[0].Account)).to.equal(beforeBalance.sub(0.42));
   });
 });
+
+describe("hardhat_setCode", function () {
+  it("Should set code at an address", async function () {
+    // Arrange
+    const wallet = new Wallet(RichAccounts[0].PrivateKey);
+    const deployer = new Deployer(hre, wallet);
+
+    const address = "0x1000000000000000000000000000000000001111";
+    const artifact = await deployer.loadArtifact("Return5");
+    const contractCode = [...ethers.utils.arrayify(artifact.deployedBytecode)];
+
+    // Act
+    await provider.send("hardhat_setCode", [address, contractCode]);
+
+    // Assert
+    const result = await provider.send("eth_call", [
+      {
+        to: address,
+        data: keccak256(ethers.utils.toUtf8Bytes("value()")).substring(0, 10),
+        from: wallet.address,
+        gas: "0x1000",
+        gasPrice: "0x0ee6b280",
+        value: "0x0",
+        nonce: "0x1",
+      },
+      "latest",
+    ]);
+    expect(BigNumber.from(result).toNumber()).to.eq(5);
+  });
+
+  it("Should update code with a different smart contract", async function () {
+    // Arrange
+    const wallet = new Wallet(RichAccounts[0].PrivateKey);
+    const deployer = new Deployer(hre, wallet);
+
+    const greeter = await deployContract(deployer, "Greeter", ["Hi"]);
+    expect(await greeter.greet()).to.eq("Hi");
+    const artifact = await deployer.loadArtifact("Return5");
+    const newContractCode = [...ethers.utils.arrayify(artifact.deployedBytecode)];
+
+    // Act
+    await provider.send("hardhat_setCode", [greeter.address, newContractCode]);
+
+    // Assert
+    const result = await provider.send("eth_call", [
+      {
+        to: greeter.address,
+        data: keccak256(ethers.utils.toUtf8Bytes("value()")).substring(0, 10),
+        from: wallet.address,
+        gas: "0x1000",
+        gasPrice: "0x0ee6b280",
+        value: "0x0",
+        nonce: "0x1",
+      },
+      "latest",
+    ]);
+    expect(BigNumber.from(result).toNumber()).to.eq(5);
+  });
+});

src/hardhat.rs

@@ -1,13 +1,17 @@
 use std::sync::{Arc, RwLock};
 
-use crate::{fork::ForkSource, node::InMemoryNodeInner, utils::mine_empty_blocks};
+use crate::{
+    fork::ForkSource,
+    node::InMemoryNodeInner,
+    utils::{bytecode_to_factory_dep, mine_empty_blocks, IntoBoxedFuture},
+};
 use jsonrpc_core::{BoxFuture, Result};
 use jsonrpc_derive::rpc;
 use zksync_basic_types::{Address, U256, U64};
 use zksync_core::api_server::web3::backend_jsonrpc::error::into_jsrpc_error;
 use zksync_state::ReadStorage;
 use zksync_types::{
-    get_nonce_key,
+    get_code_key, get_nonce_key,
     utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
 };
 use zksync_utils::{h256_to_u256, u256_to_h256};
@@ -99,6 +103,19 @@ pub trait HardhatNamespaceT {
     /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation.
     #[rpc(name = "hardhat_stopImpersonatingAccount")]
     fn stop_impersonating_account(&self, address: Address) -> BoxFuture<Result<bool>>;
+
+    /// Modifies the bytecode stored at an account's address.
+    ///
+    /// # Arguments
+    ///
+    /// * `address` - The address where the given code should be stored.
+    /// * `code` - The code to be stored.
+    ///
+    /// # Returns
+    ///
+    /// 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>) -> BoxFuture<Result<()>>;
 }
 
 impl<S: Send + Sync + 'static + ForkSource + std::fmt::Debug> HardhatNamespaceT
@@ -241,6 +258,31 @@ impl<S: Send + Sync + 'static + ForkSource + std::fmt::Debug> HardhatNamespaceT
             }
         })
     }
+
+    fn set_code(&self, address: Address, code: Vec<u8>) -> BoxFuture<Result<()>> {
+        let inner = Arc::clone(&self.node);
+        inner
+            .write()
+            .map(|mut writer| {
+                let code_key = get_code_key(&address);
+                log::info!("set code for address {address:#x}");
+                let (hash, code) = bytecode_to_factory_dep(code);
+                let hash = u256_to_h256(hash);
+                writer.fork_storage.store_factory_dep(
+                    hash,
+                    code.iter()
+                        .flat_map(|entry| {
+                            let mut bytes = vec![0u8; 32];
+                            entry.to_big_endian(&mut bytes);
+                            bytes.to_vec()
+                        })
+                        .collect(),
+                );
+                writer.fork_storage.set_value(code_key, hash);
+            })
+            .map_err(|_| into_jsrpc_error(Web3Error::InternalError))
+            .into_boxed_future()
+    }
 }
 
 #[cfg(test)]
@@ -446,4 +488,31 @@ mod tests {
         // execution should now fail again
         assert!(node.apply_txs(vec![tx]).is_err());
     }
+
+    #[tokio::test]
+    async fn test_set_code() {
+        let address = Address::repeat_byte(0x1);
+        let node = InMemoryNode::<HttpForkSource>::default();
+        let hardhat = HardhatNamespaceImpl::new(node.get_inner());
+        let new_code = vec![0x1u8; 32];
+
+        let code_before = node
+            .get_code(address, None)
+            .await
+            .expect("failed getting code")
+            .0;
+        assert_eq!(Vec::<u8>::default(), code_before);
+
+        hardhat
+            .set_code(address, new_code.clone())
+            .await
+            .expect("failed setting code");
+
+        let code_after = node
+            .get_code(address, None)
+            .await
+            .expect("failed getting code")
+            .0;
+        assert_eq!(new_code, code_after);
+    }
 }

test_endpoints.http

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