From 5f0fafd8e6d5afc48f7caf84c6c8e188850444d5 Mon Sep 17 00:00:00 2001
From: Sina Mahmoodi <itz.s1na@gmail.com>
Date: Fri, 16 Aug 2024 16:41:31 +0200
Subject: [PATCH] fix precompile moveTo

---
 internal/ethapi/api.go      | 17 +++++++
 internal/ethapi/api_test.go | 92 +++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/internal/ethapi/api.go b/internal/ethapi/api.go
index 3a465d35cd..14b7a824ef 100644
--- a/internal/ethapi/api.go
+++ b/internal/ethapi/api.go
@@ -979,12 +979,23 @@ type OverrideAccount struct {
 // StateOverride is the collection of overridden accounts.
 type StateOverride map[common.Address]OverrideAccount
 
+func (diff *StateOverride) has(address common.Address) bool {
+	_, ok := (*diff)[address]
+	return ok
+}
+
 // Apply overrides the fields of specified accounts into the given state.
 func (diff *StateOverride) Apply(statedb *state.StateDB, precompiles vm.PrecompiledContracts) error {
 	if diff == nil {
 		return nil
 	}
+	// Tracks destinations of precompiles that were moved.
+	dirtyAddrs := make(map[common.Address]struct{})
 	for addr, account := range *diff {
+		// If a precompile was moved to this address already, it can't be overridden.
+		if _, ok := dirtyAddrs[addr]; ok {
+			return fmt.Errorf("account %s has already been overridden by a precompile", addr.Hex())
+		}
 		p, isPrecompile := precompiles[addr]
 		// The MoveTo feature makes it possible to move a precompile
 		// code to another address. If the target address is another precompile
@@ -994,7 +1005,13 @@ func (diff *StateOverride) Apply(statedb *state.StateDB, precompiles vm.Precompi
 			if !isPrecompile {
 				return fmt.Errorf("account %s is not a precompile", addr.Hex())
 			}
+			// Refuse to move a precompile to an address that has been
+			// or will be overridden.
+			if diff.has(*account.MovePrecompileTo) {
+				return fmt.Errorf("account %s is already overridden", account.MovePrecompileTo.Hex())
+			}
 			precompiles[*account.MovePrecompileTo] = p
+			dirtyAddrs[*account.MovePrecompileTo] = struct{}{}
 		}
 		if isPrecompile {
 			delete(precompiles, addr)
diff --git a/internal/ethapi/api_test.go b/internal/ethapi/api_test.go
index 7ce8c29795..202b573fc0 100644
--- a/internal/ethapi/api_test.go
+++ b/internal/ethapi/api_test.go
@@ -35,6 +35,7 @@ import (
 
 	"github.com/holiman/uint256"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
 
 	"github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/accounts"
@@ -3240,6 +3241,97 @@ func TestRPCGetBlockReceipts(t *testing.T) {
 	}
 }
 
+type precompileContract struct{}
+
+func (p *precompileContract) RequiredGas(input []byte) uint64 { return 0 }
+
+func (p *precompileContract) Run(input []byte) ([]byte, error) { return nil, nil }
+
+func TestStateOverrideMovePrecompile(t *testing.T) {
+	db := state.NewDatabase(rawdb.NewMemoryDatabase())
+	statedb, err := state.New(common.Hash{}, db, nil)
+	if err != nil {
+		t.Fatalf("failed to create statedb: %v", err)
+	}
+	precompiles := map[common.Address]vm.PrecompiledContract{
+		common.BytesToAddress([]byte{0x1}): &precompileContract{},
+		common.BytesToAddress([]byte{0x2}): &precompileContract{},
+	}
+	bytes2Addr := func(b []byte) *common.Address {
+		a := common.BytesToAddress(b)
+		return &a
+	}
+	var testSuite = []struct {
+		overrides           StateOverride
+		expectedPrecompiles map[common.Address]struct{}
+		fail                bool
+	}{
+		{
+			overrides: StateOverride{
+				common.BytesToAddress([]byte{0x1}): {
+					Code:             hex2Bytes("0xff"),
+					MovePrecompileTo: bytes2Addr([]byte{0x2}),
+				},
+				common.BytesToAddress([]byte{0x2}): {
+					Code: hex2Bytes("0x00"),
+				},
+			},
+			// 0x2 has already been touched by the moveTo.
+			fail: true,
+		}, {
+			overrides: StateOverride{
+				common.BytesToAddress([]byte{0x1}): {
+					Code:             hex2Bytes("0xff"),
+					MovePrecompileTo: bytes2Addr([]byte{0xff}),
+				},
+				common.BytesToAddress([]byte{0x3}): {
+					Code:             hex2Bytes("0x00"),
+					MovePrecompileTo: bytes2Addr([]byte{0xfe}),
+				},
+			},
+			// 0x3 is not a precompile.
+			fail: true,
+		}, {
+			overrides: StateOverride{
+				common.BytesToAddress([]byte{0x1}): {
+					Code:             hex2Bytes("0xff"),
+					MovePrecompileTo: bytes2Addr([]byte{0xff}),
+				},
+				common.BytesToAddress([]byte{0x2}): {
+					Code:             hex2Bytes("0x00"),
+					MovePrecompileTo: bytes2Addr([]byte{0xfe}),
+				},
+			},
+			expectedPrecompiles: map[common.Address]struct{}{common.BytesToAddress([]byte{0xfe}): {}, common.BytesToAddress([]byte{0xff}): {}},
+		},
+	}
+
+	for i, tt := range testSuite {
+		cpy := maps.Clone(precompiles)
+		// Apply overrides
+		err := tt.overrides.Apply(statedb, cpy)
+		if tt.fail {
+			if err == nil {
+				t.Errorf("test %d: want error, have nothing", i)
+			}
+			continue
+		}
+		if err != nil {
+			t.Errorf("test %d: want no error, have %v", i, err)
+			continue
+		}
+		// Precompile keys
+		if len(cpy) != len(tt.expectedPrecompiles) {
+			t.Errorf("test %d: precompile mismatch, want %d, have %d", i, len(tt.expectedPrecompiles), len(cpy))
+		}
+		for k, _ := range tt.expectedPrecompiles {
+			if _, ok := cpy[k]; !ok {
+				t.Errorf("test %d: precompile not found: %s", i, k.String())
+			}
+		}
+	}
+}
+
 func testRPCResponseWithFile(t *testing.T, testid int, result interface{}, rpc string, file string) {
 	data, err := json.MarshalIndent(result, "", "  ")
 	if err != nil {