change query depth implementation to use context instead wasmvm to st…

…ore depth

x/wasm/keeper/contract.go

@@ -1,6 +1,7 @@
 package keeper
 
 import (
+ "context"
  "time"
 
  wasmvmtypes "github.com/CosmWasm/wasmvm/types"
@@ -415,6 +416,21 @@ func (k Keeper) queryToStore(ctx sdk.Context, contractAddress sdk.AccAddress, ke
  return prefixStore.Get(key)
 }
 
+func assertQueryDepth(ctx sdk.Context) (uint8, error) {
+ var queryDepth uint8
+ if depth := ctx.Context().Value(types.WasmVMQueryDepthContextKey); depth != nil {
+ queryDepth = depth.(uint8)
+ } else {
+ queryDepth = 1
+ }
+
+ if queryDepth > types.ContractMaxQueryDepth {
+ return queryDepth, types.ErrExceedMaxQueryDepth
+ }
+
+ return queryDepth, nil
+}
+
 func (k Keeper) queryToContract(ctx sdk.Context, contractAddress sdk.AccAddress, queryMsg []byte) ([]byte, error) {
  defer telemetry.MeasureSince(time.Now(), "wasm", "contract", "query-smart")
  ctx.GasMeter().ConsumeGas(types.InstantiateContractCosts(len(queryMsg)), "Loading CosmWasm module: query")
@@ -426,11 +442,14 @@ func (k Keeper) queryToContract(ctx sdk.Context, contractAddress sdk.AccAddress,
 
  env := types.NewEnv(ctx, contractAddress)
 
- // assert max depth to prevent stack overflow
- if err := k.wasmVM.IncreaseQueryDepth(); err != nil {
+ // assert query depth and set new query depth for the next query
+ queryDepth, err := assertQueryDepth(ctx)
+ if err != nil {
  return nil, err
  }
- defer k.wasmVM.DecreaseQueryDepth()
+
+ // set next query depth
+ ctx = ctx.WithContext(context.WithValue(ctx.Context(), types.WasmVMQueryDepthContextKey, queryDepth+1))
 
  queryResult, gasUsed, err := k.wasmVM.Query(
  codeInfo.CodeHash,

x/wasm/keeper/keeper.go

@@ -65,24 +65,23 @@ func NewKeeper(
  wasmConfig.ContractMemoryCacheSize = config.DefaultContractMemoryCacheSize
  }
 
- var writeWasmVM types.WasmerEngine
- if vm, err := wasmvm.NewVM(
+ vm, err := wasmvm.NewVM(
  filepath.Join(homePath, config.DBDir),
  supportedFeatures,
  types.ContractMemoryLimit,
  wasmConfig.ContractDebugMode,
  wasmConfig.ContractMemoryCacheSize,
- ); err != nil {
+ )
+
+ if err != nil {
  panic(err)
- } else {
- writeWasmVM = types.NewWasmerEngineWithQueryDepth(vm)
  }
 
  return Keeper{
  storeKey: storeKey,
  cdc: cdc,
  paramSpace: paramspace,
- wasmVM: writeWasmVM,
+ wasmVM: vm,
  accountKeeper: accountKeeper,
  bankKeeper: bankKeeper,
  treasuryKeeper: treasuryKeeper,

x/wasm/types/keys.go

@@ -21,8 +21,8 @@ const (
  // RouterKey is the msg router key for the wasm module
  RouterKey = ModuleName
 
- // QueryWasmVMContextKey is the context key to store wasmvm in query context
- QueryWasmVMContextKey = "wasmvm"
+ // WasmVMQueryDepthContextKey context key to keep query depth
+ WasmVMQueryDepthContextKey = "wasmvm-query-depth"
 )
 
 // Keys for wasm store

x/wasm/types/wasmer_engine.go

@@ -8,37 +8,6 @@ import (
 // ContractMaxQueryDepth maximum recursive query depth allowed
 const ContractMaxQueryDepth = 20
 
-var _ WasmerEngine = &WasmerEngineWithQueryDepth{}
-
-// WasmerEngineWithQueryDepth VM wrapper with depth counter to prevent
-// stack overflow
-type WasmerEngineWithQueryDepth struct {
- *wasmvm.VM
- QueryDepth uint8
-}
-
-// NewWasmerEngineWithQueryDepth wrap wasmer engine with query depth checker
-func NewWasmerEngineWithQueryDepth(wasmVM *wasmvm.VM) *WasmerEngineWithQueryDepth {
- wasmerEngine := WasmerEngineWithQueryDepth{VM: wasmVM}
- return &wasmerEngine
-}
-
-// IncreaseQueryDepth increase execution depth by 1 and check whether
-// the depth exceeds max one or not
-func (wasmer *WasmerEngineWithQueryDepth) IncreaseQueryDepth() error {
- if wasmer.QueryDepth >= ContractMaxQueryDepth {
- return ErrExceedMaxQueryDepth
- }
-
- wasmer.QueryDepth++
- return nil
-}
-
-// DecreaseQueryDepth decrease execution depth by 1
-func (wasmer *WasmerEngineWithQueryDepth) DecreaseQueryDepth() {
- wasmer.QueryDepth--
-}
-
 // WasmerEngine defines the WASM contract runtime engine.
 type WasmerEngine interface {
 
@@ -153,11 +122,4 @@ type WasmerEngine interface {
 
  // Cleanup should be called when no longer using this to free resources on the rust-side
  Cleanup()
-
- // IncreaseQueryDepth will increase query depth by 1 and assert the current depth
- // reached out the maximum
- IncreaseQueryDepth() error
-
- // DecreaseQueryDepth will decrease query depth by 1
- DecreaseQueryDepth()
 }