feat: add custom queries to wasm module

modules/light-clients/08-wasm/internal/ibcwasm/expected_interfaces.go

@@ -115,3 +115,21 @@ type WasmEngine interface {
 	// Unpin is idempotent.
 	Unpin(checksum wasmvm.Checksum) error
 }
+
+type WasmQuerier interface {
+	// Query takes a query request, performs the query and returns the response.
+	// It takes a gas limit measured in [CosmWasm gas] (aka. wasmvm gas) to ensure
+	// the query does not consume more gas than the contract execution has left.
+	//
+	// [CosmWasm gas]: https://github.com/CosmWasm/cosmwasm/blob/v1.3.1/docs/GAS.md
+	Query(request wasmvmtypes.QueryRequest, gasLimit uint64) ([]byte, error)
+	// GasConsumed returns the gas that was consumed by the querier during its entire
+	// lifetime or by the context in which it was executed in. The absolute gas values
+	// must not be used directly as it is undefined what is included in this value. Instead
+	// wasmvm will call GasConsumed before and after the query and use the difference
+	// as the query's gas usage.
+	// Like the gas limit above, this is measured in [CosmWasm gas] (aka. wasmvm gas).
+	//
+	// [CosmWasm gas]: https://github.com/CosmWasm/cosmwasm/blob/v1.3.1/docs/GAS.md
+	GasConsumed() uint64
+}

modules/light-clients/08-wasm/internal/ibcwasm/wasm.go

@@ -10,6 +10,8 @@ import (
 var (
 	vm WasmEngine
 
+	querier WasmQuerier
+
 	// state management
 	Schema    collections.Schema
 	Checksums collections.KeySet[[]byte]
@@ -32,6 +34,16 @@ func GetVM() WasmEngine {
 	return vm
 }
 
+// SetQuerier sets the custom wasm query handle for the 08-wasm module.
+func SetQuerier(wasmQuerier WasmQuerier) {
+	querier = wasmQuerier
+}
+
+// GetQuerier returns the custom wasm query handler for the 08-wasm module.
+func GetQuerier() WasmQuerier {
+	return querier
+}
+
 // SetupWasmStoreService sets up the 08-wasm module's collections.
 func SetupWasmStoreService(storeService storetypes.KVStoreService) {
 	sb := collections.NewSchemaBuilder(storeService)

modules/light-clients/08-wasm/keeper/keeper.go

@@ -41,6 +41,7 @@ func NewKeeperWithVM(
 	clientKeeper types.ClientKeeper,
 	authority string,
 	vm ibcwasm.WasmEngine,
+	querier ibcwasm.WasmQuerier,
 ) Keeper {
 	if clientKeeper == nil {
 		panic(errors.New("client keeper must be not nil"))
@@ -59,6 +60,7 @@ func NewKeeperWithVM(
 	}
 
 	ibcwasm.SetVM(vm)
+	ibcwasm.SetQuerier(querier)
 	ibcwasm.SetupWasmStoreService(storeService)
 
 	return Keeper{
@@ -77,13 +79,14 @@ func NewKeeperWithConfig(
 	clientKeeper types.ClientKeeper,
 	authority string,
 	wasmConfig types.WasmConfig,
+	querier ibcwasm.WasmQuerier,
 ) Keeper {
 	vm, err := wasmvm.NewVM(wasmConfig.DataDir, wasmConfig.SupportedCapabilities, types.ContractMemoryLimit, wasmConfig.ContractDebugMode, types.MemoryCacheSize)
 	if err != nil {
 		panic(fmt.Errorf("failed to instantiate new Wasm VM instance: %v", err))
 	}
 
-	return NewKeeperWithVM(cdc, storeService, clientKeeper, authority, vm)
+	return NewKeeperWithVM(cdc, storeService, clientKeeper, authority, vm, querier)
 }
 
 // GetAuthority returns the 08-wasm module's authority.

modules/light-clients/08-wasm/keeper/keeper_test.go

@@ -149,6 +149,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					ibcwasm.GetVM(),
+					nil,
 				)
 			},
 			true,
@@ -163,6 +164,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					"", // authority
 					ibcwasm.GetVM(),
+					nil,
 				)
 			},
 			false,
@@ -177,6 +179,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 					nil, // client keeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					ibcwasm.GetVM(),
+					nil,
 				)
 			},
 			false,
