Deduplicate contract names migration

cmd/util/ledger/migrations/deduplicate_contract_names_migration.go

@@ -0,0 +1,133 @@
+package migrations
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/fxamacker/cbor/v2"
+	"github.com/rs/zerolog"
+
+	"github.com/onflow/cadence/runtime/common"
+
+	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/convert"
+	"github.com/onflow/flow-go/model/flow"
+)
+
+// DeduplicateContractNamesMigration checks if the contract names have been duplicated and
+// removes the duplicate ones.
+//
+// This migration de-syncs storage used, so it should be run before the StorageUsedMigration.
+type DeduplicateContractNamesMigration struct {
+	log zerolog.Logger
+}
+
+func (d *DeduplicateContractNamesMigration) Close() error {
+	return nil
+}
+
+func (d *DeduplicateContractNamesMigration) InitMigration(
+	log zerolog.Logger,
+	_ []*ledger.Payload,
+	_ int,
+) error {
+	d.log = log.
+		With().
+		Str("migration", "DeduplicateContractNamesMigration").
+		Logger()
+
+	return nil
+}
+
+func (d *DeduplicateContractNamesMigration) MigrateAccount(
+	ctx context.Context,
+	address common.Address,
+	payloads []*ledger.Payload,
+) ([]*ledger.Payload, error) {
+	flowAddress := flow.ConvertAddress(address)
+	contractNamesID := flow.ContractNamesRegisterID(flowAddress)
+
+	var contractNamesPayload *ledger.Payload
+	contractNamesPayloadIndex := 0
+	for i, payload := range payloads {
+		key, err := payload.Key()
+		if err != nil {
+			return nil, err
+		}
+		id, err := convert.LedgerKeyToRegisterID(key)
+		if err != nil {
+			return nil, err
+		}
+		if id == contractNamesID {
+			contractNamesPayload = payload
+			contractNamesPayloadIndex = i
+			break
+		}
+	}
+	if contractNamesPayload == nil {
+		return payloads, nil
+	}
+
+	value := contractNamesPayload.Value()
+	if len(value) == 0 {
+		// Remove the empty payload
+		copy(payloads[contractNamesPayloadIndex:], payloads[contractNamesPayloadIndex+1:])
+		payloads = payloads[:len(payloads)-1]
+
+		return payloads, nil
+	}
+
+	var contractNames []string
+	err := cbor.Unmarshal(value, &contractNames)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get contract names: %w", err)
+	}
+
+	var foundDuplicate bool
+	i := 1
+	for i < len(contractNames) {
+		if contractNames[i-1] != contractNames[i] {
+
+			if contractNames[i-1] > contractNames[i] {
+				// this is not a valid state and we should fail.
+				// Contract names must be sorted by definition.
+				return nil, fmt.Errorf(
+					"contract names for account %s are not sorted: %s",
+					address.Hex(),
+					contractNames,
+				)
+			}
+
+			i++
+			continue
+		}
+		// Found duplicate (contactNames[i-1] == contactNames[i])
+		// Remove contractNames[i]
+		copy(contractNames[i:], contractNames[i+1:])
+		contractNames = contractNames[:len(contractNames)-1]
+		foundDuplicate = true
+	}
+
+	if !foundDuplicate {
+		return payloads, nil
+	}
+
+	d.log.Info().
+		Str("address", address.Hex()).
+		Strs("contract_names", contractNames).
+		Msg("removing duplicate contract names")
+
+	newContractNames, err := cbor.Marshal(contractNames)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"cannot encode contract names: %s",
+			contractNames,
+		)
+	}
+
+	payloads[contractNamesPayloadIndex] = ledger.NewPayload(convert.RegisterIDToLedgerKey(contractNamesID), newContractNames)
+	return payloads, nil
+
+}
+
+var _ AccountBasedMigration = &DeduplicateContractNamesMigration{}

cmd/util/ledger/migrations/deduplicate_contract_names_migration_test.go

