feature: add locking of staking rewards

tests/integration/distribution/keeper/msg_server_test.go

@@ -3,6 +3,7 @@ package keeper_test
 import (
 	"fmt"
 	"testing"
+	"time"
 
 	cmtabcitypes "github.com/cometbft/cometbft/abci/types"
 	"github.com/cometbft/cometbft/proto/tendermint/types"
@@ -24,6 +25,8 @@ import (
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	authsims "github.com/cosmos/cosmos-sdk/x/auth/simulation"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	"github.com/cosmos/cosmos-sdk/x/auth/vesting"
+	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
 	"github.com/cosmos/cosmos-sdk/x/bank"
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
@@ -106,6 +109,7 @@ func initFixture(t testing.TB) *fixture {
 
 	authModule := auth.NewAppModule(cdc, accountKeeper, authsims.RandomGenesisAccounts, nil)
 	bankModule := bank.NewAppModule(cdc, bankKeeper, accountKeeper, nil)
+	vestingModule := vesting.NewAppModule(accountKeeper, bankKeeper)
 	stakingModule := staking.NewAppModule(cdc, stakingKeeper, accountKeeper, bankKeeper, nil)
 	distrModule := distribution.NewAppModule(cdc, distrKeeper, accountKeeper, bankKeeper, stakingKeeper, nil)
 
@@ -122,13 +126,14 @@ func initFixture(t testing.TB) *fixture {
 			},
 			BlockIdFlag: types.BlockIDFlagCommit,
 		},
-	})
+	}).WithBlockTime(time.Now().Round(0).UTC())
 
 	integrationApp := integration.NewIntegrationApp(ctx, logger, keys, cdc, map[string]appmodule.AppModule{
 		authtypes.ModuleName:    authModule,
 		banktypes.ModuleName:    bankModule,
 		stakingtypes.ModuleName: stakingModule,
 		distrtypes.ModuleName:   distrModule,
+		vestingtypes.ModuleName: vestingModule,
 	})
 
 	sdkCtx := sdk.UnwrapSDKContext(integrationApp.Context())

tests/integration/distribution/keeper/vesting_test.go

