diff --git a/cmd/util/ledger/migrations/deduplicate_contract_names_migration.go b/cmd/util/ledger/migrations/deduplicate_contract_names_migration.go new file mode 100644 index 00000000000..ab35e04d8a3 --- /dev/null +++ b/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{} diff --git a/cmd/util/ledger/migrations/deduplicate_contract_names_migration_test.go b/cmd/util/ledger/migrations/deduplicate_contract_names_migration_test.go new file mode 100644 index 00000000000..ba81bc826cd --- /dev/null +++ b/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)) + }) + }) +}