more test cases

Author: skudasov

devbox.json

@@ -5,7 +5,7 @@
     "[email protected]",
     "[email protected]",
     "git@latest",
-    "go@latest",
+    "go@1.24.0",
     "goreleaser@latest",
     "postgresql@15",
     "python@latest",

devbox.lock

@@ -0,0 +1 @@
+- Solidity storage layout helper for anvil_setStorageAt + tests + guide

framework/.changeset/v0.8.0.md

@@ -2,20 +2,24 @@
 
 This code is used in e2e tests where we need to modify production contracts with `Anvil`'s `anvil_setStorageAt`.
 
-See a simple example where we override an array [struct](layout_api_test.go)
+See a simple example where we override different types [struct](layout_api_test.go)
 
 Run it with
 ```
 ./setup.sh
 go test -v -run TestLayoutAPI
 ./teardown.sh
 ```
-Figure out proper encoding for your `encodeFunc` and use in your e2e tests.
+This test is more like a playground for you to figure out proper encoding for `encodeFunc` and use in your e2e tests.
 
-See more [tests](layout_test.go) as examples.
+See more package [tests](layout_test.go) as examples.
 
 Layout in `testdata/layout.json` can be found in `out` after `forge build` for any contract.
 
+To double-check the layout you can also use `forge inspect <ContractName> storageLayout` in your `forge` directory.
+
+To add more types for tests use `forge build && forge inspect Counter storageLayout --json > ../testdata/layout.json`
+
 ## Useful Debug Commands
 
 ```

framework/evm_storage/README.md

@@ -2,7 +2,12 @@
 pragma solidity ^0.8.13;
 
 contract Counter {
-    uint256 public number;
+    address private _owner;
+    uint256 public number_uint256;
+    int256 public number_int256;
+    uint8 public number_uint8;
+    bool public boolean;
+    bytes32 public some_bytes;
     uint256[] public values;
     mapping(address => uint256) public scores;
 
@@ -12,11 +17,14 @@ contract Counter {
         uint8 index;
         uint8 group;
     }
+
     mapping(address => Signer) s_signers;
     Signer[] a_signers;
 
     constructor() public {
-        number = 1;
+        number_uint256 = 1;
+        number_int256 = 1;
+        number_uint8 = 1;
         values = [1, 2, 3];
         scores[address(0x5FbDB2315678afecb367f032d93F642f64180aa3)] = 1;
 
@@ -33,12 +41,17 @@ contract Counter {
         }
     }
 
+    // this function is needed to check private field mutation
+    function getOwner() external view returns (address) {
+        return _owner;
+    }
+
     function setNumber(uint256 newNumber) public {
-        number = newNumber;
+        number_uint256 = newNumber;
     }
 
     function increment() public {
-        number++;
+        number_uint256++;
     }
 
     function pushValue(uint256 value) public {

framework/evm_storage/forge/lib/forge-std/scripts/vm.py

@@ -14,11 +14,11 @@ contract CounterTest is Test {
 
     function test_Increment() public {
         counter.increment();
-        assertEq(counter.number(), 1);
+        assertEq(counter.number_uint256(), 1);
     }
 
     function testFuzz_SetNumber(uint256 x) public {
         counter.setNumber(x);
-        assertEq(counter.number(), x);
+        assertEq(counter.number_uint256(), x);
     }
 }

framework/evm_storage/forge/src/Counter.sol

@@ -3,18 +3,25 @@ package evm_storage
 import (
 	"encoding/hex"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
 	"os"
 	"strings"
 
+	"github.com/ethereum/go-ethereum/accounts/abi"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
 )
 
 const (
 	StorageSlotSizeBytes = 32
 )
 
+var (
+	ErrNoSlot = errors.New("no such slot found in layout JSON")
+)
+
 type StorageEntry struct {
 	Label  string `json:"label"`
 	Slot   string `json:"slot"`
@@ -47,19 +54,28 @@ func (s *StorageLayout) GetSlots() map[string]string {
 	return slots
 }
 
-// Slot calculates a slot in Solidity mapping for storage field and a key
-func (s *StorageLayout) Slot(label string) string {
+// MustSlot calculates a slot in Solidity mapping for storage field and a key
+func (s *StorageLayout) MustSlot(label string) string {
+	if _, ok := s.GetSlots()[label]; !ok {
+		panic(fmt.Errorf("layout label: %s, %w", label, ErrNoSlot))
+	}
 	return s.GetSlots()[label]
 }
 
-// MapSlot calculates a slot in Solidity mapping for storage field and a key
-func (s *StorageLayout) MapSlot(label, key string) string {
+// MustMapSlot calculates a slot in Solidity mapping for storage field and a key
+func (s *StorageLayout) MustMapSlot(label, key string) string {
+	if _, ok := s.GetSlots()[label]; !ok {
+		panic(fmt.Errorf("layout label: %s, %w", label, ErrNoSlot))
+	}
 	baseSlot := s.GetSlots()[label]
 	return mapSlot(baseSlot, key)
 }
 
-// ArraySlot calculates a slot in Solidity array for storage field and a key
-func (s *StorageLayout) ArraySlot(label string, index int64) string {
+// MustArraySlot calculates a slot in Solidity array for storage field and a key
+func (s *StorageLayout) MustArraySlot(label string, index int64) string {
+	if _, ok := s.GetSlots()[label]; !ok {
+		panic(fmt.Errorf("layout label: %s, %w", label, ErrNoSlot))
+	}
 	baseSlot := s.GetSlots()[label]
 	return arraySlot(baseSlot, index)
 }
@@ -84,3 +100,88 @@ func mapSlot(baseSlot string, key string) string {
 	hash := crypto.Keccak256(buf)
 	return "0x" + hex.EncodeToString(hash)
 }
+
+func ShiftHexByOffset(hexStr string, offset int) string {
+	// Strip "0x"
+	hexStr = strings.TrimPrefix(hexStr, "0x")
+
+	// Parse into big.Int
+	n := new(big.Int)
+	n.SetString(hexStr, 16)
+
+	// Shift left by offset * 8 bits
+	n.Lsh(n, uint(offset*8))
+
+	// Return as 0x-prefixed, 32-byte hex
+	return fmt.Sprintf("0x%064x", n)
+}
+
+// MustEncodeStorageSlot encodes a value for Solidity storage slots based on type
+// Panics if encoding fails
+func MustEncodeStorageSlot(solidityType string, value interface{}) string {
+	// Handle address type specially
+	if solidityType == "address" {
+		switch v := value.(type) {
+		case common.Address:
+			return fmt.Sprintf("0x%064x", v.Big())
+		case string:
+			if !common.IsHexAddress(v) {
+				panic(fmt.Sprintf("invalid address format: %v", v))
+			}
+			return fmt.Sprintf("0x%064x", common.HexToAddress(v).Big())
+		default:
+			panic(fmt.Sprintf("unsupported address type: %T", value))
+		}
+	}
+
+	// Create the ABI type
+	typ, err := abi.NewType(solidityType, "", nil)
+	if err != nil {
+		panic(fmt.Sprintf("invalid solidity type %q: %v", solidityType, err))
+	}
+
+	// Special handling for bytes32 strings
+	if solidityType == "bytes32" {
+		if s, ok := value.(string); ok {
+			if len(s) > 32 {
+				panic("string too long for bytes32")
+			}
+			var b [32]byte
+			copy(b[:], s)
+			value = b
+		}
+	}
+
+	// Encode the value
+	encoded, err := abi.Arguments{{Type: typ}}.Pack(value)
+	if err != nil {
+		panic(fmt.Sprintf("encoding failed for %v (%T) as %s: %v", value, value, solidityType, err))
+	}
+
+	// For uint256 and int256, we need to take the last 32 bytes
+	if solidityType == "uint256" || solidityType == "int256" {
+		if len(encoded) > 32 {
+			encoded = encoded[len(encoded)-32:]
+		}
+	}
+
+	return fmt.Sprintf("0x%064x", new(big.Int).SetBytes(encoded))
+}
+
+// MergeHex merges two hex strings with bitwise "OR"
+// should be used when you see values with offsets in smart contract storage layout.json file
+// example:
+//
+//	|----------------+-------------------------------------------+------+--------+-------+-------------------------|
+//	| number_uint8   | uint8                                     | 3    | 0      | 1     | src/Counter.sol:Counter |
+//	|----------------+-------------------------------------------+------+--------+-------+-------------------------|
+//	| boolean        | bool                                      | 3    | 1      | 1     | src/Counter.sol:Counter |
+//	|----------------+-------------------------------------------+------+--------+-------+-------------------------|
+func MergeHex(a, b string) string {
+	ai := new(big.Int)
+	bi := new(big.Int)
+	ai.SetString(a[2:], 16)
+	bi.SetString(b[2:], 16)
+	ai.Or(ai, bi)
+	return fmt.Sprintf("0x%064x", ai)
+}

framework/evm_storage/forge/test/Counter.t.sol

@@ -3,9 +3,11 @@ package evm_storage_test
 import (
 	"encoding/hex"
 	"fmt"
+	"math/big"
 	"strings"
 	"testing"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/stretchr/testify/require"
 
 	"github.com/smartcontractkit/chainlink-testing-framework/framework/evm_storage"
@@ -14,7 +16,7 @@ import (
 
 // TestLayoutAPI that's the example of using helpers to override storage in your contracts
 func TestLayoutAPI(t *testing.T) {
-	t.Skip("this test is for manual debugging and figuring out layout of custom structs")
+	//t.Skip("this test is for manual debugging and figuring out layout of custom structs")
 	// load contract layout file, see testdata/layout.json
 	// more docs here - https://docs.soliditylang.org/en/latest/internals/layout_in_storage.html#
 	layout, err := evm_storage.New(layoutFile)
@@ -41,12 +43,92 @@ func TestLayoutAPI(t *testing.T) {
 		return "0x" + hex.EncodeToString(buf)
 	}
 
-	slot := layout.ArraySlot("a_signers", 1)
-	data := encodeFunc("0x00000000000000000000000000000000000000a5", 255, 42)
-	fmt.Printf("setting slot: %s with data: %s\n", slot, data)
-	r := rpc.New(testRPCURL, nil)
-	err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
-	require.NoError(t, err)
-	// verify it manually
-	// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "getASigner(uint256)(address,uint8,uint8)" --rpc-url http://localhost:8545 1
+	{
+		slot := layout.MustSlot("number_uint256")
+		data := evm_storage.MustEncodeStorageSlot("uint256", big.NewInt(222))
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "number_uint256()(uint256)" --rpc-url http://localhost:8545
+	}
+	{
+		slot := layout.MustSlot("number_uint8")
+		data := evm_storage.MustEncodeStorageSlot("uint8", uint8(8))
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "number_uint8()(uint8)" --rpc-url http://localhost:8545
+	}
+	{
+		slot := layout.MustSlot("number_int256")
+		data := evm_storage.MustEncodeStorageSlot("uint256", big.NewInt(221))
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "number_int256()(int256)" --rpc-url http://localhost:8545
+	}
+	{
+		slot := layout.MustSlot("_owner")
+		data := evm_storage.MustEncodeStorageSlot("address", common.HexToAddress("0x5FbDB2315678afecb367f032d93F642f64180aa3"))
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// private fields like _owner can only be verified by getter
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "getOwner()(address)" --rpc-url http://localhost:8545
+	}
+	{
+		slot := layout.MustArraySlot("a_signers", 1)
+		data := encodeFunc("0x00000000000000000000000000000000000000a5", 255, 42)
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "getASigner(uint256)(address,uint8,uint8)" --rpc-url http://localhost:8545 1
+	}
+	{
+		slot := layout.MustMapSlot("s_signers", "0x00000000000000000000000000000000000000a5")
+		data := encodeFunc("0x00000000000000000000000000000000000000a5", 254, 40)
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "getSSigner(address)(address,uint8,uint8)" 0x00000000000000000000000000000000000000a5 --rpc-url http://localhost:8545 1
+	}
+	{
+		// offset example
+		slot := layout.MustSlot("boolean")
+		data := evm_storage.MustEncodeStorageSlot("bool", true)
+		boolValue := evm_storage.ShiftHexByOffset(data, 1)
+		uint8Value := evm_storage.MustEncodeStorageSlot("uint8", uint8(8))
+		data = evm_storage.MergeHex(uint8Value, boolValue)
+		fmt.Printf("setting slot: %s with data: %s\n", slot, data)
+		r := rpc.New(testRPCURL, nil)
+		err = r.AnvilSetStorageAt([]interface{}{contractAddr, slot, data})
+		require.NoError(t, err)
+		// Contract code:
+		// contract Counter {
+		//    address private _owner;
+		//    uint256 public number_uint256;
+		//    int256 public number_int256;
+		//    uint8 public number_uint8; <-- we need to change this
+		//    bool public boolean;       <-- and this
+		//
+		// Example layout:
+		// ╭----------------+-------------------------------------------+------+--------+-------+-------------------------╮
+		// | Name           | Type                                      | Slot | Offset | Bytes | Contract                |
+		// |----------------+-------------------------------------------+------+--------+-------+-------------------------|
+		// | number_uint8   | uint8                                     | 3    | 0      | 1     | src/Counter.sol:Counter |
+		// |----------------+-------------------------------------------+------+--------+-------+-------------------------|
+		// | boolean        | bool                                      | 3    | 1      | 1     | src/Counter.sol:Counter |
+		// |----------------+-------------------------------------------+------+--------+-------+-------------------------|
+		// Resulting value with offsets: 0x0000000000000000000000000000000000000000000000000000000000000108
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "boolean()(bool)" --rpc-url http://localhost:8545
+		// true
+		// cast call 0x5FbDB2315678afecb367f032d93F642f64180aa3 "number_uint8()(uint8)" --rpc-url http://localhost:8545
+		// 8
+	}
 }