RE: Enable Whitelist Stargate Query

CHANGELOG.md

@@ -72,6 +72,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#1667](https://github.com/osmosis-labs/osmosis/pull/1673) Move wasm-bindings code out of app package into its own root level package.
 * [#2013](https://github.com/osmosis-labs/osmosis/pull/2013) Make `SetParams`, `SetPool`, `SetTotalLiquidity`, and `SetDenomLiquidity` GAMM APIs private
 * [#1857](https://github.com/osmosis-labs/osmosis/pull/1857) x/mint rename GetLastHalvenEpochNum to GetLastReductionEpochNum
+* [#2353](https://github.com/osmosis-labs/osmosis/pull/2353) Re-enable stargate query via whitelsit
 
 ### Features
 

app/keepers/keepers.go

@@ -326,6 +326,7 @@ func (appKeepers *AppKeepers) InitNormalKeepers(
 	supportedFeatures := "iterator,staking,stargate,osmosis"
 
 	wasmOpts = append(owasm.RegisterCustomPlugins(appKeepers.GAMMKeeper, appKeepers.BankKeeper, appKeepers.TwapKeeper, appKeepers.TokenFactoryKeeper), wasmOpts...)
+	wasmOpts = append(owasm.RegisterStargateQueries(*bApp.GRPCQueryRouter(), appCodec), wasmOpts...)
 
 	wasmKeeper := wasm.NewKeeper(
 		appCodec,

wasmbinding/query_plugin.go

@@ -5,12 +5,47 @@ import (
 	"fmt"
 
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	proto "github.com/gogo/protobuf/proto"
+	abci "github.com/tendermint/tendermint/abci/types"
 
 	"github.com/osmosis-labs/osmosis/v10/wasmbinding/bindings"
 )
 
+// StargateQuerier dispatches whitelisted stargate queries
+func StargateQuerier(queryRouter baseapp.GRPCQueryRouter, codec codec.Codec) func(ctx sdk.Context, request *wasmvmtypes.StargateQuery) ([]byte, error) {
+	return func(ctx sdk.Context, request *wasmvmtypes.StargateQuery) ([]byte, error) {
+		binding, whitelisted := StargateWhitelist.Load(request.Path)
+		if !whitelisted {
+			return nil, wasmvmtypes.UnsupportedRequest{Kind: fmt.Sprintf("'%s' path is not allowed from the contract", request.Path)}
+		}
+
+		route := queryRouter.Route(request.Path)
+		if route == nil {
+			return nil, wasmvmtypes.UnsupportedRequest{Kind: fmt.Sprintf("No route to query '%s'", request.Path)}
+		}
+
+		res, err := route(ctx, abci.RequestQuery{
+			Data: request.Data,
+			Path: request.Path,
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		// normalize response to ensure backward compatibility
+		bz, err := NormalizeResponseAndJsonMarshal(binding, res.Value, codec)
+		if err != nil {
+			return nil, err
+		}
+
+		return bz, nil
+	}
+}
+
 // CustomQuerier dispatches custom CosmWasm bindings queries.
 func CustomQuerier(qp *QueryPlugin) func(ctx sdk.Context, request json.RawMessage) ([]byte, error) {
 	return func(ctx sdk.Context, request json.RawMessage) ([]byte, error) {
@@ -138,6 +173,32 @@ func CustomQuerier(qp *QueryPlugin) func(ctx sdk.Context, request json.RawMessag
 	}
 }
 
+// NormalizeReponses normalizes the responses by unmarshalling the response then marshalling them again.
+// Normalizing the response is specifically important for responses that contain type of Any.
+func NormalizeResponseAndJsonMarshal(binding interface{}, bz []byte, codec codec.Codec) ([]byte, error) {
+	// all values are proto message
+	message, ok := binding.(proto.Message)
+	if !ok {
+		return nil, wasmvmtypes.Unknown{}
+	}
+
+	// unmarshal binary into stargate response data structure
+	err := proto.Unmarshal(bz, message)
+	if err != nil {
+		return nil, wasmvmtypes.Unknown{}
+	}
+
+	bz, err = codec.MarshalJSON(message)
+	if err != nil {
+		return nil, wasmvmtypes.Unknown{}
+	}
+
+	// clear proto message
+	message.Reset()
+
+	return bz, nil
+}
+
 // ConvertSdkCoinsToWasmCoins converts sdk type coins to wasm vm type coins
 func ConvertSdkCoinsToWasmCoins(coins []sdk.Coin) wasmvmtypes.Coins {
 	var toSend wasmvmtypes.Coins

wasmbinding/stargate_whitelist.go

@@ -0,0 +1,24 @@
+package wasmbinding
+
+import (
+	"sync"
+
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+
+	epochtypes "github.com/osmosis-labs/osmosis/v10/x/epochs/types"
+)
+
+// StargateWhitelist keeps whitelist and its deterministic
+// response binding for stargate queries.
+//
+// The query can be multi-thread, so we have to use
+// thread safe sync.Map.
+var StargateWhitelist sync.Map
+
+func init() {
+	StargateWhitelist.Store("/cosmos.auth.v1beta1.Query/Account", &authtypes.QueryAccountResponse{})
+	StargateWhitelist.Store("/cosmos.bank.v1beta1.Query/AllBalances", &banktypes.QueryAllBalancesResponse{})
+
+	StargateWhitelist.Store("/osmosis.epochs.v1beta1.Query/EpochInfos", &epochtypes.QueryCurrentEpochResponse{})
+}

wasmbinding/stargate_whitelist_test.go

@@ -0,0 +1,114 @@
+package wasmbinding_test
+
+import (
+	"encoding/hex"
+	"fmt"
+	"testing"
+	"time"
+
+	proto "github.com/gogo/protobuf/proto"
+	"github.com/stretchr/testify/suite"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/query"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+
+	"github.com/osmosis-labs/osmosis/v10/app"
+	"github.com/osmosis-labs/osmosis/v10/wasmbinding"
+)
+
+type StargateTestSuite struct {
+	suite.Suite
+
+	ctx sdk.Context
+	app *app.OsmosisApp
+}
+
+func (suite *StargateTestSuite) SetupTest() {
+	suite.app = app.Setup(false)
+	suite.ctx = suite.app.BaseApp.NewContext(false, tmproto.Header{Height: 1, ChainID: "osmosis-1", Time: time.Now().UTC()})
+}
+
+func TestStargateTestSuite(t *testing.T) {
+	suite.Run(t, new(StargateTestSuite))
+}
+
+func (suite *StargateTestSuite) TestDeterministicJsonMarshal() {
+	testCases := []struct {
+		name             string
+		originalResponse string
+		updatedResponse  string
+		queryPath        string
+		expectedProto    proto.Message
+	}{
+
+		/**
+		   * Origin Response
+		   * balances:<denom:"bar" amount:"30" > pagination:<next_key:"foo" >
+		   * "0a090a036261721202333012050a03666f6f"
+		   *
+		   * New Version Response
+		   * The binary built from the proto response with additional field address
+		   * balances:<denom:"bar" amount:"30" > pagination:<next_key:"foo" > address:"cosmos1j6j5tsquq2jlw2af7l3xekyaq7zg4l8jsufu78"
+		   * "0a090a036261721202333012050a03666f6f1a2d636f736d6f73316a366a357473717571326a6c77326166376c3378656b796171377a67346c386a737566753738"
+		   // Origin proto
+		   message QueryAllBalancesResponse {
+		  	// balances is the balances of all the coins.
+		  	repeated cosmos.base.v1beta1.Coin balances = 1
+		  	[(gogoproto.nullable) = false, (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coins"];
+		  	// pagination defines the pagination in the response.
+		  	cosmos.base.query.v1beta1.PageResponse pagination = 2;
+		  }
+		  // Updated proto
+		  message QueryAllBalancesResponse {
+		  	// balances is the balances of all the coins.
+		  	repeated cosmos.base.v1beta1.Coin balances = 1
+		  	[(gogoproto.nullable) = false, (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coins"];
+		  	// pagination defines the pagination in the response.
+		  	cosmos.base.query.v1beta1.PageResponse pagination = 2;
+		  	// address is the address to query all balances for.
+		  	string address = 3;
+		  }
+		*/
+		{
+			"Query All Balances",
+			"0a090a036261721202333012050a03666f6f",
+			"0a090a036261721202333012050a03666f6f1a2d636f736d6f73316a366a357473717571326a6c77326166376c3378656b796171377a67346c386a737566753738",
+			"/cosmos.bank.v1beta1.Query/AllBalances",
+			&banktypes.QueryAllBalancesResponse{
+				Balances: sdk.NewCoins(sdk.NewCoin("bar", sdk.NewInt(30))),
+				Pagination: &query.PageResponse{
+					NextKey: []byte("foo"),
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(fmt.Sprintf("Case %s", tc.name), func() {
+			suite.SetupTest()
+
+			binding, ok := wasmbinding.StargateWhitelist.Load(tc.queryPath)
+			suite.Require().True(ok)
+
+			originVersionBz, err := hex.DecodeString(tc.originalResponse)
+			suite.Require().NoError(err)
+			jsonMarshalledOriginalBz, err := wasmbinding.NormalizeResponseAndJsonMarshal(binding, originVersionBz, suite.app.AppCodec())
+			suite.Require().NoError(err)
+
+			newVersionBz, err := hex.DecodeString(tc.updatedResponse)
+			suite.Require().NoError(err)
+			jsonMarshalledUpdatedBz, err := wasmbinding.NormalizeResponseAndJsonMarshal(binding, newVersionBz, suite.app.AppCodec())
+			suite.Require().NoError(err)
+
+			// json marshalled bytes should be the same since we use the same proto sturct for unmarshalling
+			suite.Require().Equal(jsonMarshalledOriginalBz, jsonMarshalledUpdatedBz)
+
+			// raw build also make same result
+			jsonMarshalExpectedResponse, err := suite.app.AppCodec().MarshalJSON(tc.expectedProto)
+			suite.Require().NoError(err)
+			suite.Require().Equal(jsonMarshalledUpdatedBz, jsonMarshalExpectedResponse)
+		})
+	}
+}

wasmbinding/wasm.go

@@ -2,7 +2,11 @@ package wasmbinding
 
 import (
 	"github.com/CosmWasm/wasmd/x/wasm"
+
 	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	"github.com/cosmos/cosmos-sdk/codec"
+
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 
 	gammkeeper "github.com/osmosis-labs/osmosis/v10/x/gamm/keeper"
@@ -30,3 +34,13 @@ func RegisterCustomPlugins(
 		messengerDecoratorOpt,
 	}
 }
+
+func RegisterStargateQueries(queryRouter baseapp.GRPCQueryRouter, codec codec.Codec) []wasmkeeper.Option {
+	queryPluginOpt := wasmkeeper.WithQueryPlugins(&wasmkeeper.QueryPlugins{
+		Stargate: StargateQuerier(queryRouter, codec),
+	})
+
+	return []wasm.Option{
+		queryPluginOpt,
+	}
+}