@@ -0,0 +1,139 @@
+package keeper_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gotest.tools/v3/assert"
+
+	"cosmossdk.io/math"
+
+	"github.com/cosmos/cosmos-sdk/testutil/integration"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
+	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+func TestVestingAccountRewards(t *testing.T) {
+	t.Parallel()
+	f := initFixture(t)
+
+	// Set up fee pool and parameters
+	err := f.distrKeeper.FeePool.Set(f.sdkCtx, distrtypes.FeePool{
+		CommunityPool: sdk.NewDecCoins(sdk.DecCoin{Denom: "stake", Amount: math.LegacyNewDec(10000)}),
+	})
+	require.NoError(t, err)
+	require.NoError(t, f.distrKeeper.Params.Set(f.sdkCtx, distrtypes.DefaultParams()))
+
+	// Create validator
+	validator, err := stakingtypes.NewValidator(f.valAddr.String(), PKS[0], stakingtypes.Description{})
+	assert.NilError(t, err)
+	commission := stakingtypes.NewCommission(math.LegacyZeroDec(), math.LegacyOneDec(), math.LegacyOneDec())
+	validator, err = validator.SetInitialCommission(commission)
+	assert.NilError(t, err)
+	validator.DelegatorShares = math.LegacyNewDec(100)
+	validator.Tokens = math.NewInt(1000000)
+	assert.NilError(t, f.stakingKeeper.SetValidator(f.sdkCtx, validator))
+
+	// Set module account coins
+	initTokens := f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, int64(1000))
+	err = f.bankKeeper.MintCoins(f.sdkCtx, distrtypes.ModuleName, sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, initTokens)))
+	require.NoError(t, err)
+
+	// Create vesting account
+	vestingAddr := sdk.AccAddress(PKS[1].Address())
+	baseAccount := authtypes.NewBaseAccount(vestingAddr, PKS[1], 100, 0)
+
+	// Define vesting parameters
+	vestingAmount := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(1000)))
+	vestingStartTime := f.sdkCtx.BlockTime().Unix()
+	vestingEndTime := vestingStartTime + 100000 // Long enough for the test
+
+	// Create continuous vesting account
+	vestingAcc, err := vestingtypes.NewContinuousVestingAccount(baseAccount, vestingAmount, vestingStartTime, vestingEndTime)
+	require.NoError(t, err)
+
+	// Set the vesting account in the account keeper
+	f.accountKeeper.SetAccount(f.sdkCtx, vestingAcc)
+
+	// Fund the vesting account
+	err = f.bankKeeper.MintCoins(f.sdkCtx, distrtypes.ModuleName, vestingAmount)
+	require.NoError(t, err)
+	err = f.bankKeeper.SendCoinsFromModuleToAccount(f.sdkCtx, distrtypes.ModuleName, vestingAddr, vestingAmount)
+	require.NoError(t, err)
+
+	// Delegate tokens from the vesting account
+	delTokens := sdk.TokensFromConsensusPower(1, sdk.DefaultPowerReduction)
+	validator, issuedShares := validator.AddTokensFromDel(delTokens)
+
+	valBz, err := f.stakingKeeper.ValidatorAddressCodec().StringToBytes(validator.GetOperator())
+	require.NoError(t, err)
+	delegation := stakingtypes.NewDelegation(vestingAddr.String(), validator.GetOperator(), issuedShares)
+	require.NoError(t, f.stakingKeeper.SetDelegation(f.sdkCtx, delegation))
+	require.NoError(t, f.distrKeeper.SetDelegatorStartingInfo(f.sdkCtx, valBz, vestingAddr, distrtypes.NewDelegatorStartingInfo(2, math.LegacyOneDec(), 20)))
+
+	// Setup validator rewards
+	decCoins := sdk.DecCoins{sdk.NewDecCoinFromDec(sdk.DefaultBondDenom, math.LegacyOneDec())}
+	historicalRewards := distrtypes.NewValidatorHistoricalRewards(decCoins, 2)
+	err = f.distrKeeper.SetValidatorHistoricalRewards(f.sdkCtx, valBz, 2, historicalRewards)
+	require.NoError(t, err)
+
+	// Setup current rewards and outstanding rewards
+	currentRewards := distrtypes.NewValidatorCurrentRewards(decCoins, 3)
+	err = f.distrKeeper.SetValidatorCurrentRewards(f.sdkCtx, f.valAddr, currentRewards)
+	require.NoError(t, err)
+
+	valCommission := sdk.DecCoins{
+		sdk.NewDecCoinFromDec("stake", math.LegacyNewDec(3).Quo(math.LegacyNewDec(2))),
+	}
+	err = f.distrKeeper.SetValidatorOutstandingRewards(f.sdkCtx, f.valAddr, distrtypes.ValidatorOutstandingRewards{Rewards: valCommission})
+	require.NoError(t, err)
+
+	// Store initial vesting account state
+	initialVestingAcc, ok := f.accountKeeper.GetAccount(f.sdkCtx, vestingAddr).(*vestingtypes.ContinuousVestingAccount)
+	require.True(t, ok)
+	initialOriginalVesting := initialVestingAcc.OriginalVesting
+	initialStartTime := initialVestingAcc.StartTime
+	initialEndTime := initialVestingAcc.EndTime
+
+	// Withdraw rewards
+	msg := &distrtypes.MsgWithdrawDelegatorReward{
+		DelegatorAddress: vestingAddr.String(),
+		ValidatorAddress: f.valAddr.String(),
+	}
+
+	res, err := f.app.RunMsg(
+		msg,
+		integration.WithAutomaticFinalizeBlock(),
+		integration.WithAutomaticCommit(),
+	)
+
+	assert.NilError(t, err)
+	assert.Assert(t, res != nil)
+
+	// Check the result
+	result := distrtypes.MsgWithdrawDelegatorRewardResponse{}
+	err = f.cdc.Unmarshal(res.Value, &result)
+	assert.NilError(t, err)
+
+	// Verify the vesting account after rewards
+	finalVestingAcc, ok := f.accountKeeper.GetAccount(f.sdkCtx, vestingAddr).(*vestingtypes.ContinuousVestingAccount)
+	require.True(t, ok)
+
+	// Check that original vesting times are preserved
+	assert.Equal(t, initialStartTime, finalVestingAcc.StartTime)
+	assert.Equal(t, initialEndTime, finalVestingAcc.EndTime)
+
+	// Check that rewards were received
+	finalBalance := f.bankKeeper.GetAllBalances(f.sdkCtx, vestingAddr)
+	assert.Assert(t, finalBalance.IsAllGT(vestingAmount))
+
+	// Check that the original vesting amount is unchanged
+	assert.DeepEqual(t, initialOriginalVesting, finalVestingAcc.OriginalVesting)
+
+	// Check that delegated free and delegated vesting are properly tracked
+	// The delegation should be properly split between vesting and free portions
+	assert.Assert(t, !finalVestingAcc.DelegatedVesting.IsZero() || !finalVestingAcc.DelegatedFree.IsZero())
+}

