Merge pull request #867 from CosmWasm/query-stack-size-limit

Create query stack size limit

x/wasm/keeper/keeper.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"fmt"
 	"math"
@@ -30,6 +31,13 @@ import (
 // constant value so all nodes run with the same limit.
 const contractMemoryLimit = 32
 
+type contextKey int
+
+const (
+	// private type creates an interface key for Context that cannot be accessed by any other package
+	contextKeyQueryStackSize contextKey = iota
+)
+
 // Option is an extension point to instantiate keeper with non default values
 type Option interface {
 	apply(*Keeper)
@@ -71,9 +79,10 @@ type Keeper struct {
 	wasmVMResponseHandler WasmVMResponseHandler
 	messenger             Messenger
 	// queryGasLimit is the max wasmvm gas that can be spent on executing a query with a contract
-	queryGasLimit uint64
-	paramSpace    paramtypes.Subspace
-	gasRegister   GasRegister
+	queryGasLimit     uint64
+	paramSpace        paramtypes.Subspace
+	gasRegister       GasRegister
+	maxQueryStackSize uint32
 }
 
 // NewKeeper creates a new contract Keeper instance
@@ -107,17 +116,18 @@ func NewKeeper(
 	}
 
 	keeper := &Keeper{
-		storeKey:         storeKey,
-		cdc:              cdc,
-		wasmVM:           wasmer,
-		accountKeeper:    accountKeeper,
-		bank:             NewBankCoinTransferrer(bankKeeper),
-		portKeeper:       portKeeper,
-		capabilityKeeper: capabilityKeeper,
-		messenger:        NewDefaultMessageHandler(router, channelKeeper, capabilityKeeper, bankKeeper, cdc, portSource),
-		queryGasLimit:    wasmConfig.SmartQueryGasLimit,
-		paramSpace:       paramSpace,
-		gasRegister:      NewDefaultWasmGasRegister(),
+		storeKey:          storeKey,
+		cdc:               cdc,
+		wasmVM:            wasmer,
+		accountKeeper:     accountKeeper,
+		bank:              NewBankCoinTransferrer(bankKeeper),
+		portKeeper:        portKeeper,
+		capabilityKeeper:  capabilityKeeper,
+		messenger:         NewDefaultMessageHandler(router, channelKeeper, capabilityKeeper, bankKeeper, cdc, portSource),
+		queryGasLimit:     wasmConfig.SmartQueryGasLimit,
+		paramSpace:        paramSpace,
+		gasRegister:       NewDefaultWasmGasRegister(),
+		maxQueryStackSize: types.DefaultMaxQueryStackSize,
 	}
 	keeper.wasmVMQueryHandler = DefaultQueryPlugins(bankKeeper, stakingKeeper, distKeeper, channelKeeper, queryRouter, keeper)
 	for _, o := range opts {
@@ -598,6 +608,13 @@ func (k Keeper) getLastContractHistoryEntry(ctx sdk.Context, contractAddr sdk.Ac
 // QuerySmart queries the smart contract itself.
 func (k Keeper) QuerySmart(ctx sdk.Context, contractAddr sdk.AccAddress, req []byte) ([]byte, error) {
 	defer telemetry.MeasureSince(time.Now(), "wasm", "contract", "query-smart")
+
+	// checks and increase query stack size
+	ctx, err := checkAndIncreaseQueryStackSize(ctx, k.maxQueryStackSize)
+	if err != nil {
+		return nil, err
+	}
+
 	contractInfo, codeInfo, prefixStore, err := k.contractInstance(ctx, contractAddr)
 	if err != nil {
 		return nil, err
@@ -618,6 +635,30 @@ func (k Keeper) QuerySmart(ctx sdk.Context, contractAddr sdk.AccAddress, req []b
 	return queryResult, nil
 }
 
+func checkAndIncreaseQueryStackSize(ctx sdk.Context, maxQueryStackSize uint32) (sdk.Context, error) {
+	var queryStackSize uint32
+
+	// read current value
+	if size := ctx.Context().Value(contextKeyQueryStackSize); size != nil {
+		queryStackSize = size.(uint32)
+	} else {
+		queryStackSize = 0
+	}
+
+	// increase
+	queryStackSize++
+
+	// did we go too far?
+	if queryStackSize > maxQueryStackSize {
+		return ctx, types.ErrExceedMaxQueryStackSize
+	}
+
+	// set updated stack size
+	ctx = ctx.WithContext(context.WithValue(ctx.Context(), contextKeyQueryStackSize, queryStackSize))
+
+	return ctx, nil
+}
+
 // QueryRaw returns the contract's state for give key. Returns `nil` when key is `nil`.
 func (k Keeper) QueryRaw(ctx sdk.Context, contractAddress sdk.AccAddress, key []byte) []byte {
 	defer telemetry.MeasureSince(time.Now(), "wasm", "contract", "query-raw")

x/wasm/keeper/options.go

@@ -116,3 +116,10 @@ func WithAPICosts(human, canonical uint64) Option {
 		costCanonical = canonical
 	})
 }
+
+// WithMaxQueryStackSize overwrites the default limit for maximum query stacks
+func WithMaxQueryStackSize(m uint32) Option {
+	return optsFn(func(k *Keeper) {
+		k.maxQueryStackSize = m
+	})
+}

x/wasm/keeper/options_test.go

@@ -75,6 +75,12 @@ func TestConstructorOptions(t *testing.T) {
 				assert.Equal(t, uint64(2), costCanonical)
 			},
 		},
+		"max recursion query limit": {
+			srcOpt: WithMaxQueryStackSize(1),
+			verify: func(t *testing.T, k Keeper) {
+				assert.IsType(t, uint32(1), k.maxQueryStackSize)
+			},
+		},
 	}
 	for name, spec := range specs {
 		t.Run(name, func(t *testing.T) {

x/wasm/keeper/recurse_test.go

@@ -227,6 +227,7 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 		expectQueriesFromContract int
 		expectedGas               uint64
 		expectOutOfGas            bool
+		expectError               string
 	}{
 		"no recursion, lots of work": {
 			gasLimit: 4_000_000,
@@ -259,6 +260,17 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 			expectQueriesFromContract: 4,
 			expectOutOfGas:            true,
 		},
+		"very deep recursion, hits recursion limit": {
+			gasLimit: 10_000_000,
+			msg: Recurse{
+				Depth: 100,
+				Work:  2000,
+			},
+			expectQueriesFromContract: 10,
+			expectOutOfGas:            false,
+			expectError:               "query wasm contract failed", // Error we get from the contract instance doing the failing query, not wasmd
+			expectedGas:               10*(GasWork2k+GasReturnHashed) - 264,
+		},
 	}
 
 	contractAddr, _, ctx, keeper := initRecurseContract(t)
@@ -289,7 +301,11 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 
 			// otherwise, we expect a successful call
 			_, err := keeper.QuerySmart(ctx, contractAddr, msg)
-			require.NoError(t, err)
+			if tc.expectError != "" {
+				require.ErrorContains(t, err, tc.expectError)
+			} else {
+				require.NoError(t, err)
+			}
 			if types.EnableGasVerification {
 				assert.Equal(t, tc.expectedGas, ctx.GasMeter().GasConsumed())
 			}

x/wasm/types/errors.go

@@ -84,6 +84,9 @@ var (
 
 	// ErrTopKevelKeyNotAllowed error if a JSON object has a top-level key that is not allowed
 	ErrTopKevelKeyNotAllowed = sdkErrors.Register(DefaultCodespace, 26, "top-level key is not allowed")
+
+	// ErrExceedMaxQueryStackSize error if max query stack size is exceeded
+	ErrExceedMaxQueryStackSize = sdkErrors.Register(DefaultCodespace, 27, "max query stack size exceeded")
 )
 
 type ErrNoSuchContract struct {

x/wasm/types/wasmer_engine.go

@@ -5,6 +5,9 @@ import (
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
 )
 
+// DefaultMaxQueryStackSize maximum size of the stack of contract instances doing queries
+const DefaultMaxQueryStackSize uint32 = 10
+
 // WasmerEngine defines the WASM contract runtime engine.
 type WasmerEngine interface {
 	// Create will compile the wasm code, and store the resulting pre-compile