Subaccount module part4

proto/sge/subaccount/tx.proto

@@ -12,6 +12,8 @@ service Msg {
 
   // TopUp defines a method for topping up a subaccount.
   rpc TopUp(MsgTopUp) returns (MsgTopUpResponse);
+
+  rpc WithdrawUnlockedBalances(MsgWithdrawUnlockedBalances) returns (MsgWithdrawUnlockedBalancesResponse);
 }
 
 // MsgCreateSubAccount defines the Msg/CreateSubAccount request type.
@@ -42,4 +44,13 @@ message MsgTopUp {
 }
 
 // MsgTopUpResponse defines the Msg/TopUp response type.
-message MsgTopUpResponse {}
+message MsgTopUpResponse {}
+
+// MsgWithdrawUnlockedBalances defines the Msg/WithdrawUnlockedBalances request type.
+message MsgWithdrawUnlockedBalances {
+  // sender is the subaccount owner.
+  string sender = 1;
+}
+
+// MsgWithdrawUnlockedBalancesResponse defines the Msg/WithdrawUnlockedBalances response type.
+message MsgWithdrawUnlockedBalancesResponse {}

x/subaccount/keeper/subaccount.go

@@ -77,11 +77,9 @@ func (k Keeper) SetLockedBalances(ctx sdk.Context, id uint64, lockedBalances []*
 func (k Keeper) GetLockedBalances(ctx sdk.Context, id uint64) []subaccounttypes.LockedBalance {
 	account := subaccounttypes.NewAddressFromSubaccount(id)
 	iterator := prefix.NewStore(ctx.KVStore(k.storeKey), subaccounttypes.LockedBalancePrefixKey(account)).Iterator(nil, nil)
-
-	var lockedBalances []subaccounttypes.LockedBalance
-
 	defer iterator.Close()
 
+	var lockedBalances []subaccounttypes.LockedBalance
 	for ; iterator.Valid(); iterator.Next() {
 		unlockTime, err := sdk.ParseTimeBytes(iterator.Key())
 		if err != nil {
@@ -102,6 +100,27 @@ func (k Keeper) GetLockedBalances(ctx sdk.Context, id uint64) []subaccounttypes.
 	return lockedBalances
 }
 
+// GetUnlockedBalance returns the unlocked balance of an account.
+func (k Keeper) GetUnlockedBalance(ctx sdk.Context, id uint64) sdk.Int {
+	account := subaccounttypes.NewAddressFromSubaccount(id)
+	iterator := prefix.NewStore(ctx.KVStore(k.storeKey), subaccounttypes.LockedBalancePrefixKey(account)).
+		Iterator(nil, sdk.FormatTimeBytes(ctx.BlockTime()))
+
+	unlockedBalance := sdk.ZeroInt()
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		amount := new(sdk.Int)
+		err := amount.Unmarshal(iterator.Value())
+		if err != nil {
+			panic(err)
+		}
+		unlockedBalance = unlockedBalance.Add(*amount)
+	}
+
+	return unlockedBalance
+}
+
 // SetBalance saves the balance of an account.
 func (k Keeper) SetBalance(ctx sdk.Context, subaccountID uint64, balance subaccounttypes.Balance) {
 	account := subaccounttypes.NewAddressFromSubaccount(subaccountID)

x/subaccount/keeper/subaccount_test.go

@@ -93,3 +93,39 @@ func TestSetBalances(t *testing.T) {
 	// Get balance
 	require.Equal(t, balance, k.GetBalance(ctx, 1))
 }
+
+func TestKeeper_GetLockedBalances(t *testing.T) {
+	_, k, ctx := setupKeeperAndApp(t)
+
+	beforeUnlockTime1 := time.Now().Add(-time.Hour * 24 * 365)
+	beforeUnlockTime2 := time.Now().Add(-time.Hour * 24 * 365 * 2)
+
+	afterUnlockTime1 := time.Now().Add(time.Hour * 24 * 365)
+	afterUnlockTime2 := time.Now().Add(time.Hour * 24 * 365 * 2)
+
+	// I added them unordered to make sure they are sorted
+	balanceUnlocks := []*types.LockedBalance{
+		{
+			Amount:     sdk.NewInt(10000),
+			UnlockTime: beforeUnlockTime1,
+		},
+		{
+			Amount:     sdk.NewInt(30000),
+			UnlockTime: afterUnlockTime1,
+		},
+		{
+			Amount:     sdk.NewInt(20000),
+			UnlockTime: beforeUnlockTime2,
+		},
+		{
+			Amount:     sdk.NewInt(40000),
+			UnlockTime: afterUnlockTime2,
+		},
+	}
+
+	k.SetLockedBalances(ctx, 1, balanceUnlocks)
+
+	// get unlocked balance
+	unlockedBalance := k.GetUnlockedBalance(ctx, 1)
+	require.True(t, unlockedBalance.Equal(sdk.NewInt(10000+20000)))
+}

x/subaccount/keeper/withdraw.go

@@ -0,0 +1,50 @@
+package keeper
+
+import (
+	"context"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/sge-network/sge/x/subaccount/types"
+)
+
+func (m msgServer) WithdrawUnlockedBalances(ctx context.Context, balances *types.MsgWithdrawUnlockedBalances) (*types.MsgWithdrawUnlockedBalancesResponse, error) {
+	sdkContext := sdk.UnwrapSDKContext(ctx)
+
+	sender := sdk.MustAccAddressFromBech32(balances.Sender)
+	if !m.keeper.HasSubAccount(sdkContext, sender) {
+		return nil, types.ErrSubaccountDoesNotExist
+	}
+
+	params := m.keeper.GetParams(sdkContext)
+	subaccountID := m.keeper.GetSubAccountByOwner(sdkContext, sender)
+	subaccountAddr := types.NewAddressFromSubaccount(subaccountID)
+
+	balance, unlockedBalance, bankBalance := m.getBalances(sdkContext, subaccountID, subaccountAddr, params)
+
+	// calculate withdrawable balance, which is the minimum between the available balance, and
+	// what has been unlocked so far. Also, it cannot be greater than the bank balance.
+	// Available reports the deposited amount - spent amount - lost amount - withdrawn amount.
+	withdrawableBalance := sdk.MinInt(sdk.MinInt(balance.Available(), unlockedBalance), bankBalance.Amount)
+	if withdrawableBalance.IsZero() {
+		return nil, types.ErrNothingToWithdraw
+	}
+
+	balance.WithdrawmAmount = balance.WithdrawmAmount.Add(withdrawableBalance)
+	m.keeper.SetBalance(sdkContext, subaccountID, balance)
+
+	err := m.bankKeeper.SendCoins(sdkContext, subaccountAddr, sender, sdk.NewCoins(sdk.NewCoin(params.LockedBalanceDenom, withdrawableBalance)))
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.MsgWithdrawUnlockedBalancesResponse{}, nil
+}
+
+// getBalances returns the balance, unlocked balance and bank balance of a subaccount
+func (m msgServer) getBalances(sdkContext sdk.Context, subaccountID uint64, subaccountAddr sdk.AccAddress, params types.Params) (types.Balance, sdk.Int, sdk.Coin) {
+	balance := m.keeper.GetBalance(sdkContext, subaccountID)
+	unlockedBalance := m.keeper.GetUnlockedBalance(sdkContext, subaccountID)
+	bankBalance := m.bankKeeper.GetBalance(sdkContext, subaccountAddr, params.LockedBalanceDenom)
+
+	return balance, unlockedBalance, bankBalance
+}

x/subaccount/keeper/withdraw_test.go

@@ -0,0 +1,151 @@
+package keeper_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/cosmos/cosmos-sdk/simapp"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/sge-network/sge/testutil/sample"
+	"github.com/sge-network/sge/x/subaccount/keeper"
+	"github.com/sge-network/sge/x/subaccount/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMsgServer_WithdrawUnlockedBalances(t *testing.T) {
+	sender := sample.NativeAccAddress()
+	subaccountOwner := sample.NativeAccAddress()
+	lockedTime := time.Now().Add(time.Hour * 24 * 365)
+	lockedTime2 := time.Now().Add(time.Hour * 24 * 365 * 2)
+
+	app, _, msgServer, ctx := setupMsgServerAndApp(t)
+
+	t.Log("fund sender account")
+	err := simapp.FundAccount(app.BankKeeper, ctx, sender, sdk.NewCoins(sdk.NewInt64Coin("usge", 1000)))
+
+	t.Log("Create sub account")
+	_, err = msgServer.CreateSubAccount(sdk.WrapSDKContext(ctx), &types.MsgCreateSubAccount{
+		Sender:          sender.String(),
+		SubAccountOwner: subaccountOwner.String(),
+		LockedBalances: []*types.LockedBalance{
+			{
+				Amount:     sdk.NewInt(100),
+				UnlockTime: lockedTime,
+			},
+			{
+				Amount:     sdk.NewInt(200),
+				UnlockTime: lockedTime2,
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	t.Log("check balance of sub account")
+	subAccountAddr := types.NewAddressFromSubaccount(1)
+	balance := app.BankKeeper.GetBalance(ctx, subAccountAddr, "usge")
+	require.Equal(t, sdk.NewInt(300), balance.Amount)
+
+	t.Log("check balance of subaccount owner")
+	balance = app.BankKeeper.GetBalance(ctx, subaccountOwner, "usge")
+	require.Equal(t, sdk.NewInt(0), balance.Amount)
+
+	t.Log("Withdraw unlocked balances, with 0 expires")
+	_, err = msgServer.WithdrawUnlockedBalances(sdk.WrapSDKContext(ctx), &types.MsgWithdrawUnlockedBalances{
+		Sender: subaccountOwner.String(),
+	})
+	require.ErrorContains(t, err, types.ErrNothingToWithdraw.Error())
+
+	t.Log("balance of subaccount owner should be zero")
+	balance = app.BankKeeper.GetBalance(ctx, subaccountOwner, "usge")
+	require.True(t, balance.IsZero())
+
+	t.Log("expire first locked balance")
+	ctx = ctx.WithBlockTime(lockedTime.Add(1 * time.Second))
+	t.Log("Withdraw unlocked balances, with 1 expires")
+	_, err = msgServer.WithdrawUnlockedBalances(sdk.WrapSDKContext(ctx), &types.MsgWithdrawUnlockedBalances{
+		Sender: subaccountOwner.String(),
+	})
+	require.NoError(t, err)
+
+	t.Log("balance of subaccount owner should be the same as first locked balance")
+	balance = app.BankKeeper.GetBalance(ctx, subaccountOwner, "usge")
+	require.True(t, balance.Amount.Equal(sdk.NewInt(100)), balance.Amount.String())
+
+	t.Log("expire second locked balance, also force money to be spent")
+	// we force some money to be spent on the subaccount to correctly test
+	// that if the amount is unlocked but spent, it will not be withdrawable.
+	subaccountBalance := app.SubaccountKeeper.GetBalance(ctx, 1)
+	require.NoError(t, subaccountBalance.Spend(sdk.NewInt(100)))
+	app.SubaccountKeeper.SetBalance(ctx, 1, subaccountBalance)
+
+	ctx = ctx.WithBlockTime(lockedTime2.Add(1 * time.Second))
+	t.Log("Withdraw unlocked balances, with 2 expires")
+	_, err = msgServer.WithdrawUnlockedBalances(sdk.WrapSDKContext(ctx), &types.MsgWithdrawUnlockedBalances{
+		Sender: subaccountOwner.String(),
+	})
+	require.NoError(t, err)
+
+	t.Log("balance of subaccount owner should be the same as both expired locked balances minus spent money")
+	balance = app.BankKeeper.GetBalance(ctx, subaccountOwner, "usge")
+	require.Equal(t, sdk.NewInt(200), balance.Amount)
+
+	t.Log("check bank balance of sub account address")
+	balance = app.BankKeeper.GetBalance(ctx, subAccountAddr, "usge")
+	require.Equal(t, sdk.NewInt(100), balance.Amount)
+
+	t.Log("after unspending the money of the subaccount, the owner will be able to get the money back when withdrawing")
+	subaccountBalance = app.SubaccountKeeper.GetBalance(ctx, 1)
+	require.NoError(t, subaccountBalance.Unspend(sdk.NewInt(100)))
+	app.SubaccountKeeper.SetBalance(ctx, 1, subaccountBalance)
+	_, err = msgServer.WithdrawUnlockedBalances(sdk.WrapSDKContext(ctx), &types.MsgWithdrawUnlockedBalances{
+		Sender: subaccountOwner.String(),
+	})
+	require.NoError(t, err)
+
+	// check balances
+	balance = app.BankKeeper.GetBalance(ctx, subAccountAddr, "usge")
+	require.Equal(t, sdk.NewInt(0), balance.Amount)
+	subaccountBalance = app.SubaccountKeeper.GetBalance(ctx, 1)
+	require.Equal(t, sdk.NewInt(300), subaccountBalance.WithdrawmAmount)
+
+	// check that the owner has received the last money
+	balance = app.BankKeeper.GetBalance(ctx, subaccountOwner, "usge")
+	require.Equal(t, sdk.NewInt(300), balance.Amount)
+
+	// check that the owner can't withdraw again
+	_, err = msgServer.WithdrawUnlockedBalances(sdk.WrapSDKContext(ctx), &types.MsgWithdrawUnlockedBalances{
+		Sender: subaccountOwner.String(),
+	})
+	require.ErrorContains(t, err, types.ErrNothingToWithdraw.Error())
+}
+
+func TestMsgServer_WithdrawUnlockedBalances_Errors(t *testing.T) {
+	sender := sample.AccAddress()
+	tests := []struct {
+		name        string
+		msg         types.MsgWithdrawUnlockedBalances
+		prepare     func(ctx sdk.Context, keeper keeper.Keeper)
+		expectedErr string
+	}{
+		{
+			name: "sub account does not exist",
+			msg: types.MsgWithdrawUnlockedBalances{
+				Sender: sender,
+			},
+			prepare:     func(ctx sdk.Context, keeper keeper.Keeper) {},
+			expectedErr: types.ErrSubaccountDoesNotExist.Error(),
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			_, k, msgServer, ctx := setupMsgServerAndApp(t)
+
+			tt.prepare(ctx, k)
+
+			_, err := msgServer.WithdrawUnlockedBalances(sdk.WrapSDKContext(ctx), &tt.msg)
+			require.ErrorContains(t, err, tt.expectedErr)
+		})
+	}
+}

x/subaccount/types/codec.go

@@ -14,6 +14,7 @@ func RegisterCodec(_ *codec.LegacyAmino) {}
 func RegisterInterfaces(registry cdctypes.InterfaceRegistry) {
 	registry.RegisterImplementations((*sdk.Msg)(nil),
 		&MsgCreateSubAccount{},
+		&MsgWithdrawUnlockedBalances{},
 	)
 
 	msgservice.RegisterMsgServiceDesc(registry, &_Msg_serviceDesc)

x/subaccount/types/errors.go

@@ -13,4 +13,7 @@ var (
 
 	// ErrSubaccountDoesNotExist is the error when sub account does not exist
 	ErrSubaccountDoesNotExist = sdkerrors.Register(ModuleName, 3, "sub account does not exist")
+
+	// ErrNothingToWithdraw is the error returned when there is nothing to withdraw
+	ErrNothingToWithdraw = sdkerrors.Register(ModuleName, 4, "nothing to withdraw")
 )

x/subaccount/types/msg.go

@@ -8,6 +8,7 @@ import (
 var (
 	_ sdk.Msg = &MsgCreateSubAccount{}
 	_ sdk.Msg = &MsgTopUp{}
+	_ sdk.Msg = &MsgWithdrawUnlockedBalances{}
 )
 
 func (msg *MsgCreateSubAccount) GetSigners() []sdk.AccAddress {
@@ -67,3 +68,20 @@ func (msg *MsgTopUp) GetSigners() []sdk.AccAddress {
 	}
 	return []sdk.AccAddress{signer}
 }
+
+func (msg *MsgWithdrawUnlockedBalances) ValidateBasic() error {
+	_, err := sdk.AccAddressFromBech32(msg.Sender)
+	if err != nil {
+		return errors.ErrInvalidAddress
+	}
+
+	return nil
+}
+
+func (msg *MsgWithdrawUnlockedBalances) GetSigners() []sdk.AccAddress {
+	signer, err := sdk.AccAddressFromBech32(msg.Sender)
+	if err != nil {
+		panic(err)
+	}
+	return []sdk.AccAddress{signer}
+}

x/subaccount/types/subaccount.go

@@ -2,6 +2,8 @@ package types
 
 import (
 	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
 // Validate performs a basic validation of the LockedBalance fields.
@@ -20,3 +22,33 @@ func (lb *LockedBalance) Validate() error {
 
 	return nil
 }
+
+// Available reports the coins that are available in the subaccount.
+func (m *Balance) Available() sdk.Int {
+	return m.DepositedAmount.
+		Sub(m.WithdrawmAmount).
+		Sub(m.SpentAmount).
+		Sub(m.LostAmount)
+}
+
+func (m *Balance) Spend(amt sdk.Int) error {
+	if !amt.IsPositive() {
+		return fmt.Errorf("amount is not positive")
+	}
+	if amt.GT(m.Available()) {
+		return fmt.Errorf("amount is greater than available")
+	}
+	m.SpentAmount = m.SpentAmount.Add(amt)
+	return nil
+}
+
+func (m *Balance) Unspend(amt sdk.Int) error {
+	if !amt.IsPositive() {
+		return fmt.Errorf("amount is not positive")
+	}
+	if amt.GT(m.SpentAmount) {
+		return fmt.Errorf("amount is greater than spent")
+	}
+	m.SpentAmount = m.SpentAmount.Sub(amt)
+	return nil
+}