fix: donate all vesting token to work with dust delegations

Author: javiersuweijie

x/auth/vesting/handler_test.go

@@ -1,7 +1,9 @@
 package vesting_test
 
 import (
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/suite"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
@@ -108,8 +110,22 @@ func (suite *HandlerTestSuite) TestMsgDonateVestingToken() {
 	addr1 := sdk.AccAddress([]byte("addr1_______________"))
 	addr2 := sdk.AccAddress([]byte("addr2_______________"))
 	addr3 := sdk.AccAddress([]byte("addr3_______________"))
+	addr4 := sdk.AccAddress([]byte("addr4_______________"))
 
 	valAddr := sdk.ValAddress([]byte("validator___________"))
+	suite.app.StakingKeeper.SetValidator(ctx, stakingtypes.Validator{
+		OperatorAddress:   valAddr.String(),
+		ConsensusPubkey:   nil,
+		Jailed:            false,
+		Status:            0,
+		Tokens:            sdk.NewInt(2),
+		DelegatorShares:   sdk.MustNewDecFromStr("1.1"),
+		Description:       stakingtypes.Description{},
+		UnbondingHeight:   0,
+		UnbondingTime:     time.Time{},
+		Commission:        stakingtypes.Commission{},
+		MinSelfDelegation: sdk.NewInt(1),
+	})
 
 	acc1 := suite.app.AccountKeeper.NewAccountWithAddress(ctx, addr1)
 	suite.app.AccountKeeper.SetAccount(ctx, acc1)
@@ -133,6 +149,18 @@ func (suite *HandlerTestSuite) TestMsgDonateVestingToken() {
 	suite.app.AccountKeeper.SetAccount(ctx, acc3)
 	suite.Require().NoError(simapp.FundAccount(suite.app.BankKeeper, ctx, addr3, balances))
 
+	acc4 := types.NewPermanentLockedAccount(
+		suite.app.AccountKeeper.NewAccountWithAddress(ctx, addr4).(*authtypes.BaseAccount), balances,
+	)
+	acc4.DelegatedVesting = balances
+	suite.app.AccountKeeper.SetAccount(ctx, acc4)
+	suite.app.StakingKeeper.SetDelegation(ctx, stakingtypes.Delegation{
+		DelegatorAddress: addr4.String(),
+		ValidatorAddress: valAddr.String(),
+		Shares:           sdk.MustNewDecFromStr("0.1"),
+	})
+	suite.Require().NoError(simapp.FundAccount(suite.app.BankKeeper, ctx, addr4, balances))
+
 	testCases := []struct {
 		name      string
 		msg       *types.MsgDonateAllVestingTokens
@@ -153,12 +181,19 @@ func (suite *HandlerTestSuite) TestMsgDonateVestingToken() {
 			msg:       types.NewMsgDonateAllVestingTokens(addr3),
 			expectErr: false,
 		},
+		{
+			name:      "donate from vesting account with dust delegation",
+			msg:       types.NewMsgDonateAllVestingTokens(addr4),
+			expectErr: false,
+		},
 	}
 
 	for _, tc := range testCases {
 		tc := tc
 
 		suite.Run(tc.name, func() {
+			// Rollback context after every test case
+			ctx, _ := ctx.CacheContext()
 			res, err := suite.handler(ctx, tc.msg)
 			if tc.expectErr {
 				suite.Require().Error(err)
@@ -178,6 +213,9 @@ func (suite *HandlerTestSuite) TestMsgDonateVestingToken() {
 				suite.Require().True(ok)
 				balance := suite.app.BankKeeper.GetAllBalances(ctx, fromAddr)
 				suite.Require().Empty(balance)
+
+				_, broken := stakingkeeper.DelegatorSharesInvariant(suite.app.StakingKeeper)(ctx)
+				suite.Require().False(broken)
 			}
 		})
 	}

x/auth/vesting/msg_server.go

@@ -2,6 +2,8 @@ package vesting
 
 import (
 	"context"
+	"fmt"
+	"math"
 
 	"github.com/armon/go-metrics"
 
@@ -219,9 +221,33 @@ func (s msgServer) DonateAllVestingTokens(goCtx context.Context, msg *types.MsgD
 		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, "account %s not exists", msg.FromAddress)
 	}
 
-	// check whether an account has any type of staking entry
-	if len(sk.GetDelegatorDelegations(ctx, acc.GetAddress(), 1)) != 0 ||
-		len(sk.GetUnbondingDelegations(ctx, acc.GetAddress(), 1)) != 0 ||
+	// get all delegations of an account and undust those that have less than 1 uluna
+	delegations := sk.GetDelegatorDelegations(ctx, acc.GetAddress(), math.MaxUint16)
+	for _, delegation := range delegations {
+		validatorAddr, err := sdk.ValAddressFromBech32(delegation.ValidatorAddress)
+		if err != nil {
+			return nil, err
+		}
+		validator, found := sk.GetValidator(ctx, validatorAddr)
+		if !found {
+			return nil, fmt.Errorf("validator not found")
+		}
+		// Try to delete the dust delegation
+		_, removedTokens := sk.RemoveValidatorTokensAndShares(ctx, validator, delegation.Shares)
+		// If the delegation is not dust, return an error and stop the donation flow
+		if !removedTokens.IsZero() {
+			return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, "account %s has a non-zero staking entry", msg.FromAddress)
+		}
+
+		// Remove the dust delegation shares from the validator
+		err = sk.RemoveDelegation(ctx, delegation)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// check whether an account has any other type of staking entries
+	if len(sk.GetUnbondingDelegations(ctx, acc.GetAddress(), 1)) != 0 ||
 		len(sk.GetRedelegations(ctx, acc.GetAddress(), 1)) != 0 {
 		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, "account %s has staking entry", msg.FromAddress)
 	}

x/auth/vesting/types/expected_keepers.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 )
@@ -23,4 +24,7 @@ type StakingKeeper interface {
 	GetDelegatorDelegations(ctx sdk.Context, delegator sdk.AccAddress, maxRetrieve uint16) (delegations []stakingtypes.Delegation)
 	GetUnbondingDelegations(ctx sdk.Context, delegator sdk.AccAddress, maxRetrieve uint16) (unbondingDelegations []stakingtypes.UnbondingDelegation)
 	GetRedelegations(ctx sdk.Context, delegator sdk.AccAddress, maxRetrieve uint16) (redelegations []stakingtypes.Redelegation)
+	RemoveValidatorTokensAndShares(ctx sdk.Context, validator stakingtypes.Validator, sharesToRemove sdk.Dec) (valOut stakingtypes.Validator, removedTokens math.Int)
+	GetValidator(ctx sdk.Context, addr sdk.ValAddress) (validator stakingtypes.Validator, found bool)
+	RemoveDelegation(ctx sdk.Context, delegation stakingtypes.Delegation) error
 }