Skip to content

Commit

Permalink
Fix
Browse files Browse the repository at this point in the history
  • Loading branch information
chipshort committed Aug 21, 2024
1 parent 456af0e commit 71cf6a8
Show file tree
Hide file tree
Showing 8 changed files with 73 additions and 5 deletions.
15 changes: 15 additions & 0 deletions x/wasm/keeper/handler_plugin.go
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,21 @@ func (h SDKMessageHandler) handleSdkMessage(ctx sdk.Context, contractAddr sdk.Ad
return nil, errorsmod.Wrapf(sdkerrors.ErrUnknownRequest, "can't route message %+v", msg)
}

// callDepthMessageHandler is a wrapper around a Messenger that checks the call depth before dispatching a message.
type callDepthMessageHandler struct {
Messenger
MaxCallDepth uint32
}

func (h callDepthMessageHandler) DispatchMsg(ctx sdk.Context, contractAddr sdk.AccAddress, contractIBCPortID string, msg wasmvmtypes.CosmosMsg) (events []sdk.Event, data [][]byte, err error) {
ctx, err = checkAndIncreaseCallDepth(ctx, h.MaxCallDepth)
if err != nil {
return nil, nil, err
}

return h.Messenger.DispatchMsg(ctx, contractAddr, contractIBCPortID, msg)
}

