Update storage health check to add account key validation, etc.

cmd/util/cmd/check-storage/cmd.go

@@ -3,20 +3,25 @@ package check_storage
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/onflow/atree"
 	"github.com/onflow/cadence/interpreter"
+	"github.com/onflow/crypto/hash"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/onflow/cadence/common"
 	"github.com/onflow/cadence/runtime"
 
+	"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"
 	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
+	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/fvm/evm/emulator/state"
+	storageState "github.com/onflow/flow-go/fvm/storage/state"
 	"github.com/onflow/flow-go/fvm/systemcontracts"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/model/flow"
@@ -32,6 +37,8 @@ var (
 	flagNWorker            int
 	flagHasAccountFormatV1 bool
 	flagHasAccountFormatV2 bool
+	flagIsAccountStatusV4  bool
+	flagCheckContract      bool
 )
 
 var (
@@ -109,9 +116,23 @@ func init() {
 	Cmd.Flags().BoolVar(
 		&flagHasAccountFormatV2,
 		"account-format-v2",
-		false,
+		true,
 		"State contains accounts in v2 format",
 	)
+
+	Cmd.Flags().BoolVar(
+		&flagIsAccountStatusV4,
+		"account-status-v4",
+		false,
+		"State is migrated to account status v4 format",
+	)
+
+	Cmd.Flags().BoolVar(
+		&flagCheckContract,
+		"check-contract",
+		false,
+		"check stored contract",
+	)
 }
 
 func run(*cobra.Command, []string) {
@@ -397,7 +418,48 @@ func checkAccountStorageHealth(accountRegisters *registers.AccountRegisters, nWo
 		)
 	}
 
-	// TODO: check health of non-atree registers
+	if flagIsAccountStatusV4 {
+		// Validate account public key storage
+		err = migrations.ValidateAccountPublicKeyV4(address, accountRegisters)
+		if err != nil {
+			issues = append(
+				issues,
+				accountStorageIssue{
+					Address: address.Hex(),
+					Kind:    storageErrorKindString[accountKeyErrorKind],
+					Msg:     err.Error(),
+				},
+			)
+		}
+
+		// Check account public keys
+		err = checkAccountPublicKeys(address, accountRegisters)
+		if err != nil {
+			issues = append(
+				issues,
+				accountStorageIssue{
+					Address: address.Hex(),
+					Kind:    storageErrorKindString[accountKeyErrorKind],
+					Msg:     err.Error(),
+				},
+			)
+		}
+	}
+
+	if flagCheckContract {
+		// Check contracts
+		err = checkContracts(address, accountRegisters)
+		if err != nil {
+			issues = append(
+				issues,
+				accountStorageIssue{
+					Address: address.Hex(),
+					Kind:    storageErrorKindString[contractErrorKind],
+					Msg:     err.Error(),
+				},
+			)
+		}
+	}
 
 	return issues
 }
