Add distribution query

2023-07-18 15:29:16 +02:00
parent 5d05b85875
commit a17f5f2fbf
3 changed files with 133 additions and 32 deletions
--- a/x/wasm/keeper/query_plugins.go
+++ b/x/wasm/keeper/query_plugins.go
@@ -6,7 +6,6 @@ import (
 	"fmt"

 	errorsmod "cosmossdk.io/errors"
-	"github.com/CosmWasm/wasmd/x/wasm/types"
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
 	abci "github.com/cometbft/cometbft/abci/types"
 	"github.com/cosmos/cosmos-sdk/baseapp"
@@ -18,6 +17,8 @@ import (
 	distributiontypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	channeltypes "github.com/cosmos/ibc-go/v7/modules/core/04-channel/types"
+
+	"github.com/CosmWasm/wasmd/x/wasm/types"
 )

 type QueryHandler struct {
@@ -78,12 +79,13 @@ func (q QueryHandler) GasConsumed() uint64 {
 type CustomQuerier func(ctx sdk.Context, request json.RawMessage) ([]byte, error)

 type QueryPlugins struct {
-	Bank     func(ctx sdk.Context, request *wasmvmtypes.BankQuery) ([]byte, error)
-	Custom   CustomQuerier
-	IBC      func(ctx sdk.Context, caller sdk.AccAddress, request *wasmvmtypes.IBCQuery) ([]byte, error)
-	Staking  func(ctx sdk.Context, request *wasmvmtypes.StakingQuery) ([]byte, error)
-	Stargate func(ctx sdk.Context, request *wasmvmtypes.StargateQuery) ([]byte, error)
-	Wasm     func(ctx sdk.Context, request *wasmvmtypes.WasmQuery) ([]byte, error)
+	Bank         func(ctx sdk.Context, request *wasmvmtypes.BankQuery) ([]byte, error)
+	Custom       CustomQuerier
+	IBC          func(ctx sdk.Context, caller sdk.AccAddress, request *wasmvmtypes.IBCQuery) ([]byte, error)
+	Staking      func(ctx sdk.Context, request *wasmvmtypes.StakingQuery) ([]byte, error)
+	Stargate     func(ctx sdk.Context, request *wasmvmtypes.StargateQuery) ([]byte, error)
+	Wasm         func(ctx sdk.Context, request *wasmvmtypes.WasmQuery) ([]byte, error)
+	Distribution func(ctx sdk.Context, request *wasmvmtypes.DistributionQuery) ([]byte, error)
 }

 type contractMetaDataSource interface {
@@ -106,12 +108,13 @@ func DefaultQueryPlugins(
 	wasm wasmQueryKeeper,
 ) QueryPlugins {
 	return QueryPlugins{
-		Bank:     BankQuerier(bank),
-		Custom:   NoCustomQuerier,
-		IBC:      IBCQuerier(wasm, channelKeeper),
-		Staking:  StakingQuerier(staking, distKeeper),
-		Stargate: RejectStargateQuerier(),
-		Wasm:     WasmQuerier(wasm),
+		Bank:         BankQuerier(bank),
+		Custom:       NoCustomQuerier,
+		IBC:          IBCQuerier(wasm, channelKeeper),
+		Staking:      StakingQuerier(staking, distKeeper),
+		Stargate:     RejectStargateQuerier(),
+		Wasm:         WasmQuerier(wasm),
+		Distribution: DistributionQuerier(distKeeper),
 	}
 }

@@ -138,29 +141,30 @@ func (e QueryPlugins) Merge(o *QueryPlugins) QueryPlugins {
 	if o.Wasm != nil {
 		e.Wasm = o.Wasm
 	}
+	if o.Distribution != nil {
+		e.Distribution = o.Distribution
+	}
 	return e
 }

 // HandleQuery executes the requested query
-func (e QueryPlugins) HandleQuery(ctx sdk.Context, caller sdk.AccAddress, request wasmvmtypes.QueryRequest) ([]byte, error) {
+func (e QueryPlugins) HandleQuery(ctx sdk.Context, caller sdk.AccAddress, req wasmvmtypes.QueryRequest) ([]byte, error) {
 	// do the query
-	if request.Bank != nil {
-		return e.Bank(ctx, request.Bank)
-	}
-	if request.Custom != nil {
-		return e.Custom(ctx, request.Custom)
-	}
-	if request.IBC != nil {
-		return e.IBC(ctx, caller, request.IBC)
-	}
-	if request.Staking != nil {
-		return e.Staking(ctx, request.Staking)
-	}
-	if request.Stargate != nil {
-		return e.Stargate(ctx, request.Stargate)
-	}
-	if request.Wasm != nil {
-		return e.Wasm(ctx, request.Wasm)
+	switch {
+	case req.Bank != nil:
+		return e.Bank(ctx, req.Bank)
+	case req.Custom != nil:
+		return e.Custom(ctx, req.Custom)
+	case req.IBC != nil:
+		return e.IBC(ctx, caller, req.IBC)
+	case req.Staking != nil:
+		return e.Staking(ctx, req.Staking)
+	case req.Stargate != nil:
+		return e.Stargate(ctx, req.Stargate)
+	case req.Wasm != nil:
+		return e.Wasm(ctx, req.Wasm)
+	case req.Distribution != nil:
+		return e.Distribution(ctx, req.Distribution)
 	}
 	return nil, wasmvmtypes.Unknown{}
 }
@@ -589,6 +593,22 @@ func WasmQuerier(k wasmQueryKeeper) func(ctx sdk.Context, request *wasmvmtypes.W
 	}
 }

+func DistributionQuerier(k types.DistributionKeeper) func(ctx sdk.Context, request *wasmvmtypes.DistributionQuery) ([]byte, error) {
+	return func(ctx sdk.Context, req *wasmvmtypes.DistributionQuery) ([]byte, error) {
+		if req.DelegatorWithdrawAddress == nil {
+			return nil, wasmvmtypes.UnsupportedRequest{Kind: "unknown distribution query"}
+		}
+		addr, err := sdk.AccAddressFromBech32(req.DelegatorWithdrawAddress.DelegatorAddress)
+		if err != nil {
+			return nil, sdkerrors.ErrInvalidAddress.Wrap("delegator address")
+		}
+		res := wasmvmtypes.DelegatorWithdrawAddressResponse{
+			WithdrawAddress: k.GetDelegatorWithdrawAddr(ctx, addr).String(),
+		}
+		return json.Marshal(res)
+	}
+}
+
 // ConvertSdkCoinsToWasmCoins covert sdk type to wasmvm coins type
 func ConvertSdkCoinsToWasmCoins(coins []sdk.Coin) wasmvmtypes.Coins {
 	converted := make(wasmvmtypes.Coins, len(coins))
--- a/x/wasm/keeper/query_plugins_test.go
+++ b/x/wasm/keeper/query_plugins_test.go
@@ -8,6 +8,10 @@ import (
 	"testing"
 	"time"

+	"github.com/cometbft/cometbft/libs/rand"
+	"github.com/cosmos/cosmos-sdk/types/address"
+	distributiontypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
+
 	errorsmod "cosmossdk.io/errors"
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
 	dbm "github.com/cometbft/cometbft-db"
@@ -757,6 +761,82 @@ func TestAcceptListStargateQuerier(t *testing.T) {
 	}
 }

+func TestDistributionQuerier(t *testing.T) {
+	ctx := sdk.Context{}
+	var myAddr sdk.AccAddress = rand.Bytes(address.Len)
+	var myOtherAddr sdk.AccAddress = rand.Bytes(address.Len)
+	specs := map[string]struct {
+		q       wasmvmtypes.DistributionQuery
+		mockFn  func(ctx sdk.Context, delAddr sdk.AccAddress) sdk.AccAddress
+		expAddr string
+		expErr  bool
+	}{
+		"withdrawal override": {
+			q: wasmvmtypes.DistributionQuery{
+				DelegatorWithdrawAddress: &wasmvmtypes.DelegatorWithdrawAddressQuery{DelegatorAddress: myAddr.String()},
+			},
+			mockFn: func(_ sdk.Context, delAddr sdk.AccAddress) sdk.AccAddress {
+				return myOtherAddr
+			},
+			expAddr: myOtherAddr.String(),
+		},
+		"no withdrawal override": {
+			q: wasmvmtypes.DistributionQuery{
+				DelegatorWithdrawAddress: &wasmvmtypes.DelegatorWithdrawAddressQuery{DelegatorAddress: myAddr.String()},
+			},
+			mockFn: func(_ sdk.Context, delAddr sdk.AccAddress) sdk.AccAddress {
+				return delAddr
+			},
+			expAddr: myAddr.String(),
+		},
+		"empty address": {
+			q: wasmvmtypes.DistributionQuery{
+				DelegatorWithdrawAddress: &wasmvmtypes.DelegatorWithdrawAddressQuery{},
+			},
+			expErr: true,
+		},
+		"unknown query": {
+			q:      wasmvmtypes.DistributionQuery{},
+			expErr: true,
+		},
+	}
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			mock := distrKeeperMock{GetDelegatorWithdrawAddrFn: spec.mockFn}
+			q := keeper.DistributionQuerier(mock)
+
+			gotBz, gotErr := q(ctx, &spec.q)
+			if spec.expErr {
+				require.Error(t, gotErr)
+				return
+			}
+			require.NoError(t, gotErr)
+			var rsp wasmvmtypes.DelegatorWithdrawAddressResponse
+			require.NoError(t, json.Unmarshal(gotBz, &rsp))
+			assert.Equal(t, spec.expAddr, rsp.WithdrawAddress)
+		})
+	}
+}
+
+type distrKeeperMock struct {
+	DelegationRewardsFn        func(c context.Context, req *distributiontypes.QueryDelegationRewardsRequest) (*distributiontypes.QueryDelegationRewardsResponse, error)
+	GetDelegatorWithdrawAddrFn func(ctx sdk.Context, delAddr sdk.AccAddress) sdk.AccAddress
+}
+
+func (m distrKeeperMock) DelegationRewards(ctx context.Context, req *distributiontypes.QueryDelegationRewardsRequest) (*distributiontypes.QueryDelegationRewardsResponse, error) {
+	if m.DelegationRewardsFn == nil {
+		panic("not expected to be called")
+	}
+	return m.DelegationRewardsFn(ctx, req)
+}
+
+func (m distrKeeperMock) GetDelegatorWithdrawAddr(ctx sdk.Context, delAddr sdk.AccAddress) sdk.AccAddress {
+	if m.GetDelegatorWithdrawAddrFn == nil {
+		panic("not expected to be called")
+	}
+	return m.GetDelegatorWithdrawAddrFn(ctx, delAddr)
+}
+
 type mockWasmQueryKeeper struct {
 	GetContractInfoFn func(ctx sdk.Context, contractAddress sdk.AccAddress) *types.ContractInfo
 	QueryRawFn        func(ctx sdk.Context, contractAddress sdk.AccAddress, key []byte) []byte
--- a/x/wasm/types/expected_keepers.go
+++ b/x/wasm/types/expected_keepers.go
@@ -52,7 +52,8 @@ type AccountKeeper interface {

 // DistributionKeeper defines a subset of methods implemented by the cosmos-sdk distribution keeper
 type DistributionKeeper interface {
-	DelegationRewards(c context.Context, req *distrtypes.QueryDelegationRewardsRequest) (*distrtypes.QueryDelegationRewardsResponse, error)
+	DelegationRewards(ctx context.Context, req *distrtypes.QueryDelegationRewardsRequest) (*distrtypes.QueryDelegationRewardsResponse, error)
+	GetDelegatorWithdrawAddr(ctx sdk.Context, delAddr sdk.AccAddress) sdk.AccAddress
 }

 // StakingKeeper defines a subset of methods implemented by the cosmos-sdk staking keeper