Add FT/NFT type change rules for contract update validator

cmd/util/ledger/migrations/cadence.go

@@ -15,30 +15,35 @@ import (
 	"github.com/onflow/flow-go/model/flow"
 )
 
-func NewCadence1InterfaceStaticTypeConverter(chainID flow.ChainID) statictypes.InterfaceTypeConverterFunc {
+func NewInterfaceTypeConversionRules(chainID flow.ChainID) StaticTypeMigrationRules {
 	systemContracts := systemcontracts.SystemContractsForChain(chainID)
 
 	oldFungibleTokenResolverType, newFungibleTokenResolverType := fungibleTokenResolverRule(systemContracts)
 
-	rules := StaticTypeMigrationRules{
+	return StaticTypeMigrationRules{
 		oldFungibleTokenResolverType.ID(): newFungibleTokenResolverType,
 	}
-
-	return NewStaticTypeMigrator[*interpreter.InterfaceStaticType](rules)
 }
 
-func NewCadence1CompositeStaticTypeConverter(chainID flow.ChainID) statictypes.CompositeTypeConverterFunc {
-
+func NewCompositeTypeConversionRules(chainID flow.ChainID) StaticTypeMigrationRules {
 	systemContracts := systemcontracts.SystemContractsForChain(chainID)
 
 	oldFungibleTokenVaultCompositeType, newFungibleTokenVaultType := fungibleTokenVaultRule(systemContracts)
 	oldNonFungibleTokenNFTCompositeType, newNonFungibleTokenNFTType := nonFungibleTokenNFTRule(systemContracts)
 
-	rules := StaticTypeMigrationRules{
+	return StaticTypeMigrationRules{
 		oldFungibleTokenVaultCompositeType.ID():  newFungibleTokenVaultType,
 		oldNonFungibleTokenNFTCompositeType.ID(): newNonFungibleTokenNFTType,
 	}
+}
 
+func NewCadence1InterfaceStaticTypeConverter(chainID flow.ChainID) statictypes.InterfaceTypeConverterFunc {
+	rules := NewInterfaceTypeConversionRules(chainID)
+	return NewStaticTypeMigrator[*interpreter.InterfaceStaticType](rules)
+}
+
+func NewCadence1CompositeStaticTypeConverter(chainID flow.ChainID) statictypes.CompositeTypeConverterFunc {
+	rules := NewCompositeTypeConversionRules(chainID)
 	return NewStaticTypeMigrator[*interpreter.CompositeStaticType](rules)
 }
 
@@ -193,7 +198,9 @@ func NewCadence1ContractsMigrations(
 	stagedContracts []StagedContract,
 ) []ledger.Migration {
 
-	stagedContractsMigration := NewStagedContractsMigration(chainID)
+	stagedContractsMigration := NewStagedContractsMigration(chainID).
+		WithContractUpdateValidation()
+
 	stagedContractsMigration.RegisterContractUpdates(stagedContracts)
 
 	return []ledger.Migration{

cmd/util/ledger/migrations/staged_contracts_migration.go

@@ -9,13 +9,13 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/onflow/cadence/runtime/sema"
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/cadence/runtime"
 	"github.com/onflow/cadence/runtime/common"
 	"github.com/onflow/cadence/runtime/interpreter"
 	"github.com/onflow/cadence/runtime/old_parser"
+	"github.com/onflow/cadence/runtime/sema"
 	"github.com/onflow/cadence/runtime/stdlib"
 
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
@@ -27,13 +27,14 @@ import (
 )
 
 type StagedContractsMigration struct {
-	name                   string
-	chainID                flow.ChainID
-	log                    zerolog.Logger
-	mutex                  sync.RWMutex
-	stagedContracts        map[common.Address]map[flow.RegisterID]Contract
-	contractsByLocation    map[common.Location][]byte
-	enableUpdateValidation bool
+	name                           string
+	chainID                        flow.ChainID
+	log                            zerolog.Logger
+	mutex                          sync.RWMutex
+	stagedContracts                map[common.Address]map[flow.RegisterID]Contract
+	contractsByLocation            map[common.Location][]byte
+	enableUpdateValidation         bool
+	userDefinedTypeChangeCheckFunc func(oldTypeID common.TypeID, newTypeID common.TypeID) (checked bool, valid bool)
 }
 
 type StagedContract struct {
@@ -59,6 +60,7 @@ func NewStagedContractsMigration(chainID flow.ChainID) *StagedContractsMigration
 
 func (m *StagedContractsMigration) WithContractUpdateValidation() *StagedContractsMigration {
 	m.enableUpdateValidation = true
+	m.userDefinedTypeChangeCheckFunc = newUserDefinedTypeChangeCheckerFunc(m.chainID)
 	return m
 }
 
@@ -215,7 +217,7 @@ func (m *StagedContractsMigration) MigrateAccount(
 		oldCode := payload.Value()
 
 		if m.enableUpdateValidation {
-			err = CheckContractUpdateValidity(
+			err = m.checkContractUpdateValidity(
 				mr,
 				address,
 				name,
@@ -257,7 +259,7 @@ func (m *StagedContractsMigration) MigrateAccount(
 	return payloads, nil
 }
 
-func CheckContractUpdateValidity(
+func (m *StagedContractsMigration) checkContractUpdateValidity(
 	mr *migratorRuntime,
 	address common.Address,
 	contractName string,
@@ -295,6 +297,10 @@ func CheckContractUpdateValidity(
 		elaborations,
 	)
 
+	validator.WithUserDefinedTypeChangeChecker(
+		m.userDefinedTypeChangeCheckFunc,
+	)
+
 	return validator.Validate()
 }
 
@@ -347,3 +353,28 @@ func StagedContractsFromCSV(path string) ([]StagedContract, error) {
 
 	return contracts, nil
 }
+
+func newUserDefinedTypeChangeCheckerFunc(
+	chainID flow.ChainID,
+) func(oldTypeID common.TypeID, newTypeID common.TypeID) (checked, valid bool) {
+
+	typeChangeRules := map[common.TypeID]common.TypeID{}
+
+	compositeTypeRules := NewCompositeTypeConversionRules(chainID)
+	for typeID, newStaticType := range compositeTypeRules {
+		typeChangeRules[typeID] = newStaticType.ID()
+	}
+
+	interfaceTypeRules := NewInterfaceTypeConversionRules(chainID)
+	for typeID, newStaticType := range interfaceTypeRules {
+		typeChangeRules[typeID] = newStaticType.ID()
+	}
+
+	return func(oldTypeID common.TypeID, newTypeID common.TypeID) (checked, valid bool) {
+		expectedNewTypeID, found := typeChangeRules[oldTypeID]
+		if found {
+			return true, expectedNewTypeID == newTypeID
+		}
+		return false, false
+	}
+}

cmd/util/ledger/migrations/staged_contracts_migration_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/rs/zerolog"
 
+	"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"
@@ -669,3 +670,161 @@ contract MultilineContract{
 		require.Empty(t, contracts)
 	})
 }
+
+func TestStagedContractsWithUpdateValidator(t *testing.T) {
+	t.Parallel()
+
+	chainID := flow.Emulator
+	systemContracts := systemcontracts.SystemContractsForChain(chainID)
+
+	address, err := common.HexToAddress("0x1")
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	t.Run("FungibleToken.Vault", func(t *testing.T) {
+		t.Parallel()
+
+		ftAddress := common.Address(systemContracts.FungibleToken.Address)
+
+		oldCodeA := fmt.Sprintf(`
+            import FungibleToken from %s
+
+            pub contract A {
+                pub var vault: @FungibleToken.Vault?
+                init() {
+                    self.vault <- nil
+                }
+            }
+        `,
+			ftAddress.HexWithPrefix(),
+		)
+
+		newCodeA := fmt.Sprintf(`
+            import FungibleToken from %s
+            access(all) contract A {
+                access(all) var vault: @{FungibleToken.Vault}?
+                init() {
+                    self.vault <- nil
+                }
+            }
+        `,
+			ftAddress.HexWithPrefix(),
+		)
+
+		ftContract := `
+            access(all) contract FungibleToken {
+                access(all) resource interface Vault {}
+		    }
+        `
+
+		stagedContracts := []StagedContract{
+			{
+				Contract: Contract{
+					Name: "A",
+					Code: []byte(newCodeA),
+				},
+				Address: address,
+			},
+		}
+
+		migration := NewStagedContractsMigration(chainID)
+		migration.RegisterContractUpdates(stagedContracts)
+		migration.WithContractUpdateValidation()
+
+		logWriter := &logWriter{}
+		log := zerolog.New(logWriter)
+		err = migration.InitMigration(log, nil, 0)
+		require.NoError(t, err)
+
+		payloads := []*ledger.Payload{
+			newContractPayload(address, "A", []byte(oldCodeA)),
+			newContractPayload(ftAddress, "FungibleToken", []byte(ftContract)),
+		}
+
+		payloads, err = migration.MigrateAccount(ctx, address, payloads)
+		require.NoError(t, err)
+
+		err = migration.Close()
+		require.NoError(t, err)
+
+		require.Empty(t, logWriter.logs)
+
+		require.Len(t, payloads, 2)
+		require.Equal(t, newCodeA, string(payloads[0].Value()))
+	})
+
+	t.Run("other type", func(t *testing.T) {
+		t.Parallel()
+
+		otherAddress, err := common.HexToAddress("0x2")
+		require.NoError(t, err)
+
+		oldCodeA := fmt.Sprintf(`
+            import FungibleToken from %s
+
+            pub contract A {
+                pub var vault: @FungibleToken.Vault?
+                init() {
+                    self.vault <- nil
+                }
+            }
+        `,
+			otherAddress.HexWithPrefix(), // Importing from some other address
+		)
+
+		newCodeA := fmt.Sprintf(`
+            import FungibleToken from %s
+            access(all) contract A {
+                access(all) var vault: @{FungibleToken.Vault}?
+                init() {
+                    self.vault <- nil
+                }
+            }
+        `,
+			otherAddress.HexWithPrefix(), // Importing from some other address
+		)
+
+		ftContract := `
+            access(all) contract FungibleToken {
+                access(all) resource interface Vault {}
+		    }
+        `
+
+		stagedContracts := []StagedContract{
+			{
+				Contract: Contract{
+					Name: "A",
+					Code: []byte(newCodeA),
+				},
+				Address: address,
+			},
+		}
+
+		migration := NewStagedContractsMigration(chainID)
+		migration.RegisterContractUpdates(stagedContracts)
+		migration.WithContractUpdateValidation()
+
+		logWriter := &logWriter{}
+		log := zerolog.New(logWriter)
+		err = migration.InitMigration(log, nil, 0)
+		require.NoError(t, err)
+
+		payloads := []*ledger.Payload{
+			newContractPayload(address, "A", []byte(oldCodeA)),
+			newContractPayload(otherAddress, "FungibleToken", []byte(ftContract)),
+		}
+
+		payloads, err = migration.MigrateAccount(ctx, address, payloads)
+		require.NoError(t, err)
+
+		err = migration.Close()
+		require.NoError(t, err)
+
+		require.Len(t, logWriter.logs, 1)
+		assert.Contains(t, logWriter.logs[0], "cannot update contract `A`")
+
+		require.Len(t, payloads, 2)
+		require.Equal(t, oldCodeA, string(payloads[0].Value()))
+	})
+}