@@ -0,0 +1,221 @@
+package migrations_test
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"sort"
+	"testing"
+
+	"github.com/fxamacker/cbor/v2"
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/require"
+
+	"github.com/onflow/cadence/runtime/common"
+
+	"github.com/onflow/flow-go/cmd/util/ledger/migrations"
+	"github.com/onflow/flow-go/fvm/environment"
+	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/convert"
+	"github.com/onflow/flow-go/model/flow"
+)
+
+func TestDeduplicateContractNamesMigration(t *testing.T) {
+	migration := migrations.DeduplicateContractNamesMigration{}
+	log := zerolog.New(zerolog.NewTestWriter(t))
+	err := migration.InitMigration(log, nil, 0)
+	require.NoError(t, err)
+
+	address, err := common.HexToAddress("0x1")
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	accountStatus := environment.NewAccountStatus()
+	accountStatus.SetStorageUsed(1000)
+	accountStatusPayload := ledger.NewPayload(
+		convert.RegisterIDToLedgerKey(
+			flow.AccountStatusRegisterID(flow.ConvertAddress(address)),
+		),
+		accountStatus.ToBytes(),
+	)
+
+	contractNamesPayload := func(contractNames []byte) *ledger.Payload {
+		return ledger.NewPayload(
+			convert.RegisterIDToLedgerKey(
+				flow.RegisterID{
+					Owner: string(address.Bytes()),
+					Key:   flow.ContractNamesKey,
+				},
+			),
+			contractNames,
+		)
+	}
+
+	requireContractNames := func(payloads []*ledger.Payload, f func([]string)) {
+		for _, payload := range payloads {
+			key, err := payload.Key()
+			require.NoError(t, err)
+			id, err := convert.LedgerKeyToRegisterID(key)
+			require.NoError(t, err)
+
+			if id.Key != flow.ContractNamesKey {
+				continue
+			}
+
+			contracts := make([]string, 0)
+			err = cbor.Unmarshal(payload.Value(), &contracts)
+			require.NoError(t, err)
+
+			f(contracts)
+
+		}
+	}
+
+	t.Run("no contract names", func(t *testing.T) {
+		payloads, err := migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+			},
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, 1, len(payloads))
+	})
+
+	t.Run("one contract", func(t *testing.T) {
+		contractNames := []string{"test"}
+		newContractNames, err := cbor.Marshal(contractNames)
+		require.NoError(t, err)
+
+		payloads, err := migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+				contractNamesPayload(newContractNames),
+			},
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, 2, len(payloads))
+
+		requireContractNames(payloads, func(contracts []string) {
+			require.Equal(t, 1, len(contracts))
+			require.Equal(t, "test", contracts[0])
+		})
+	})
+
+	t.Run("two unique contracts", func(t *testing.T) {
+		contractNames := []string{"test", "test2"}
+		newContractNames, err := cbor.Marshal(contractNames)
+		require.NoError(t, err)
+
+		payloads, err := migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+				contractNamesPayload(newContractNames),
+			},
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, 2, len(payloads))
+
+		requireContractNames(payloads, func(contracts []string) {
+			require.Equal(t, 2, len(contracts))
+			require.Equal(t, "test", contracts[0])
+			require.Equal(t, "test2", contracts[1])
+		})
+	})
+
+	t.Run("two contracts", func(t *testing.T) {
+		contractNames := []string{"test", "test"}
+		newContractNames, err := cbor.Marshal(contractNames)
+		require.NoError(t, err)
+
+		payloads, err := migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+				contractNamesPayload(newContractNames),
+			},
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, 2, len(payloads))
+
+		requireContractNames(payloads, func(contracts []string) {
+			require.Equal(t, 1, len(contracts))
+			require.Equal(t, "test", contracts[0])
+		})
+	})
+
+	t.Run("not sorted contracts", func(t *testing.T) {
+		contractNames := []string{"test2", "test"}
+		newContractNames, err := cbor.Marshal(contractNames)
+		require.NoError(t, err)
+
+		_, err = migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+				contractNamesPayload(newContractNames),
+			},
+		)
+
+		require.Error(t, err)
+	})
+
+	t.Run("duplicate contracts", func(t *testing.T) {
+		contractNames := []string{"test", "test", "test2", "test3", "test3"}
+		newContractNames, err := cbor.Marshal(contractNames)
+		require.NoError(t, err)
+
+		payloads, err := migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+				contractNamesPayload(newContractNames),
+			},
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, 2, len(payloads))
+
+		requireContractNames(payloads, func(contracts []string) {
+			require.Equal(t, 3, len(contracts))
+			require.Equal(t, "test", contracts[0])
+			require.Equal(t, "test2", contracts[1])
+			require.Equal(t, "test3", contracts[2])
+		})
+	})
+
+	t.Run("random contracts", func(t *testing.T) {
+		contractNames := make([]string, 1000)
+		uniqueContracts := 1
+		for i := 0; i < 1000; i++ {
+			// i > 0 so it's easier to know how many unique contracts there are
+			if i > 0 && rand.Float32() < 0.5 {
+				uniqueContracts++
+			}
+			contractNames[i] = fmt.Sprintf("test%d", uniqueContracts)
+		}
+
+		// sort contractNames alphabetically, because they are not sorted
+		sort.Slice(contractNames, func(i, j int) bool {
+			return contractNames[i] < contractNames[j]
+		})
+
+		newContractNames, err := cbor.Marshal(contractNames)
+		require.NoError(t, err)
+
+		payloads, err := migration.MigrateAccount(ctx, address,
+			[]*ledger.Payload{
+				accountStatusPayload,
+				contractNamesPayload(newContractNames),
+			},
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, 2, len(payloads))
+
+		requireContractNames(payloads, func(contracts []string) {
+			require.Equal(t, uniqueContracts, len(contracts))
+		})
+	})
+}