Merge pull request #309 from scorpioborn/fix/subacc-withdraw

Fix / `subaccount` withdraw locked balances

x/subaccount/keeper/balance.go

@@ -5,6 +5,7 @@ import (
 	sdkmath "cosmossdk.io/math"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrtypes "github.com/cosmos/cosmos-sdk/types/errors"
 
 	"github.com/sge-network/sge/app/params"
 	"github.com/sge-network/sge/utils"
@@ -136,20 +137,20 @@ func (k Keeper) getAccountSummary(sdkContext sdk.Context, subaccountAddr sdk.Acc
 	return balance, unlockedBalance, bankBalance
 }
 
-// getBalances returns the balance, unlocked balance and bank balance of a subaccount
-func (k Keeper) withdraw(ctx sdk.Context, subAccAddr sdk.AccAddress, ownerAddr sdk.AccAddress) error {
-	summary, unlockedBalance, bankBalance := k.getAccountSummary(ctx, subAccAddr)
+// withdrawUnlocked returns the balance, unlocked balance and bank balance of a subaccount
+func (k Keeper) withdrawUnlocked(ctx sdk.Context, subAccAddr sdk.AccAddress, ownerAddr sdk.AccAddress) error {
+	accSummary, unlockedBalance, bankBalance := k.getAccountSummary(ctx, subAccAddr)
 
-	withdrawableBalance := summary.WithdrawableBalance(unlockedBalance, bankBalance.Amount)
+	withdrawableBalance := accSummary.WithdrawableBalance(unlockedBalance, bankBalance.Amount)
 	if withdrawableBalance.IsZero() {
 		return types.ErrNothingToWithdraw
 	}
 
-	if err := summary.Withdraw(withdrawableBalance); err != nil {
+	if err := accSummary.Withdraw(withdrawableBalance); err != nil {
 		return err
 	}
 
-	k.SetAccountSummary(ctx, subAccAddr, summary)
+	k.SetAccountSummary(ctx, subAccAddr, accSummary)
 
 	err := k.bankKeeper.SendCoins(ctx, subAccAddr, ownerAddr, sdk.NewCoins(sdk.NewCoin(params.DefaultBondDenom, withdrawableBalance)))
 	if err != nil {
@@ -158,3 +159,62 @@ func (k Keeper) withdraw(ctx sdk.Context, subAccAddr sdk.AccAddress, ownerAddr s
 
 	return nil
 }
+
+// withdrawLockedAndUnlocked withdraws unlocked balance first, if more balance is needed to be deducted,
+// modifies the locked balances accordingly
+func (k Keeper) withdrawLockedAndUnlocked(ctx sdk.Context, subAccAddr sdk.AccAddress, ownerAddr sdk.AccAddress, subAmountDeduct sdkmath.Int,
+) error {
+	accSummary, unlockedBalance, bankBalance := k.getAccountSummary(ctx, subAccAddr)
+	withdrawableBalance := accSummary.WithdrawableBalance(unlockedBalance, bankBalance.Amount)
+
+	if withdrawableBalance.GT(sdkmath.ZeroInt()) {
+		if err := k.bankKeeper.SendCoins(ctx,
+			subAccAddr,
+			ownerAddr,
+			sdk.NewCoins(sdk.NewCoin(params.DefaultBondDenom, withdrawableBalance))); err != nil {
+			return sdkerrors.Wrapf(types.ErrSendCoinError, "error sending coin from subaccount to main account %s", err)
+		}
+	}
+
+	lockedAmountToWithdraw := subAmountDeduct.Sub(withdrawableBalance)
+
+	if lockedAmountToWithdraw.GT(sdkmath.ZeroInt()) {
+		if err := k.bankKeeper.SendCoins(ctx, subAccAddr, ownerAddr,
+			sdk.NewCoins(sdk.NewCoin(params.DefaultBondDenom, lockedAmountToWithdraw))); err != nil {
+			return sdkerrors.Wrapf(types.ErrSendCoinError, "error sending coin from subaccount to main account %s", err)
+		}
+
+		// calculate locked balances
+		lockedBalances := k.GetLockedBalances(ctx, subAccAddr)
+		updatedLockedBalances := []types.LockedBalance{}
+		for _, lb := range lockedBalances {
+			if lb.Amount.GTE(lockedAmountToWithdraw) {
+				lb.Amount = lb.Amount.Sub(lockedAmountToWithdraw)
+				updatedLockedBalances = append(updatedLockedBalances, lb)
+
+				lockedAmountToWithdraw = sdkmath.ZeroInt()
+				break
+			} else {
+				lb.Amount = sdkmath.ZeroInt()
+				updatedLockedBalances = append(updatedLockedBalances, lb)
+
+				lockedAmountToWithdraw = lockedAmountToWithdraw.Sub(lb.Amount)
+			}
+		}
+
+		if lockedAmountToWithdraw.GT(sdkmath.ZeroInt()) {
+			return sdkerrors.Wrapf(sdkerrtypes.ErrInvalidRequest,
+				"not enough balance in sub account locked balances, need more %s tokens",
+				lockedAmountToWithdraw)
+		}
+
+		k.SetLockedBalances(ctx, subAccAddr, updatedLockedBalances)
+	}
+
+	if err := accSummary.Withdraw(subAmountDeduct); err != nil {
+		return sdkerrors.Wrapf(types.ErrWithdrawLocked, "%s", err)
+	}
+	k.SetAccountSummary(ctx, subAccAddr, accSummary)
+
+	return nil
+}

x/subaccount/keeper/msg_server_balance.go

@@ -32,7 +32,7 @@ func (k msgServer) WithdrawUnlockedBalances(goCtx context.Context, msg *types.Ms
 		return nil, types.ErrSubaccountDoesNotExist
 	}
 
-	err := k.keeper.withdraw(ctx, subAccAddr, ownerAddr)
+	err := k.keeper.withdrawUnlocked(ctx, subAccAddr, ownerAddr)
 	if err != nil {
 		return nil, err
 	}

x/subaccount/keeper/msg_server_bet.go

@@ -7,7 +7,6 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
-	sdkmath "cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrtypes "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/sge-network/sge/app/params"
@@ -57,61 +56,14 @@ func (k msgServer) Wager(goCtx context.Context, msg *types.MsgWager) (*types.Msg
 		return nil, sdkerrors.Wrapf(sdkerrtypes.ErrInvalidRequest, "not enough balance in main account")
 	}
 
-	accSummary, unlockedBalance, _ := k.keeper.getAccountSummary(ctx, subAccAddr)
-	if unlockedBalance.GTE(payload.SubaccDeductAmount) {
-		if err := k.keeper.bankKeeper.SendCoins(ctx,
-			subAccAddr,
-			sdk.MustAccAddressFromBech32(msg.Creator),
-			sdk.NewCoins(sdk.NewCoin(params.DefaultBondDenom, payload.SubaccDeductAmount))); err != nil {
-			return nil, sdkerrors.Wrapf(types.ErrSendCoinError, "error sending coin from subaccount to main account %s", err)
-		}
-	} else {
-		lockedAmountToWithdraw := payload.SubaccDeductAmount.Sub(unlockedBalance)
-
-		if err := accSummary.Withdraw(lockedAmountToWithdraw); err != nil {
-			return nil, sdkerrors.Wrapf(types.ErrWithdrawLocked, "%s", err)
-		}
-
-		if err := k.keeper.bankKeeper.SendCoins(ctx, subAccAddr, subAccOwner,
-			sdk.NewCoins(sdk.NewCoin(params.DefaultBondDenom, lockedAmountToWithdraw))); err != nil {
-			return nil, sdkerrors.Wrapf(types.ErrSendCoinError, "error sending coin from subaccount to main account %s", err)
-		}
-
-		// calculate locked balances
-		lockedBalances := k.keeper.GetLockedBalances(ctx, subAccAddr)
-		updatedLockedBalances := []types.LockedBalance{}
-		for _, lb := range lockedBalances {
-			if lockedAmountToWithdraw.LTE(sdkmath.ZeroInt()) {
-				break
-			}
-
-			if lb.Amount.GT(lockedAmountToWithdraw) {
-				lb.Amount = lb.Amount.Sub(lockedAmountToWithdraw)
-				updatedLockedBalances = append(updatedLockedBalances, lb)
-
-				lockedAmountToWithdraw = sdkmath.ZeroInt()
-				break
-			} else {
-				lb.Amount = sdkmath.ZeroInt()
-				updatedLockedBalances = append(updatedLockedBalances, lb)
-
-				lockedAmountToWithdraw = lockedAmountToWithdraw.Sub(lb.Amount)
-			}
-		}
-
-		if lockedAmountToWithdraw.GT(sdkmath.ZeroInt()) {
-			return nil, sdkerrors.Wrapf(sdkerrtypes.ErrInvalidRequest, "not enough balance in sub account")
-		}
-
-		k.keeper.SetLockedBalances(ctx, subAccAddr, updatedLockedBalances)
+	if err := k.keeper.withdrawLockedAndUnlocked(ctx, subAccAddr, subAccOwner, payload.SubaccDeductAmount); err != nil {
+		return nil, err
 	}
 
 	if err := k.keeper.betKeeper.Wager(ctx, bet, oddsMap); err != nil {
 		return nil, err
 	}
 
-	k.keeper.SetAccountSummary(ctx, subAccAddr, accSummary)
-
 	msg.EmitEvent(&ctx, payload.Msg, subAccOwner.String())
 
 	return &types.MsgWagerResponse{