From 4debd3d31cdbca35994237b9fe780918959fd8b1 Mon Sep 17 00:00:00 2001 From: Damilola Edwards Date: Mon, 8 Jan 2024 15:41:20 +0100 Subject: [PATCH 1/5] add snapshot and revertTo logic --- chain/standard_cheat_code_contract.go | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/chain/standard_cheat_code_contract.go b/chain/standard_cheat_code_contract.go index 88b8f6d2..fdb15b58 100644 --- a/chain/standard_cheat_code_contract.go +++ b/chain/standard_cheat_code_contract.go @@ -282,6 +282,27 @@ func getStandardCheatCodeContract(tracer *cheatCodeTracer) (*CheatCodeContract, return nil, nil }, ) + + // snapshot: Takes a snapshot of the current state of the evm and returns the id associated with the snapshot + contract.addMethod( + "snapshot", abi.Arguments{}, abi.Arguments{{Type: typeUint256}}, + func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) { + snapshotID := tracer.evm.StateDB.Snapshot() + + return []any{snapshotID}, nil + }, + ) + + // revertTo(uint256): Revert the state of the evm to a previous snapshot. Takes the snapshot id to revert to. + contract.addMethod( + "revertTo", abi.Arguments{{Type: typeUint256}}, abi.Arguments{{Type: typeBool}}, + func(tracer *cheatCodeTracer, inputs []any) ([]any, *cheatCodeRawReturnData) { + snapshotID := inputs[0].(*big.Int) + tracer.evm.StateDB.RevertToSnapshot(int(snapshotID.Int64())) + + return []any{true}, nil + }, + ) // FFI: Run arbitrary command on base OS contract.addMethod( From 4bb231aefbfd66b9a9707d317fdd7fcfebfb04c7 Mon Sep 17 00:00:00 2001 From: Damilola Edwards Date: Mon, 8 Jan 2024 15:49:59 +0100 Subject: [PATCH 2/5] added test files --- fuzzing/fuzzer_test.go | 1 + .../cheat_codes/vm/snapshot_and_revert_to.sol | 54 +++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol diff --git a/fuzzing/fuzzer_test.go b/fuzzing/fuzzer_test.go index 7fdc4f47..d8acc6de 100644 --- a/fuzzing/fuzzer_test.go +++ b/fuzzing/fuzzer_test.go @@ -209,6 +209,7 @@ func TestCheatCodes(t *testing.T) { "testdata/contracts/cheat_codes/utils/to_string.sol", "testdata/contracts/cheat_codes/utils/sign.sol", "testdata/contracts/cheat_codes/utils/parse.sol", + "testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol", "testdata/contracts/cheat_codes/vm/coinbase.sol", "testdata/contracts/cheat_codes/vm/chain_id.sol", "testdata/contracts/cheat_codes/vm/deal.sol", diff --git a/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol b/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol new file mode 100644 index 00000000..f19fc731 --- /dev/null +++ b/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol @@ -0,0 +1,54 @@ +// This test ensures that we can take a snapshot of the current state of the testchain and revert to the state at that snapshot using the snapshot and revertTo cheatcodes +pragma solidity ^0.8.0; + +interface CheatCodes { + function warp(uint256) external; + + function deal(address, uint256) external; + + function snapshot() external returns (uint256); + + function revertTo(uint256) external returns (bool); +} + + struct Storage { + uint slot0; + uint slot1; + } + +contract TestContract { + Storage store; + uint256 timestamp; + + function test() public { + // Obtain our cheat code contract reference. + CheatCodes cheats = CheatCodes( + 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D + ); + + store.slot0 = 10; + store.slot1 = 20; + timestamp = block.timestamp; + cheats.deal(address(this), 5 ether); + + uint256 snapshot = cheats.snapshot(); // saves the state + + // let's change the state + store.slot0 = 300; + store.slot1 = 400; + cheats.deal(address(this), 500 ether); + cheats.warp(12345); // block.timestamp = 12345 + + assert(store.slot0 == 300); + assert(store.slot1 == 400); + assert(address(this).balance == 500 ether); + assert(block.timestamp == 12345); + + cheats.revertTo(snapshot); // restores the state + + assert(store.slot0 == 10); + assert(store.slot1 == 20); + assert(address(this).balance == 5 ether); + assert(block.timestamp == timestamp); + } +} \ No newline at end of file From 7d03603edcf06291c2ad696d945f66df050a9e52 Mon Sep 17 00:00:00 2001 From: Damilola Edwards Date: Mon, 8 Jan 2024 15:57:17 +0100 Subject: [PATCH 3/5] go fmt --- chain/standard_cheat_code_contract.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/chain/standard_cheat_code_contract.go b/chain/standard_cheat_code_contract.go index fdb15b58..07fb5071 100644 --- a/chain/standard_cheat_code_contract.go +++ b/chain/standard_cheat_code_contract.go @@ -282,7 +282,7 @@ func getStandardCheatCodeContract(tracer *cheatCodeTracer) (*CheatCodeContract, return nil, nil }, ) - + // snapshot: Takes a snapshot of the current state of the evm and returns the id associated with the snapshot contract.addMethod( "snapshot", abi.Arguments{}, abi.Arguments{{Type: typeUint256}}, From a1cc0f825257c936e5650dbebf6347ede31c65dd Mon Sep 17 00:00:00 2001 From: Anish Naik Date: Mon, 25 Mar 2024 13:30:24 -0400 Subject: [PATCH 4/5] fix formatting --- .../cheat_codes/vm/snapshot_and_revert_to.sol | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol b/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol index f19fc731..ff6a8feb 100644 --- a/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol +++ b/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol @@ -11,10 +11,10 @@ interface CheatCodes { function revertTo(uint256) external returns (bool); } - struct Storage { - uint slot0; - uint slot1; - } +struct Storage { + uint slot0; + uint slot1; +} contract TestContract { Storage store; @@ -31,21 +31,25 @@ contract TestContract { timestamp = block.timestamp; cheats.deal(address(this), 5 ether); - uint256 snapshot = cheats.snapshot(); // saves the state + // Save state + uint256 snapshot = cheats.snapshot(); - // let's change the state + // Change state store.slot0 = 300; store.slot1 = 400; cheats.deal(address(this), 500 ether); - cheats.warp(12345); // block.timestamp = 12345 + cheats.warp(12345); + // Assert that state has been changed assert(store.slot0 == 300); assert(store.slot1 == 400); assert(address(this).balance == 500 ether); assert(block.timestamp == 12345); - cheats.revertTo(snapshot); // restores the state + // Revert to snapshot + cheats.revertTo(snapshot); + // Ensure state has been reset assert(store.slot0 == 10); assert(store.slot1 == 20); assert(address(this).balance == 5 ether); From d0a9955708ac7dddb9eb221b41b5e9c26123c378 Mon Sep 17 00:00:00 2001 From: Anish Naik Date: Mon, 25 Mar 2024 13:31:03 -0400 Subject: [PATCH 5/5] fix formatting again --- .../contracts/cheat_codes/vm/snapshot_and_revert_to.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol b/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol index ff6a8feb..577ff194 100644 --- a/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol +++ b/fuzzing/testdata/contracts/cheat_codes/vm/snapshot_and_revert_to.sol @@ -55,4 +55,4 @@ contract TestContract { assert(address(this).balance == 5 ether); assert(block.timestamp == timestamp); } -} \ No newline at end of file +}