Add migration to add a key to all service accounts

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,
+					),
 				},
 			),
 		},

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
+
+}

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()
+}