@@ -409,12 +471,16 @@ const (
 	cadenceAtreeStorageErrorKind
 	evmAtreeStorageErrorKind
 	storageFormatErrorKind
+	accountKeyErrorKind
+	contractErrorKind
 )
 
 var storageErrorKindString = map[storageErrorKind]string{
 	otherErrorKind:               "error_check_storage_failed",
 	cadenceAtreeStorageErrorKind: "error_cadence_atree_storage",
 	evmAtreeStorageErrorKind:     "error_evm_atree_storage",
+	accountKeyErrorKind:          "error_account_public_key",
+	contractErrorKind:            "error_contract",
 }
 
 type accountStorageIssue struct {
@@ -484,3 +550,95 @@ func checkAccountFormat(
 
 	return nil
 }
+
+func checkAccountPublicKeys(
+	address common.Address,
+	accountRegisters *registers.AccountRegisters,
+) error {
+	// Skip empty address because it doesn't have any account status registers.
+	if len(address) == 0 || address == common.ZeroAddress {
+		return nil
+	}
+
+	accounts := newAccounts(accountRegisters)
+
+	keyCount, err := accounts.GetAccountPublicKeyCount(flow.BytesToAddress(address.Bytes()))
+	if err != nil {
+		return err
+	}
+
+	// Check keys
+	for keyIndex := range keyCount {
+		_, err = accounts.GetAccountPublicKey(flow.BytesToAddress(address.Bytes()), keyIndex)
+		if err != nil {
+			return err
+		}
+	}
+
+	// NOTE: don't need to check unreachable keys because it is checked in ValidateAccountPublicKeyV4().
+
+	return nil
+}
+
+func checkContracts(
+	address common.Address,
+	accountRegisters *registers.AccountRegisters,
+) error {
+	if len(address) == 0 || address == common.ZeroAddress {
+		return nil
+	}
+
+	accounts := newAccounts(accountRegisters)
+
+	contractNames, err := accounts.GetContractNames(flow.BytesToAddress(address.Bytes()))
+	if err != nil {
+		return err
+	}
+
+	// Check contract
+	contractRegisterKeys := make(map[string]bool)
+	for _, contractName := range contractNames {
+		contractRegisterKeys[flow.ContractKey(contractName)] = true
+
+		_, err = accounts.GetContract(contractName, flow.BytesToAddress(address.Bytes()))
+		if err != nil {
+			return err
+		}
+	}
+
+	// Check unreachable contract registers
+	err = accountRegisters.ForEachKey(func(key string) error {
+		if strings.HasPrefix(key, flow.CodeKeyPrefix) && !contractRegisterKeys[key] {
+			return fmt.Errorf("found unreachable contract register %s, contract names %v", key, contractNames)
+		}
+		return nil
+	})
+
+	return err
+}
+
+func newAccounts(accountRegisters *registers.AccountRegisters) environment.Accounts {
+	// Create a new transaction state with a dummy hasher
+	// because we do not need spock proofs for migrations.
+	transactionState := storageState.NewTransactionStateFromExecutionState(
+		storageState.NewExecutionStateWithSpockStateHasher(
+			registers.StorageSnapshot{
+				Registers: accountRegisters,
+			},
+			storageState.DefaultParameters(),
+			func() hash.Hasher {
+				return dummyHasher{}
+			},
+		),
+	)
+	return environment.NewAccounts(transactionState)
+}
+
+type dummyHasher struct{}
+
+func (d dummyHasher) Algorithm() hash.HashingAlgorithm { return hash.UnknownHashingAlgorithm }
+func (d dummyHasher) Size() int                        { return 0 }
+func (d dummyHasher) ComputeHash([]byte) hash.Hash     { return nil }
+func (d dummyHasher) Write([]byte) (int, error)        { return 0, nil }
+func (d dummyHasher) SumHash() hash.Hash               { return nil }
+func (d dummyHasher) Reset()                           {}

cmd/util/ledger/migrations/account_key_deduplication_migration.go

@@ -98,6 +98,7 @@ func NewAccountPublicKeyDeduplicationMigration(
 		chainID:   chainID,
 		reporter:  rwf.ReportWriter("account-public-key-deduplication-migration_summary"),
 		outputDir: outputDir,
+		validate:  validate,
 	}
 
 	if validate {
@@ -131,7 +132,7 @@ func (m *AccountPublicKeyDeduplicationMigration) MigrateAccount(
 	}
 
 	if m.validate {
-		err := validateAccountPublicKeyV4(address, accountRegisters)
+		err := ValidateAccountPublicKeyV4(address, accountRegisters)
 		if err != nil {
 			m.validationReporter.Write(validationError{
 				Address: address.Hex(),

cmd/util/ledger/migrations/account_key_deduplication_migration_test.go

@@ -232,7 +232,7 @@ func TestMigration(t *testing.T) {
 		require.Equal(t, byte(0x40), encodedAccountStatusV4[0])
 		require.Equal(t, encodedAccountStatusV3[1:], encodedAccountStatusV4[1:])
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 
@@ -276,7 +276,7 @@ func TestMigration(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, encodedPk1, encodedAccountPublicKey0)
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 
@@ -321,7 +321,7 @@ func TestMigration(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, encodedPk1, encodedAccountPublicKey0)
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 
@@ -399,7 +399,7 @@ func TestMigration(t *testing.T) {
 		require.NoError(t, err)
 		require.ElementsMatch(t, [][]byte{{}, encodedSpk2}, encodedPks)
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 
@@ -488,7 +488,7 @@ func TestMigration(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, uint64(2), seqNum)
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 
@@ -548,7 +548,7 @@ func TestMigration(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, encodedPk1, encodedAccountPublicKey0)
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 
@@ -618,7 +618,7 @@ func TestMigration(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, uint64(1), seqNum)
 
-		err = validateAccountPublicKeyV4(common.Address(owner), accountRegisters)
+		err = ValidateAccountPublicKeyV4(common.Address(owner), accountRegisters)
 		require.NoError(t, err)
 	})
 }

cmd/util/ledger/migrations/account_key_deduplication_migration_validation.go

@@ -13,10 +13,14 @@ import (
 	"github.com/onflow/flow-go/model/flow"
 )
 
-func validateAccountPublicKeyV4(
+func ValidateAccountPublicKeyV4(
 	address common.Address,
 	accountRegisters *registers.AccountRegisters,
 ) error {
+	// Skip empty address because it doesn't have account status register.
+	if len(address) == 0 || address == common.ZeroAddress {
+		return nil
+	}
 
 	// Validate account status register
 	accountPublicKeyCount, storedKeyCount, deduplicated, err := validateAccountStatusV4Register(address, accountRegisters)
@@ -90,7 +94,7 @@ func validateAccountStatusV4Register(
 	deduplicated bool,
 	err error,
 ) {
-	owner := string(address[:])
+	owner := flow.AddressToRegisterOwner(flow.Address(address[:]))
 
 	encodedAccountStatus, err := accountRegisters.Get(owner, flow.AccountStatusKey)
 	if err != nil {