Enable deterministic Cosmwasm stargate queries

app/keepers/keepers.go

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

wasmbinding/query_plugin.go

@@ -5,12 +5,46 @@ import (
 	"fmt"
 
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+	"github.com/cosmos/cosmos-sdk/baseapp"
 	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) 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 := NormalizeReponse(binding, res.Value)
+		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) {
@@ -110,6 +144,33 @@ 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 NormalizeReponse(binding interface{}, bz []byte) ([]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{}
+	}
+
+	// build new deterministic response
+	bz, err = proto.Marshal(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,122 @@
+package wasmbinding_test
+
+import (
+	"encoding/hex"
+	"testing"
+
+	proto "github.com/gogo/protobuf/proto"
+	"github.com/stretchr/testify/require"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/query"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+
+	"github.com/osmosis-labs/osmosis/v10/wasmbinding"
+)
+
+/**
+ * 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"
+// 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;
+}
+// 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;
+}
+*/
+
+func TestDeterministic_AllBalances(t *testing.T) {
+	originVersionBz, err := hex.DecodeString("0a090a036261721202333012050a03666f6f")
+	require.NoError(t, err)
+
+	newVersionBz, err := hex.DecodeString("0a090a036261721202333012050a03666f6f1a2d636f736d6f73316a366a357473717571326a6c77326166376c3378656b796171377a67346c386a737566753738")
+	require.NoError(t, err)
+
+	binding, ok := wasmbinding.StargateWhitelist.Load("/cosmos.bank.v1beta1.Query/AllBalances")
+	require.True(t, ok)
+
+	// new version response should be changed into origin version response
+	normalizedBz, err := wasmbinding.NormalizeReponse(binding, newVersionBz)
+	require.NoError(t, err)
+
+	require.Equal(t, originVersionBz, normalizedBz)
+	require.NotEqual(t, newVersionBz, normalizedBz)
+
+	// raw build also make same result
+	expectedResponse := banktypes.QueryAllBalancesResponse{
+		Balances: sdk.NewCoins(sdk.NewCoin("bar", sdk.NewInt(30))),
+		Pagination: &query.PageResponse{
+			NextKey: []byte("foo"),
+		},
+	}
+	expectedResponseBz, err := proto.Marshal(&expectedResponse)
+	require.NoError(t, err)
+	require.Equal(t, expectedResponseBz, normalizedBz)
+
+	// should be cleared
+	data := binding.(*banktypes.QueryAllBalancesResponse)
+	require.Empty(t, data.Balances)
+	require.Empty(t, data.Pagination)
+}
+
+/**
+ *
+ * Origin Response
+ * 0a530a202f636f736d6f732e617574682e763162657461312e426173654163636f756e74122f0a2d636f736d6f7331346c3268686a6e676c3939367772703935673867646a6871653038326375367a7732706c686b
+ *
+ * Updated Response
+ * 0a530a202f636f736d6f732e617574682e763162657461312e426173654163636f756e74122f0a2d636f736d6f7331646a783375676866736d6b6135386676673076616a6e6533766c72776b7a6a346e6377747271122d636f736d6f7331646a783375676866736d6b6135386676673076616a6e6533766c72776b7a6a346e6377747271
+// Origin proto
+message QueryAccountResponse {
+  // account defines the account of the corresponding address.
+  google.protobuf.Any account = 1 [(cosmos_proto.accepts_interface) = "AccountI"];
+}
+// Updated proto
+message QueryAccountResponse {
+  // account defines the account of the corresponding address.
+  google.protobuf.Any account = 1 [(cosmos_proto.accepts_interface) = "AccountI"];
+  // address is the address to query for.
+	string address = 2;
+}
+*/
+
+func TestDeterministic_Account(t *testing.T) {
+	originVersionBz, err := hex.DecodeString("0a530a202f636f736d6f732e617574682e763162657461312e426173654163636f756e74122f0a2d636f736d6f733166387578756c746e3873717a687a6e72737a3371373778776171756867727367366a79766679")
+	require.NoError(t, err)
+
+	newVersionBz, err := hex.DecodeString("0a530a202f636f736d6f732e617574682e763162657461312e426173654163636f756e74122f0a2d636f736d6f733166387578756c746e3873717a687a6e72737a3371373778776171756867727367366a79766679122d636f736d6f733166387578756c746e3873717a687a6e72737a3371373778776171756867727367366a79766679")
+	require.NoError(t, err)
+
+	binding, ok := wasmbinding.StargateWhitelist.Load("/cosmos.auth.v1beta1.Query/Account")
+	require.True(t, ok)
+
+	// new version response should be changed into origin version response
+	normalizedBz, err := wasmbinding.NormalizeReponse(binding, newVersionBz)
+	require.NoError(t, err)
+
+	require.Equal(t, originVersionBz, normalizedBz)
+	require.NotEqual(t, newVersionBz, normalizedBz)
+
+	// should be cleared
+	data := binding.(*authtypes.QueryAccountResponse)
+	require.Empty(t, data.Account)
+}

wasmbinding/wasm.go

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