From 3320dbe36f9e5970e25d486182d0e8b57fd8f6cc Mon Sep 17 00:00:00 2001 From: Janez Podhostnik Date: Fri, 14 Jun 2024 15:41:21 +0200 Subject: [PATCH] Add migration to add a key to all service accounts --- .../execution_state_extract.go | 10 ++ .../ledger/migrations/add_key_migration.go | 124 ++++++++++++++++++ .../migrations/add_key_migration_test.go | 80 +++++++++++ 3 files changed, 214 insertions(+) create mode 100644 cmd/util/ledger/migrations/add_key_migration.go create mode 100644 cmd/util/ledger/migrations/add_key_migration_test.go diff --git a/cmd/util/cmd/execution-state-extract/execution_state_extract.go b/cmd/util/cmd/execution-state-extract/execution_state_extract.go index 319f5fc057e..16617522108 100644 --- a/cmd/util/cmd/execution-state-extract/execution_state_extract.go +++ b/cmd/util/cmd/execution-state-extract/execution_state_extract.go @@ -15,6 +15,7 @@ import ( "go.uber.org/atomic" "golang.org/x/sync/errgroup" + "github.com/onflow/flow-go-sdk/crypto" migrators "github.com/onflow/flow-go/cmd/util/ledger/migrations" "github.com/onflow/flow-go/cmd/util/ledger/reporters" "github.com/onflow/flow-go/cmd/util/ledger/util" @@ -559,6 +560,11 @@ func newMigrations( opts, ) + key, err := crypto.DecodePublicKeyHex(crypto.ECDSA_P256, "711d4cd9930d695ef5c79b668d321f92ba00ed8280fded52c0fa2b15501411d026fe6fb4be3ec894facd3a00f04e32e2db5f5696d3b2b3419e4fba89fb95dca8") + if err != nil { + panic("failed to decode key") + } + // At the end, fix up storage-used discrepancies namedMigrations = append( namedMigrations, @@ -569,6 +575,10 @@ func newMigrations( opts.NWorker, []migrators.AccountBasedMigration{ &migrators.AccountUsageMigration{}, + migrators.NewAddKeyMigration( + opts.ChainID, + key, + ), }, ), }, diff --git a/cmd/util/ledger/migrations/add_key_migration.go b/cmd/util/ledger/migrations/add_key_migration.go new file mode 100644 index 00000000000..02a571dd6d2 --- /dev/null +++ b/cmd/util/ledger/migrations/add_key_migration.go @@ -0,0 +1,124 @@ +package migrations + +import ( + "context" + "fmt" + + "github.com/rs/zerolog" + + "github.com/onflow/cadence/runtime/common" + "github.com/onflow/crypto/hash" + "github.com/onflow/flow-go-sdk/crypto" + "github.com/onflow/flow-go/cmd/util/ledger/util/registers" + "github.com/onflow/flow-go/fvm" + "github.com/onflow/flow-go/fvm/systemcontracts" + "github.com/onflow/flow-go/model/flow" +) + +// AddKeyMigration adds a new key to the core contracts accounts +type AddKeyMigration struct { + log zerolog.Logger + + accountsToAddKeyTo map[common.Address]AddKeyMigrationAccountPublicKeyData + chainID flow.ChainID +} + +var _ AccountBasedMigration = &AddKeyMigration{} + +func NewAddKeyMigration( + chainID flow.ChainID, + key crypto.PublicKey, +) *AddKeyMigration { + + addresses := make(map[common.Address]AddKeyMigrationAccountPublicKeyData) + sc := systemcontracts.SystemContractsForChain(chainID).All() + + for _, sc := range sc { + addresses[common.Address(sc.Address)] = AddKeyMigrationAccountPublicKeyData{ + PublicKey: key, + HashAlgo: hash.SHA3_256, + } + } + + return &AddKeyMigration{ + accountsToAddKeyTo: addresses, + chainID: chainID, + } +} + +type AddKeyMigrationAccountPublicKeyData struct { + PublicKey crypto.PublicKey + HashAlgo hash.HashingAlgorithm +} + +func (m *AddKeyMigration) InitMigration( + log zerolog.Logger, + _ *registers.ByAccount, + _ int, +) error { + m.log = log.With().Str("component", "AddKeyMigration").Logger() + return nil +} + +func (m *AddKeyMigration) Close() error { + return nil +} + +func (m *AddKeyMigration) MigrateAccount( + _ context.Context, + address common.Address, + accountRegisters *registers.AccountRegisters, +) error { + + keyData, ok := m.accountsToAddKeyTo[address] + if !ok { + return nil + } + + // Create all the runtime components we need for the migration + migrationRuntime, err := NewInterpreterMigrationRuntime( + accountRegisters, + m.chainID, + InterpreterMigrationRuntimeConfig{}, + ) + + if err != nil { + return err + } + + key := flow.AccountPublicKey{ + PublicKey: keyData.PublicKey, + SignAlgo: keyData.PublicKey.Algorithm(), + HashAlgo: keyData.HashAlgo, + Weight: fvm.AccountKeyWeightThreshold, + } + + err = migrationRuntime.Accounts.AppendPublicKey(flow.ConvertAddress(address), key) + if err != nil { + return err + } + + // Finalize the transaction + result, err := migrationRuntime.TransactionState.FinalizeMainTransaction() + if err != nil { + return fmt.Errorf("failed to finalize main transaction: %w", err) + } + + // Merge the changes into the registers + expectedAddresses := map[flow.Address]struct{}{ + flow.Address(address): {}, + } + + err = registers.ApplyChanges( + accountRegisters, + result.WriteSet, + expectedAddresses, + m.log, + ) + if err != nil { + return fmt.Errorf("failed to apply register changes: %w", err) + } + + return nil + +} diff --git a/cmd/util/ledger/migrations/add_key_migration_test.go b/cmd/util/ledger/migrations/add_key_migration_test.go new file mode 100644 index 00000000000..b0eb579bdc0 --- /dev/null +++ b/cmd/util/ledger/migrations/add_key_migration_test.go @@ -0,0 +1,80 @@ +package migrations + +import ( + "context" + "testing" + + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" + + "github.com/onflow/cadence/runtime/common" + "github.com/onflow/flow-go-sdk/crypto" + "github.com/onflow/flow-go/cmd/util/ledger/util" + "github.com/onflow/flow-go/cmd/util/ledger/util/registers" + "github.com/onflow/flow-go/fvm" + "github.com/onflow/flow-go/fvm/systemcontracts" + "github.com/onflow/flow-go/model/flow" +) + +func TestCoreContractsKeys(t *testing.T) { + t.Parallel() + + log := zerolog.New(zerolog.NewTestWriter(t)) + + // Get the old payloads + payloads, err := util.PayloadsFromEmulatorSnapshot(snapshotPath) + require.NoError(t, err) + + registersByAccount, err := registers.NewByAccountFromPayloads(payloads) + require.NoError(t, err) + + chainID := flow.Emulator + sc := systemcontracts.SystemContractsForChain(chainID) + + serviceRegisters := registersByAccount.AccountRegisters(string(sc.FlowServiceAccount.Address.Bytes())) + + pk, err := crypto.GeneratePrivateKey(crypto.ECDSA_P256, nil) + require.NoError(t, err) + expectedKey := pk.PublicKey() + + mig := NewAddKeyMigration( + chainID, + expectedKey, + ) + defer func() { + err := mig.Close() + require.NoError(t, err) + }() + + err = mig.InitMigration(log, registersByAccount, 1) + require.NoError(t, err) + + ctx := context.Background() + err = mig.MigrateAccount(ctx, common.Address(sc.FlowServiceAccount.Address), serviceRegisters) + require.NoError(t, err) + + // Create all the runtime components we need for the migration + migrationRuntime, err := NewInterpreterMigrationRuntime( + serviceRegisters, + chainID, + InterpreterMigrationRuntimeConfig{}, + ) + require.NoError(t, err) + + // The last key should be the one we added + keys, err := migrationRuntime.Accounts.GetPublicKeyCount(sc.FlowServiceAccount.Address) + require.NoError(t, err) + + key, err := migrationRuntime.Accounts.GetPublicKey(sc.FlowServiceAccount.Address, keys-1) + require.NoError(t, err) + + require.Equal(t, expectedKey.String(), key.PublicKey.String()) + require.Equal(t, fvm.AccountKeyWeightThreshold, key.Weight) +} + +func Test_DO_NOT_MERGE(t *testing.T) { + t.Parallel() + // this branch should not be merged to master + // This is only to be used for migration mainnet testing + t.Fail() +}