refactor: use IterateClientStates in migration code

docs/migrations/v6-to-v7.md

@@ -30,7 +30,7 @@ app.UpgradeKeeper.SetUpgradeHandler(
     upgradeName,
     func(ctx sdk.Context, _ upgradetypes.Plan, _ module.VersionMap) (module.VersionMap, error) {
         // prune expired tendermint consensus states to save storage space
-        ibctm.PruneTendermintConsensusStates(ctx, app.Codec, appCodec, keys[ibchost.StoreKey])
+        ibctm.PruneTendermintConsensusStates(ctx, app.Codec, app.IBCKeeper.ClientKeeper)
 
         return app.mm.RunMigrations(ctx, app.configurator, fromVM)
     },

modules/core/02-client/keeper/migrations.go

@@ -34,5 +34,5 @@ func (m Migrator) Migrate1to2(ctx sdk.Context) error {
 // - removes the localhost client
 // - asserts that existing tendermint clients are properly registered on the chain codec
 func (m Migrator) Migrate2to3(ctx sdk.Context) error {
-	return v7.MigrateStore(ctx, m.keeper.storeKey, m.keeper.cdc)
+	return v7.MigrateStore(ctx, m.keeper.cdc, m.keeper)
 }

modules/core/02-client/migrations/v7/expected_keepers.go

@@ -0,0 +1,15 @@
+package v7
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/cosmos/ibc-go/v6/modules/core/exported"
+)
+
+// ClientKeeper expected IBC client keeper
+type ClientKeeper interface {
+	GetClientState(ctx sdk.Context, clientID string) (exported.ClientState, bool)
+	SetClientState(ctx sdk.Context, clientID string, clientState exported.ClientState)
+	IterateClientStates(ctx sdk.Context, prefix []byte, cb func(string, exported.ClientState) bool)
+	ClientStore(ctx sdk.Context, clientID string) sdk.KVStore
+}

modules/core/02-client/migrations/v7/store.go

@@ -6,8 +6,6 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
-	"github.com/cosmos/cosmos-sdk/store/prefix"
-	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
@@ -29,18 +27,16 @@ const Localhost string = "09-localhost"
 // - Pruning all solo machine consensus states
 // - Removing the localhost client
 // - Asserting existing tendermint clients are properly registered on the chain codec
-func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec) error {
-	store := ctx.KVStore(storeKey)
-
-	if err := handleSolomachineMigration(ctx, store, cdc); err != nil {
+func MigrateStore(ctx sdk.Context, cdc codec.BinaryCodec, keeper ClientKeeper) error {
+	if err := handleSolomachineMigration(ctx, cdc, keeper); err != nil {
 		return err
 	}
 
-	if err := handleTendermintMigration(ctx, store, cdc); err != nil {
+	if err := handleTendermintMigration(ctx, keeper); err != nil {
 		return err
 	}
 
-	if err := handleLocalhostMigration(ctx, store, cdc); err != nil {
+	if err := handleLocalhostMigration(ctx, keeper); err != nil {
 		return err
 	}
 
@@ -49,15 +45,14 @@ func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.Binar
 
 // handleSolomachineMigration iterates over the solo machine clients and migrates client state from
 // protobuf definition v2 to v3. All consensus states stored outside of the client state are pruned.
-func handleSolomachineMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.BinaryCodec) error {
-	clients, err := collectClients(ctx, store, exported.Solomachine)
+func handleSolomachineMigration(ctx sdk.Context, cdc codec.BinaryCodec, keeper ClientKeeper) error {
+	clients, err := collectClients(ctx, keeper, exported.Solomachine)
 	if err != nil {
 		return err
 	}
 
 	for _, clientID := range clients {
-		clientPrefix := []byte(fmt.Sprintf("%s/%s/", host.KeyClientStorePrefix, clientID))
-		clientStore := prefix.NewStore(store, clientPrefix)
+		clientStore := keeper.ClientStore(ctx, clientID)
 
 		bz := clientStore.Get(host.ClientStateKey())
 		if bz == nil {
@@ -81,8 +76,7 @@ func handleSolomachineMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.Bi
 			return sdkerrors.Wrap(err, "failed to unmarshal client state bytes into solo machine client state")
 		}
 
-		// update solomachine in store
-		clientStore.Set(host.ClientStateKey(), bz)
+		keeper.SetClientState(ctx, clientID, &updatedClientState)
 
 		removeAllClientConsensusStates(clientStore)
 	}
@@ -92,32 +86,24 @@ func handleSolomachineMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.Bi
 
 // handlerTendermintMigration asserts that the tendermint client in state can be decoded properly.
 // This ensures the upgrading chain properly registered the tendermint client types on the chain codec.
-func handleTendermintMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.BinaryCodec) error {
-	clients, err := collectClients(ctx, store, exported.Tendermint)
+func handleTendermintMigration(ctx sdk.Context, keeper ClientKeeper) error {
+	clients, err := collectClients(ctx, keeper, exported.Tendermint)
 	if err != nil {
 		return err
 	}
+	fmt.Println(clients)
 
 	if len(clients) > 1 {
 		return sdkerrors.Wrap(sdkerrors.ErrLogic, "more than one Tendermint client collected")
 	}
 
 	clientID := clients[0]
-
-	clientPrefix := []byte(fmt.Sprintf("%s/%s/", host.KeyClientStorePrefix, clientID))
-	clientStore := prefix.NewStore(store, clientPrefix)
-
-	bz := clientStore.Get(host.ClientStateKey())
-	if bz == nil {
-		return clienttypes.ErrClientNotFound
-	}
-
-	var clientState exported.ClientState
-	if err := cdc.UnmarshalInterface(bz, &clientState); err != nil {
-		return sdkerrors.Wrap(err, "failed to unmarshal client state bytes into tendermint client state")
+	clientState, ok := keeper.GetClientState(ctx, clientID)
+	if !ok {
+		return sdkerrors.Wrapf(clienttypes.ErrClientNotFound, "clientID %s", clientID)
 	}
 
-	_, ok := clientState.(*ibctm.ClientState)
+	_, ok = clientState.(*ibctm.ClientState)
 	if !ok {
 		return sdkerrors.Wrap(clienttypes.ErrInvalidClient, "client state is not tendermint even though client id contains 07-tendermint")
 	}
@@ -126,15 +112,14 @@ func handleTendermintMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.Bin
 }
 
 // handleLocalhostMigration removes all client and consensus states associated with the localhost client type.
-func handleLocalhostMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.BinaryCodec) error {
-	clients, err := collectClients(ctx, store, Localhost)
+func handleLocalhostMigration(ctx sdk.Context, keeper ClientKeeper) error {
+	clients, err := collectClients(ctx, keeper, Localhost)
 	if err != nil {
 		return err
 	}
 
 	for _, clientID := range clients {
-		clientPrefix := []byte(fmt.Sprintf("%s/%s/", host.KeyClientStorePrefix, clientID))
-		clientStore := prefix.NewStore(store, clientPrefix)
+		clientStore := keeper.ClientStore(ctx, clientID)
 
 		// delete the client state
 		clientStore.Delete(host.ClientStateKey())
@@ -150,9 +135,11 @@ func handleLocalhostMigration(ctx sdk.Context, store sdk.KVStore, cdc codec.Bina
 // avoid state corruption as modifying state during iteration is unsafe. A special case
 // for tendermint clients is included as only one tendermint clientID is required for
 // v7 migrations.
-func collectClients(ctx sdk.Context, store sdk.KVStore, clientType string) (clients []string, err error) {
-	clientPrefix := []byte(fmt.Sprintf("%s/%s", host.KeyClientStorePrefix, clientType))
-	iterator := sdk.KVStorePrefixIterator(store, clientPrefix)
+func collectClients(ctx sdk.Context, keeper ClientKeeper, clientType string) ([]string, error) {
+	store := ctx.KVStore(k.storeKey)
+	iterator := sdk.KVStorePrefixIterator(store, host.PrefixedClientStoreKey(prefix))
+
+	var clientIDs []string
 
 	defer iterator.Close()
 	for ; iterator.Valid(); iterator.Next() {
@@ -163,15 +150,15 @@ func collectClients(ctx sdk.Context, store sdk.KVStore, clientType string) (clie
 		}
 
 		clientID := host.MustParseClientStatePath(path)
-		clients = append(clients, clientID)
+		clientIDs = append(clientIDs, clientID)
 
 		// optimization: exit after a single tendermint client iteration
 		if strings.Contains(clientID, exported.Tendermint) {
-			return clients, nil
+			return clientIDs, nil
 		}
 	}
 
-	return clients, nil
+	return clientIDs, nil
 }
 
 // removeAllClientConsensusStates removes all client consensus states from the associated

modules/core/02-client/migrations/v7/store_test.go

@@ -1,6 +1,7 @@
 package v7_test
 
 import (
+	"fmt"
 	"strconv"
 	"testing"
 
@@ -44,23 +45,25 @@ func TestIBCTestSuite(t *testing.T) {
 func (suite *MigrationsV7TestSuite) TestMigrateStoreTendermint() {
 	path := ibctesting.NewPath(suite.chainA, suite.chainB)
 	suite.coordinator.SetupClients(path)
+	fmt.Println("here")
 
 	registry := codectypes.NewInterfaceRegistry()
 	cdc := codec.NewProtoCodec(registry)
 
 	solomachine.RegisterInterfaces(registry)
-	err := v7.MigrateStore(suite.chainA.GetContext(), suite.chainA.GetSimApp().GetKey(host.StoreKey), cdc)
+	err := v7.MigrateStore(suite.chainA.GetContext(), cdc, suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().Error(err)
 
 	ibctm.RegisterInterfaces(registry)
-	err = v7.MigrateStore(suite.chainA.GetContext(), suite.chainA.GetSimApp().GetKey(host.StoreKey), cdc)
+	err = v7.MigrateStore(suite.chainA.GetContext(), cdc, suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().NoError(err)
 }
 
 // create multiple solo machine clients, tendermint and localhost clients
 // ensure that solo machine clients are migrated and their consensus states are removed
 // ensure the localhost is deleted entirely.
 func (suite *MigrationsV7TestSuite) TestMigrateStore() {
+	fmt.Println("there")
 	paths := []*ibctesting.Path{
 		ibctesting.NewPath(suite.chainA, suite.chainB),
 		ibctesting.NewPath(suite.chainA, suite.chainB),
@@ -79,7 +82,7 @@ func (suite *MigrationsV7TestSuite) TestMigrateStore() {
 	suite.createSolomachineClients(solomachines)
 	suite.createLocalhostClients()
 
-	err := v7.MigrateStore(suite.chainA.GetContext(), suite.chainA.GetSimApp().GetKey(host.StoreKey), suite.chainA.App.AppCodec())
+	err := v7.MigrateStore(suite.chainA.GetContext(), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().NoError(err)
 
 	suite.assertSolomachineClients(solomachines)

modules/light-clients/07-tendermint/expected_keepers.go

@@ -0,0 +1,16 @@
+package tendermint
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/tendermint/tendermint/libs/log"
+
+	"github.com/cosmos/ibc-go/v6/modules/core/exported"
+)
+
+// ClientKeeper expected account IBC client keeper
+type ClientKeeper interface {
+	GetClientState(ctx sdk.Context, clientID string) (exported.ClientState, bool)
+	IterateClientStates(ctx sdk.Context, prefix []byte, cb func(string, exported.ClientState) bool)
+	ClientStore(ctx sdk.Context, clientID string) sdk.KVStore
+	Logger(ctx sdk.Context) log.Logger
+}

modules/light-clients/07-tendermint/migrations.go

@@ -1,60 +1,33 @@
 package tendermint
 
 import (
-	"fmt"
-	"strings"
-
 	"github.com/cosmos/cosmos-sdk/codec"
-	"github.com/cosmos/cosmos-sdk/store/prefix"
-	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
 	clienttypes "github.com/cosmos/ibc-go/v6/modules/core/02-client/types"
-	host "github.com/cosmos/ibc-go/v6/modules/core/24-host"
 	"github.com/cosmos/ibc-go/v6/modules/core/exported"
 )
 
 // PruneTendermintConsensusStates prunes all expired tendermint consensus states. This function
 // may optionally be called during in-place store migrations. The ibc store key must be provided.
-func PruneTendermintConsensusStates(ctx sdk.Context, cdc codec.BinaryCodec, storeKey storetypes.StoreKey) error {
-	store := ctx.KVStore(storeKey)
-
-	// iterate over ibc store with prefix: clients/07-tendermint,
-	tendermintClientPrefix := []byte(fmt.Sprintf("%s/%s", host.KeyClientStorePrefix, exported.Tendermint))
-	iterator := sdk.KVStorePrefixIterator(store, tendermintClientPrefix)
-
+func PruneTendermintConsensusStates(ctx sdk.Context, cdc codec.BinaryCodec, clientKeeper ClientKeeper) error {
 	var clientIDs []string
-
-	// collect all clients to avoid performing store state changes during iteration
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		path := string(iterator.Key())
-		if !strings.Contains(path, host.KeyClientState) {
-			// skip non client state keys
-			continue
-		}
-
-		clientID := host.MustParseClientStatePath(path)
+	clientKeeper.IterateClientStates(ctx, []byte(exported.Tendermint), func(clientID string, _ exported.ClientState) bool {
 		clientIDs = append(clientIDs, clientID)
-	}
+		return false
+	})
 
 	// keep track of the total consensus states pruned so chains can
 	// understand how much space is saved when the migration is run
 	var totalPruned int
 
 	for _, clientID := range clientIDs {
-		clientPrefix := []byte(fmt.Sprintf("%s/%s/", host.KeyClientStorePrefix, clientID))
-		clientStore := prefix.NewStore(ctx.KVStore(storeKey), clientPrefix)
+		clientStore := clientKeeper.ClientStore(ctx, clientID)
 
-		bz := clientStore.Get(host.ClientStateKey())
-		if bz == nil {
-			return clienttypes.ErrClientNotFound
-		}
-
-		var clientState exported.ClientState
-		if err := cdc.UnmarshalInterface(bz, &clientState); err != nil {
-			return sdkerrors.Wrap(err, "failed to unmarshal client state bytes into tendermint client state")
+		clientState, ok := clientKeeper.GetClientState(ctx, clientID)
+		if !ok {
+			return sdkerrors.Wrapf(clienttypes.ErrClientNotFound, "clientID %s", clientID)
 		}
 
 		tmClientState, ok := clientState.(*ClientState)
@@ -65,7 +38,7 @@ func PruneTendermintConsensusStates(ctx sdk.Context, cdc codec.BinaryCodec, stor
 		totalPruned += PruneAllExpiredConsensusStates(ctx, clientStore, cdc, tmClientState)
 	}
 
-	clientLogger := ctx.Logger().With("module", "x/"+host.ModuleName+"/"+clienttypes.SubModuleName)
+	clientLogger := clientKeeper.Logger(ctx)
 	clientLogger.Info("pruned expired tendermint consensus states", "total", totalPruned)
 
 	return nil

modules/light-clients/07-tendermint/migrations_test.go

@@ -99,7 +99,7 @@ func (suite *TendermintTestSuite) TestPruneTendermintConsensusStates() {
 	// This will cause the consensus states created before the first time increment
 	// to be expired
 	suite.coordinator.IncrementTimeBy(7 * 24 * time.Hour)
-	err = ibctm.PruneTendermintConsensusStates(suite.chainA.GetContext(), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().GetKey(host.StoreKey))
+	err = ibctm.PruneTendermintConsensusStates(suite.chainA.GetContext(), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().NoError(err)
 
 	for _, path := range paths {

testing/simapp/app.go

@@ -899,7 +899,7 @@ func (app *SimApp) setupUpgradeHandlers() {
 			app.mm,
 			app.configurator,
 			app.appCodec,
-			app.keys[ibchost.StoreKey],
+			app.IBCKeeper.ClientKeeper,
 		),
 	)
 }

testing/simapp/upgrades/v7/upgrades.go

@@ -2,11 +2,11 @@ package v7
 
 import (
 	"github.com/cosmos/cosmos-sdk/codec"
-	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
 
+	clientkeeper "github.com/cosmos/ibc-go/v6/modules/core/02-client/keeper"
 	ibctm "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint"
 )
 
@@ -20,13 +20,14 @@ func CreateUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
 	cdc codec.BinaryCodec,
-	hostStoreKey *storetypes.KVStoreKey,
+	clientKeeper clientkeeper.Keeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx sdk.Context, _ upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		// OPTIONAL: prune expired tendermint consensus states to save storage space
-		if err := ibctm.PruneTendermintConsensusStates(ctx, cdc, hostStoreKey); err != nil {
+		if err := ibctm.PruneTendermintConsensusStates(ctx, cdc, clientKeeper); err != nil {
 			return nil, err
 		}
+
 		return mm.RunMigrations(ctx, configurator, vm)
 	}
 }