ethclient/gethclient: fix bugs in override object encoding

ethclient/gethclient/gethclient.go

@@ -19,6 +19,7 @@ package gethclient
 
 import (
 	"context"
+	"encoding/json"
 	"math/big"
 	"runtime"
 	"runtime/debug"
@@ -118,15 +119,6 @@ func (ec *Client) GetProof(ctx context.Context, account common.Address, keys []s
 	return &result, err
 }
 
-// OverrideAccount specifies the state of an account to be overridden.
-type OverrideAccount struct {
-	Nonce     uint64                      `json:"nonce"`
-	Code      []byte                      `json:"code"`
-	Balance   *big.Int                    `json:"balance"`
-	State     map[common.Hash]common.Hash `json:"state"`
-	StateDiff map[common.Hash]common.Hash `json:"stateDiff"`
-}
-
 // CallContract executes a message call transaction, which is directly executed in the VM
 // of the node, but never mined into the blockchain.
 //
@@ -141,7 +133,7 @@ func (ec *Client) CallContract(ctx context.Context, msg ethereum.CallMsg, blockN
 	var hex hexutil.Bytes
 	err := ec.c.CallContext(
 		ctx, &hex, "eth_call", toCallArg(msg),
-		toBlockNumArg(blockNumber), toOverrideMap(overrides),
+		toBlockNumArg(blockNumber), overrides,
 	)
 	return hex, err
 }
@@ -210,26 +202,48 @@ func toCallArg(msg ethereum.CallMsg) interface{} {
 	return arg
 }
 
-func toOverrideMap(overrides *map[common.Address]OverrideAccount) interface{} {
-	if overrides == nil {
-		return nil
-	}
-	type overrideAccount struct {
-		Nonce     hexutil.Uint64              `json:"nonce"`
-		Code      hexutil.Bytes               `json:"code"`
-		Balance   *hexutil.Big                `json:"balance"`
-		State     map[common.Hash]common.Hash `json:"state"`
-		StateDiff map[common.Hash]common.Hash `json:"stateDiff"`
-	}
-	result := make(map[common.Address]overrideAccount)
-	for addr, override := range *overrides {
-		result[addr] = overrideAccount{
-			Nonce:     hexutil.Uint64(override.Nonce),
-			Code:      override.Code,
-			Balance:   (*hexutil.Big)(override.Balance),
-			State:     override.State,
-			StateDiff: override.StateDiff,
-		}
-	}
-	return &result
+// OverrideAccount specifies the state of an account to be overridden.
+type OverrideAccount struct {
+	// Nonce sets nonce of the account. Note: the nonce override will only
+	// be applied when it is set to a non-zero value.
+	Nonce uint64
+
+	// Code sets the contract code. The override will be applied
+	// when the code is non-nil, i.e. setting empty code is possible
+	// using an empty slice.
+	Code []byte
+
+	// Balance sets the account balance.
+	Balance *big.Int
+
+	// State sets the complete storage. The override will be applied
+	// when the given map is non-nil. Using an empty map wipes the
+	// entire contract storage during the call.
+	State map[common.Hash]common.Hash
+
+	// StateDiff allows overriding individual storage slots.
+	StateDiff 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"`
+	}
+
+	output := acc{
+		Nonce:     hexutil.Uint64(a.Nonce),
+		Balance:   (*hexutil.Big)(a.Balance),
+		StateDiff: a.StateDiff,
+	}
+	if a.Code != nil {
+		output.Code = hexutil.Encode(a.Code)
+	}
+	if a.State != nil {
+		output.State = a.State
+	}
+	return json.Marshal(output)
 }

ethclient/gethclient/gethclient_test.go

@@ -19,6 +19,7 @@ package gethclient
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"math/big"
 	"testing"
 
@@ -327,3 +328,53 @@ func testCallContract(t *testing.T, client *rpc.Client) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }
+
+func TestOverrideAccountMarshal(t *testing.T) {
+	om := map[common.Address]OverrideAccount{
+		common.Address{0x11}: OverrideAccount{
+			// Zero-valued nonce is not overriddden, but simply dropped by the encoder.
+			Nonce: 0,
+		},
+		common.Address{0xaa}: OverrideAccount{
+			Nonce: 5,
+		},
+		common.Address{0xbb}: OverrideAccount{
+			Code: []byte{1},
+		},
+		common.Address{0xcc}: OverrideAccount{
+			// 'code', 'balance', 'state' should be set when input is
+			// a non-nil but empty value.
+			Code:    []byte{},
+			Balance: big.NewInt(0),
+			State:   map[common.Hash]common.Hash{},
+			// For 'stateDiff' the behavior is different, empty map
+			// is ignored because it makes no difference.
+			StateDiff: map[common.Hash]common.Hash{},
+		},
+	}
+
+	marshalled, err := json.MarshalIndent(&om, "", "  ")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	expected := `{
+  "0x1100000000000000000000000000000000000000": {},
+  "0xaa00000000000000000000000000000000000000": {
+    "nonce": "0x5"
+  },
+  "0xbb00000000000000000000000000000000000000": {
+    "code": "0x01"
+  },
+  "0xcc00000000000000000000000000000000000000": {
+    "code": "0x",
+    "balance": "0x0",
+    "state": {}
+  }
+}`
+
+	if string(marshalled) != expected {
+		t.Error("wrong output:", string(marshalled))
+		t.Error("want:", expected)
+	}
+}