feat: types.HeaderHooks

Author: ARR4N

core/state/state.libevm.go

@@ -19,46 +19,47 @@ package state
 import (
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/core/types"
+	"github.com/ava-labs/libevm/libevm/pseudo"
 )
 
 // GetExtra returns the extra payload from the [types.StateAccount] associated
-// with the address, or a zero-value `SA` if not found. The
-// [types.ExtraPayloads] MUST be sourced from [types.RegisterExtras].
-func GetExtra[SA any](s *StateDB, p types.ExtraPayloads[SA], addr common.Address) SA {
+// with the address, or a zero-value `SA` if not found. The [pseudo.Accessor]
+// MUST be sourced from [types.RegisterExtras].
+func GetExtra[SA any](s *StateDB, a pseudo.Accessor[types.StateOrSlimAccount, SA], addr common.Address) SA {
 	stateObject := s.getStateObject(addr)
 	if stateObject != nil {
-		return p.StateAccount.Get(&stateObject.data)
+		return a.Get(&stateObject.data)
 	}
 	var zero SA
 	return zero
 }
 
 // SetExtra sets the extra payload for the address. See [GetExtra] for details.
-func SetExtra[SA any](s *StateDB, p types.ExtraPayloads[SA], addr common.Address, extra SA) {
+func SetExtra[SA any](s *StateDB, a pseudo.Accessor[types.StateOrSlimAccount, SA], addr common.Address, extra SA) {
 	stateObject := s.getOrNewStateObject(addr)
 	if stateObject != nil {
-		setExtraOnObject(stateObject, p, addr, extra)
+		setExtraOnObject(stateObject, a, addr, extra)
 	}
 }
 
-func setExtraOnObject[SA any](s *stateObject, p types.ExtraPayloads[SA], addr common.Address, extra SA) {
+func setExtraOnObject[SA any](s *stateObject, a pseudo.Accessor[types.StateOrSlimAccount, SA], addr common.Address, extra SA) {
 	s.db.journal.append(extraChange[SA]{
-		payloads: p,
+		accessor: a,
 		account:  &addr,
-		prev:     p.StateAccount.Get(&s.data),
+		prev:     a.Get(&s.data),
 	})
-	p.StateAccount.Set(&s.data, extra)
+	a.Set(&s.data, extra)
 }
 
 // extraChange is a [journalEntry] for [SetExtra] / [setExtraOnObject].
 type extraChange[SA any] struct {
-	payloads types.ExtraPayloads[SA]
+	accessor pseudo.Accessor[types.StateOrSlimAccount, SA]
 	account  *common.Address
 	prev     SA
 }
 
 func (e extraChange[SA]) dirtied() *common.Address { return e.account }
 
 func (e extraChange[SA]) revert(s *StateDB) {
-	e.payloads.StateAccount.Set(&s.getStateObject(*e.account).data, e.prev)
+	e.accessor.Set(&s.getStateObject(*e.account).data, e.prev)
 }

core/state/state.libevm_test.go

@@ -45,7 +45,7 @@ func TestGetSetExtra(t *testing.T) {
 	t.Cleanup(types.TestOnlyClearRegisteredExtras)
 	// Just as its Data field is a pointer, the registered type is a pointer to
 	// test deep copying.
-	payloads := types.RegisterExtras[*accountExtra]()
+	payloads := types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, *accountExtra]().StateAccount
 
 	rng := ethtest.NewPseudoRand(42)
 	addr := rng.Address()
@@ -87,7 +87,7 @@ func TestGetSetExtra(t *testing.T) {
 				Root:     types.EmptyRootHash,
 				CodeHash: types.EmptyCodeHash[:],
 			}
-			payloads.StateAccount.Set(want, extra)
+			payloads.Set(want, extra)
 
 			if diff := cmp.Diff(want, got); diff != "" {
 				t.Errorf("types.FullAccount(%T.Account()) diff (-want +got):\n%s", iter, diff)

core/state/state_object.libevm_test.go

@@ -46,21 +46,21 @@ func TestStateObjectEmpty(t *testing.T) {
 		{
 			name: "explicit false bool",
 			registerAndSet: func(acc *types.StateAccount) {
-				types.RegisterExtras[bool]().StateAccount.Set(acc, false)
+				types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, bool]().StateAccount.Set(acc, false)
 			},
 			wantEmpty: true,
 		},
 		{
 			name: "implicit false bool",
 			registerAndSet: func(*types.StateAccount) {
-				types.RegisterExtras[bool]()
+				types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, bool]()
 			},
 			wantEmpty: true,
 		},
 		{
 			name: "true bool",
 			registerAndSet: func(acc *types.StateAccount) {
-				types.RegisterExtras[bool]().StateAccount.Set(acc, true)
+				types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, bool]().StateAccount.Set(acc, true)
 			},
 			wantEmpty: false,
 		},

core/types/block.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/common/hexutil"
+	"github.com/ava-labs/libevm/libevm/pseudo"
 	"github.com/ava-labs/libevm/rlp"
 )
 
@@ -93,6 +94,8 @@ type Header struct {
 
 	// ParentBeaconRoot was added by EIP-4788 and is ignored in legacy headers.
 	ParentBeaconRoot *common.Hash `json:"parentBeaconBlockRoot" rlp:"optional"`
+
+	extra *pseudo.Type // See RegisterExtras()
 }
 
 // field type overrides for gencodec

core/types/block.libevm.go

@@ -17,13 +17,74 @@
 package types
 
 import (
+	"fmt"
 	"io"
 
+	"github.com/ava-labs/libevm/libevm/pseudo"
 	"github.com/ava-labs/libevm/rlp"
 )
 
+// HeaderHooks are required for all types registered with [RegisterExtras] for
+// [Header] payloads.
+type HeaderHooks interface {
+	EncodeRLP(*Header, io.Writer) error
+	DecodeRLP(*Header, *rlp.Stream) error
+}
+
+var _ interface {
+	rlp.Encoder
+	rlp.Decoder
+} = (*Header)(nil)
+
+// EncodeRLP implements the [rlp.Encoder] interface.
 func (h *Header) EncodeRLP(w io.Writer) error {
+	if r := &registeredExtras; r.Registered() {
+		return r.Get().hooks.hooksFromHeader(h).EncodeRLP(h, w)
+	}
 	return h.encodeRLP(w)
 }
 
-var _ rlp.Encoder = (*Header)(nil)
+// decodeHeaderRLPDirectly bypasses the [Header.DecodeRLP] method to avoid
+// infinite recursion.
+func decodeHeaderRLPDirectly(h *Header, s *rlp.Stream) error {
+	type withoutMethods Header
+	return s.Decode((*withoutMethods)(h))
+}
+
+// DecodeRLP implements the [rlp.Decoder] interface.
+func (h *Header) DecodeRLP(s *rlp.Stream) error {
+	if r := &registeredExtras; r.Registered() {
+		return r.Get().hooks.hooksFromHeader(h).DecodeRLP(h, s)
+	}
+	return decodeHeaderRLPDirectly(h, s)
+}
+
+func (e ExtraPayloads[HPtr, SA]) hooksFromHeader(h *Header) HeaderHooks {
+	return e.Header.Get(h)
+}
+
+func (h *Header) extraPayload() *pseudo.Type {
+	r := &registeredExtras
+	if !r.Registered() {
+		// See params.ChainConfig.extraPayload() for panic rationale.
+		panic(fmt.Sprintf("%T.extraPayload() called before RegisterExtras()", r))
+	}
+	if h.extra == nil {
+		h.extra = r.Get().newHeader()
+	}
+	return h.extra
+}
+
+// NOOPHeaderHooks implements [HeaderHooks] such that they are equivalent to
+// no type having been registered.
+type NOOPHeaderHooks struct{}
+
+var _ HeaderHooks = (*NOOPHeaderHooks)(nil)
+
+func (*NOOPHeaderHooks) EncodeRLP(h *Header, w io.Writer) error {
+	return h.encodeRLP(w)
+}
+
+func (*NOOPHeaderHooks) DecodeRLP(h *Header, s *rlp.Stream) error {
+	return decodeHeaderRLPDirectly(h, s)
+}

core/types/block.libevm_test.go

@@ -0,0 +1,112 @@
+// Copyright 2024 the libevm authors.
+//
+// The libevm additions to go-ethereum are free software: you can redistribute
+// them and/or modify them under the terms of the GNU Lesser General Public License
+// as published by the Free Software Foundation, either version 3 of the License,
+// or (at your option) any later version.
+//
+// The libevm additions are distributed in the hope that they will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
+// General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see
+// <http://www.gnu.org/licenses/>.
+
+package types_test
+
+import (
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	. "github.com/ava-labs/libevm/core/types"
+	"github.com/ava-labs/libevm/crypto"
+	"github.com/ava-labs/libevm/libevm/ethtest"
+	"github.com/ava-labs/libevm/rlp"
+)
+
+type stubHeaderHooks struct {
+	rlpSuffix           []byte
+	gotRawRLPToDecode   []byte
+	setHeaderToOnDecode Header
+
+	errEncode, errDecode error
+}
+
+func fakeHeaderRLP(h *Header, suffix []byte) []byte {
+	return append(crypto.Keccak256(h.ParentHash[:]), suffix...)
+}
+
+func (hh *stubHeaderHooks) EncodeRLP(h *Header, w io.Writer) error {
+	if _, err := w.Write(fakeHeaderRLP(h, hh.rlpSuffix)); err != nil {
+		return err
+	}
+	return hh.errEncode
+}
+
+func (hh *stubHeaderHooks) DecodeRLP(h *Header, s *rlp.Stream) error {
+	r, err := s.Raw()
+	if err != nil {
+		return err
+	}
+	hh.gotRawRLPToDecode = r
+	*h = hh.setHeaderToOnDecode
+	return hh.errDecode
+}
+
+func TestHeaderHooks(t *testing.T) {
+	TestOnlyClearRegisteredExtras()
+	defer TestOnlyClearRegisteredExtras()
+
+	extras := RegisterExtras[stubHeaderHooks, *stubHeaderHooks, struct{}]()
+	rng := ethtest.NewPseudoRand(13579)
+
+	t.Run("EncodeRLP", func(t *testing.T) {
+		suffix := rng.Bytes(8)
+
+		hdr := &Header{
+			ParentHash: rng.Hash(),
+		}
+		extras.Header.Get(hdr).rlpSuffix = append([]byte{}, suffix...)
+
+		got, err := rlp.EncodeToBytes(hdr)
+		require.NoError(t, err, "rlp.EncodeToBytes(%T)", hdr)
+		assert.Equal(t, fakeHeaderRLP(hdr, suffix), got)
+	})
+
+	t.Run("DecodeRLP", func(t *testing.T) {
+		input, err := rlp.EncodeToBytes(rng.Bytes(8))
+		require.NoError(t, err)
+
+		hdr := new(Header)
+		stub := &stubHeaderHooks{
+			setHeaderToOnDecode: Header{
+				Extra: []byte("arr4n was here"),
+			},
+		}
+		extras.Header.Set(hdr, stub)
+		require.NoErrorf(t, rlp.DecodeBytes(input, hdr), "rlp.DecodeBytes(%#x)", input)
+
+		assert.Equal(t, input, stub.gotRawRLPToDecode, "raw RLP received by hooks")
+		assert.Equalf(t, &stub.setHeaderToOnDecode, hdr, "%T after RLP decoding with hook")
+	})
+
+	t.Run("error propagation", func(t *testing.T) {
+		errEncode := errors.New("uh oh")
+		errDecode := errors.New("something bad happened")
+
+		hdr := new(Header)
+		extras.Header.Set(hdr, &stubHeaderHooks{
+			errEncode: errEncode,
+			errDecode: errDecode,
+		})
+
+		assert.Equal(t, errEncode, rlp.Encode(io.Discard, hdr), "via rlp.Encode()")
+		assert.Equal(t, errDecode, rlp.DecodeBytes([]byte{0}, hdr), "via rlp.DecodeBytes()")
+	})
+}