From e9dd9d536f0492a706a1ee9c888ec31bdf3a11ea Mon Sep 17 00:00:00 2001 From: aon <21188659+aon@users.noreply.github.com> Date: Tue, 5 Dec 2023 11:09:07 -0300 Subject: [PATCH 1/2] feat: add vm.toString cheatcode --- e2e-tests/contracts/TestCheatcodes.sol | 81 +++++++++++++++++++++++ e2e-tests/test/cheatcodes.test.ts | 90 ++++++++++++++++++++++++++ src/cheatcodes.rs | 58 ++++++++++++++++- 3 files changed, 228 insertions(+), 1 deletion(-) diff --git a/e2e-tests/contracts/TestCheatcodes.sol b/e2e-tests/contracts/TestCheatcodes.sol index d6737496..7a150f65 100644 --- a/e2e-tests/contracts/TestCheatcodes.sol +++ b/e2e-tests/contracts/TestCheatcodes.sol @@ -126,6 +126,73 @@ contract TestCheatcodes { ); } + function testToStringFromAddress() external { + address testAddress = 0x413D15117be7a498e68A64FcfdB22C6e2AaE1808; + (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call( + abi.encodeWithSignature("toString(address)", testAddress) + ); + bytes memory data = trimReturnBytes(rawData); + string memory testString = string(abi.encodePacked(data)); + require( + keccak256(bytes(testString)) == keccak256(bytes("0x413D15117be7a498e68A64FcfdB22C6e2AaE1808")), + "toString mismatch" + ); + } + + function testToStringFromBool() external { + (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("toString(bool)", false)); + bytes memory data = trimReturnBytes(rawData); + string memory testString = string(abi.encodePacked(data)); + require(keccak256(bytes(testString)) == keccak256(bytes("false")), "toString mismatch"); + + (success, rawData) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("toString(bool)", true)); + data = trimReturnBytes(rawData); + testString = string(abi.encodePacked(data)); + require(keccak256(bytes(testString)) == keccak256(bytes("true")), "toString mismatch"); + } + + function testToStringFromUint256(uint256 value, string memory stringValue) external { + (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("toString(uint256)", value)); + bytes memory data = trimReturnBytes(rawData); + string memory testString = string(abi.encodePacked(data)); + require(keccak256(bytes(testString)) == keccak256(bytes(stringValue)), "toString mismatch"); + } + + function testToStringFromInt256(int256 value, string memory stringValue) external { + (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature("toString(int256)", value)); + bytes memory data = trimReturnBytes(rawData); + string memory testString = string(abi.encodePacked(data)); + require(keccak256(bytes(testString)) == keccak256(bytes(stringValue)), "toString mismatch"); + } + + function testToStringFromBytes32() external { + bytes32 testBytes = hex"4ec893b0a778b562e893cee722869c3e924e9ee46ec897cabda6b765a6624324"; + (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call( + abi.encodeWithSignature("toString(bytes32)", testBytes) + ); + bytes memory data = trimReturnBytes(rawData); + string memory testString = string(abi.encodePacked(data)); + require( + keccak256(bytes(testString)) == + keccak256(bytes("0x4ec893b0a778b562e893cee722869c3e924e9ee46ec897cabda6b765a6624324")), + "toString mismatch" + ); + } + + function testToStringFromBytes() external { + bytes memory testBytes = hex"89987299ea14decf0e11d068474a6e459439802edca8aacf9644222e490d8ef6db"; + (bool success, bytes memory rawData) = CHEATCODE_ADDRESS.call( + abi.encodeWithSignature("toString(bytes)", testBytes) + ); + bytes memory data = trimReturnBytes(rawData); + string memory testString = string(abi.encodePacked(data)); + require( + keccak256(bytes(testString)) == + keccak256(bytes("0x89987299ea14decf0e11d068474a6e459439802edca8aacf9644222e490d8ef6db")), + "toString mismatch" + ); + } + function testWarp(uint256 timestamp) external { uint256 initialTimestamp = block.timestamp; require(timestamp != initialTimestamp, "timestamp must be different than current block timestamp"); @@ -148,6 +215,20 @@ contract TestCheatcodes { testStoreInstance.testStoredValue(value); } + + function trimReturnBytes(bytes memory rawData) internal pure returns (bytes memory) { + uint256 lengthStartingPos = rawData.length - 32; + bytes memory lengthSlice = new bytes(32); + for (uint256 i = 0; i < 32; i++) { + lengthSlice[i] = rawData[lengthStartingPos + i]; + } + uint256 length = abi.decode(lengthSlice, (uint256)); + bytes memory data = new bytes(length); + for (uint256 i = 0; i < length; i++) { + data[i] = rawData[i]; + } + return data; + } } contract testStoreTarget { diff --git a/e2e-tests/test/cheatcodes.test.ts b/e2e-tests/test/cheatcodes.test.ts index 4ccbc2cd..2909d823 100644 --- a/e2e-tests/test/cheatcodes.test.ts +++ b/e2e-tests/test/cheatcodes.test.ts @@ -140,6 +140,96 @@ describe("Cheatcodes", function () { expect(receipt2.status).to.eq(1); }); + it("Should test vm.toString from address", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + const tx = await cheatcodes.testToStringFromAddress({ gasLimit: 100000000 }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + + it("Should test vm.toString from bool", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + const tx = await cheatcodes.testToStringFromBool({ gasLimit: 100000000 }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + + it("Should test vm.toString from uint256", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + const value = hre.ethers.BigNumber.from(hre.ethers.utils.randomBytes(32)); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + + const tx = await cheatcodes.testToStringFromUint256(value, value.toString(), { gasLimit: 100000000 }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + + it("Should test vm.toString from int256", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + const value = hre.ethers.BigNumber.from(hre.ethers.utils.randomBytes(32)); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + + const tx = await cheatcodes.testToStringFromUint256(value, value.toString(), { gasLimit: 100000000 }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + + it("Should test vm.toString from bytes32", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + + const tx = await cheatcodes.testToStringFromBytes32({ gasLimit: 100000000 }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + + it("Should test vm.toString from bytes", async function () { + // Arrange + const wallet = new Wallet(RichAccounts[0].PrivateKey); + const deployer = new Deployer(hre, wallet); + + // Act + const cheatcodes = await deployContract(deployer, "TestCheatcodes", []); + + const tx = await cheatcodes.testToStringFromBytes({ gasLimit: 100000000 }); + const receipt = await tx.wait(); + + // Assert + expect(receipt.status).to.eq(1); + }); + it("Should test vm.warp", async function () { // Arrange const wallet = new Wallet(RichAccounts[0].PrivateKey); diff --git a/src/cheatcodes.rs b/src/cheatcodes.rs index d6b3610a..cc27c737 100644 --- a/src/cheatcodes.rs +++ b/src/cheatcodes.rs @@ -2,7 +2,7 @@ use crate::{ node::{BlockContext, InMemoryNodeInner}, utils::bytecode_to_factory_dep, }; -use ethers::{abi::AbiDecode, prelude::abigen}; +use ethers::{abi::AbiDecode, prelude::abigen, utils::to_checksum}; use itertools::Itertools; use multivm::zk_evm_1_3_3::tracing::AfterExecutionData; use multivm::zk_evm_1_3_3::vm_state::PrimitiveValue; @@ -99,6 +99,12 @@ abigen!( function startPrank(address sender) function startPrank(address sender, address origin) function stopPrank() + function toString(address value) + function toString(bool value) + function toString(uint256 value) + function toString(int256 value) + function toString(bytes32 value) + function toString(bytes value) function warp(uint256 timestamp) function store(address account, bytes32 slot, bytes32 value) ]"# @@ -370,6 +376,37 @@ impl CheatcodeTracer { self.start_prank_opts = None; } + ToString0(ToString0Call { value }) => { + tracing::info!("Converting address into string"); + let address = Address::from(value); + let address_with_checksum = to_checksum(&address, None); + self.add_return_data(address_with_checksum.as_bytes()); + } + ToString1(ToString1Call { value }) => { + tracing::info!("Converting bool into string"); + let bool_value = value.to_string(); + self.add_return_data(bool_value.as_bytes()); + } + ToString2(ToString2Call { value }) => { + tracing::info!("Converting uint256 into string"); + let uint_value = value.to_string(); + self.add_return_data(uint_value.as_bytes()); + } + ToString3(ToString3Call { value }) => { + tracing::info!("Converting int256 into string"); + let int_value = value.to_string(); + self.add_return_data(int_value.as_bytes()); + } + ToString4(ToString4Call { value }) => { + tracing::info!("Converting bytes32 into string"); + let bytes_value = format!("0x{}", hex::encode(value)); + self.add_return_data(bytes_value.as_bytes()); + } + ToString5(ToString5Call { value }) => { + tracing::info!("Converting bytes into string"); + let bytes_value = format!("0x{}", hex::encode(value)); + self.add_return_data(bytes_value.as_bytes()); + } Warp(WarpCall { timestamp }) => { tracing::info!("Setting block timestamp {}", timestamp); self.node_ctx.set_time(timestamp.as_u64()); @@ -412,6 +449,25 @@ impl CheatcodeTracer { write_value: value, }); } + + fn add_return_data(&mut self, data: &[u8]) { + let data_length = data.len(); + let mut data: Vec = data + .chunks(32) + .map(|b| { + // Copies the bytes into a 32 byte array + // padding with zeros to the right if necessary + let mut bytes = [0u8; 32]; + bytes[..b.len()].copy_from_slice(b); + bytes.into() + }) + .collect_vec(); + + // Add the length of the data to the end of the return data + data.push(data_length.into()); + + self.returndata = Some(data); + } } pub struct CheatcodeNodeContext { From 5187aecc95a04313113923926b81f2fb93327882 Mon Sep 17 00:00:00 2001 From: aon <21188659+aon@users.noreply.github.com> Date: Wed, 6 Dec 2023 14:53:50 -0300 Subject: [PATCH 2/2] fix: reorder cheatcodes --- src/cheatcodes.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/cheatcodes.rs b/src/cheatcodes.rs index 4da41485..f37df115 100644 --- a/src/cheatcodes.rs +++ b/src/cheatcodes.rs @@ -90,6 +90,8 @@ pub trait NodeCtx { abigen!( CheatcodeContract, r#"[ + + function addr(uint256 privateKey) function deal(address who, uint256 newBalance) function etch(address who, bytes calldata code) @@ -100,6 +102,7 @@ abigen!( function startPrank(address sender) function startPrank(address sender, address origin) function stopPrank() + function store(address account, bytes32 slot, bytes32 value) function toString(address value) function toString(bool value) function toString(uint256 value) @@ -107,7 +110,6 @@ abigen!( function toString(bytes32 value) function toString(bytes value) function warp(uint256 timestamp) - function store(address account, bytes32 slot, bytes32 value) ]"# );