internal/ethapi: add transient storage override support for eth_call

Author: hero5512

ethclient/gethclient/gethclient.go

@@ -301,21 +301,26 @@ type OverrideAccount struct {
 
 	// StateDiff allows overriding individual storage slots.
 	StateDiff map[common.Hash]common.Hash
+
+	// TransientStorage allows overriding transient storage slots.
+	TransientStorage map[common.Hash]common.Hash
 }
 
 func (a OverrideAccount) MarshalJSON() ([]byte, error) {
 	type acc struct {
-		Nonce     hexutil.Uint64              `json:"nonce,omitempty"`
-		Code      string                      `json:"code,omitempty"`
-		Balance   *hexutil.Big                `json:"balance,omitempty"`
-		State     interface{}                 `json:"state,omitempty"`
-		StateDiff map[common.Hash]common.Hash `json:"stateDiff,omitempty"`
+		Nonce            hexutil.Uint64              `json:"nonce,omitempty"`
+		Code             string                      `json:"code,omitempty"`
+		Balance          *hexutil.Big                `json:"balance,omitempty"`
+		State            interface{}                 `json:"state,omitempty"`
+		StateDiff        map[common.Hash]common.Hash `json:"stateDiff,omitempty"`
+		TransientStorage map[common.Hash]common.Hash `json:"transientStorage,omitempty"`
 	}
 
 	output := acc{
-		Nonce:     hexutil.Uint64(a.Nonce),
-		Balance:   (*hexutil.Big)(a.Balance),
-		StateDiff: a.StateDiff,
+		Nonce:            hexutil.Uint64(a.Nonce),
+		Balance:          (*hexutil.Big)(a.Balance),
+		StateDiff:        a.StateDiff,
+		TransientStorage: a.TransientStorage,
 	}
 	if a.Code != nil {
 		output.Code = hexutil.Encode(a.Code)

internal/ethapi/api_test.go

@@ -1226,6 +1226,24 @@ func TestCall(t *testing.T) {
 			},
 			expectErr: errors.New(`block override "withdrawals" is not supported for this RPC method`),
 		},
+		// Test transient storage override
+		{
+			name: "transient storage override takes effect",
+			call: TransactionArgs{
+				From: &accounts[1].addr,
+				To:   &randomAccounts[2].addr,
+			},
+			overrides: override.StateOverride{
+				randomAccounts[2].addr: override.OverrideAccount{
+					// PUSH1 0x00 TLOAD PUSH1 0x00 MSTORE PUSH1 0x20 PUSH1 0x00 RETURN
+					Code: hex2Bytes("0x60005c60005260206000f3"),
+					TransientStorage: map[common.Hash]common.Hash{
+						common.Hash{}: common.HexToHash("0xabcd"),
+					},
+				},
+			},
+			want: "0x000000000000000000000000000000000000000000000000000000000000abcd",
+		},
 	}
 	for _, tc := range testSuite {
 		result, err := api.Call(context.Background(), tc.call, &rpc.BlockNumberOrHash{BlockNumber: &tc.blockNumber}, &tc.overrides, &tc.blockOverrides)

internal/ethapi/override/override.go

@@ -33,7 +33,7 @@ import (
 // OverrideAccount indicates the overriding fields of account during the execution
 // of a message call.
 // Note, state and stateDiff can't be specified at the same time. If state is
-// set, message execution will only use the data in the given state. Otherwise
+// set, message execution will only use the data in the given state. Otherwise,
 // if stateDiff is set, all diff will be applied first and then execute the call
 // message.
 type OverrideAccount struct {
@@ -42,6 +42,7 @@ type OverrideAccount struct {
 	Balance          *hexutil.Big                `json:"balance"`
 	State            map[common.Hash]common.Hash `json:"state"`
 	StateDiff        map[common.Hash]common.Hash `json:"stateDiff"`
+	TransientStorage map[common.Hash]common.Hash `json:"transientStorage"`
 	MovePrecompileTo *common.Address             `json:"movePrecompileToAddress"`
 }
 
@@ -111,6 +112,12 @@ func (diff *StateOverride) Apply(statedb *state.StateDB, precompiles vm.Precompi
 				statedb.SetState(addr, key, value)
 			}
 		}
+		// Apply transient storage overrides.
+		if account.TransientStorage != nil {
+			for key, value := range account.TransientStorage {
+				statedb.SetTransientState(addr, key, value)
+			}
+		}
 	}
 	// Now finalize the changes. Finalize is normally performed between transactions.
 	// By using finalize, the overrides are semantically behaving as

internal/ethapi/override/override_test.go

@@ -128,3 +128,58 @@ func hex2Bytes(str string) *hexutil.Bytes {
 	rpcBytes := hexutil.Bytes(common.FromHex(str))
 	return &rpcBytes
 }
+
+func TestStateOverrideTransientStorage(t *testing.T) {
+	db := state.NewDatabase(triedb.NewDatabase(rawdb.NewMemoryDatabase(), nil), nil)
+	statedb, err := state.New(types.EmptyRootHash, db)
+	if err != nil {
+		t.Fatalf("failed to create statedb: %v", err)
+	}
+
+	addr := common.BytesToAddress([]byte{0x1})
+	key1 := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000001")
+	key2 := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000002")
+	value1 := common.HexToHash("0x1111111111111111111111111111111111111111111111111111111111111111")
+	value2 := common.HexToHash("0x2222222222222222222222222222222222222222222222222222222222222222")
+
+	// Verify initial state is empty
+	if got := statedb.GetTransientState(addr, key1); got != (common.Hash{}) {
+		t.Fatalf("expected initial transient state to be empty, got %s", got.Hex())
+	}
+	if got := statedb.GetTransientState(addr, key2); got != (common.Hash{}) {
+		t.Fatalf("expected initial transient state to be empty, got %s", got.Hex())
+	}
+
+	// Apply override with transient storage
+	override := StateOverride{
+		addr: OverrideAccount{
+			TransientStorage: map[common.Hash]common.Hash{
+				key1: value1,
+				key2: value2,
+			},
+		},
+	}
+
+	if err := override.Apply(statedb, nil); err != nil {
+		t.Fatalf("failed to apply override: %v", err)
+	}
+
+	// Verify transient storage was set
+	if got := statedb.GetTransientState(addr, key1); got != value1 {
+		t.Errorf("expected transient state for key1 to be %s, got %s", value1.Hex(), got.Hex())
+	}
+	if got := statedb.GetTransientState(addr, key2); got != value2 {
+		t.Errorf("expected transient state for key2 to be %s, got %s", value2.Hex(), got.Hex())
+	}
+
+	// Verify other addresses/keys remain empty
+	otherAddr := common.BytesToAddress([]byte{0x2})
+	if got := statedb.GetTransientState(otherAddr, key1); got != (common.Hash{}) {
+		t.Errorf("expected transient state for different address to be empty, got %s", got.Hex())
+	}
+
+	otherKey := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000003")
+	if got := statedb.GetTransientState(addr, otherKey); got != (common.Hash{}) {
+		t.Errorf("expected transient state for different key to be empty, got %s", got.Hex())
+	}
+}