feat(ADR-036): custom withdrawal address (#309)

Adds the ability to specify custom address to separate operational and
rewards address

[References ADR-036](babylonlabs-io/pm#133)

CHANGELOG.md

@@ -39,6 +39,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Improvements
 
+- [#309](https://github.com/babylonlabs-io/babylon/pull/309) feat(adr-036): custom withdrawal address
 - [#305](https://github.com/babylonlabs-io/babylon/pull/305) chore: add more error logs to `VerifyInclusionProofAndGetHeight`
 - [#304](https://github.com/babylonlabs-io/babylon/pull/304) Add highest voted height to finality provider
 - [#311](https://github.com/babylonlabs-io/babylon/pull/311) Enforce version 2

proto/babylon/incentive/query.proto

@@ -6,6 +6,7 @@ import "google/api/annotations.proto";
 import "babylon/incentive/params.proto";
 import "babylon/incentive/incentive.proto";
 import "cosmos/base/v1beta1/coin.proto";
+import "cosmos_proto/cosmos.proto";
 
 option go_package = "github.com/babylonlabs-io/babylon/x/incentive/types";
 
@@ -23,6 +24,11 @@ service Query {
     rpc BTCStakingGauge(QueryBTCStakingGaugeRequest) returns (QueryBTCStakingGaugeResponse) {
         option (google.api.http).get = "/babylon/incentive/btc_staking_gauge/{height}";
     }
+
+    // DelegatorWithdrawAddress queries withdraw address of a delegator.
+    rpc DelegatorWithdrawAddress(QueryDelegatorWithdrawAddressRequest) returns (QueryDelegatorWithdrawAddressResponse) {
+        option (google.api.http).get = "/babylon/incentive/delegators/{delegator_address}/withdraw_address";
+    }
 }
 
 // QueryParamsRequest is request type for the Query/Params RPC method.
@@ -84,3 +90,23 @@ message QueryBTCStakingGaugeResponse {
     // gauge is the BTC staking gauge at the queried height 
     BTCStakingGaugeResponse gauge = 1;
 }
+
+// QueryDelegatorWithdrawAddressRequest is the request type for the
+// Query/DelegatorWithdrawAddress RPC method.
+message QueryDelegatorWithdrawAddressRequest {
+    option (gogoproto.equal)           = false;
+    option (gogoproto.goproto_getters) = false;
+
+    // delegator_address defines the delegator address to query for.
+    string delegator_address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+}
+
+// QueryDelegatorWithdrawAddressResponse is the response type for the
+// Query/DelegatorWithdrawAddress RPC method.
+message QueryDelegatorWithdrawAddressResponse {
+    option (gogoproto.equal)           = false;
+    option (gogoproto.goproto_getters) = false;
+
+    // withdraw_address defines the delegator address to query for.
+    string withdraw_address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+}

proto/babylon/incentive/tx.proto

@@ -17,6 +17,8 @@ service Msg {
     rpc WithdrawReward(MsgWithdrawReward) returns (MsgWithdrawRewardResponse);
     // UpdateParams updates the incentive module parameters.
     rpc UpdateParams(MsgUpdateParams) returns (MsgUpdateParamsResponse);
+    // SetWithdrawAddress defines a method to change the withdraw address of a stakeholder
+    rpc SetWithdrawAddress(MsgSetWithdrawAddress) returns (MsgSetWithdrawAddressResponse);
 }
 
 
@@ -56,3 +58,14 @@ message MsgUpdateParams {
 }
 // MsgUpdateParamsResponse is the response to the MsgUpdateParams message.
 message MsgUpdateParamsResponse {}
+
+// MsgSetWithdrawAddress sets the withdraw address
+message MsgSetWithdrawAddress {
+    option (cosmos.msg.v1.signer) = "delegator_address";
+
+    string delegator_address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+    string withdraw_address  = 2 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+}
+// MsgSetWithdrawAddressResponse defines the Msg/SetWithdrawAddress response
+// type.
+message MsgSetWithdrawAddressResponse {}

x/incentive/client/cli/tx.go

@@ -2,7 +2,7 @@ package cli
 
 import (
 	"fmt"
-
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/spf13/cobra"
 
 	"github.com/babylonlabs-io/babylon/x/incentive/types"
@@ -23,6 +23,7 @@ func GetTxCmd() *cobra.Command {
 
 	cmd.AddCommand(
 		NewWithdrawRewardCmd(),
+		NewSetWithdrawAddressCmd(),
 	)
 
 	return cmd
@@ -52,3 +53,34 @@ func NewWithdrawRewardCmd() *cobra.Command {
 
 	return cmd
 }
+
+func NewSetWithdrawAddressCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "set-withdraw-addr [withdraw-addr]",
+		Short: "change the default withdraw address for rewards",
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			clientCtx, err := client.GetClientTxContext(cmd)
+			if err != nil {
+				return err
+			}
+
+			delAddr := clientCtx.GetFromAddress()
+			withdrawAddr, err := sdk.AccAddressFromBech32(args[0])
+			if err != nil {
+				return err
+			}
+
+			msg := &types.MsgSetWithdrawAddress{
+				DelegatorAddress: delAddr.String(),
+				WithdrawAddress:  withdrawAddr.String(),
+			}
+
+			return tx.GenerateOrBroadcastTxCLI(clientCtx, cmd.Flags(), msg)
+		},
+	}
+
+	flags.AddTxFlagsToCmd(cmd)
+
+	return cmd
+}

x/incentive/keeper/msg_server.go

@@ -2,7 +2,6 @@ package keeper
 
 import (
 	"context"
-
 	errorsmod "cosmossdk.io/errors"
 	"github.com/babylonlabs-io/babylon/x/incentive/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -65,3 +64,22 @@ func (ms msgServer) WithdrawReward(goCtx context.Context, req *types.MsgWithdraw
 		Coins: withdrawnCoins,
 	}, nil
 }
+
+func (ms msgServer) SetWithdrawAddress(ctx context.Context, msg *types.MsgSetWithdrawAddress) (*types.MsgSetWithdrawAddressResponse, error) {
+	delegatorAddress, err := sdk.AccAddressFromBech32(msg.DelegatorAddress)
+	if err != nil {
+		return nil, status.Error(codes.InvalidArgument, err.Error())
+	}
+
+	withdrawAddress, err := sdk.AccAddressFromBech32(msg.WithdrawAddress)
+	if err != nil {
+		return nil, status.Error(codes.InvalidArgument, err.Error())
+	}
+
+	err = ms.SetWithdrawAddr(ctx, delegatorAddress, withdrawAddress)
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.MsgSetWithdrawAddressResponse{}, nil
+}

x/incentive/keeper/msg_server_test.go

@@ -63,3 +63,55 @@ func FuzzWithdrawReward(f *testing.F) {
 		require.True(t, newRg.IsFullyWithdrawn())
 	})
 }
+
+func FuzzSetWithdrawAddr(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+	f.Fuzz(func(t *testing.T, seed int64) {
+		r := rand.New(rand.NewSource(seed))
+
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+
+		// mock bank keeper
+		bk := types.NewMockBankKeeper(ctrl)
+
+		ik, ctx := testkeeper.IncentiveKeeper(t, bk, nil, nil)
+		ms := keeper.NewMsgServerImpl(*ik)
+
+		// generate and set a random reward gauge with a random set of withdrawable coins
+		rg := datagen.GenRandomRewardGauge(r)
+		rg.WithdrawnCoins = datagen.GenRandomWithdrawnCoins(r, rg.Coins)
+		sType := datagen.GenRandomStakeholderType(r)
+		sAddr := datagen.GenRandomAccount().GetAddress()
+		withdrawalAddr := datagen.GenRandomAccount().GetAddress()
+
+		ik.SetRewardGauge(ctx, sType, sAddr, rg)
+
+		// mock transfer of withdrawable coins
+		withdrawableCoins := rg.GetWithdrawableCoins()
+		bk.EXPECT().SendCoinsFromModuleToAccount(gomock.Any(), gomock.Eq(types.ModuleName), gomock.Eq(withdrawalAddr), gomock.Eq(withdrawableCoins)).Times(1)
+
+		_, err := ms.SetWithdrawAddress(ctx, &types.MsgSetWithdrawAddress{
+			DelegatorAddress: sAddr.String(),
+			WithdrawAddress:  withdrawalAddr.String(),
+		})
+		require.NoError(t, err)
+
+		rgauge := ik.GetRewardGauge(ctx, sType, sAddr)
+		require.NotNil(t, rgauge)
+		require.False(t, rgauge.IsFullyWithdrawn())
+
+		// invoke withdraw and assert consistency
+		resp, err := ms.WithdrawReward(ctx, &types.MsgWithdrawReward{
+			Type:    sType.String(),
+			Address: sAddr.String(),
+		})
+		require.NoError(t, err)
+		require.Equal(t, withdrawableCoins, resp.Coins)
+
+		// ensure reward gauge is now empty
+		newRg := ik.GetRewardGauge(ctx, sType, sAddr)
+		require.NotNil(t, newRg)
+		require.True(t, newRg.IsFullyWithdrawn())
+	})
+}

x/incentive/keeper/query_delegator.go

@@ -0,0 +1,28 @@
+package keeper
+
+import (
+	"context"
+	"github.com/babylonlabs-io/babylon/x/incentive/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func (k Keeper) DelegatorWithdrawAddress(goCtx context.Context, req *types.QueryDelegatorWithdrawAddressRequest) (*types.QueryDelegatorWithdrawAddressResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	}
+
+	ctx := sdk.UnwrapSDKContext(goCtx)
+	delegatorAddress, err := sdk.AccAddressFromBech32(req.DelegatorAddress)
+	if err != nil {
+		return nil, status.Error(codes.InvalidArgument, err.Error())
+	}
+
+	withdrawAddr, err := k.GetWithdrawAddr(ctx, delegatorAddress)
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.QueryDelegatorWithdrawAddressResponse{WithdrawAddress: withdrawAddr.String()}, nil
+}

x/incentive/keeper/query_delegator_test.go

@@ -0,0 +1,21 @@
+package keeper_test
+
+import (
+	"github.com/babylonlabs-io/babylon/testutil/datagen"
+	testkeeper "github.com/babylonlabs-io/babylon/testutil/keeper"
+	"github.com/babylonlabs-io/babylon/x/incentive/types"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestDelegatorAddressQuery(t *testing.T) {
+	keeper, ctx := testkeeper.IncentiveKeeper(t, nil, nil, nil)
+	withdrawalAddr := datagen.GenRandomAccount().GetAddress()
+	delegatorAddr := datagen.GenRandomAccount().GetAddress()
+	err := keeper.SetWithdrawAddr(ctx, delegatorAddr, withdrawalAddr)
+	require.NoError(t, err)
+
+	response, err := keeper.DelegatorWithdrawAddress(ctx, &types.QueryDelegatorWithdrawAddressRequest{DelegatorAddress: delegatorAddr.String()})
+	require.NoError(t, err)
+	require.Equal(t, &types.QueryDelegatorWithdrawAddressResponse{WithdrawAddress: withdrawalAddr.String()}, response)
+}

x/incentive/keeper/reward_gauge.go

@@ -20,8 +20,19 @@ func (k Keeper) withdrawReward(ctx context.Context, sType types.StakeholderType,
 	if !withdrawableCoins.IsAllPositive() {
 		return nil, types.ErrNoWithdrawableCoins
 	}
+
+	withdrawAddr, err := k.GetWithdrawAddr(ctx, addr)
+	if err != nil {
+		return nil, err
+	}
+
+	// Fallback to the stakeholder's address if no specific withdrawal address is set
+	if withdrawAddr == nil {
+		withdrawAddr = addr
+	}
+
 	// transfer withdrawable coins from incentive module account to the stakeholder's address
-	if err := k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, addr, withdrawableCoins); err != nil {
+	if err := k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, withdrawAddr, withdrawableCoins); err != nil {
 		return nil, err
 	}
 	// empty reward gauge

x/incentive/keeper/store.go

@@ -0,0 +1,23 @@
+package keeper
+
+import (
+	"context"
+	"github.com/babylonlabs-io/babylon/x/incentive/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func (k Keeper) GetWithdrawAddr(ctx context.Context, addr sdk.AccAddress) (sdk.AccAddress, error) {
+	store := k.storeService.OpenKVStore(ctx)
+	b, err := store.Get(types.GetWithdrawAddrKey(addr))
+	if b == nil {
+		return addr, err
+	}
+
+	return b, nil
+}
+
+func (k Keeper) SetWithdrawAddr(ctx context.Context, addr, withdrawAddr sdk.AccAddress) error {
+	store := k.storeService.OpenKVStore(ctx)
+
+	return store.Set(types.GetWithdrawAddrKey(addr), withdrawAddr.Bytes())
+}

x/incentive/types/keys.go

@@ -1,6 +1,10 @@
 package types
 
-import "cosmossdk.io/collections"
+import (
+	"cosmossdk.io/collections"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
+)
 
 const (
 	// ModuleName defines the module name
@@ -17,9 +21,14 @@ const (
 )
 
 var (
-	ParamsKey                 = []byte{0x01}             // key prefix for the parameters
-	BTCStakingGaugeKey        = []byte{0x02}             // key prefix for BTC staking gauge at each height
-	ReservedKey               = []byte{0x03}             // reserved //nolint:unused
-	RewardGaugeKey            = []byte{0x04}             // key prefix for reward gauge for a given stakeholder in a given type
-	RefundableMsgKeySetPrefix = collections.NewPrefix(5) // key prefix for refundable msg key set
+	ParamsKey                   = []byte{0x01}             // key prefix for the parameters
+	BTCStakingGaugeKey          = []byte{0x02}             // key prefix for BTC staking gauge at each height
+	DelegatorWithdrawAddrPrefix = []byte{0x03}             // key for delegator withdraw address
+	RewardGaugeKey              = []byte{0x04}             // key prefix for reward gauge for a given stakeholder in a given type
+	RefundableMsgKeySetPrefix   = collections.NewPrefix(5) // key prefix for refundable msg key set
 )
+
+// GetWithdrawAddrKey creates the key for a delegator's withdraw addr.
+func GetWithdrawAddrKey(delAddr sdk.AccAddress) []byte {
+	return append(DelegatorWithdrawAddrPrefix, address.MustLengthPrefix(delAddr.Bytes())...)
+}