@@ -191,6 +194,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					nil,
+					nil,
 				)
 			},
 			false,
@@ -205,6 +209,7 @@ func (suite *KeeperTestSuite) TestNewKeeper() {
 					GetSimApp(suite.chainA).IBCKeeper.ClientKeeper,
 					GetSimApp(suite.chainA).WasmClientKeeper.GetAuthority(),
 					ibcwasm.GetVM(),
+					nil,
 				)
 			},
 			false,

modules/light-clients/08-wasm/testing/simapp/app.go

@@ -471,12 +471,12 @@ func NewSimApp(
 		// NOTE: mockVM is used for testing purposes only!
 		app.WasmClientKeeper = wasmkeeper.NewKeeperWithVM(
 			appCodec, runtime.NewKVStoreService(keys[wasmtypes.StoreKey]), app.IBCKeeper.ClientKeeper,
-			authtypes.NewModuleAddress(govtypes.ModuleName).String(), mockVM,
+			authtypes.NewModuleAddress(govtypes.ModuleName).String(), mockVM, nil,
 		)
 	} else {
 		app.WasmClientKeeper = wasmkeeper.NewKeeperWithConfig(
 			appCodec, runtime.NewKVStoreService(keys[wasmtypes.StoreKey]), app.IBCKeeper.ClientKeeper,
-			authtypes.NewModuleAddress(govtypes.ModuleName).String(), wasmConfig,
+			authtypes.NewModuleAddress(govtypes.ModuleName).String(), wasmConfig, nil,
 		)
 	}
 

modules/light-clients/08-wasm/types/client_state_test.go

@@ -2,6 +2,7 @@ package types_test
 
 import (
 	"encoding/json"
+	"math"
 	"time"
 
 	wasmvm "github.com/CosmWasm/wasmvm"
@@ -617,3 +618,87 @@ func (suite *TypesTestSuite) TestVerifyNonMembership() {
 		})
 	}
 }
