added key assignment

x/ccv/provider/keeper/key_assignment.go

@@ -1,6 +1,7 @@
 package keeper
 
 import (
+	"encoding/base64"
 	"fmt"
 
 	errorsmod "cosmossdk.io/errors"
@@ -15,6 +16,53 @@ import (
 	ccvtypes "github.com/cosmos/interchain-security/v4/x/ccv/types"
 )
 
+// ParseConsumerKey parses the ED25519 PubKey`consumerKey` from a JSON string
+// and constructs its corresponding `tmprotocrypto.PublicKey`
+func (k Keeper) ParseConsumerKey(consumerKey string) (tmprotocrypto.PublicKey, error) {
+	// parse consumer key as long as it's in the right format
+	pkType, keyStr, err := types.ParseConsumerKeyFromJson(consumerKey)
+	if err != nil {
+		return tmprotocrypto.PublicKey{}, err
+	}
+
+	// Note: the correct way to decide if a key type is supported is to check the
+	// consensus params. However this functionality was disabled in https://github.com/cosmos/interchain-security/pull/916
+	// as a quick way to get ed25519 working, avoiding amino/proto-any marshalling issues.
+
+	// make sure the consumer key type is supported
+	// cp := ctx.ConsensusParams()
+	// if cp != nil && cp.Validator != nil {
+	// 	if !tmstrings.StringInSlice(pkType, cp.Validator.PubKeyTypes) {
+	// 		return nil, errorsmod.Wrapf(
+	// 			stakingtypes.ErrValidatorPubKeyTypeNotSupported,
+	// 			"got: %s, expected one of: %s", pkType, cp.Validator.PubKeyTypes,
+	// 		)
+	// 	}
+	// }
+
+	// For now, only accept ed25519.
+	// TODO: decide what types should be supported.
+	if pkType != "/cosmos.crypto.ed25519.PubKey" {
+		return tmprotocrypto.PublicKey{}, errorsmod.Wrapf(
+			stakingtypes.ErrValidatorPubKeyTypeNotSupported,
+			"got: %s, expected: %s", pkType, "/cosmos.crypto.ed25519.PubKey",
+		)
+	}
+
+	pubKeyBytes, err := base64.StdEncoding.DecodeString(keyStr)
+	if err != nil {
+		return tmprotocrypto.PublicKey{}, err
+	}
+
+	consumerTMPublicKey := tmprotocrypto.PublicKey{
+		Sum: &tmprotocrypto.PublicKey_Ed25519{
+			Ed25519: pubKeyBytes,
+		},
+	}
+
+	return consumerTMPublicKey, nil
+}
+
 // GetValidatorConsumerPubKey returns a validator's public key assigned for a consumer chain
 func (k Keeper) GetValidatorConsumerPubKey(
 	ctx sdk.Context,

x/ccv/provider/keeper/msg_server.go

@@ -2,15 +2,12 @@ package keeper
 
 import (
 	"context"
-	"encoding/base64"
-
 	errorsmod "cosmossdk.io/errors"
 
 	cryptocodec "github.com/cosmos/cosmos-sdk/crypto/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
-	tmprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 	tmtypes "github.com/cometbft/cometbft/types"
 
 	"github.com/cosmos/interchain-security/v4/x/ccv/provider/types"
@@ -45,47 +42,11 @@ func (k msgServer) AssignConsumerKey(goCtx context.Context, msg *types.MsgAssign
 		return nil, stakingtypes.ErrNoValidatorFound
 	}
 
-	// parse consumer key as long as it's in the right format
-	pkType, keyStr, err := types.ParseConsumerKeyFromJson(msg.ConsumerKey)
-	if err != nil {
-		return nil, err
-	}
-
-	// Note: the correct way to decide if a key type is supported is to check the
-	// consensus params. However this functionality was disabled in https://github.com/cosmos/interchain-security/pull/916
-	// as a quick way to get ed25519 working, avoiding amino/proto-any marshalling issues.
-
-	// make sure the consumer key type is supported
-	// cp := ctx.ConsensusParams()
-	// if cp != nil && cp.Validator != nil {
-	// 	if !tmstrings.StringInSlice(pkType, cp.Validator.PubKeyTypes) {
-	// 		return nil, errorsmod.Wrapf(
-	// 			stakingtypes.ErrValidatorPubKeyTypeNotSupported,
-	// 			"got: %s, expected one of: %s", pkType, cp.Validator.PubKeyTypes,
-	// 		)
-	// 	}
-	// }
-
-	// For now, only accept ed25519.
-	// TODO: decide what types should be supported.
-	if pkType != "/cosmos.crypto.ed25519.PubKey" {
-		return nil, errorsmod.Wrapf(
-			stakingtypes.ErrValidatorPubKeyTypeNotSupported,
-			"got: %s, expected: %s", pkType, "/cosmos.crypto.ed25519.PubKey",
-		)
-	}
-
-	pubKeyBytes, err := base64.StdEncoding.DecodeString(keyStr)
+	consumerTMPublicKey, err := k.ParseConsumerKey(msg.ConsumerKey)
 	if err != nil {
 		return nil, err
 	}
 
-	consumerTMPublicKey := tmprotocrypto.PublicKey{
-		Sum: &tmprotocrypto.PublicKey_Ed25519{
-			Ed25519: pubKeyBytes,
-		},
-	}
-
 	if err := k.Keeper.AssignConsumerKey(ctx, msg.ChainId, validator, consumerTMPublicKey); err != nil {
 		return nil, err
 	}

x/ccv/provider/keeper/partial_set_security.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	errorsmod "cosmossdk.io/errors"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/cosmos/interchain-security/v4/x/ccv/provider/types"
 )
 
@@ -28,7 +29,20 @@ func (k Keeper) HandleOptIn(ctx sdk.Context, chainID string, providerAddr types.
 	}
 
 	if consumerKey != nil {
-		// TODO (PR 1586): assign consumer key in this case
+		consumerTMPublicKey, err := k.ParseConsumerKey(*consumerKey)
+		if err != nil {
+			return err
+		}
+
+		validator, found := k.stakingKeeper.GetValidatorByConsAddr(ctx, providerAddr.Address)
+		if !found {
+			return stakingtypes.ErrNoValidatorFound
+		}
+
+		err = k.AssignConsumerKey(ctx, chainID, validator, consumerTMPublicKey)
+		if err != nil {
+			return err
+		}
 	}
 
 	return nil

x/ccv/provider/keeper/partial_set_security_test.go

@@ -1,8 +1,14 @@
 package keeper_test
 
 import (
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/crypto/keys/ed25519"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	testkeeper "github.com/cosmos/interchain-security/v4/testutil/keeper"
 	"github.com/cosmos/interchain-security/v4/x/ccv/provider/types"
+	ccvtypes "github.com/cosmos/interchain-security/v4/x/ccv/types"
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 	"testing"
 )
@@ -31,6 +37,53 @@ func TestHandleOptIn(t *testing.T) {
 	require.False(t, providerKeeper.IsToBeOptedIn(ctx, "chainID", providerAddr))
 }
 
+func TestHandleOptInWithConsumerKey(t *testing.T) {
+	providerKeeper, ctx, ctrl, mocks := testkeeper.GetProviderKeeperAndCtx(t, testkeeper.NewInMemKeeperParams(t))
+	defer ctrl.Finish()
+
+	// generate a consensus public key for the provider
+	providerConsPubKey := ed25519.GenPrivKeyFromSecret([]byte{1}).PubKey()
+	consAddr := sdk.ConsAddress(providerConsPubKey.Address())
+	providerAddr := types.NewProviderConsAddress(consAddr)
+
+	calls := []*gomock.Call{
+		mocks.MockStakingKeeper.EXPECT().
+			GetValidatorByConsAddr(gomock.Any(), gomock.Any()).
+			DoAndReturn(func(ctx sdk.Context, addr sdk.ConsAddress) (stakingtypes.Validator, bool) {
+				if addr.Equals(providerAddr.Address) {
+					// Given `providerAddr`, `GetValidatorByConsAddr` returns a validator with the
+					// exact same `ConsensusPubkey`
+					pkAny, _ := codectypes.NewAnyWithValue(providerConsPubKey)
+					return stakingtypes.Validator{ConsensusPubkey: pkAny}, true
+				} else {
+					// for any other consensus address, we cannot find a validator
+					return stakingtypes.Validator{}, false
+				}
+			}).Times(2),
+	}
+
+	gomock.InOrder(calls...)
+	providerKeeper.SetProposedConsumerChain(ctx, "chainID", 1)
+
+	// create a sample consumer key to assign to the `providerAddr` validator
+	// on the consumer chain with id `chainID`
+	consumerKey := "{\"@type\":\"/cosmos.crypto.ed25519.PubKey\",\"key\":\"Ui5Gf1+mtWUdH8u3xlmzdKID+F3PK0sfXZ73GZ6q6is=\"}"
+	expectedConsumerPubKey, _ := providerKeeper.ParseConsumerKey(consumerKey)
+
+	err := providerKeeper.HandleOptIn(ctx, "chainID", providerAddr, &consumerKey)
+	require.NoError(t, err)
+
+	// assert that the consumeKey was assigned to `providerAddr` validator on chain with id `chainID`
+	actualConsumerPubKey, found := providerKeeper.GetValidatorConsumerPubKey(ctx, "chainID", providerAddr)
+	require.True(t, found)
+	require.Equal(t, expectedConsumerPubKey, actualConsumerPubKey)
+
+	// assert that the `consumerAddr` to `providerAddr` association was set as well
+	consumerAddr, _ := ccvtypes.TMCryptoPublicKeyToConsAddr(actualConsumerPubKey)
+	actualProviderConsAddr, found := providerKeeper.GetValidatorByConsumerAddr(ctx, "chainID", types.NewConsumerConsAddress(consumerAddr))
+	require.Equal(t, providerAddr, actualProviderConsAddr)
+}
+
 func TestHandleOptOut(t *testing.T) {
 	providerKeeper, ctx, ctrl, _ := testkeeper.GetProviderKeeperAndCtx(t, testkeeper.NewInMemKeeperParams(t))
 	defer ctrl.Finish()