From 6b6f9c73d6c89a611837e57503a3cb063870848f Mon Sep 17 00:00:00 2001
From: Joowon Yun <40225835+JoowonYun@users.noreply.github.com>
Date: Thu, 31 Oct 2024 16:03:36 +0900
Subject: [PATCH] fix: dust share (#12)

---
 x/staking/keeper/delegation.go      |  7 ++++-
 x/staking/keeper/delegation_test.go | 48 +++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go
index 606cb1e51c89..b3aa947f2bdf 100644
--- a/x/staking/keeper/delegation.go
+++ b/x/staking/keeper/delegation.go
@@ -1379,7 +1379,12 @@ func (k Keeper) ValidateUnbondAmount(
 	// due to rounding, however we don't want to truncate the shares or take the
 	// minimum because we want to allow for the full withdraw of shares from a
 	// delegation.
-	if shares.GT(delShares) {
+	tolerance, err := validator.SharesFromTokens(math.OneInt())
+	if err != nil {
+		return shares, err
+	}
+
+	if delShares.Sub(shares).LT(tolerance) {
 		shares = delShares
 	}
 
diff --git a/x/staking/keeper/delegation_test.go b/x/staking/keeper/delegation_test.go
index 1404b18cd5d0..062e5c61d4ac 100644
--- a/x/staking/keeper/delegation_test.go
+++ b/x/staking/keeper/delegation_test.go
@@ -1168,3 +1168,51 @@ func (s *KeeperTestSuite) TestSetUnbondingDelegationEntry() {
 	// unbondingID comes from a global counter -> gaps in unbondingIDs are OK as long as every unbondingID is unique
 	require.Equal(uint64(2), resUnbonding.Entries[1].UnbondingId)
 }
+
+func (s *KeeperTestSuite) TestUndelegateWithDustShare() {
+	ctx, keeper := s.ctx, s.stakingKeeper
+	require := s.Require()
+
+	addrDels, valAddrs := createValAddrs(2)
+
+	s.accountKeeper.EXPECT().AddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
+
+	// construct the validators[0] & slash 1stake
+	amt := math.NewInt(100)
+	validator := testutil.NewValidator(s.T(), valAddrs[0], PKs[0])
+	validator, _ = validator.AddTokensFromDel(amt)
+	validator = validator.RemoveTokens(math.NewInt(1))
+	validator = stakingkeeper.TestingUpdateValidator(keeper, ctx, validator, true)
+
+	// first add a validators[0] to delegate too
+	bond1to1 := stakingtypes.NewDelegation(addrDels[0].String(), valAddrs[0].String(), math.LegacyNewDec(100))
+	require.NoError(keeper.SetDelegation(ctx, bond1to1))
+	resBond, err := keeper.GetDelegation(ctx, addrDels[0], valAddrs[0])
+	require.NoError(err)
+	require.Equal(bond1to1, resBond)
+
+	// second delegators[1] add a validators[0] to delegate
+	bond2to1 := stakingtypes.NewDelegation(addrDels[1].String(), valAddrs[0].String(), math.LegacyNewDec(1))
+	validator, delegatorShare := validator.AddTokensFromDel(math.NewInt(1))
+	bond2to1.Shares = delegatorShare
+	_ = stakingkeeper.TestingUpdateValidator(keeper, ctx, validator, true)
+	require.NoError(keeper.SetDelegation(ctx, bond2to1))
+	resBond, err = keeper.GetDelegation(ctx, addrDels[1], valAddrs[0])
+	require.NoError(err)
+	require.Equal(bond2to1, resBond)
+
+	// check delegation state
+	delegations, err := keeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.NoError(err)
+	require.Equal(2, len(delegations))
+
+	// undelegate all delegator[0]'s delegate
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[0].String(), valAddrs[0].String(), sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(99))))
+	require.NoError(err)
+
+	// remain only delegator[1]'s delegate
+	delegations, err = keeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.NoError(err)
+	require.Equal(1, len(delegations))
+	require.Equal(delegations[0].DelegatorAddress, addrDels[1].String())
+}