refactor!: all temp-registration functions require a lock

Author: ARR4N

core/state/statedb.libevm.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/core/state/snapshot"
+	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/libevm/register"
 	"github.com/ava-labs/libevm/libevm/stateconf"
 )
@@ -91,8 +92,11 @@ func RegisterExtras(s StateDBHooks) {
 // consumers that require access to extras. Said consumers SHOULD NOT, however
 // call this function directly. Use the libevm/temporary.WithRegisteredExtras()
 // function instead as it atomically overrides all possible packages.
-func WithTempRegisteredExtras(s StateDBHooks, fn func()) {
-	registeredExtras.TempOverride(s, fn)
+func WithTempRegisteredExtras(lock libevm.ExtrasLock, s StateDBHooks, fn func() error) error {
+	if err := lock.Verify(); err != nil {
+		return err
+	}
+	return registeredExtras.TempOverride(s, fn)
 }
 
 // TestOnlyClearRegisteredExtras clears the arguments previously passed to

core/state/statedb.libevm_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/ava-labs/libevm/core/state/snapshot"
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/ethdb"
+	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/trie"
 	"github.com/ava-labs/libevm/trie/trienode"
@@ -216,13 +217,17 @@ func TestTransformStateKey(t *testing.T) {
 	assertCommittedEq(t, flippedKey, flippedVal, noTransform)
 
 	t.Run("WithTempRegisteredExtras", func(t *testing.T) {
-		WithTempRegisteredExtras(noopHooks{}, func() {
-			// No-op hooks are equivalent to using the `noTransform` option.
-			// NOTE this is NOT the intended usage of [WithTempRegisteredExtras]
-			// and is simply an easy way to test the temporary registration.
-			assertEq(t, regularKey, regularVal)
-			assertEq(t, flippedKey, flippedVal)
+		err := libevm.WithTemporaryExtrasLock(func(lock libevm.ExtrasLock) error {
+			return WithTempRegisteredExtras(lock, noopHooks{}, func() error {
+				// No-op hooks are equivalent to using the `noTransform` option.
+				// NOTE this is NOT the intended usage of [WithTempRegisteredExtras]
+				// and is simply an easy way to test the temporary registration.
+				assertEq(t, regularKey, regularVal)
+				assertEq(t, flippedKey, flippedVal)
+				return nil
+			})
 		})
+		require.NoError(t, err)
 	})
 
 	updatedVal := common.Hash{'u', 'p', 'd', 'a', 't', 'e', 'd'}

core/types/rlp_payload.libevm.go

@@ -21,6 +21,7 @@ import (
 	"io"
 
 	"github.com/ava-labs/libevm/common"
+	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/libevm/pseudo"
 	"github.com/ava-labs/libevm/libevm/register"
 	"github.com/ava-labs/libevm/libevm/testonly"
@@ -114,9 +115,12 @@ func WithTempRegisteredExtras[
 	H, B, SA any,
 	HPtr HeaderHooksPointer[H],
 	BPtr BlockBodyHooksPointer[B, BPtr],
-](fn func(ExtraPayloads[HPtr, BPtr, SA])) {
+](lock libevm.ExtrasLock, fn func(ExtraPayloads[HPtr, BPtr, SA]) error) error {
+	if err := lock.Verify(); err != nil {
+		return err
+	}
 	payloads, ctors := payloadsAndConstructors[H, HPtr, B, BPtr, SA]()
-	registeredExtras.TempOverride(ctors, func() { fn(payloads) })
+	return registeredExtras.TempOverride(ctors, func() error { return fn(payloads) })
 }
 
 // A HeaderHooksPointer is a type constraint for an implementation of

core/types/tempextras.libevm_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/rlp"
 )
 
@@ -58,19 +59,23 @@ func TestTempRegisteredExtras(t *testing.T) {
 
 	t.Run("before_temp", testPrimaryExtras)
 	t.Run("WithTempRegisteredExtras", func(t *testing.T) {
-		WithTempRegisteredExtras(func(extras ExtraPayloads[*NOOPHeaderHooks, *tempBlockBodyHooks, bool]) {
-			const val = "Hello, world"
-			b := new(Block)
-			payload := &tempBlockBodyHooks{X: val}
-			extras.Block.Set(b, payload)
+		err := libevm.WithTemporaryExtrasLock(func(lock libevm.ExtrasLock) error {
+			return WithTempRegisteredExtras(lock, func(extras ExtraPayloads[*NOOPHeaderHooks, *tempBlockBodyHooks, bool]) error {
+				const val = "Hello, world"
+				b := new(Block)
+				payload := &tempBlockBodyHooks{X: val}
+				extras.Block.Set(b, payload)
 
-			got, err := rlp.EncodeToBytes(b)
-			require.NoErrorf(t, err, "rlp.EncodeToBytes(%T) with %T hooks", b, extras.Block.Get(b))
-			want, err := rlp.EncodeToBytes([]string{val})
-			require.NoErrorf(t, err, "rlp.EncodeToBytes(%T{%[1]v})", []string{val})
+				got, err := rlp.EncodeToBytes(b)
+				require.NoErrorf(t, err, "rlp.EncodeToBytes(%T) with %T hooks", b, extras.Block.Get(b))
+				want, err := rlp.EncodeToBytes([]string{val})
+				require.NoErrorf(t, err, "rlp.EncodeToBytes(%T{%[1]v})", []string{val})
 
-			assert.Equalf(t, want, got, "rlp.EncodeToBytes(%T) with %T hooks", b, payload)
+				assert.Equalf(t, want, got, "rlp.EncodeToBytes(%T) with %T hooks", b, payload)
+				return nil
+			})
 		})
+		require.NoError(t, err)
 	})
 	t.Run("after_temp", testPrimaryExtras)
 }

core/vm/evm.libevm_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/params"
 )
 
@@ -72,10 +73,14 @@ func TestOverrideNewEVMArgs(t *testing.T) {
 	assertChainID(t, chainID)
 
 	t.Run("WithTempRegisteredHooks", func(t *testing.T) {
-		override := evmArgOverrider{newEVMchainID: 24680}
-		WithTempRegisteredHooks(&override, func() {
-			assertChainID(t, override.newEVMchainID)
+		err := libevm.WithTemporaryExtrasLock(func(lock libevm.ExtrasLock) error {
+			override := evmArgOverrider{newEVMchainID: 24680}
+			return WithTempRegisteredHooks(lock, &override, func() error {
+				assertChainID(t, override.newEVMchainID)
+				return nil
+			})
 		})
+		require.NoError(t, err)
 		t.Run("after", func(t *testing.T) {
 			assertChainID(t, chainID)
 		})

core/vm/hooks.libevm.go

@@ -17,6 +17,7 @@
 package vm
 
 import (
+	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/libevm/register"
 	"github.com/ava-labs/libevm/params"
 )
@@ -36,8 +37,11 @@ func RegisterHooks(h Hooks) {
 // consumers that require access to extras. Said consumers SHOULD NOT, however
 // call this function directly. Use the libevm/temporary.WithRegisteredExtras()
 // function instead as it atomically overrides all possible packages.
-func WithTempRegisteredHooks(h Hooks, fn func()) {
-	libevmHooks.TempOverride(h, fn)
+func WithTempRegisteredHooks(lock libevm.ExtrasLock, h Hooks, fn func() error) error {
+	if err := lock.Verify(); err != nil {
+		return err
+	}
+	return libevmHooks.TempOverride(h, fn)
 }
 
 // TestOnlyClearRegisteredHooks clears the [Hooks] previously passed to

libevm/register/register.go

@@ -72,22 +72,23 @@ func (o *AtMostOnce[T]) TestOnlyClear() {
 //
 // It is valid to call this method with or without a prior call to
 // [AtMostOnce.Register].
-func (o *AtMostOnce[T]) TempOverride(with T, fn func()) {
-	o.temp(&with, fn)
+func (o *AtMostOnce[T]) TempOverride(with T, fn func() error) error {
+	return o.temp(&with, fn)
 }
 
 // TempClear calls `fn`, clearing any registered `T`, but only for the life of
 // the call. It is not threadsafe.
 //
 // It is valid to call this method with or without a prior call to
 // [AtMostOnce.Register].
-func (o *AtMostOnce[T]) TempClear(fn func()) {
-	o.temp(nil, fn)
+func (o *AtMostOnce[T]) TempClear(fn func() error) error {
+	return o.temp(nil, fn)
 }
 
-func (o *AtMostOnce[T]) temp(with *T, fn func()) {
+func (o *AtMostOnce[T]) temp(with *T, fn func() error) error {
 	old := o.v
 	o.v = with
-	fn()
+	err := fn()
 	o.v = old
+	return err
 }

libevm/register/register_test.go

@@ -17,6 +17,7 @@
 package register
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -56,9 +57,10 @@ func TestAtMostOnce(t *testing.T) {
 
 	t.Run("TempOverride", func(t *testing.T) {
 		t.Run("during", func(t *testing.T) {
-			sut.TempOverride(val+1, func() {
+			require.NoError(t, sut.TempOverride(val+1, func() error {
 				assertRegistered(t, val+1)
-			})
+				return nil
+			}))
 		})
 		t.Run("after", func(t *testing.T) {
 			assertRegistered(t, val)
@@ -67,12 +69,20 @@ func TestAtMostOnce(t *testing.T) {
 
 	t.Run("TempClear", func(t *testing.T) {
 		t.Run("during", func(t *testing.T) {
-			sut.TempClear(func() {
+			require.NoError(t, sut.TempClear(func() error {
 				assert.False(t, sut.Registered(), "Registered()")
-			})
+				return nil
+			}))
 		})
 		t.Run("after", func(t *testing.T) {
 			assertRegistered(t, val)
 		})
 	})
+
+	t.Run("error_propagation", func(t *testing.T) {
+		errFoo := errors.New("foo")
+		fn := func() error { return errFoo }
+		assert.ErrorIs(t, sut.TempOverride(0, fn), errFoo, "TempOverride()")
+		assert.ErrorIs(t, sut.TempClear(fn), errFoo, "TempClear()")
+	})
 }

params/config.libevm.go

@@ -114,11 +114,8 @@ func WithTempRegisteredExtras[C ChainConfigHooks, R RulesHooks](
 	if err := lock.Verify(); err != nil {
 		return err
 	}
-
 	payloads, ctors := payloadsAndConstructors(e)
-	var err error
-	registeredExtras.TempOverride(ctors, func() { err = fn(payloads) })
-	return err
+	return registeredExtras.TempOverride(ctors, func() error { return fn(payloads) })
 }
 
 // TestOnlyClearRegisteredExtras clears the [Extras] previously passed to