refactor!: store key (WIP)

x/ccv/provider/keeper/hooks.go

@@ -100,7 +100,7 @@ func ValidatorConsensusKeyInUse(k *Keeper, ctx sdk.Context, valAddr sdk.ValAddre
 
 	inUse := false
 
-	for _, validatorConsumerAddrs := range k.GetAllValidatorsByConsumerAddr(ctx, nil) {
+	for _, validatorConsumerAddrs := range k.GetValidatorsByConsumerAddr(ctx, consensusAddr) {
 		if sdk.ConsAddress(validatorConsumerAddrs.ConsumerAddr).Equals(consensusAddr) {
 			inUse = true
 			break

x/ccv/provider/keeper/key_assignment.go

@@ -1,10 +1,8 @@
 package keeper
 
 import (
-	"fmt"
-
 	errorsmod "cosmossdk.io/errors"
-
+	"fmt"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
@@ -135,30 +133,23 @@
 	store.Set(types.ValidatorsByConsumerAddrKey(chainID, consumerAddr), bz)
 }
 
-// GetValidatorsByConsumerAddrs gets all the mappings from consensus addresses
+// GetValidatorsByConsumerAddr gets all the mappings from consensus addresses
 // on a given consumer chain to consensus addresses on the provider chain.
 // If chainID is nil, it returns all the mappings from consensus addresses on all consumer chains.
 //
 // Note that the mappings for a consumer chain are stored under keys with the following format:
-// ValidatorsByConsumerAddrBytePrefix | len(chainID) | chainID | consumerAddress
+// ValidatorsByConsumerAddrBytePrefix | consumerAddress | len(chainID) | chainID
 // Thus, the returned array is
 //   - in ascending order of consumerAddresses, if chainID is not nil;
 //   - in undetermined order, if chainID is nil.
-func (k Keeper) GetAllValidatorsByConsumerAddr(ctx sdk.Context, chainID *string) (validatorConsumerAddrs []types.ValidatorByConsumerAddr) {
+
+func (k Keeper) GetValidatorsByConsumerAddr(ctx sdk.Context, consumerConsAddr sdk.ConsAddress) (validatorConsumerAddrs []types.ValidatorByConsumerAddr) {
 	store := ctx.KVStore(k.storeKey)
-	var prefix []byte
-	if chainID == nil {
-		// iterate over the mappings from consensus addresses on all consumer chains
-		prefix = []byte{types.ValidatorsByConsumerAddrBytePrefix}
-	} else {
-		// iterate over the mappings from consensus addresses on chainID
-		prefix = types.ChainIdWithLenKey(types.ValidatorsByConsumerAddrBytePrefix, *chainID)
-	}
+	prefix := types.PrefixWithLenConsAddress(types.ValidatorsByConsumerAddrBytePrefix, consumerConsAddr)
 	iterator := sdk.KVStorePrefixIterator(store, prefix)
 	defer iterator.Close()
 	for ; iterator.Valid(); iterator.Next() {
-		// TODO: store chainID and consumer cons address in value bytes, marshaled as protobuf type
-		chainID, consumerAddrTmp, err := types.ParseChainIdAndConsAddrKey(types.ValidatorsByConsumerAddrBytePrefix, iterator.Key())
+		chainID, consumerAddrTmp, err := types.ParseConsAddrKeyAndChainID(types.ValidatorsByConsumerAddrBytePrefix, iterator.Key())
 		if err != nil {
 			// An error here would indicate something is very wrong,
 			// store keys are assumed to be correctly serialized in SetValidatorByConsumerAddr.
@@ -177,6 +168,41 @@
 	return validatorConsumerAddrs
 }
 
+// ValidatorsByConsumerAddrBytePrefix | consumerAddress | len(chainID) | chainID
+func (k Keeper) GetAllValidatorsByConsumerAddr(ctx sdk.Context, chainID *string) (validatorConsumerAddrs []types.ValidatorByConsumerAddr) {
+	store := ctx.KVStore(k.storeKey)
+	prefix := []byte{types.ValidatorsByConsumerAddrBytePrefix}
+	iterator := sdk.KVStorePrefixIterator(store, prefix)
+	for ; iterator.Valid(); iterator.Next() {
+		cID, consumerAddrTmp, err := types.ParseConsAddrKeyAndChainID(types.ValidatorsByConsumerAddrBytePrefix, iterator.Key())
+		if err != nil {
+			// An error here would indicate something is very wrong,
+			// store keys are assumed to be correctly serialized in SetValidatorByConsumerAddr.
+			panic(fmt.Sprintf("failed to parse chainID and consumer address: %v", err))
+		}
+
+		if cID == *chainID || chainID == nil {
+			consumerAddr := types.NewConsumerConsAddress(consumerAddrTmp)
+			providerAddr := types.NewProviderConsAddress(iterator.Value())
+
+			validatorConsumerAddrs = append(validatorConsumerAddrs, types.ValidatorByConsumerAddr{
+				ConsumerAddr: consumerAddr.ToSdkConsAddr(),
+				ProviderAddr: providerAddr.ToSdkConsAddr(),
+				ChainId:      cID,
+			})
+		}
+	}
+
+	return validatorConsumerAddrs
+}
+
+// DeleteValidatorByConsumerAddrLegacy deletes the mapping from a validator's consensus address on a consumer
+// to the validator's consensus address on the provider
+func (k Keeper) DeleteValidatorByConsumerAddrLegacy(ctx sdk.Context, chainID string, consumerAddr types.ConsumerConsAddress) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(types.ValidatorsByConsumerAddrKeyLegacy(chainID, consumerAddr))
+}
+
 // DeleteValidatorByConsumerAddr deletes the mapping from a validator's consensus address on a consumer
 // to the validator's consensus address on the provider
 func (k Keeper) DeleteValidatorByConsumerAddr(ctx sdk.Context, chainID string, consumerAddr types.ConsumerConsAddress) {
@@ -610,10 +636,22 @@
 	for _, validatorConsumerAddr := range k.GetAllValidatorConsumerPubKeys(ctx, &chainID) {
 		providerAddr := types.NewProviderConsAddress(validatorConsumerAddr.ProviderAddr)
 		k.DeleteValidatorConsumerPubKey(ctx, chainID, providerAddr)
+
+		//consAddrTmp, err := ccvtypes.TMCryptoPublicKeyToConsAddr(*validatorConsumerAddr.ConsumerKey)
+		//if err != nil {
+		//	// todo
+		//	panic(err)
+		//}
+		//consumerAddr := types.NewConsumerConsAddress(consAddrTmp)
+		//
+		//k.DeleteValidatorByConsumerAddr(ctx, chainID, consumerAddr)
 	}
 
-	// delete ValidatorsByConsumerAddr
+	//// delete ValidatorsByConsumerAddr
 	for _, validatorConsumerAddr := range k.GetAllValidatorsByConsumerAddr(ctx, &chainID) {
+		if validatorConsumerAddr.ChainId != chainID {
+			continue
+		}
 		consumerAddr := types.NewConsumerConsAddress(validatorConsumerAddr.ConsumerAddr)
 		k.DeleteValidatorByConsumerAddr(ctx, chainID, consumerAddr)
 	}

x/ccv/provider/keeper/key_assignment_test.go

@@ -118,7 +118,7 @@ func TestValidatorByConsumerAddrCRUD(t *testing.T) {
 	require.NotEmpty(t, providerAddrResult, "provider address is empty")
 	require.Equal(t, providerAddr, providerAddrResult)
 
-	keeper.DeleteValidatorByConsumerAddr(ctx, chainID, consumerAddr)
+	keeper.DeleteValidatorByConsumerAddrLegacy(ctx, chainID, consumerAddr)
 	providerAddrResult, found = keeper.GetValidatorByConsumerAddr(ctx, chainID, consumerAddr)
 	require.False(t, found, "provider address was found")
 	require.Empty(t, providerAddrResult, "provider address not empty")

x/ccv/provider/keeper/migration.go

@@ -0,0 +1,49 @@
+package keeper
+
+import (
+	"github.com/coinbase/rosetta-sdk-go/storage/errors"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/interchain-security/v3/x/ccv/provider/types"
+)
+
+type Migrator struct {
+	ccvProviderKeeper Keeper
+}
+
+func NewMigrator(ccvProviderKeeper Keeper) Migrator {
+	return Migrator{ccvProviderKeeper: ccvProviderKeeper}
+}
+
+func (m Migrator) Migrate2to3(ctx sdk.Context) error {
+	return m.ccvProviderKeeper.MigrateConsumerAddrStoreKey(ctx)
+}
+
+func (k Keeper) MigrateConsumerAddrStoreKey(ctx sdk.Context) error {
+	store := ctx.KVStore(k.storeKey)
+	var prefixCon []byte
+
+	prefixCon = []byte{types.ValidatorsByConsumerAddrBytePrefix}
+
+	iteratorCon := sdk.KVStorePrefixIterator(store, prefixCon)
+	defer iteratorCon.Close()
+	for ; iteratorCon.Valid(); iteratorCon.Next() {
+		// ParseChainIdAndConsAddrKey is used in other places, so it was not renamed to ParseChainIdAndConsAddrKeyLegacy
+		chainID, consumerAddrTmp, err := types.ParseChainIdAndConsAddrKey(types.ValidatorsByConsumerAddrBytePrefix, iteratorCon.Key())
+		if err != nil {
+			// An error here would indicate something is very wrong,
+			// store keys are assumed to be correctly serialized in SetValidatorByConsumerAddr.
+			return errors.ErrParseKeyPairFailed
+		}
+		consumerAddr := types.NewConsumerConsAddress(consumerAddrTmp)
+		providerAddr := types.NewProviderConsAddress(iteratorCon.Value())
+
+		// bytePrefix | len(ConsAddress) | ConsAddress | chainID
+		k.SetValidatorByConsumerAddr(ctx, chainID, consumerAddr, providerAddr)
+
+		// delete old kv
+		k.DeleteValidatorByConsumerAddrLegacy(ctx, chainID, consumerAddr)
+
+	}
+
+	return nil
+}

x/ccv/provider/module.go

@@ -107,6 +107,10 @@ func (AppModule) RegisterInvariants(ir sdk.InvariantRegistry) {
 func (am AppModule) RegisterServices(cfg module.Configurator) {
 	providertypes.RegisterMsgServer(cfg.MsgServer(), keeper.NewMsgServerImpl(am.keeper))
 	providertypes.RegisterQueryServer(cfg.QueryServer(), am.keeper)
+	m := keeper.NewMigrator(*am.keeper)
+	if err := cfg.RegisterMigration(providertypes.ModuleName, 1, m.Migrate2to3); err != nil {
+		panic(fmt.Sprintf("failed to migrate provider from version 1 to 2: %v", err))
+	}
 }
 
 // InitGenesis performs genesis initialization for the provider module. It returns no validator updates.
@@ -129,7 +133,7 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 }
 
 // ConsensusVersion implements AppModule/ConsensusVersion.
-func (AppModule) ConsensusVersion() uint64 { return 2 }
+func (AppModule) ConsensusVersion() uint64 { return 3 }
 
 // BeginBlock implements the AppModule interface
 func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {

x/ccv/provider/types/keys.go

@@ -343,22 +343,34 @@ func MustParseGlobalSlashEntryKey(bz []byte) (
 	return recvTime, chainID, ibcSeqNum
 }
 
+// ConsumerValidatorsKeyLegacy returns the key under which the
+// validator assigned keys for every consumer chain are stored
+// key: bytePrefix | len(chainID) | chainID | ConsAddress
+func ConsumerValidatorsKeyLegacy(chainID string, addr ProviderConsAddress) []byte {
+	return ChainIdAndConsAddrKeyLegacy(ConsumerValidatorsBytePrefix, chainID, addr.ToSdkConsAddr())
+}
+
 // ConsumerValidatorsKey returns the key under which the
 // validator assigned keys for every consumer chain are stored
+// key bytePrefix | len(ConsAddress) | ConsAddress | chainID
 func ConsumerValidatorsKey(chainID string, addr ProviderConsAddress) []byte {
 	return ChainIdAndConsAddrKey(ConsumerValidatorsBytePrefix, chainID, addr.ToSdkConsAddr())
 }
 
-// ValidatorsByConsumerAddrKey returns the key under which the mapping from validator addresses
+// ValidatorsByConsumerAddrKeyLegacy returns the key under which the mapping from validator addresses
 // on consumer chains to validator addresses on the provider chain is stored
+func ValidatorsByConsumerAddrKeyLegacy(chainID string, addr ConsumerConsAddress) []byte {
+	return ChainIdAndConsAddrKeyLegacy(ValidatorsByConsumerAddrBytePrefix, chainID, addr.ToSdkConsAddr())
+}
+
 func ValidatorsByConsumerAddrKey(chainID string, addr ConsumerConsAddress) []byte {
 	return ChainIdAndConsAddrKey(ValidatorsByConsumerAddrBytePrefix, chainID, addr.ToSdkConsAddr())
 }
 
 // KeyAssignmentReplacementsKey returns the key under which the
 // key assignments that need to be replaced in the current block are stored
 func KeyAssignmentReplacementsKey(chainID string, addr ProviderConsAddress) []byte {
-	return ChainIdAndConsAddrKey(KeyAssignmentReplacementsBytePrefix, chainID, addr.ToSdkConsAddr())
+	return ChainIdAndConsAddrKeyLegacy(KeyAssignmentReplacementsBytePrefix, chainID, addr.ToSdkConsAddr())
 }
 
 // ConsumerAddrsToPruneKey returns the key under which the
@@ -414,6 +426,17 @@ func ChainIdWithLenKey(prefix byte, chainID string) []byte {
 	)
 }
 
+// PrefixWithLenConsAddress returns the key with the following format:
+// bytePrefix | Address
+func PrefixWithLenConsAddress(prefix byte, addr sdk.ConsAddress) []byte {
+	addrL := len(addr)
+	return ccvtypes.AppendMany(
+		[]byte{prefix},
+		sdk.Uint64ToBigEndian(uint64(addrL)),
+		addr,
+	)
+}
+
 // ParseChainIdAndTsKey returns the chain ID and time for a ChainIdAndTs key
 func ParseChainIdAndTsKey(prefix byte, bz []byte) (string, time.Time, error) {
 	expectedPrefix := []byte{prefix}
@@ -455,9 +478,9 @@ func ParseChainIdAndUintIdKey(prefix byte, bz []byte) (string, uint64, error) {
 	return chainID, uintID, nil
 }
 
-// ChainIdAndConsAddrKey returns the key with the following format:
+// ChainIdAndConsAddrKeyLegacy returns the key with the following format:
 // bytePrefix | len(chainID) | chainID | ConsAddress
-func ChainIdAndConsAddrKey(prefix byte, chainID string, addr sdk.ConsAddress) []byte {
+func ChainIdAndConsAddrKeyLegacy(prefix byte, chainID string, addr sdk.ConsAddress) []byte {
 	partialKey := ChainIdWithLenKey(prefix, chainID)
 	return ccvtypes.AppendMany(
 		// Append the partialKey
@@ -467,7 +490,18 @@ func ChainIdAndConsAddrKey(prefix byte, chainID string, addr sdk.ConsAddress) []
 	)
 }
 
-// ParseChainIdAndConsAddrKey returns the chain ID and ConsAddress for a ChainIdAndConsAddrKey key
+// ChainIdAndConsAddrKey returns the key with the following format:
+// bytePrefix | ConsAddress | len(chainID) | chainID
+func ChainIdAndConsAddrKey(prefix byte, chainID string, addr sdk.ConsAddress) []byte {
+	partialKey := PrefixWithLenConsAddress(prefix, addr)
+	return ccvtypes.AppendMany(
+		partialKey,
+		// Append the chainID
+		[]byte(chainID),
+	)
+}
+
+// ParseChainIdAndConsAddrKey returns the chain ID and ConsAddress for a ChainIdAndConsAddrKey key: bytePrefix | len(chainID) | chainID | ConsAddress
 func ParseChainIdAndConsAddrKey(prefix byte, bz []byte) (string, sdk.ConsAddress, error) {
 	expectedPrefix := []byte{prefix}
 	prefixL := len(expectedPrefix)
@@ -480,6 +514,20 @@ func ParseChainIdAndConsAddrKey(prefix byte, bz []byte) (string, sdk.ConsAddress
 	return chainID, addr, nil
 }
 
+// ParseConsAddrKeyAndChainID returns the chain ID and ConsAddress for a ConsAddrKeyAndChainID key:  bytePrefix | len(ConsAddress) | ConsAddress | chainID
+func ParseConsAddrKeyAndChainID(prefix byte, bz []byte) (string, sdk.ConsAddress, error) {
+	expectedPrefix := []byte{prefix}
+	prefixL := len(expectedPrefix)
+	if prefix := bz[:prefixL]; !bytes.Equal(prefix, expectedPrefix) {
+		return "", nil, fmt.Errorf("invalid prefix; expected: %X, got: %X", expectedPrefix, prefix)
+	}
+	addrL := sdk.BigEndianToUint64(bz[prefixL : prefixL+8])
+	addr := bz[prefixL+8 : prefixL+8+int(addrL)]
+	chainID := string(bz[prefixL+8+int(addrL):])
+
+	return chainID, addr, nil
+}
+
 func VSCMaturedHandledThisBlockKey() []byte {
 	return []byte{VSCMaturedHandledThisBlockBytePrefix}
 }

x/ccv/provider/types/keys_test.go

@@ -92,8 +92,8 @@ func getAllFullyDefinedKeys() [][]byte {
 		providertypes.ThrottledPacketDataSizeKey("chainID"),
 		providertypes.ThrottledPacketDataKey("chainID", 88),
 		providertypes.GlobalSlashEntryKey(providertypes.GlobalSlashEntry{}),
-		providertypes.ConsumerValidatorsKey("chainID", providertypes.NewProviderConsAddress([]byte{0x05})),
-		providertypes.ValidatorsByConsumerAddrKey("chainID", providertypes.NewConsumerConsAddress([]byte{0x05})),
+		providertypes.ConsumerValidatorsKeyLegacy("chainID", providertypes.NewProviderConsAddress([]byte{0x05})),
+		providertypes.ValidatorsByConsumerAddrKeyLegacy("chainID", providertypes.NewConsumerConsAddress([]byte{0x05})),
 		providertypes.KeyAssignmentReplacementsKey("chainID", providertypes.NewProviderConsAddress([]byte{0x05})),
 		providertypes.ConsumerAddrsToPruneKey("chainID", 88),
 		providertypes.SlashLogKey(providertypes.NewProviderConsAddress([]byte{0x05})),
@@ -228,7 +228,7 @@ func TestChainIdAndConsAddrAndParse(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		key := providertypes.ChainIdAndConsAddrKey(test.prefix, test.chainID, test.addr)
+		key := providertypes.ChainIdAndConsAddrKeyLegacy(test.prefix, test.chainID, test.addr)
 		require.NotEmpty(t, key)
 		// Expected bytes = prefix + chainID length + chainID + consAddr bytes
 		expectedLen := 1 + 8 + len(test.chainID) + len(test.addr)