feat: include rotate keys logic in abci

tests/integration/staking/keeper/unbonding_test.go

@@ -44,6 +44,7 @@ func SetupUnbondingTests(t *testing.T, f *fixture, hookCalled *bool, ubdeID *uin
 	mockStackingHooks.EXPECT().BeforeDelegationSharesModified(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockStackingHooks.EXPECT().BeforeValidatorModified(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	mockStackingHooks.EXPECT().BeforeValidatorSlashed(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	mockStackingHooks.EXPECT().AfterConsensusPubKeyUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 	f.stakingKeeper.SetHooks(types.NewMultiStakingHooks(mockStackingHooks))
 
 	addrDels = simtestutil.AddTestAddrsIncremental(f.bankKeeper, f.stakingKeeper, f.sdkCtx, 2, math.NewInt(10000))

x/distribution/keeper/hooks.go

@@ -9,6 +9,7 @@ import (
 	"cosmossdk.io/x/distribution/types"
 	stakingtypes "cosmossdk.io/x/staking/types"
 
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -187,3 +188,7 @@ func (h Hooks) BeforeDelegationRemoved(_ context.Context, _ sdk.AccAddress, _ sd
 func (h Hooks) AfterUnbondingInitiated(_ context.Context, _ uint64) error {
 	return nil
 }
+
+func (h Hooks) AfterConsensusPubKeyUpdate(_ context.Context, _, _ cryptotypes.PubKey, _ sdk.Coin) error {
+	return nil
+}

x/evidence/keeper/infraction.go

@@ -40,6 +40,13 @@ func (k Keeper) handleEquivocationEvidence(ctx context.Context, evidence *types.
 	}
 
 	if len(validator.GetOperator()) != 0 {
+		// get the consAddr again, this is because validator might've rotated it's key.
+		valConsAddr, err := validator.GetConsAddr()
+		if err != nil {
+			return err
+		}
+		consAddr = valConsAddr
+
 		if _, err := k.slashingKeeper.GetPubkey(ctx, consAddr.Bytes()); err != nil {
 			// Ignore evidence that cannot be handled.
 			//

x/slashing/keeper/hooks.go

@@ -9,6 +9,7 @@ import (
 	sdkmath "cosmossdk.io/math"
 	"cosmossdk.io/x/slashing/types"
 
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -99,3 +100,15 @@ func (h Hooks) BeforeValidatorSlashed(_ context.Context, _ sdk.ValAddress, _ sdk
 func (h Hooks) AfterUnbondingInitiated(_ context.Context, _ uint64) error {
 	return nil
 }
+
+func (h Hooks) AfterConsensusPubKeyUpdate(ctx context.Context, oldPubKey, newPubKey cryptotypes.PubKey, _ sdk.Coin) error {
+	if err := h.k.PerformConsensusPubKeyUpdate(ctx, oldPubKey, newPubKey); err != nil {
+		return err
+	}
+
+	if err := h.k.AddrPubkeyRelation.Remove(ctx, oldPubKey.Address()); err != nil {
+		return err
+	}
+
+	return nil
+}

x/slashing/keeper/signing_info.go

@@ -11,6 +11,7 @@ import (
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/x/slashing/types"
 
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -209,3 +210,44 @@ func (k Keeper) GetValidatorMissedBlocks(ctx context.Context, addr sdk.ConsAddre
 
 	return missedBlocks, err
 }
+
+// PerformConsensusPubKeyUpdate updates cons address to its pub key relation
+// Updates signing info, missed blocks (removes old one, and sets new one)
+func (k Keeper) PerformConsensusPubKeyUpdate(ctx context.Context, oldPubKey, newPubKey cryptotypes.PubKey) error {
+	// Connect new consensus address with PubKey
+	if err := k.AddrPubkeyRelation.Set(ctx, newPubKey.Address(), newPubKey); err != nil {
+		return err
+	}
+
+	// Migrate ValidatorSigningInfo from oldPubKey to newPubKey
+	signingInfo, err := k.ValidatorSigningInfo.Get(ctx, sdk.ConsAddress(oldPubKey.Address()))
+	if err != nil {
+		return types.ErrInvalidConsPubKey.Wrap("failed to get signing info for old public key")
+	}
+
+	if err := k.ValidatorSigningInfo.Set(ctx, sdk.ConsAddress(newPubKey.Address()), signingInfo); err != nil {
+		return err
+	}
+
+	if err := k.ValidatorSigningInfo.Remove(ctx, sdk.ConsAddress(oldPubKey.Address())); err != nil {
+		return err
+	}
+
+	// Migrate ValidatorMissedBlockBitArray from oldPubKey to newPubKey
+	missedBlocks, err := k.GetValidatorMissedBlocks(ctx, sdk.ConsAddress(oldPubKey.Address()))
+	if err != nil {
+		return err
+	}
+
+	if err := k.DeleteMissedBlockBitmap(ctx, sdk.ConsAddress(oldPubKey.Address())); err != nil {
+		return err
+	}
+
+	for _, missed := range missedBlocks {
+		if err := k.SetMissedBlockBitmapValue(ctx, sdk.ConsAddress(newPubKey.Address()), missed.Index, missed.Missed); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

x/slashing/types/errors.go

@@ -13,4 +13,5 @@ var (
 	ErrNoSigningInfoFound           = errors.Register(ModuleName, 8, "no validator signing info found")
 	ErrValidatorTombstoned          = errors.Register(ModuleName, 9, "validator already tombstoned")
 	ErrInvalidSigner                = errors.Register(ModuleName, 10, "expected authority account as only signer for proposal message")
+	ErrInvalidConsPubKey            = errors.Register(ModuleName, 11, "invalid consensus pubkey")
 )

x/staking/keeper/cons_pubkey.go

@@ -6,10 +6,14 @@ import (
 	"time"
 
 	"cosmossdk.io/collections"
+	"cosmossdk.io/collections/indexes"
+	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/x/staking/types"
 
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 )
 
 // maxRotations is the value of max rotations can be made in unbonding period for a validator.
@@ -47,6 +51,54 @@ func (k Keeper) setConsPubKeyRotationHistory(
 	return k.setConsKeyQueue(ctx, queueTime, valAddr)
 }
 
+// This method gets called from the `ApplyAndReturnValidatorSetUpdates`(from endblocker) method.
+//
+// This method makes the relative state changes to update the keys,
+// also maintains a map with old to new conskey rotation which is needed to retrieve the old conskey.
+// And also triggers the hook to make changes required in slashing and distribution modules.
+func (k Keeper) updateToNewPubkey(ctx context.Context, val types.Validator, oldPubKey, newPubKey *codectypes.Any, fee sdk.Coin) error {
+	consAddr, err := val.GetConsAddr()
+	if err != nil {
+		return err
+	}
+
+	if err := k.ValidatorByConsensusAddress.Remove(ctx, consAddr); err != nil {
+		return err
+	}
+
+	if err := k.DeleteValidatorByPowerIndex(ctx, val); err != nil {
+		return err
+	}
+
+	val.ConsensusPubkey = newPubKey
+	if err := k.SetValidator(ctx, val); err != nil {
+		return err
+	}
+	if err := k.SetValidatorByConsAddr(ctx, val); err != nil {
+		return err
+	}
+	if err := k.SetValidatorByPowerIndex(ctx, val); err != nil {
+		return err
+	}
+
+	oldPk, ok := oldPubKey.GetCachedValue().(cryptotypes.PubKey)
+	if !ok {
+		return errorsmod.Wrapf(sdkerrors.ErrInvalidType, "Expecting cryptotypes.PubKey, got %T", oldPk)
+	}
+
+	newPk, ok := newPubKey.GetCachedValue().(cryptotypes.PubKey)
+	if !ok {
+		return errorsmod.Wrapf(sdkerrors.ErrInvalidType, "Expecting cryptotypes.PubKey, got %T", newPk)
+	}
+
+	// Sets a map to newly rotated consensus key with old consensus key
+	if err := k.RotatedConsKeyMapIndex.Set(ctx, oldPk.Address(), newPk.Address()); err != nil {
+		return err
+	}
+
+	return k.Hooks().AfterConsensusPubKeyUpdate(ctx, oldPk, newPk, fee)
+}
+
 // exceedsMaxRotations returns true if the key rotations exceed the limit, currently we are limiting one rotation for unbonding period.
 func (k Keeper) exceedsMaxRotations(ctx context.Context, valAddr sdk.ValAddress) error {
 	count := 0
@@ -90,3 +142,61 @@ func bytesSliceExists(sliceList [][]byte, targetBytes []byte) bool {
 	}
 	return false
 }
+
+// UpdateAllMaturedConsKeyRotatedKeys udpates all the matured key rotations.
+func (k Keeper) UpdateAllMaturedConsKeyRotatedKeys(ctx sdk.Context, maturedTime time.Time) error {
+	maturedRotatedValAddrs, err := k.GetAllMaturedRotatedKeys(ctx, maturedTime)
+	if err != nil {
+		return err
+	}
+
+	for _, valAddr := range maturedRotatedValAddrs {
+		err := k.deleteConsKeyIndexKey(ctx, valAddr, maturedTime)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// deleteConsKeyIndexKey deletes the key which is formed with the given valAddr, time.
+func (k Keeper) deleteConsKeyIndexKey(ctx sdk.Context, valAddr sdk.ValAddress, ts time.Time) error {
+	rng := new(collections.Range[collections.Pair[[]byte, time.Time]]).
+		StartInclusive(collections.Join(valAddr.Bytes(), time.Time{})).
+		EndInclusive(collections.Join(valAddr.Bytes(), ts))
+
+	return k.ValidatorConsensusKeyRotationRecordIndexKey.Walk(ctx, rng, func(key collections.Pair[[]byte, time.Time]) (stop bool, err error) {
+		return false, k.ValidatorConsensusKeyRotationRecordIndexKey.Remove(ctx, key)
+	})
+}
+
+// GetAllMaturedRotatedKeys returns all matured valaddresses .
+func (k Keeper) GetAllMaturedRotatedKeys(ctx sdk.Context, matureTime time.Time) ([][]byte, error) {
+	valAddrs := [][]byte{}
+
+	// get an iterator for all timeslices from time 0 until the current HeaderInfo time
+	rng := new(collections.Range[time.Time]).EndInclusive(matureTime)
+	err := k.ValidatorConsensusKeyRotationRecordQueue.Walk(ctx, rng, func(key time.Time, value types.ValAddrsOfRotatedConsKeys) (stop bool, err error) {
+		valAddrs = append(valAddrs, value.Addresses...)
+		return false, k.ValidatorConsensusKeyRotationRecordQueue.Remove(ctx, key)
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return valAddrs, nil
+}
+
+// GetBlockConsPubKeyRotationHistory iterator over the rotation history for the given height.
+func (k Keeper) GetBlockConsPubKeyRotationHistory(ctx context.Context) ([]types.ConsPubKeyRotationHistory, error) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+	iterator, err := k.RotationHistory.Indexes.Block.MatchExact(ctx, uint64(sdkCtx.BlockHeight()))
+	if err != nil {
+		return nil, err
+	}
+	defer iterator.Close()
+
+	return indexes.CollectValues(ctx, k.RotationHistory, iterator)
+}

x/staking/keeper/hooks_test.go

@@ -0,0 +1,15 @@
+package keeper_test
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func (s *KeeperTestSuite) TestHookAfterConsensusPubKeyUpdate() {
+	stKeeper := s.stakingKeeper
+	ctx := s.ctx
+	require := s.Require()
+
+	rotationFee := sdk.NewInt64Coin("stake", 1000000)
+	err := stKeeper.Hooks().AfterConsensusPubKeyUpdate(ctx, PKs[0], PKs[1], rotationFee)
+	require.NoError(err)
+}

x/staking/keeper/msg_server.go

@@ -661,7 +661,7 @@ func (k msgServer) RotateConsPubKey(ctx context.Context, msg *types.MsgRotateCon
 		return nil, err
 	}
 
-	err = k.Keeper.bankKeeper.SendCoinsFromAccountToModule(ctx, sdk.AccAddress(valAddr), types.DistributionModuleName, sdk.NewCoins(params.KeyRotationFee))
+	err = k.Keeper.bankKeeper.SendCoinsFromAccountToModule(ctx, sdk.AccAddress(valAddr), types.PoolModuleName, sdk.NewCoins(params.KeyRotationFee))
 	if err != nil {
 		return nil, err
 	}

x/staking/keeper/val_state_change.go

@@ -13,6 +13,8 @@ import (
 	"cosmossdk.io/math"
 	"cosmossdk.io/x/staking/types"
 
+	cryptocodec "github.com/cosmos/cosmos-sdk/crypto/codec"
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -112,6 +114,11 @@ func (k Keeper) BlockValidatorUpdates(ctx context.Context) ([]abci.ValidatorUpda
 		)
 	}
 
+	err = k.UpdateAllMaturedConsKeyRotatedKeys(sdkCtx, sdkCtx.HeaderInfo().Time)
+	if err != nil {
+		return nil, err
+	}
+
 	return validatorUpdates, nil
 }
 
@@ -235,6 +242,54 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx context.Context) (updates
 		updates = append(updates, validator.ABCIValidatorUpdateZero())
 	}
 
+	// ApplyAndReturnValidatorSetUpdates checks if there is ConsPubKeyRotationHistory
+	// with ConsPubKeyRotationHistory.RotatedHeight == ctx.BlockHeight() and if so, generates 2 ValidatorUpdate,
+	// one for a remove validator and one for create new validator
+	historyObjects, err := k.GetBlockConsPubKeyRotationHistory(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, history := range historyObjects {
+		valAddr := history.OperatorAddress
+		if err != nil {
+			return nil, err
+		}
+
+		validator := k.mustGetValidator(ctx, valAddr)
+
+		oldPk := history.OldConsPubkey.GetCachedValue().(cryptotypes.PubKey)
+		oldTmPk, err := cryptocodec.ToCmtProtoPublicKey(oldPk)
+		if err != nil {
+			return nil, err
+		}
+
+		newPk := history.NewConsPubkey.GetCachedValue().(cryptotypes.PubKey)
+		newTmPk, err := cryptocodec.ToCmtProtoPublicKey(newPk)
+		if err != nil {
+			return nil, err
+		}
+
+		if !(validator.Jailed || validator.Status != types.Bonded) {
+			updates = append(updates, abci.ValidatorUpdate{
+				PubKey: oldTmPk,
+				Power:  0,
+			})
+
+			updates = append(updates, abci.ValidatorUpdate{
+				PubKey: newTmPk,
+				Power:  validator.ConsensusPower(powerReduction),
+			})
+
+			if err := k.updateToNewPubkey(ctx, validator, history.OldConsPubkey, history.NewConsPubkey, history.Fee); err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	// TODO: at previousVotes Iteration logic of AllocateTokens, previousVote using OldConsPubKey
+	// match up with ConsPubKeyRotationHistory, and replace validator for token allocation
+
 	// Update the pools based on the recent updates in the validator set:
 	// - The tokens from the non-bonded candidates that enter the new validator set need to be transferred
 	// to the Bonded pool.