refactor(x/distribution): audit QA (#21316)

Co-authored-by: Randy Grok <[email protected]> Co-authored-by: marbar3778 <[email protected]> Co-authored-by: Cosmos SDK <[email protected]> Co-authored-by: github-merge-queue <[email protected]> Co-authored-by: Julián Toledano <[email protected]> Co-authored-by: Julien Robert <[email protected]> Co-authored-by: Aaron Craelius <[email protected]>
cosmos · Aug 17, 2024 · ee04cee · ee04cee
1 parent 651868a
commit ee04cee
Show file tree

Hide file tree

Showing 5 changed files with 149 additions and 17 deletions.
diff --git a/x/distribution/README.md b/x/distribution/README.md
@@ -186,7 +186,7 @@ it can be updated with governance or the address with authority.
 * Params: `0x09 | ProtocolBuffer(Params)`
 
 ```protobuf reference
-https://github.com/cosmos/cosmos-sdk/blob/v0.47.0-rc1/proto/cosmos/distribution/v1beta1/distribution.proto#L12-L42
+https://github.com/cosmos/cosmos-sdk/blob/release/v0.52.x/x/distribution/proto/cosmos/distribution/v1beta1/distribution.proto#L12-L44
 ```
 
 ## Begin Block
@@ -272,7 +272,7 @@ The withdraw address cannot be any of the module accounts. These accounts are bl
 Response:
 
 ```protobuf reference
-https://github.com/cosmos/cosmos-sdk/blob/v0.47.0-rc1/proto/cosmos/distribution/v1beta1/tx.proto#L49-L60
+https://github.com/cosmos/cosmos-sdk/blob/release/v0.52.x/x/distribution/proto/cosmos/distribution/v1beta1/tx.proto#L62-L73
 ```
 
 ```go
@@ -324,7 +324,7 @@ The final calculated stake is equivalent to the actual staked coins in the deleg
 Response:
 
 ```protobuf reference
-https://github.com/cosmos/cosmos-sdk/blob/v0.47.0-rc1/proto/cosmos/distribution/v1beta1/tx.proto#L66-L77
+https://github.com/cosmos/cosmos-sdk/blob/release/v0.52.x/x/distribution/proto/cosmos/distribution/v1beta1/tx.proto#L79-L90
 ```
 
 ### WithdrawValidatorCommission
@@ -368,7 +368,7 @@ func (k Keeper) initializeDelegation(ctx context.Context, val sdk.ValAddress, de
 Distribution module params can be updated through `MsgUpdateParams`, which can be done using governance proposal and the signer will always be gov module account address.
 
 ```protobuf reference
-https://github.com/cosmos/cosmos-sdk/blob/v0.47.0-rc1/proto/cosmos/distribution/v1beta1/tx.proto#L133-L147
+https://github.com/cosmos/cosmos-sdk/blob/release/v0.52.x/x/distribution/proto/cosmos/distribution/v1beta1/tx.proto#L159:L172
 ```
 
 The message handling can fail if:

diff --git a/x/distribution/keeper/allocation_test.go b/x/distribution/keeper/allocation_test.go
@@ -375,3 +375,131 @@ func TestAllocateTokensTruncation(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, val2OutstandingRewards.Rewards.IsValid())
 }
+
+func TestAllocateTokensToValidatorWithoutCommission(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	key := storetypes.NewKVStoreKey(disttypes.StoreKey)
+	testCtx := testutil.DefaultContextWithDB(t, key, storetypes.NewTransientStoreKey("transient_test"))
+	cdcOpts := codectestutil.CodecOptions{}
+	encCfg := moduletestutil.MakeTestEncodingConfig(cdcOpts, distribution.AppModule{})
+	ctx := testCtx.Ctx.WithHeaderInfo(header.Info{Time: time.Now()})
+
+	bankKeeper := distrtestutil.NewMockBankKeeper(ctrl)
+	stakingKeeper := distrtestutil.NewMockStakingKeeper(ctrl)
+	accountKeeper := distrtestutil.NewMockAccountKeeper(ctrl)
+
+	env := runtime.NewEnvironment(runtime.NewKVStoreService(key), coretesting.NewNopLogger())
+
+	valCodec := address.NewBech32Codec("cosmosvaloper")
+
+	accountKeeper.EXPECT().GetModuleAddress("distribution").Return(distrAcc.GetAddress())
+	stakingKeeper.EXPECT().ValidatorAddressCodec().Return(valCodec).AnyTimes()
+
+	authorityAddr, err := cdcOpts.GetAddressCodec().BytesToString(authtypes.NewModuleAddress("gov"))
+	require.NoError(t, err)
+
+	distrKeeper := keeper.NewKeeper(
+		encCfg.Codec,
+		env,
+		accountKeeper,
+		bankKeeper,
+		stakingKeeper,
+		testCometService,
+		"fee_collector",
+		authorityAddr,
+	)
+
+	// create validator with 0% commission
+	operatorAddr, err := stakingKeeper.ValidatorAddressCodec().BytesToString(valConsPk0.Address())
+	require.NoError(t, err)
+	val, err := distrtestutil.CreateValidator(valConsPk0, operatorAddr, math.NewInt(100))
+	require.NoError(t, err)
+	val.Commission = stakingtypes.NewCommission(math.LegacyNewDec(0), math.LegacyNewDec(0), math.LegacyNewDec(0))
+	stakingKeeper.EXPECT().ValidatorByConsAddr(gomock.Any(), sdk.GetConsAddress(valConsPk0)).Return(val, nil).AnyTimes()
+
+	// allocate tokens
+	tokens := sdk.DecCoins{
+		{Denom: sdk.DefaultBondDenom, Amount: math.LegacyNewDec(10)},
+	}
+	require.NoError(t, distrKeeper.AllocateTokensToValidator(ctx, val, tokens))
+
+	// check commission
+	var expectedCommission sdk.DecCoins = nil
+
+	valBz, err := valCodec.StringToBytes(val.GetOperator())
+	require.NoError(t, err)
+
+	valCommission, err := distrKeeper.ValidatorsAccumulatedCommission.Get(ctx, valBz)
+	require.NoError(t, err)
+	require.Equal(t, expectedCommission, valCommission.Commission)
+
+	// check current rewards
+	expectedRewards := tokens
+
+	currentRewards, err := distrKeeper.ValidatorCurrentRewards.Get(ctx, valBz)
+	require.NoError(t, err)
+	require.Equal(t, expectedRewards, currentRewards.Rewards)
+}
+
+func TestAllocateTokensWithZeroTokens(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	key := storetypes.NewKVStoreKey(disttypes.StoreKey)
+	testCtx := testutil.DefaultContextWithDB(t, key, storetypes.NewTransientStoreKey("transient_test"))
+	cdcOpts := codectestutil.CodecOptions{}
+	encCfg := moduletestutil.MakeTestEncodingConfig(cdcOpts, distribution.AppModule{})
+	ctx := testCtx.Ctx.WithHeaderInfo(header.Info{Time: time.Now()})
+
+	bankKeeper := distrtestutil.NewMockBankKeeper(ctrl)
+	stakingKeeper := distrtestutil.NewMockStakingKeeper(ctrl)
+	accountKeeper := distrtestutil.NewMockAccountKeeper(ctrl)
+
+	env := runtime.NewEnvironment(runtime.NewKVStoreService(key), coretesting.NewNopLogger())
+
+	valCodec := address.NewBech32Codec("cosmosvaloper")
+
+	accountKeeper.EXPECT().GetModuleAddress("distribution").Return(distrAcc.GetAddress())
+	stakingKeeper.EXPECT().ValidatorAddressCodec().Return(valCodec).AnyTimes()
+
+	authorityAddr, err := cdcOpts.GetAddressCodec().BytesToString(authtypes.NewModuleAddress("gov"))
+	require.NoError(t, err)
+
+	distrKeeper := keeper.NewKeeper(
+		encCfg.Codec,
+		env,
+		accountKeeper,
+		bankKeeper,
+		stakingKeeper,
+		testCometService,
+		"fee_collector",
+		authorityAddr,
+	)
+
+	// create validator with 50% commission
+	operatorAddr, err := stakingKeeper.ValidatorAddressCodec().BytesToString(valConsPk0.Address())
+	require.NoError(t, err)
+	val, err := distrtestutil.CreateValidator(valConsPk0, operatorAddr, math.NewInt(100))
+	require.NoError(t, err)
+	val.Commission = stakingtypes.NewCommission(math.LegacyNewDecWithPrec(5, 1), math.LegacyNewDecWithPrec(5, 1), math.LegacyNewDec(0))
+	stakingKeeper.EXPECT().ValidatorByConsAddr(gomock.Any(), sdk.GetConsAddress(valConsPk0)).Return(val, nil).AnyTimes()
+
+	// allocate zero tokens
+	tokens := sdk.DecCoins{}
+	require.NoError(t, distrKeeper.AllocateTokensToValidator(ctx, val, tokens))
+
+	// check commission
+	var expected sdk.DecCoins = nil
+
+	valBz, err := valCodec.StringToBytes(val.GetOperator())
+	require.NoError(t, err)
+
+	valCommission, err := distrKeeper.ValidatorsAccumulatedCommission.Get(ctx, valBz)
+	require.NoError(t, err)
+	require.Equal(t, expected, valCommission.Commission)
+
+	// check current rewards
+	currentRewards, err := distrKeeper.ValidatorCurrentRewards.Get(ctx, valBz)
+	require.NoError(t, err)
+	require.Equal(t, expected, currentRewards.Rewards)
+}
diff --git a/x/distribution/keeper/keeper.go b/x/distribution/keeper/keeper.go
@@ -46,7 +46,7 @@ type Keeper struct {
 	DelegatorStartingInfo collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], types.DelegatorStartingInfo]
 	// ValidatorsAccumulatedCommission key: valAddr | value: ValidatorAccumulatedCommission
 	ValidatorsAccumulatedCommission collections.Map[sdk.ValAddress, types.ValidatorAccumulatedCommission]
-	// ValidatorOutstandingRewards key: valAddr | value: ValidatorOustandingRewards
+	// ValidatorOutstandingRewards key: valAddr | value: ValidatorOutstandingRewards
 	ValidatorOutstandingRewards collections.Map[sdk.ValAddress, types.ValidatorOutstandingRewards]
 	// ValidatorHistoricalRewards key: valAddr+period | value: ValidatorHistoricalRewards
 	ValidatorHistoricalRewards collections.Map[collections.Pair[sdk.ValAddress, uint64], types.ValidatorHistoricalRewards]

diff --git a/x/distribution/keeper/validator.go b/x/distribution/keeper/validator.go
@@ -13,6 +13,10 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
+const (
+	maxReferenceCount = 2
+)
+
 // initialize rewards for a new validator
 func (k Keeper) initializeValidator(ctx context.Context, val sdk.ValidatorI) error {
 	valBz, err := k.stakingKeeper.ValidatorAddressCodec().StringToBytes(val.GetOperator())
@@ -126,8 +130,8 @@ func (k Keeper) incrementReferenceCount(ctx context.Context, valAddr sdk.ValAddr
 	}
 
 	historical.ReferenceCount++
-	if historical.ReferenceCount > 2 {
-		panic("reference count should never exceed 2")
+	if historical.ReferenceCount > maxReferenceCount {
+		return fmt.Errorf("reference count should never exceed %d", maxReferenceCount)
 	}
 	return k.ValidatorHistoricalRewards.Set(ctx, collections.Join(valAddr, period), historical)
 }
@@ -140,7 +144,7 @@ func (k Keeper) decrementReferenceCount(ctx context.Context, valAddr sdk.ValAddr
 	}
 
 	if historical.ReferenceCount == 0 {
-		panic("cannot set negative reference count")
+		return fmt.Errorf("cannot set negative reference count")
 	}
 	historical.ReferenceCount--
 	if historical.ReferenceCount == 0 {
@@ -152,7 +156,7 @@ func (k Keeper) decrementReferenceCount(ctx context.Context, valAddr sdk.ValAddr
 
 func (k Keeper) updateValidatorSlashFraction(ctx context.Context, valAddr sdk.ValAddress, fraction math.LegacyDec) error {
 	if fraction.GT(math.LegacyOneDec()) || fraction.IsNegative() {
-		panic(fmt.Sprintf("fraction must be >=0 and <=1, current fraction: %v", fraction))
+		return fmt.Errorf("fraction must be >=0 and <=1, current fraction: %v", fraction)
 	}
 
 	headerinfo := k.HeaderService.HeaderInfo(ctx)

diff --git a/x/distribution/migrations/v4/migrate.go b/x/distribution/migrations/v4/migrate.go
@@ -16,16 +16,16 @@ import (
 var OldProposerKey = []byte{0x01}
 
 // MigrateStore removes the last proposer from store.
-func MigrateStore(ctx context.Context, env appmodule.Environment, cdc codec.BinaryCodec) error {
-	store := env.KVStoreService.OpenKVStore(ctx)
-	return store.Delete(OldProposerKey)
+func MigrateStore(ctx context.Context, env appmodule.Environment, _ codec.BinaryCodec) error {
+	kvStore := env.KVStoreService.OpenKVStore(ctx)
+	return kvStore.Delete(OldProposerKey)
 }
 
 // GetPreviousProposerConsAddr returns the proposer consensus address for the
 // current block.
 func GetPreviousProposerConsAddr(ctx context.Context, storeService store.KVStoreService, cdc codec.BinaryCodec) (sdk.ConsAddress, error) {
-	store := storeService.OpenKVStore(ctx)
-	bz, err := store.Get(OldProposerKey)
+	kvStore := storeService.OpenKVStore(ctx)
+	bz, err := kvStore.Get(OldProposerKey)
 	if err != nil {
 		return nil, err
 	}
@@ -43,9 +43,9 @@ func GetPreviousProposerConsAddr(ctx context.Context, storeService store.KVStore
 	return addrValue.GetValue(), nil
 }
 
-// set the proposer public key for this block
+// SetPreviousProposerConsAddr set the proposer public key for this block.
 func SetPreviousProposerConsAddr(ctx context.Context, storeService store.KVStoreService, cdc codec.BinaryCodec, consAddr sdk.ConsAddress) error {
-	store := storeService.OpenKVStore(ctx)
+	kvStore := storeService.OpenKVStore(ctx)
 	bz := cdc.MustMarshal(&gogotypes.BytesValue{Value: consAddr})
-	return store.Set(OldProposerKey, bz)
+	return kvStore.Set(OldProposerKey, bz)
 }