Merge pull request #318 from terra-money/fix/v12/burn-after-unbond

fix: upgrade function to burn after unbond

app/upgrades/v2.12/upgrade.go

@@ -1,12 +1,16 @@
 package v2_12
 
 import (
+	"cosmossdk.io/math"
+	"github.com/terra-money/core/v2/app/keepers"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
+	distributionkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
-	"github.com/terra-money/core/v2/app/keepers"
 )
 
 type EscrowUpdate struct {
@@ -22,105 +26,29 @@ func CreateUpgradeHandler(
 	cfg module.Configurator,
 	k keepers.TerraAppKeepers,
 ) upgradetypes.UpgradeHandler {
-	return func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {		
-		// Included in the start to also run on testnet
+	return func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+
 		if err := updateValidatorsMinCommissionRate(ctx, k.StakingKeeper); err != nil {
 			return nil, err
 		}
 
+		var addr sdk.AccAddress
+		var multisigAddr sdk.AccAddress
+
 		if ctx.ChainID() != "phoenix-1" {
-			return mm.RunMigrations(ctx, cfg, vm)
+			addr = sdk.MustAccAddressFromBech32("")
+			multisigAddr = sdk.MustAccAddressFromBech32("")
+		} else {
+			addr = sdk.MustAccAddressFromBech32("")
+			multisigAddr = sdk.MustAccAddressFromBech32("")
+		}
+
+		if err := burnTokensFromAccount(ctx, k.StakingKeeper, k.BankKeeper, k.DistrKeeper, addr); err != nil {
+			return nil, err
+		}
+		if err := burnTokensFromAccount(ctx, k.StakingKeeper, k.BankKeeper, k.DistrKeeper, multisigAddr); err != nil {
+			return nil, err
 		}
-		addr := sdk.MustAccAddressFromBech32("")
-		multisigAddr := sdk.MustAccAddressFromBech32("")
-
-		// Iterate delegations and unbond all shares
-		// burning the coins immediately
-		k.StakingKeeper.IterateDelegatorDelegations(ctx, addr, func(d stakingtypes.Delegation) (stop bool) {
-			valAddr, err := sdk.ValAddressFromBech32(d.ValidatorAddress)
-			if err != nil {
-				panic(err)
-			}
-			// Use this method without adding unbonding to the unbondings queue
-			// because it's not necessary to wait for the unbonding period
-			// (basically burn the shares and coins immediately)
-			_, err = k.StakingKeeper.Unbond(ctx, addr, valAddr, d.Shares)
-			if err != nil {
-				panic(err)
-			}
-			return false
-		})
-
-		// Given one of the states can be undelegating, we need to iterate over all unbonding delegations
-		// and remove them manually to ensure that the undelegated coins are burned.
-		bondDenom := k.StakingKeeper.GetParams(ctx).BondDenom
-		k.StakingKeeper.IterateDelegatorUnbondingDelegations(ctx, addr, func(ubd stakingtypes.UnbondingDelegation) (stop bool) {
-			balances := sdk.NewCoins()
-			for i := 0; i < len(ubd.Entries); i++ {
-				entry := ubd.Entries[i]
-				ubd.RemoveEntry(int64(i))
-				i--
-				k.StakingKeeper.DeleteUnbondingIndex(ctx, entry.UnbondingId)
-
-				// track undelegation only when remaining or truncated shares are non-zero
-				if !entry.Balance.IsZero() {
-					amt := sdk.NewCoin(bondDenom, entry.Balance)
-					if err := k.BankKeeper.UndelegateCoinsFromModuleToAccount(
-						ctx, stakingtypes.NotBondedPoolName, addr, sdk.NewCoins(amt),
-					); err != nil {
-						panic(err)
-					}
-
-					balances = balances.Add(amt)
-				}
-			}
-			k.StakingKeeper.RemoveUnbondingDelegation(ctx, ubd)
-			return false
-		})
-
-		// Redelegations are two queues but no coins are custodied in any "redelegations_pool",
-		// so we can just iterate over all redelegations and remove the indices to prevent issues.
-		k.StakingKeeper.IterateDelegatorRedelegations(ctx, addr, func(red stakingtypes.Redelegation) (stop bool) {
-			balances := sdk.NewCoins()
-			for i := 0; i < len(red.Entries); i++ {
-				entry := red.Entries[i]
-				red.RemoveEntry(int64(i))
-				i--
-				k.StakingKeeper.DeleteUnbondingIndex(ctx, entry.UnbondingId)
-
-				if !entry.InitialBalance.IsZero() {
-					balances = balances.Add(sdk.NewCoin(bondDenom, entry.InitialBalance))
-				}
-			}
-			k.StakingKeeper.RemoveRedelegation(ctx, red)
-			return false
-		})
-
-		// Burn all coins in the addr
-		k.BankKeeper.IterateAccountBalances(ctx, addr, func(balance sdk.Coin) bool {
-			err := k.BankKeeper.SendCoinsFromAccountToModule(ctx, addr, stakingtypes.NotBondedPoolName, sdk.NewCoins(balance))
-			if err != nil {
-				panic(err)
-			}
-			err = k.BankKeeper.BurnCoins(ctx, stakingtypes.NotBondedPoolName, sdk.NewCoins(balance))
-			if err != nil {
-				panic(err)
-			}
-			return false
-		})
-
-		// Burn all coins from the multisig account
-		k.BankKeeper.IterateAccountBalances(ctx, multisigAddr, func(balance sdk.Coin) bool {
-			err := k.BankKeeper.SendCoinsFromAccountToModule(ctx, multisigAddr, stakingtypes.NotBondedPoolName, sdk.NewCoins(balance))
-			if err != nil {
-				panic(err)
-			}
-			err = k.BankKeeper.BurnCoins(ctx, stakingtypes.NotBondedPoolName, sdk.NewCoins(balance))
-			if err != nil {
-				panic(err)
-			}
-			return false
-		})
 
 		return mm.RunMigrations(ctx, cfg, vm)
 	}
@@ -131,30 +59,127 @@ func updateValidatorsMinCommissionRate(ctx sdk.Context, sk *stakingkeeper.Keeper
 	stakingParams := sk.GetParams(ctx)
 	stakingParams.MinCommissionRate = sdk.MustNewDecFromStr("0.05")
 	if err := sk.SetParams(ctx, stakingParams); err != nil {
-	 return err
+		return err
 	}
-   
+
 	// Update all validators to have a min commission rate of 5%
 	validators := sk.GetAllValidators(ctx)
 	for _, validator := range validators {
-	 update := false
-	 commission := validator.Commission
-	 if commission.MaxRate.LT(sdk.MustNewDecFromStr("0.05")) {
-	  commission.MaxRate = sdk.MustNewDecFromStr("0.05")
-	  update = true
-	 }
-	 if commission.Rate.LT(sdk.MustNewDecFromStr("0.05")) {
-	  // force update without checking the <24h restriction and the max update rate
-	  commission.Rate = sdk.MustNewDecFromStr("0.05")
-	  update = true
-	 }
-	 if update {
-	  validator.Commission.UpdateTime = ctx.BlockTime()
-	  if err := sk.Hooks().BeforeValidatorModified(ctx, validator.GetOperator()); err != nil {
-	   return err
-	  }
-	  sk.SetValidator(ctx, validator)
-	 }
+		update := false
+		commission := validator.Commission
+		if commission.MaxRate.LT(sdk.MustNewDecFromStr("0.05")) {
+			commission.MaxRate = sdk.MustNewDecFromStr("0.05")
+			update = true
+		}
+		if commission.Rate.LT(sdk.MustNewDecFromStr("0.05")) {
+			// force update without checking the <24h restriction and the max update rate
+			commission.Rate = sdk.MustNewDecFromStr("0.05")
+			update = true
+		}
+		if update {
+			validator.Commission.UpdateTime = ctx.BlockTime()
+			if err := sk.Hooks().BeforeValidatorModified(ctx, validator.GetOperator()); err != nil {
+				return err
+			}
+			sk.SetValidator(ctx, validator)
+		}
 	}
 	return nil
-   }
+}
+
+func burnTokensFromAccount(ctx sdk.Context, sk *stakingkeeper.Keeper, bk bankkeeper.Keeper, dk distributionkeeper.Keeper, addr sdk.AccAddress) error {
+	// Iterate delegations and unbond all shares
+	// burning the coins immediately
+	bondDenom := sk.GetParams(ctx).BondDenom
+	var err error
+	sk.IterateDelegatorDelegations(ctx, addr, func(d stakingtypes.Delegation) (stop bool) {
+		var valAddr sdk.ValAddress
+		valAddr, err = sdk.ValAddressFromBech32(d.ValidatorAddress)
+		if err != nil {
+			return true
+		}
+
+		// Withdraw delegation rewards first
+		_, err = dk.WithdrawDelegationRewards(ctx, addr, valAddr)
+		if err != nil {
+			return true
+		}
+		// Use this method without adding unbonding to the unbondings queue
+		// because it's not necessary to wait for the unbonding period
+		var unbondedAmount math.Int
+		unbondedAmount, err = sk.Unbond(ctx, addr, valAddr, d.Shares)
+		if err != nil {
+			return true
+		}
+
+		// After unbonding, burn the coins depending on the validator's status
+		validator := sk.Validator(ctx, valAddr)
+		if validator.IsBonded() {
+			if err = bk.BurnCoins(ctx, stakingtypes.BondedPoolName, sdk.NewCoins(sdk.NewCoin(bondDenom, unbondedAmount))); err != nil {
+				return true
+			}
+		} else {
+			if err = bk.BurnCoins(ctx, stakingtypes.NotBondedPoolName, sdk.NewCoins(sdk.NewCoin(bondDenom, unbondedAmount))); err != nil {
+				return true
+			}
+		}
+
+		return false
+	})
+	if err != nil {
+		return err
+	}
+
+	// Given one of the states can be undelegating, we need to iterate over all unbonding delegations
+	// and remove them manually to ensure that the undelegated coins are burned.
+	sk.IterateDelegatorUnbondingDelegations(ctx, addr, func(ubd stakingtypes.UnbondingDelegation) (stop bool) {
+		for i := 0; i < len(ubd.Entries); i++ {
+			entry := ubd.Entries[i]
+			ubd.RemoveEntry(int64(i))
+			i--
+			sk.DeleteUnbondingIndex(ctx, entry.UnbondingId)
+
+			// track undelegation only when remaining or truncated shares are non-zero
+			if !entry.Balance.IsZero() {
+				amt := sdk.NewCoin(bondDenom, entry.Balance)
+				if err = bk.BurnCoins(
+					ctx, stakingtypes.NotBondedPoolName, sdk.NewCoins(amt),
+				); err != nil {
+					return true
+				}
+			}
+		}
+		sk.RemoveUnbondingDelegation(ctx, ubd)
+		return false
+	})
+	if err != nil {
+		return err
+	}
+
+	// Redelegations are two queues but no coins are custodied in any "redelegations_pool",
+	// so we can just iterate over all redelegations and remove the indices to prevent issues.
+	sk.IterateDelegatorRedelegations(ctx, addr, func(red stakingtypes.Redelegation) (stop bool) {
+		for i := 0; i < len(red.Entries); i++ {
+			entry := red.Entries[i]
+			red.RemoveEntry(int64(i))
+			i--
+			sk.DeleteUnbondingIndex(ctx, entry.UnbondingId)
+		}
+		sk.RemoveRedelegation(ctx, red)
+		return false
+	})
+
+	// Burn all coins in the addr
+	bk.IterateAccountBalances(ctx, addr, func(balance sdk.Coin) bool {
+		err = bk.SendCoinsFromAccountToModule(ctx, addr, stakingtypes.NotBondedPoolName, sdk.NewCoins(balance))
+		if err != nil {
+			return true
+		}
+		err = bk.BurnCoins(ctx, stakingtypes.NotBondedPoolName, sdk.NewCoins(balance))
+		if err != nil {
+			return true
+		}
+		return false
+	})
+	return err
+}