Merge pull request #5405 from onflow/supun/contract-update-validator

cmd/util/ledger/migrations/cadence.go

@@ -193,6 +193,9 @@ func NewCadence1ContractsMigrations(
 	stagedContracts []StagedContract,
 ) []ledger.Migration {
 
+	stagedContractsMigration := NewStagedContractsMigration(chainID)
+	stagedContractsMigration.RegisterContractUpdates(stagedContracts)
+
 	return []ledger.Migration{
 		NewAccountBasedMigration(
 			log,
@@ -211,7 +214,7 @@ func NewCadence1ContractsMigrations(
 			log,
 			nWorker,
 			[]AccountBasedMigration{
-				NewStagedContractsMigration(stagedContracts),
+				stagedContractsMigration,
 			},
 		),
 	}

cmd/util/ledger/migrations/cadence_values_migration_test.go

@@ -7,8 +7,6 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/onflow/flow-go/fvm/environment"
-
 	"github.com/rs/zerolog"
 
 	_ "github.com/glebarez/go-sqlite"
@@ -22,6 +20,7 @@ import (
 
 	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
+	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/model/flow"
 )
@@ -99,7 +98,7 @@ func TestCadenceValuesMigration(t *testing.T) {
 	checkReporters(t, rwf, address)
 
 	// Check error logs.
-	require.Nil(t, logWriter.logs)
+	require.Empty(t, logWriter.logs)
 }
 
 // TODO:

cmd/util/ledger/migrations/change_contract_code_migration.go

@@ -1,175 +1,46 @@
 package migrations
 
 import (
-	"context"
 	"fmt"
-	"strings"
-	"sync"
 
 	coreContracts "github.com/onflow/flow-core-contracts/lib/go/contracts"
 	ftContracts "github.com/onflow/flow-ft/lib/go/contracts"
 	nftContracts "github.com/onflow/flow-nft/lib/go/contracts"
-	"github.com/rs/zerolog"
 
 	sdk "github.com/onflow/flow-go-sdk"
 
 	"github.com/onflow/cadence/runtime/common"
 
 	evm "github.com/onflow/flow-go/fvm/evm/stdlib"
 	"github.com/onflow/flow-go/fvm/systemcontracts"
-	"github.com/onflow/flow-go/ledger"
-	"github.com/onflow/flow-go/ledger/common/convert"
 	"github.com/onflow/flow-go/model/flow"
 )
 
 type ChangeContractCodeMigration struct {
-	log       zerolog.Logger
-	mutex     sync.RWMutex
-	contracts map[common.Address]map[flow.RegisterID]string
+	*StagedContractsMigration
 }
 
 var _ AccountBasedMigration = (*ChangeContractCodeMigration)(nil)
 
-func (d *ChangeContractCodeMigration) Close() error {
-	d.mutex.RLock()
-	defer d.mutex.RUnlock()
-
-	if len(d.contracts) > 0 {
-		var sb strings.Builder
-		sb.WriteString("failed to find all contract registers that need to be changed:\n")
-		for address, contracts := range d.contracts {
-			_, _ = fmt.Fprintf(&sb, "- address: %s\n", address)
-			for registerID := range contracts {
-				_, _ = fmt.Fprintf(&sb, "  - %s\n", flow.RegisterIDContractName(registerID))
-			}
-		}
-		return fmt.Errorf(sb.String())
-	}
-
-	return nil
-}
-
-func (d *ChangeContractCodeMigration) InitMigration(
-	log zerolog.Logger,
-	_ []*ledger.Payload,
-	_ int,
-) error {
-	d.log = log.
-		With().
-		Str("migration", "ChangeContractCodeMigration").
-		Logger()
-
-	return nil
-}
-
-func (d *ChangeContractCodeMigration) MigrateAccount(
-	_ context.Context,
-	address common.Address,
-	payloads []*ledger.Payload,
-) ([]*ledger.Payload, error) {
-
-	contracts, ok := (func() (map[flow.RegisterID]string, bool) {
-		d.mutex.Lock()
-		defer d.mutex.Unlock()
-
-		contracts, ok := d.contracts[address]
-
-		// remove address from set of addresses
-		// to keep track of which addresses are left to change
-		delete(d.contracts, address)
-
-		return contracts, ok
-	})()
-
-	if !ok {
-		// no contracts to change on this address
-		return payloads, nil
-	}
-
-	for payloadIndex, payload := range payloads {
-		key, err := payload.Key()
-		if err != nil {
-			return nil, err
-		}
-
-		registerID, err := convert.LedgerKeyToRegisterID(key)
-		if err != nil {
-			return nil, err
-		}
-
-		newContract, ok := contracts[registerID]
-		if !ok {
-			// not a contract register, or
-			// not interested in this contract
-			continue
-		}
-
-		// change contract code
-		payloads[payloadIndex] = ledger.NewPayload(
-			key,
-			[]byte(newContract),
-		)
-
-		// TODO: maybe log diff between old and new
-
-		// remove contract from list of contracts to change
-		// to keep track of which contracts are left to change
-		delete(contracts, registerID)
-	}
-
-	if len(contracts) > 0 {
-		var sb strings.Builder
-		_, _ = fmt.Fprintf(&sb, "failed to find all contract registers that need to be changed for address %s:\n", address)
-		for registerID := range contracts {
-			_, _ = fmt.Fprintf(&sb, "- %s\n", flow.RegisterIDContractName(registerID))
-		}
-		return nil, fmt.Errorf(sb.String())
-	}
-
-	return payloads, nil
-}
-
-func (d *ChangeContractCodeMigration) RegisterContractChange(
-	address common.Address,
-	contractName string,
-	newContractCode string,
-) (
-	previousNewContractCode string,
-) {
-	d.mutex.Lock()
-	defer d.mutex.Unlock()
-
-	if d.contracts == nil {
-		d.contracts = map[common.Address]map[flow.RegisterID]string{}
-	}
-
-	if _, ok := d.contracts[address]; !ok {
-		d.contracts[address] = map[flow.RegisterID]string{}
+func NewChangeContractCodeMigration(chainID flow.ChainID) *ChangeContractCodeMigration {
+	return &ChangeContractCodeMigration{
+		StagedContractsMigration: NewStagedContractsMigration(chainID).
+			// TODO:
+			//WithContractUpdateValidation().
+			WithName("ChangeContractCodeMigration"),
 	}
-
-	registerID := flow.ContractRegisterID(flow.ConvertAddress(address), contractName)
-
-	previousNewContractCode = d.contracts[address][registerID]
-
-	d.contracts[address][registerID] = newContractCode
-
-	return
-}
-
-type SystemContractChange struct {
-	Address         common.Address
-	ContractName    string
-	NewContractCode string
 }
 
 func NewSystemContractChange(
 	systemContract systemcontracts.SystemContract,
 	newContractCode []byte,
-) SystemContractChange {
-	return SystemContractChange{
-		Address:         common.Address(systemContract.Address),
-		ContractName:    systemContract.Name,
-		NewContractCode: string(newContractCode),
+) StagedContract {
+	return StagedContract{
+		Address: common.Address(systemContract.Address),
+		Contract: Contract{
+			Name: systemContract.Name,
+			Code: newContractCode,
+		},
 	}
 }
 
@@ -185,7 +56,7 @@ type SystemContractChangesOptions struct {
 	EVM EVMContractChange
 }
 
-func SystemContractChanges(chainID flow.ChainID, options SystemContractChangesOptions) []SystemContractChange {
+func SystemContractChanges(chainID flow.ChainID, options SystemContractChangesOptions) []StagedContract {
 	systemContracts := systemcontracts.SystemContractsForChain(chainID)
 
 	var stakingCollectionAddress, stakingProxyAddress common.Address
@@ -211,7 +82,7 @@ func SystemContractChanges(chainID flow.ChainID, options SystemContractChangesOp
 	fungibleTokenMetadataViewsAddress := common.Address(systemContracts.FungibleToken.Address)
 	fungibleTokenSwitchboardAddress := common.Address(systemContracts.FungibleToken.Address)
 
-	contractChanges := []SystemContractChange{
+	contractChanges := []StagedContract{
 		// epoch related contracts
 		NewSystemContractChange(
 			systemContracts.Epoch,
@@ -269,35 +140,41 @@ func SystemContractChanges(chainID flow.ChainID, options SystemContractChangesOp
 			),
 		),
 		{
-			Address:      stakingCollectionAddress,
-			ContractName: "FlowStakingCollection",
-			NewContractCode: string(coreContracts.FlowStakingCollection(
-				systemContracts.FungibleToken.Address.HexWithPrefix(),
-				systemContracts.FlowToken.Address.HexWithPrefix(),
-				systemContracts.IDTableStaking.Address.HexWithPrefix(),
-				stakingProxyAddress.HexWithPrefix(),
-				lockedTokensAddress.HexWithPrefix(),
-				systemContracts.FlowStorageFees.Address.HexWithPrefix(),
-				systemContracts.ClusterQC.Address.HexWithPrefix(),
-				systemContracts.DKG.Address.HexWithPrefix(),
-				systemContracts.Epoch.Address.HexWithPrefix(),
-			)),
+			Address: stakingCollectionAddress,
+			Contract: Contract{
+				Name: "FlowStakingCollection",
+				Code: coreContracts.FlowStakingCollection(
+					systemContracts.FungibleToken.Address.HexWithPrefix(),
+					systemContracts.FlowToken.Address.HexWithPrefix(),
+					systemContracts.IDTableStaking.Address.HexWithPrefix(),
+					stakingProxyAddress.HexWithPrefix(),
+					lockedTokensAddress.HexWithPrefix(),
+					systemContracts.FlowStorageFees.Address.HexWithPrefix(),
+					systemContracts.ClusterQC.Address.HexWithPrefix(),
+					systemContracts.DKG.Address.HexWithPrefix(),
+					systemContracts.Epoch.Address.HexWithPrefix(),
+				),
+			},
 		},
 		{
-			Address:         stakingProxyAddress,
-			ContractName:    "StakingProxy",
-			NewContractCode: string(coreContracts.FlowStakingProxy()),
+			Address: stakingProxyAddress,
+			Contract: Contract{
+				Name: "StakingProxy",
+				Code: coreContracts.FlowStakingProxy(),
+			},
 		},
 		{
-			Address:      lockedTokensAddress,
-			ContractName: "LockedTokens",
-			NewContractCode: string(coreContracts.FlowLockedTokens(
-				systemContracts.FungibleToken.Address.HexWithPrefix(),
-				systemContracts.FlowToken.Address.HexWithPrefix(),
-				systemContracts.IDTableStaking.Address.HexWithPrefix(),
-				stakingProxyAddress.HexWithPrefix(),
-				systemContracts.FlowStorageFees.Address.HexWithPrefix(),
-			)),
+			Address: lockedTokensAddress,
+			Contract: Contract{
+				Name: "LockedTokens",
+				Code: coreContracts.FlowLockedTokens(
+					systemContracts.FungibleToken.Address.HexWithPrefix(),
+					systemContracts.FlowToken.Address.HexWithPrefix(),
+					systemContracts.IDTableStaking.Address.HexWithPrefix(),
+					stakingProxyAddress.HexWithPrefix(),
+					systemContracts.FlowStorageFees.Address.HexWithPrefix(),
+				),
+			},
 		},
 
 		// token related contracts
@@ -327,14 +204,16 @@ func SystemContractChanges(chainID flow.ChainID, options SystemContractChangesOp
 			),
 		),
 		{
-			Address:      fungibleTokenMetadataViewsAddress,
-			ContractName: "FungibleTokenMetadataViews",
-			NewContractCode: string(ftContracts.FungibleTokenMetadataViews(
-				// Use `Hex()`, since this method adds the prefix.
-				systemContracts.FungibleToken.Address.Hex(),
-				systemContracts.MetadataViews.Address.Hex(),
-				systemContracts.ViewResolver.Address.Hex(),
-			)),
+			Address: fungibleTokenMetadataViewsAddress,
+			Contract: Contract{
+				Name: "FungibleTokenMetadataViews",
+				Code: ftContracts.FungibleTokenMetadataViews(
+					// Use `Hex()`, since this method adds the prefix.
+					systemContracts.FungibleToken.Address.Hex(),
+					systemContracts.MetadataViews.Address.Hex(),
+					systemContracts.ViewResolver.Address.Hex(),
+				),
+			},
 		},
 
 		// NFT related contracts
@@ -361,12 +240,14 @@ func SystemContractChanges(chainID flow.ChainID, options SystemContractChangesOp
 	if chainID != flow.Emulator {
 		contractChanges = append(
 			contractChanges,
-			SystemContractChange{
-				Address:      fungibleTokenSwitchboardAddress,
-				ContractName: "FungibleTokenSwitchboard",
-				NewContractCode: string(ftContracts.FungibleTokenSwitchboard(
-					systemContracts.FungibleToken.Address.HexWithPrefix(),
-				)),
+			StagedContract{
+				Address: fungibleTokenSwitchboardAddress,
+				Contract: Contract{
+					Name: "FungibleTokenSwitchboard",
+					Code: ftContracts.FungibleTokenSwitchboard(
+						systemContracts.FungibleToken.Address.HexWithPrefix(),
+					),
+				},
 			},
 		)
 	}
@@ -406,13 +287,9 @@ func NewSystemContactsMigration(
 	chainID flow.ChainID,
 	options SystemContractChangesOptions,
 ) *ChangeContractCodeMigration {
-	migration := &ChangeContractCodeMigration{}
+	migration := NewChangeContractCodeMigration(chainID)
 	for _, change := range SystemContractChanges(chainID, options) {
-		migration.RegisterContractChange(
-			change.Address,
-			change.ContractName,
-			change.NewContractCode,
-		)
+		migration.RegisterContractChange(change)
 	}
 	return migration
 }