+
+type CustomQuery struct {
+	Echo *QueryEcho `json:"echo,omitempty"`
+}
+
+type QueryEcho struct {
+	Data string `json:"data"`
+}
+
+type CustomQueryHandler struct{}
+
+func (*CustomQueryHandler) GasConsumed() uint64 {
+	return 0
+}
+
+func (*CustomQueryHandler) Query(request wasmvmtypes.QueryRequest, gasLimit uint64) ([]byte, error) {
+	var customQuery CustomQuery
+	err := json.Unmarshal([]byte(request.Custom), &customQuery)
+	if err != nil {
+		return nil, wasmtesting.ErrMockContract
+	}
+
+	if customQuery.Echo != nil {
+		data, err := json.Marshal(customQuery.Echo.Data)
+		return data, err
+	}
+
+	return nil, wasmtesting.ErrMockContract
+}
+
+func (suite *TypesTestSuite) TestCustomQuery() {
+	testCases := []struct {
+		name     string
+		malleate func()
+	}{
+		{
+			"success: custom query",
+			func() {
+				suite.mockVM.RegisterQueryCallback(types.StatusMsg{}, func(checksum wasmvm.Checksum, env wasmvmtypes.Env, queryMsg []byte, store wasmvm.KVStore, goapi wasmvm.GoAPI, querier wasmvm.Querier, gasMeter wasmvm.GasMeter, gasLimit uint64, deserCost wasmvmtypes.UFraction) ([]byte, uint64, error) {
+					echo := CustomQuery{
+						Echo: &QueryEcho{
+							Data: "hello world",
+						},
+					}
+					echoJSON, err := json.Marshal(echo)
+					suite.Require().NoError(err)
+					resp, err := querier.Query(wasmvmtypes.QueryRequest{
+						Custom: json.RawMessage(echoJSON),
+					}, math.MaxUint64)
+					suite.Require().NoError(err)
+
+					var respData string
+					err = json.Unmarshal(resp, &respData)
+
+					suite.Require().NoError(err)
+					suite.Require().Equal("hello world", respData)
+
+					resp, err = json.Marshal(types.StatusResult{Status: exported.Active.String()})
+					suite.Require().NoError(err)
+					return resp, wasmtesting.DefaultGasUsed, nil
+				})
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(tc.name, func() {
+			suite.SetupWasmWithMockVM()
+
+			endpoint := wasmtesting.NewWasmEndpoint(suite.chainA)
+			err := endpoint.CreateClient()
+			suite.Require().NoError(err)
+
+			tc.malleate()
+
+			clientStore := suite.chainA.App.GetIBCKeeper().ClientKeeper.ClientStore(suite.chainA.GetContext(), endpoint.ClientID)
+			clientState := endpoint.GetClientState()
+
+			ibcwasm.SetQuerier(&CustomQueryHandler{})
+			clientState.Status(suite.chainA.GetContext(), clientStore, suite.chainA.App.AppCodec())
+			ibcwasm.SetQuerier(nil)
+		})
+	}
+}

modules/light-clients/08-wasm/types/vm.go

@@ -49,7 +49,7 @@ func instantiateContract(ctx sdk.Context, clientStore storetypes.KVStore, checks
 	}
 
 	ctx.GasMeter().ConsumeGas(VMGasRegister.NewContractInstanceCosts(true, len(msg)), "Loading CosmWasm module: instantiate")
-	response, gasUsed, err := ibcwasm.GetVM().Instantiate(checksum, env, msgInfo, msg, newStoreAdapter(clientStore), wasmvmAPI, nil, multipliedGasMeter, gasLimit, costJSONDeserialization)
+	response, gasUsed, err := ibcwasm.GetVM().Instantiate(checksum, env, msgInfo, msg, newStoreAdapter(clientStore), wasmvmAPI, ibcwasm.GetQuerier(), multipliedGasMeter, gasLimit, costJSONDeserialization)
 	VMGasRegister.consumeRuntimeGas(ctx, gasUsed)
 	return response, err
 }
@@ -67,7 +67,7 @@ func callContract(ctx sdk.Context, clientStore storetypes.KVStore, checksum Chec
 	env := getEnv(ctx, clientID)
 
 	ctx.GasMeter().ConsumeGas(VMGasRegister.InstantiateContractCosts(true, len(msg)), "Loading CosmWasm module: sudo")
-	resp, gasUsed, err := ibcwasm.GetVM().Sudo(checksum, env, msg, newStoreAdapter(clientStore), wasmvmAPI, nil, multipliedGasMeter, gasLimit, costJSONDeserialization)
+	resp, gasUsed, err := ibcwasm.GetVM().Sudo(checksum, env, msg, newStoreAdapter(clientStore), wasmvmAPI, ibcwasm.GetQuerier(), multipliedGasMeter, gasLimit, costJSONDeserialization)
 	VMGasRegister.consumeRuntimeGas(ctx, gasUsed)
 	return resp, err
 }
@@ -81,7 +81,7 @@ func migrateContract(ctx sdk.Context, clientID string, clientStore storetypes.KV
 	env := getEnv(ctx, clientID)
 
 	ctx.GasMeter().ConsumeGas(VMGasRegister.InstantiateContractCosts(true, len(msg)), "Loading CosmWasm module: migrate")
-	resp, gasUsed, err := ibcwasm.GetVM().Migrate(checksum, env, msg, newStoreAdapter(clientStore), wasmvmAPI, nil, multipliedGasMeter, gasLimit, costJSONDeserialization)
+	resp, gasUsed, err := ibcwasm.GetVM().Migrate(checksum, env, msg, newStoreAdapter(clientStore), wasmvmAPI, ibcwasm.GetQuerier(), multipliedGasMeter, gasLimit, costJSONDeserialization)
 	VMGasRegister.consumeRuntimeGas(ctx, gasUsed)
 	return resp, err
 }
@@ -99,7 +99,7 @@ func queryContract(ctx sdk.Context, clientStore storetypes.KVStore, checksum Che
 	env := getEnv(ctx, clientID)
 
 	ctx.GasMeter().ConsumeGas(VMGasRegister.InstantiateContractCosts(true, len(msg)), "Loading CosmWasm module: query")
-	resp, gasUsed, err := ibcwasm.GetVM().Query(checksum, env, msg, newStoreAdapter(clientStore), wasmvmAPI, nil, multipliedGasMeter, gasLimit, costJSONDeserialization)
+	resp, gasUsed, err := ibcwasm.GetVM().Query(checksum, env, msg, newStoreAdapter(clientStore), wasmvmAPI, ibcwasm.GetQuerier(), multipliedGasMeter, gasLimit, costJSONDeserialization)
 	VMGasRegister.consumeRuntimeGas(ctx, gasUsed)
 	return resp, err
 }