refactor!: temporary extras require proof of global lock

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/extraslock.go

@@ -0,0 +1,63 @@
+// Copyright 2025 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 libevm
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+)
+
+var (
+	extrasMu     sync.Mutex
+	extrasHandle atomic.Uint64
+)
+
+// An ExtrasLock is a handle that proves a current call to
+// [WithTemporaryExtrasLock].
+type ExtrasLock struct {
+	handle *uint64
+}
+
+// WithTemporaryExtrasLock takes a global lock and calls `fn` with a handle that
+// can be used to prove that the lock is held. All package-specific temporary
+// overrides require this proof.
+//
+// WithTemporaryExtrasLock MUST NOT be used on a live chain. It is solely
+// intended for off-chain consumers that require access to extras.
+func WithTemporaryExtrasLock(fn func(lock ExtrasLock) error) error {
+	extrasMu.Lock()
+	defer func() {
+		extrasHandle.Add(1)
+		extrasMu.Unlock()
+	}()
+
+	v := extrasHandle.Load()
+	return fn(ExtrasLock{&v})
+}
+
+// ErrExpiredExtrasLock is returned by [ExtrasLock.Verify] if the lock has been
+// persisted beyond the call to [WithTemporaryExtrasLock] that created it.
+var ErrExpiredExtrasLock = errors.New("libevm.ExtrasLock expired")
+
+// Verify verifies that the lock is valid.
+func (l ExtrasLock) Verify() error {
+	if *l.handle != extrasHandle.Load() {
+		return ErrExpiredExtrasLock
+	}
+	return nil
+}

libevm/extraslock_test.go

@@ -0,0 +1,52 @@
+// Copyright 2025 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 libevm_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	// Testing from outside the package to guarantee usage of the public API only.
+	. "github.com/ava-labs/libevm/libevm"
+)
+
+func TestExtrasLock(t *testing.T) {
+	var zero ExtrasLock
+	assert.Panics(t, func() { _ = zero.Verify() }, "Verify() method of zero-value ExtrasLock{}")
+
+	testIntegration := func(t *testing.T) {
+		t.Helper()
+		require.NoError(t,
+			WithTemporaryExtrasLock((ExtrasLock).Verify),
+			"WithTemporaryExtrasLock((ExtrasLock).Verify)",
+		)
+	}
+	t.Run("initial_usage", testIntegration)
+
+	t.Run("lock_expiration", func(t *testing.T) {
+		var persisted ExtrasLock
+		require.NoError(t, WithTemporaryExtrasLock(func(l ExtrasLock) error {
+			persisted = l
+			return l.Verify()
+		}))
+		assert.ErrorIs(t, persisted.Verify(), ErrExpiredExtrasLock, "Verify() of persisted ExtrasLock")
+	})
+
+	t.Run("repeat_usage", testIntegration)
+}

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()") //nolint:testifylint // Blindly using require is an anti-pattern!!!
+		assert.ErrorIs(t, sut.TempClear(fn), errFoo, "TempClear()")
+	})
 }