From ccf62d22b328fbb79e10ef49233167317285fb98 Mon Sep 17 00:00:00 2001 From: Jrigada Date: Fri, 1 Dec 2023 15:53:06 -0300 Subject: [PATCH] feat: added vm.load functionality and tests --- e2e-tests/contracts/TestCheatcodes.sol | 12 +++++++ e2e-tests/test/cheatcodes.test.ts | 47 ++++++++++++++++++-------- src/cheatcodes.rs | 38 +++++++++++++-------- 3 files changed, 67 insertions(+), 30 deletions(-) diff --git a/e2e-tests/contracts/TestCheatcodes.sol b/e2e-tests/contracts/TestCheatcodes.sol index 171d3d2a..2694c5da 100644 --- a/e2e-tests/contracts/TestCheatcodes.sol +++ b/e2e-tests/contracts/TestCheatcodes.sol @@ -147,6 +147,18 @@ contract TestCheatcodes { testStoreInstance.testStoredValue(value); require(success, "store failed"); } + + function testLoad(bytes32 slot) external { + TestLoadTarget testLoadTarget = new TestLoadTarget(); + (bool success, bytes memory data) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("load(address,bytes32)", address(testLoadTarget), slot)); + require(success, "load failed"); + bytes32 loadedValue = abi.decode(data, (bytes32)); + require(loadedValue == bytes32(uint256(1337)), "address mismatch"); + } +} + +contract TestLoadTarget { + bytes32 public testValue = bytes32(uint256(1337)); //slot 0 } contract testStoreTarget { diff --git a/e2e-tests/test/cheatcodes.test.ts b/e2e-tests/test/cheatcodes.test.ts index 6c9b2f25..faf91a66 100644 --- a/e2e-tests/test/cheatcodes.test.ts +++ b/e2e-tests/test/cheatcodes.test.ts @@ -64,6 +64,23 @@ describe("Cheatcodes", function () { expect(finalRandomWalletCode).to.not.eq(initialRandomWalletCode); }); + it("Should test vm.load", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + const slot = hre.ethers.constants.HashZero; + const tx = await cheatcodes.testLoad(slot, { + gasLimit: 10000000, + }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + it("Should test vm.roll", async function () { // Arrange const wallet = new Wallet(RichAccounts[0].PrivateKey); @@ -140,44 +157,44 @@ describe("Cheatcodes", function () { expect(receipt2.status).to.eq(1); }); - it("Should test vm.warp", async function () { + it("Should test vm.store", async function () { // Arrange const wallet = new Wallet(RichAccounts[0].PrivateKey); const deployer = new Deployer(hre, wallet); - const timeIncreaseInMS = 123; - let expectedTimestamp: number = await provider.send("config_getCurrentTimestamp", []); - expectedTimestamp += timeIncreaseInMS; - // Act const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); - const tx = await cheatcodes.testWarp(expectedTimestamp, { - gasLimit: 1000000, + const slot = hre.ethers.constants.HashZero; + const value = hre.ethers.constants.MaxUint256; + const tx = await cheatcodes.testStore(slot, value, { + gasLimit: 10000000, }); - expectedTimestamp += 2; // New transaction will add two blocks const receipt = await tx.wait(); // Assert expect(receipt.status).to.eq(1); - const newBlockTimestamp = (await provider.getBlock("latest")).timestamp; - expect(newBlockTimestamp).to.equal(expectedTimestamp); }); - it("Should test vm.store", async function () { + it("Should test vm.warp", async function () { // Arrange const wallet = new Wallet(RichAccounts[0].PrivateKey); const deployer = new Deployer(hre, wallet); + const timeIncreaseInMS = 123; + let expectedTimestamp: number = await provider.send("config_getCurrentTimestamp", []); + expectedTimestamp += timeIncreaseInMS; + // Act const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); - const slot = hre.ethers.constants.HashZero; - const value = hre.ethers.constants.MaxUint256; - const tx = await cheatcodes.testStore(slot, value, { - gasLimit: 10000000, + const tx = await cheatcodes.testWarp(expectedTimestamp, { + gasLimit: 1000000, }); + expectedTimestamp += 2; // New transaction will add two blocks const receipt = await tx.wait(); // Assert expect(receipt.status).to.eq(1); + const newBlockTimestamp = (await provider.getBlock("latest")).timestamp; + expect(newBlockTimestamp).to.equal(expectedTimestamp); }); }); diff --git a/src/cheatcodes.rs b/src/cheatcodes.rs index 87029674..caf1fc27 100644 --- a/src/cheatcodes.rs +++ b/src/cheatcodes.rs @@ -86,6 +86,7 @@ abigen!( function deal(address who, uint256 newBalance) function etch(address who, bytes calldata code) function getNonce(address account) + function load(address account, bytes32 slot) function roll(uint256 blockNumber) function setNonce(address account, uint64 nonce) function startPrank(address sender) @@ -276,6 +277,13 @@ impl CheatcodeTracer { tracing::info!("👷 Setting returndata",); self.returndata = Some(vec![account_nonce]); } + Load(LoadCall { account, slot }) => { + tracing::info!("Getting storage slot {:?} for account {:?}", slot, account); + let key = StorageKey::new(AccountTreeId::new(account), H256(slot)); + let mut storage = storage.borrow_mut(); + let value = storage.read_value(&key); + self.returndata = Some(vec![h256_to_u256(value)]); + } SetNonce(SetNonceCall { account, nonce }) => { tracing::info!("Setting nonce for {account:?} to {nonce}"); let mut storage = storage.borrow_mut(); @@ -346,21 +354,6 @@ impl CheatcodeTracer { self.start_prank_opts = None; } - Warp(WarpCall { timestamp }) => { - tracing::info!("Setting block timestamp {}", timestamp); - self.node_ctx.set_time(timestamp.as_u64()); - - let key = StorageKey::new( - AccountTreeId::new(zksync_types::SYSTEM_CONTEXT_ADDRESS), - zksync_types::CURRENT_VIRTUAL_BLOCK_INFO_POSITION, - ); - let mut storage = storage.borrow_mut(); - let (block_number, _) = unpack_block_info(h256_to_u256(storage.read_value(&key))); - storage.set_value( - key, - u256_to_h256(pack_block_info(block_number, timestamp.as_u64())), - ); - } Store(StoreCall { account, slot, @@ -376,6 +369,21 @@ impl CheatcodeTracer { let mut storage = storage.borrow_mut(); storage.set_value(key, H256(value)); } + Warp(WarpCall { timestamp }) => { + tracing::info!("Setting block timestamp {}", timestamp); + self.node_ctx.set_time(timestamp.as_u64()); + + let key = StorageKey::new( + AccountTreeId::new(zksync_types::SYSTEM_CONTEXT_ADDRESS), + zksync_types::CURRENT_VIRTUAL_BLOCK_INFO_POSITION, + ); + let mut storage = storage.borrow_mut(); + let (block_number, _) = unpack_block_info(h256_to_u256(storage.read_value(&key))); + storage.set_value( + key, + u256_to_h256(pack_block_info(block_number, timestamp.as_u64())), + ); + } }; } }