// MessageHandlerChain defines a chain of handlers that are called one by one until it can be handled.
type MessageHandlerChain struct {
handlers []Messenger
Expand Down
19 changes: 19 additions & 0 deletions x/wasm/keeper/keeper.go
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ type Keeper struct {
queryGasLimit uint64
gasRegister types.GasRegister
maxQueryStackSize uint32
maxCallDepth uint32
acceptedAccountTypes map[reflect.Type]struct{}
accountPruner AccountPruner
// propagate gov authZ to sub-messages
Expand Down Expand Up @@ -754,6 +755,24 @@ func checkAndIncreaseQueryStackSize(ctx sdk.Context, maxQueryStackSize uint32) (
return types.WithQueryStackSize(ctx, queryStackSize), nil
}

func checkAndIncreaseCallDepth(ctx sdk.Context, maxCallDepth uint32) (sdk.Context, error) {
var callDepth uint32 = 0
if size, ok := types.CallDepth(ctx); ok {
callDepth = size
}

// increase
callDepth++

// did we go too far?
if callDepth > maxCallDepth {
return sdk.Context{}, types.ErrExceedMaxCallDepth
}

// set updated stack size
return types.WithCallDepth(sdk.UnwrapSDKContext(ctx), callDepth), 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")
Expand Down
3 changes: 3 additions & 0 deletions x/wasm/keeper/keeper_cgo.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ func NewKeeper(
queryGasLimit: wasmConfig.SmartQueryGasLimit,
gasRegister: types.NewDefaultWasmGasRegister(),
maxQueryStackSize: types.DefaultMaxQueryStackSize,
maxCallDepth: types.DefaultMaxCallDepth,
acceptedAccountTypes: defaultAcceptedAccountTypes,
propagateGovAuthorization: map[types.AuthorizationPolicyAction]struct{}{
types.AuthZActionInstantiate: {},
Expand All @@ -59,6 +60,8 @@ func NewKeeper(
for _, o := range preOpts {
o.apply(keeper)
}
// always wrap the messenger, even if it was replaced by an option
keeper.messenger = callDepthMessageHandler{keeper.messenger, keeper.maxCallDepth}
// only set the wasmvm if no one set this in the options
// NewVM does a lot, so better not to create it and silently drop it.
if keeper.wasmVM == nil {
Expand Down
6 changes: 6 additions & 0 deletions x/wasm/keeper/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,12 @@ func WithMaxQueryStackSize(m uint32) Option {
})
}

func WithMaxCallDepth(m uint32) Option {
return optsFn(func(k *Keeper) {
k.maxCallDepth = m
})
}

// WithAcceptedAccountTypesOnContractInstantiation sets the accepted account types. Account types of this list won't be overwritten or cause a failure
// when they exist for an address on contract instantiation.
//
Expand Down
15 changes: 11 additions & 4 deletions x/wasm/keeper/options_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,9 @@ func TestConstructorOptions(t *testing.T) {
"message handler": {
srcOpt: WithMessageHandler(&wasmtesting.MockMessageHandler{}),
verify: func(t *testing.T, k Keeper) {
t.Helper()
assert.IsType(t, &wasmtesting.MockMessageHandler{}, k.messenger)
require.IsType(t, callDepthMessageHandler{}, k.messenger)
messenger, _ := k.messenger.(callDepthMessageHandler)
assert.IsType(t, &wasmtesting.MockMessageHandler{}, messenger.Messenger)
},
},
"query plugins": {
Expand All @@ -68,7 +69,7 @@ func TestConstructorOptions(t *testing.T) {
},
"message handler decorator": {
srcOpt: WithMessageHandlerDecorator(func(old Messenger) Messenger {
require.IsType(t, &MessageHandlerChain{}, old)
require.IsType(t, callDepthMessageHandler{}, old)
return &wasmtesting.MockMessageHandler{}
}),
verify: func(t *testing.T, k Keeper) {
Expand Down Expand Up @@ -111,13 +112,19 @@ func TestConstructorOptions(t *testing.T) {
assert.Equal(t, uint64(2), costCanonical)
},
},
"max recursion query limit": {
"max query recursion limit": {
srcOpt: WithMaxQueryStackSize(1),
verify: func(t *testing.T, k Keeper) {
t.Helper()
assert.IsType(t, uint32(1), k.maxQueryStackSize)
},
},
"max message recursion limit": {
srcOpt: WithMaxCallDepth(1),
verify: func(t *testing.T, k Keeper) {
assert.IsType(t, uint32(1), k.maxCallDepth)
},
},
"accepted account types": {
srcOpt: WithAcceptedAccountTypesOnContractInstantiation(&authtypes.BaseAccount{}, &vestingtypes.ContinuousVestingAccount{}),
verify: func(t *testing.T, k Keeper) {
Expand Down
11 changes: 11 additions & 0 deletions x/wasm/types/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ const (
contextKeySubMsgAuthzPolicy = iota
// gas register
contextKeyGasRegister = iota

contextKeyCallDepth contextKey = iota
)

// WithTXCounter stores a transaction counter value in the context
Expand All @@ -41,6 +43,15 @@ func QueryStackSize(ctx sdk.Context) (uint32, bool) {
return val, ok
}

func WithCallDepth(ctx sdk.Context, counter uint32) sdk.Context {
return ctx.WithValue(contextKeyCallDepth, counter)
}

func CallDepth(ctx sdk.Context) (uint32, bool) {
val, ok := ctx.Value(contextKeyCallDepth).(uint32)
return val, ok
}

// WithSubMsgAuthzPolicy stores the authorization policy for submessages into the context returned
func WithSubMsgAuthzPolicy(ctx sdk.Context, policy AuthorizationPolicy) sdk.Context {
if policy == nil {
Expand Down
5 changes: 5 additions & 0 deletions x/wasm/types/errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,11 @@ var (
ErrNoSuchCodeFn = WasmVMFlavouredErrorFactory(errorsmod.Register(DefaultCodespace, 28, "no such code"),
func(id uint64) error { return wasmvmtypes.NoSuchCode{CodeID: id} },
)

// code 29 reserved for wasmd 0.50+

// ErrExceedMaxCallDepth error if max message stack size is exceeded
ErrExceedMaxCallDepth = errorsmod.Register(DefaultCodespace, 30, "max call depth exceeded")
)

// WasmVMErrorable mapped error type in wasmvm and are not redacted
Expand Down
4 changes: 3 additions & 1 deletion x/wasm/types/wasmer_engine.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,11 @@ import (
sdk "github.com/cosmos/cosmos-sdk/types"
)

// DefaultMaxQueryStackSize maximum size of the stack of contract instances doing queries
// DefaultMaxQueryStackSize maximum size of the stack of recursive queries a contract can make
const DefaultMaxQueryStackSize uint32 = 10

const DefaultMaxCallDepth uint32 = 500

// WasmEngine defines the WASM contract runtime engine.
type WasmEngine interface {
// Create will compile the wasm code, and store the resulting pre-compile
Expand Down

0 comments on commit 71cf6a8

Please sign in to comment.