migration: Add tests

op-chain-ops/cmd/celo-migrate/state.go

@@ -34,11 +34,11 @@ var (
 	alfajoresChainId uint64 = 44787
 	mainnetChainId   uint64 = 42220
 
-	// Whitelist of accounts that are allowed to be overwritten
+	// Allowlist of accounts that are allowed to be overwritten
 	// If the value for an account is set to true, the nonce and storage will be overwritten
 	// This must be checked for each account, as this might create issues with contracts
 	// calling `CREATE` or `CREATE2`
-	accountOverwriteWhitelist = map[uint64]map[common.Address]bool{
+	accountOverwriteAllowlist = map[uint64]map[common.Address]bool{
 		// Add any addresses that should be allowed to overwrite existing accounts here.
 		alfajoresChainId: {
 			// Create2Deployer
@@ -96,7 +96,7 @@ func applyStateMigrationChanges(config *genesis.DeployConfig, genesis *core.Gene
 	}
 
 	// Apply the changes to the state DB.
-	err = applyAllocsToState(db, genesis, cfg)
+	err = applyAllocsToState(db, genesis, accountOverwriteAllowlist[cfg.ChainID.Uint64()])
 	if err != nil {
 		return nil, fmt.Errorf("cannot apply allocations to state: %w", err)
 	}
@@ -246,43 +246,48 @@ func applyStateMigrationChanges(config *genesis.DeployConfig, genesis *core.Gene
 // If an account already exists, it adds the balance of the new account to the existing balance.
 // If the code of an existing account is different from the code in the genesis block, it logs a warning.
 // This changes the state root, so `Commit` needs to be called after this function.
-func applyAllocsToState(db vm.StateDB, genesis *core.Genesis, config *params.ChainConfig) error {
+func applyAllocsToState(db vm.StateDB, genesis *core.Genesis, allowlist map[common.Address]bool) error {
 	log.Info("Starting to migrate OP contracts into state DB")
 
 	copyCounter := 0
 	overwriteCounter := 0
-	whitelist := accountOverwriteWhitelist[config.ChainID.Uint64()]
 
 	for k, v := range genesis.Alloc {
 		// Check that the balance of the account to written is zero,
-		// as we must not create new CELo tokens
-		if v.Balance.Cmp(big.NewInt(0)) != 0 {
-			log.Error("Account balance is not zero, would change celo supply", "address", k.Hex())
+		// as we must not create new CELO tokens
+		if v.Balance != nil && v.Balance.Cmp(big.NewInt(0)) != 0 {
 			return fmt.Errorf("account balance is not zero, would change celo supply: %s", k.Hex())
 		}
 
-		overwriteNonceAndState := true
+		overwrite := true
 		if db.Exist(k) {
-			var whitelisted bool
-			overwriteNonceAndState, whitelisted = whitelist[k]
+			var allowed bool
+			overwrite, allowed = allowlist[k]
 
-			// If the account is not whitelisted and has a non zero nonce or code size, bail out we will need to manually investigate how to handle this.
-			if !whitelisted && (db.GetCodeSize(k) > 0 || db.GetNonce(k) > 0) {
-				return fmt.Errorf("account exists and is not whitelisted, account: %s, nonce: %d, code: %d", k.Hex(), db.GetNonce(k), db.GetCode(k))
+			// If the account is not allowed and has a non zero nonce or code size, bail out we will need to manually investigate how to handle this.
+			if !allowed && (db.GetCodeSize(k) > 0 || db.GetNonce(k) > 0) {
+				return fmt.Errorf("account exists and is not allowed, account: %s, nonce: %d, code: %d", k.Hex(), db.GetNonce(k), db.GetCode(k))
+			}
+
+			// This means that the account just has balance, in that case we wan to copy over the account
+			if db.GetCodeSize(k) == 0 && db.GetNonce(k) == 0 {
+				overwrite = true
 			}
-			overwriteCounter++
 		}
 
 		// This carries over any existing balance
 		db.CreateAccount(k)
-		db.SetCode(k, v.Code)
 
-		if overwriteNonceAndState {
+		if overwrite {
+			overwriteCounter++
+
+			db.SetCode(k, v.Code)
 			db.SetNonce(k, v.Nonce)
 			for key, value := range v.Storage {
 				db.SetState(k, key, value)
 			}
 		}
+
 		copyCounter++
 		log.Info("Copied account", "address", k.Hex())
 	}

op-chain-ops/cmd/celo-migrate/state_test.go

@@ -0,0 +1,151 @@
+package main
+
+import (
+	"bytes"
+	"math/big"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	contractCode         = []byte{0x01, 0x02}
+	defaultBalance int64 = 123
+)
+
+func TestApplyAllocsToState(t *testing.T) {
+	tests := []struct {
+		name             string
+		addr             common.Address
+		existingAccount  *types.Account
+		newAccount       types.Account
+		allowlist        map[common.Address]bool
+		balanceInAccount bool
+		wantErr          bool
+	}{
+		{
+			name: "Write to empty account",
+			addr: common.HexToAddress("01"),
+			newAccount: types.Account{
+				Code:  contractCode,
+				Nonce: 1,
+			},
+			balanceInAccount: false,
+			wantErr:          false,
+		},
+		{
+			name: "Copy account with non-zero balance fails",
+			addr: common.HexToAddress("a"),
+			existingAccount: &types.Account{
+				Balance: big.NewInt(defaultBalance),
+			},
+			newAccount: types.Account{
+				Balance: big.NewInt(1),
+			},
+			wantErr: true,
+		},
+		{
+			name: "Write to account with only balance should overwrite and keep balance",
+			addr: common.HexToAddress("a"),
+			existingAccount: &types.Account{
+				Balance: big.NewInt(defaultBalance),
+			},
+			newAccount: types.Account{
+				Code:  contractCode,
+				Nonce: 5,
+			},
+			balanceInAccount: true,
+			wantErr:          false,
+		},
+		{
+			name: "Write to account with existing nonce fails",
+			addr: common.HexToAddress("c"),
+			existingAccount: &types.Account{
+				Balance: big.NewInt(defaultBalance),
+				Nonce:   5,
+			},
+			newAccount: types.Account{
+				Code:  contractCode,
+				Nonce: 5,
+			},
+			wantErr: true,
+		},
+		{
+			name: "Write to account with contract code fails",
+			addr: common.HexToAddress("b"),
+			existingAccount: &types.Account{
+				Balance: big.NewInt(defaultBalance),
+				Code:    bytes.Repeat([]byte{0x01}, 10),
+			},
+			newAccount: types.Account{
+				Code:  contractCode,
+				Nonce: 5,
+			},
+			wantErr: true,
+		},
+		{
+			name: "Write account with allowlist overwrites",
+			addr: common.HexToAddress("d"),
+			existingAccount: &types.Account{
+				Balance: big.NewInt(defaultBalance),
+				Code:    bytes.Repeat([]byte{0x01}, 10),
+			},
+			newAccount: types.Account{
+				Code:  contractCode,
+				Nonce: 5,
+			},
+			allowlist:        map[common.Address]bool{common.HexToAddress("d"): true},
+			balanceInAccount: true,
+			wantErr:          false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db := rawdb.NewMemoryDatabase()
+			tdb := state.NewDatabase(db)
+			sdb, _ := state.New(types.EmptyRootHash, tdb, nil)
+
+			if tt.existingAccount != nil {
+				sdb.CreateAccount(tt.addr)
+
+				if tt.existingAccount.Balance != nil {
+					sdb.SetBalance(tt.addr, uint256.MustFromBig(tt.existingAccount.Balance))
+				}
+				if tt.existingAccount.Nonce != 0 {
+					sdb.SetNonce(tt.addr, tt.existingAccount.Nonce)
+				}
+				if tt.existingAccount.Code != nil {
+					sdb.SetCode(tt.addr, tt.existingAccount.Code)
+				}
+			}
+
+			if err := applyAllocsToState(sdb, &core.Genesis{Alloc: types.GenesisAlloc{tt.addr: tt.newAccount}}, tt.allowlist); (err != nil) != tt.wantErr {
+				t.Errorf("applyAllocsToState() error = %v, wantErr %v", err, tt.wantErr)
+			}
+
+			// Don't check account state if an error was thrown
+			if tt.wantErr {
+				return
+			}
+
+			if !sdb.Exist(tt.addr) {
+				t.Errorf("account does not exists as expected: %v", tt.addr.Hex())
+			}
+
+			assert.Equal(t, tt.newAccount.Nonce, sdb.GetNonce(tt.addr))
+			assert.Equal(t, tt.newAccount.Code, sdb.GetCode(tt.addr))
+
+			if tt.balanceInAccount {
+				assert.True(t, big.NewInt(defaultBalance).Cmp(sdb.GetBalance(tt.addr).ToBig()) == 0)
+			} else {
+				assert.True(t, big.NewInt(0).Cmp(sdb.GetBalance(tt.addr).ToBig()) == 0)
+			}
+		})
+	}
+}