testutil/integration/router.go

@@ -3,6 +3,7 @@
 import (
 	"context"
 	"fmt"
+	"time"
 
 	cmtabcitypes "github.com/cometbft/cometbft/abci/types"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
@@ -105,7 +106,7 @@
 
 	bApp.Commit()
 
-	ctx := sdkCtx.WithBlockHeader(cmtproto.Header{ChainID: appName}).WithIsCheckTx(true)
+	ctx := sdkCtx.WithBlockHeader(cmtproto.Header{ChainID: appName, Time: time.Now().Round(0).UTC()}).WithIsCheckTx(true)
 
 	return &App{
 		BaseApp:       bApp,

x/auth/vesting/types/vesting_account.go

@@ -162,6 +162,38 @@
 	return bva.BaseAccount.Validate()
 }
 
+// UpdateSchedule updates the vesting schedule for a base vesting account.
+// It calculates the proportion of the passed amount that should be used for updating
+// based on the ratio of delegated vesting to total delegation.
+func (bva *BaseVestingAccount) UpdateSchedule(amount sdk.Coins) error {
+	totalDelegated := bva.DelegatedFree.Add(bva.DelegatedVesting...)
+	if totalDelegated.IsZero() {
+		return nil // No delegations, nothing to update
+	}
+
+	// Calculate what portion of the amount should be applied based on delegated vesting ratio
+	var updatedAmount sdk.Coins
+	for _, coin := range amount {
+		denom := coin.Denom
+		totalDelegatedForDenom := totalDelegated.AmountOf(denom)
+		if totalDelegatedForDenom.IsZero() {
+			continue
+		}
+
+		delegatedVestingForDenom := bva.DelegatedVesting.AmountOf(denom)
+		ratio := math.LegacyNewDecFromInt(delegatedVestingForDenom).Quo(math.LegacyNewDecFromInt(totalDelegatedForDenom))
+		amountToUse := math.LegacyNewDecFromInt(coin.Amount).Mul(ratio).RoundInt()
+
+		if !amountToUse.IsZero() {
+			updatedAmount = updatedAmount.Add(sdk.NewCoin(denom, amountToUse))
+		}
+	}
+
+	// Add the calculated amount to original vesting
+	bva.OriginalVesting = bva.OriginalVesting.Add(updatedAmount...)
+	return nil
+}
+
 // Continuous Vesting Account
 
 var (
@@ -254,6 +286,16 @@
 	return cva.BaseVestingAccount.Validate()
 }
 
+// UpdateSchedule updates the vesting schedule for a continuous vesting account.
+// It delegates to the base vesting account implementation and adjusts end time if needed.
+func (cva *ContinuousVestingAccount) UpdateSchedule(amount sdk.Coins) error {
+	if err := cva.BaseVestingAccount.UpdateSchedule(amount); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // Periodic Vesting Account
 
 var (
@@ -387,6 +429,96 @@
 	return pva.BaseVestingAccount.Validate()
 }
 
+// UpdateSchedule updates the vesting schedule for a periodic vesting account.
+// It takes in the amount of coins from the rewards and updates the vesting schedule
+// based on the ratio of delegated vesting to total delegated coins.
+func (pva *PeriodicVestingAccount) UpdateSchedule(amount sdk.Coins) error {
+	if err := pva.BaseVestingAccount.UpdateSchedule(amount); err != nil {
+		return err
+	}
+
+	// If there are no periods or amount is zero, nothing to do
+	if len(pva.VestingPeriods) == 0 || amount.IsZero() {
+		return nil
+	}
+
+	// For periodic vesting, distribute the new amount proportionally across existing periods
+	// or add a new period if the account's end time has passed
+	currentTime := time.Now().Unix()
+
+	if currentTime > pva.EndTime {
+		// Account has completed vesting, add a new period
+		// Use the average length of existing periods as a reference
+		totalLength := int64(0)
+		for _, period := range pva.VestingPeriods {
+			totalLength += period.Length
+		}
+
+		avgPeriodLength := totalLength / int64(len(pva.VestingPeriods))
+		if avgPeriodLength <= 0 {
+			avgPeriodLength = 30 * 24 * 60 * 60 // Default to 30 days if can't determine
+		}
+
+		newPeriod := Period{
+			Length: avgPeriodLength,
+			Amount: amount,
+		}
+
+		pva.VestingPeriods = append(pva.VestingPeriods, newPeriod)
+		pva.EndTime += avgPeriodLength
+	} else {
+		// Account is still vesting, distribute proportionally across remaining periods
+		remainingPeriods := 0
+		for i := range pva.VestingPeriods {
+			periodEndTime := pva.StartTime
+			for j := 0; j <= i; j++ {
+				periodEndTime += pva.VestingPeriods[j].Length
+			}
+
+			if periodEndTime > currentTime {
+				remainingPeriods++
+			}
+		}
+
+		if remainingPeriods == 0 {
+			remainingPeriods = 1 // At least one period should remain
+		}
+
+		// Distribute amount evenly across remaining periods
+		amountPerPeriod := sdk.NewCoins()
+		for _, coin := range amount {
+			amtPerPeriod := coin.Amount.Quo(math.NewInt(int64(remainingPeriods)))
+			if !amtPerPeriod.IsZero() {
+				amountPerPeriod = amountPerPeriod.Add(sdk.NewCoin(coin.Denom, amtPerPeriod))
+			}
+		}
+
+		periodCount := 0
+		for i := range pva.VestingPeriods {
+			periodEndTime := pva.StartTime
+			for j := 0; j <= i; j++ {
+				periodEndTime += pva.VestingPeriods[j].Length
+			}
+
+			if periodEndTime > currentTime {
+				if periodCount < remainingPeriods-1 {
+					pva.VestingPeriods[i].Amount = pva.VestingPeriods[i].Amount.Add(amountPerPeriod...)
+					periodCount++
+				} else {
+					// Last period gets any remaining amount
+					remaining := amount
+					for d := 0; d < periodCount; d++ {
+						remaining = remaining.Sub(amountPerPeriod...)
+					}
+					pva.VestingPeriods[i].Amount = pva.VestingPeriods[i].Amount.Add(remaining...)
+				}
+			}
+		}
+	}
+
+	return nil
+}
+
 // Delayed Vesting Account
 
 var (
@@ -453,6 +585,11 @@
 	return dva.BaseVestingAccount.Validate()
 }
 
+// UpdateSchedule updates the vesting schedule for a delayed vesting account.
+func (dva *DelayedVestingAccount) UpdateSchedule(amount sdk.Coins) error {
+	return dva.BaseVestingAccount.UpdateSchedule(amount)
+}
+
 //-----------------------------------------------------------------------------
 // Permanent Locked Vesting Account
 
@@ -518,3 +655,8 @@
 
 	return plva.BaseVestingAccount.Validate()
 }
+
+// UpdateSchedule updates the vesting schedule for a permanent locked account.
+func (plva *PermanentLockedAccount) UpdateSchedule(amount sdk.Coins) error {
+	return plva.BaseVestingAccount.UpdateSchedule(amount)
+}