feat: PrecompileEnvironment.ReentrancyGuard()

Author: ARR4N

core/vm/contracts.libevm.go

@@ -190,6 +190,9 @@ type PrecompileEnvironment interface {
 	// Invalidate invalidates the transaction calling this precompile.
 	InvalidateExecution(error)
 
+	// ReentrancyGuard ...
+	ReentrancyGuard(key []byte) error
+
 	// Call is equivalent to [EVM.Call] except that the `caller` argument is
 	// removed and automatically determined according to the type of call that
 	// invoked the precompile.

core/vm/contracts.libevm_test.go

@@ -839,3 +839,45 @@ func ExamplePrecompileEnvironment() {
 	// variable to include it in this example function.
 	_ = actualCaller
 }
+
+func TestReentrancyGuard(t *testing.T) {
+	sut := common.HexToAddress("7E57ED")
+	eve := common.HexToAddress("BAD")
+	eveCalled := false
+
+	zero := func() *uint256.Int {
+		return uint256.NewInt(0)
+	}
+
+	returnIfGuarded := []byte("guarded")
+
+	hooks := &hookstest.Stub{
+		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
+			eve: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
+				eveCalled = true
+				return env.Call(sut, []byte{}, env.Gas(), zero()) // i.e. reenter
+			}),
+			sut: vm.NewStatefulPrecompile(func(env vm.PrecompileEnvironment, input []byte) (ret []byte, err error) {
+				// The argument is optional and used only to allow more than one
+				// guard in a contract.
+				if err := env.ReentrancyGuard(nil); err != nil {
+					return returnIfGuarded, err
+				}
+				if env.Addresses().Caller == eve {
+					// A real precompile MUST NOT panic under any circumstances.
+					// It is done here to avoid a loop should the guard not
+					// work.
+					panic("reentrancy")
+				}
+				return env.Call(eve, []byte{}, env.Gas(), zero())
+			}),
+		},
+	}
+	hooks.Register(t)
+
+	_, evm := ethtest.NewZeroEVM(t)
+	got, _, err := evm.Call(vm.AccountRef{}, sut, []byte{}, 1e6, zero())
+	require.True(t, eveCalled, "Malicious contract called")
+	assert.Equal(t, err, vm.ErrExecutionReverted, "Precompile reverted")
+	assert.Equal(t, returnIfGuarded, got, "Precompile reverted with expected data")
+}

core/vm/environment.libevm.go

@@ -25,6 +25,7 @@ import (
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/common/math"
 	"github.com/ava-labs/libevm/core/types"
+	"github.com/ava-labs/libevm/crypto"
 	"github.com/ava-labs/libevm/libevm"
 	"github.com/ava-labs/libevm/libevm/options"
 	"github.com/ava-labs/libevm/params"
@@ -103,6 +104,21 @@ func (e *environment) BlockHeader() (types.Header, error) {
 	return *hdr, nil
 }
 
+func reentrancyGuardSlot(key []byte) common.Hash {
+	return crypto.Keccak256Hash([]byte("libevm-reentrancy-guard"), key)
+}
+
+func (e *environment) ReentrancyGuard(key []byte) error {
+	self := e.Addresses().Self
+	slot := reentrancyGuardSlot(key)
+
+	if e.evm.StateDB.GetTransientState(self, slot) != (common.Hash{}) {
+		return ErrExecutionReverted
+	}
+	e.evm.StateDB.SetTransientState(e.Addresses().Self, slot, common.Hash{1})
+	return nil
+}
+
 func (e *environment) Call(addr common.Address, input []byte, gas uint64, value *uint256.Int, opts ...CallOption) ([]byte, error) {
 	return e.callContract(Call, addr, input, gas, value, opts...)
 }