Allow better override of wasmVM in x/wasm keeper (#1494)

* Allow better override of wasmVM in x/wasm keeper

* Add post vm options

---------

Co-authored-by: Alex Peters <alpe@users.noreply.github.com>

x/wasm/keeper/keeper_cgo.go

@@ -35,15 +35,10 @@ func NewKeeper(
 	authority string,
 	opts ...Option,
 ) Keeper {
-	wasmer, err := wasmvm.NewVM(filepath.Join(homeDir, "wasm"), availableCapabilities, contractMemoryLimit, wasmConfig.ContractDebugMode, wasmConfig.MemoryCacheSize)
-	if err != nil {
-		panic(err)
-	}
-
 	keeper := &Keeper{
 		storeKey:             storeKey,
 		cdc:                  cdc,
-		wasmVM:               wasmer,
+		wasmVM:               nil,
 		accountKeeper:        accountKeeper,
 		bank:                 NewBankCoinTransferrer(bankKeeper),
 		accountPruner:        NewVestingCoinBurner(bankKeeper),
@@ -60,10 +55,24 @@ func NewKeeper(
 		authority: authority,
 	}
 	keeper.wasmVMQueryHandler = DefaultQueryPlugins(bankKeeper, stakingKeeper, distrKeeper, channelKeeper, keeper)
-	for _, o := range opts {
+	preOpts, postOpts := splitOpts(opts)
+	for _, o := range preOpts {
+		o.apply(keeper)
+	}
+	// 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 {
+		var err error
+		keeper.wasmVM, err = wasmvm.NewVM(filepath.Join(homeDir, "wasm"), availableCapabilities, contractMemoryLimit, wasmConfig.ContractDebugMode, wasmConfig.MemoryCacheSize)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	for _, o := range postOpts {
 		o.apply(keeper)
 	}
-	// not updateable, yet
+	// not updatable, yet
 	keeper.wasmVMResponseHandler = NewDefaultWasmVMContractResponseHandler(NewMessageDispatcher(keeper.messenger, keeper))
 	return *keeper
 }

x/wasm/keeper/options.go

@@ -16,6 +16,13 @@ func (f optsFn) apply(keeper *Keeper) {
 	f(keeper)
 }
 
+// option that is applied after keeper is setup with the VM. Used for decorators mainly.
+type postOptsFn func(*Keeper)
+
+func (f postOptsFn) apply(keeper *Keeper) {
+	f(keeper)
+}
+
 // WithWasmEngine is an optional constructor parameter to replace the default wasmVM engine with the
 // given one.
 func WithWasmEngine(x types.WasmerEngine) Option {
@@ -26,7 +33,7 @@ func WithWasmEngine(x types.WasmerEngine) Option {
 
 // WithWasmEngineDecorator is an optional constructor parameter to decorate the default wasmVM engine.
 func WithWasmEngineDecorator(d func(old types.WasmerEngine) types.WasmerEngine) Option {
-	return optsFn(func(k *Keeper) {
+	return postOptsFn(func(k *Keeper) {
 		k.wasmVM = d(k.wasmVM)
 	})
 }
@@ -42,7 +49,7 @@ func WithMessageHandler(x Messenger) Option {
 // WithMessageHandlerDecorator is an optional constructor parameter to decorate the wasm handler for wasmVM messages.
 // This option should not be combined with Option `WithMessageEncoders` or `WithMessageHandler`
 func WithMessageHandlerDecorator(d func(old Messenger) Messenger) Option {
-	return optsFn(func(k *Keeper) {
+	return postOptsFn(func(k *Keeper) {
 		k.messenger = d(k.messenger)
 	})
 }
@@ -58,7 +65,7 @@ func WithQueryHandler(x WasmVMQueryHandler) Option {
 // WithQueryHandlerDecorator is an optional constructor parameter to decorate the default wasm query handler for wasmVM requests.
 // This option should not be combined with Option `WithQueryPlugins` or `WithQueryHandler`
 func WithQueryHandlerDecorator(d func(old WasmVMQueryHandler) WasmVMQueryHandler) Option {
-	return optsFn(func(k *Keeper) {
+	return postOptsFn(func(k *Keeper) {
 		k.wasmVMQueryHandler = d(k.wasmVMQueryHandler)
 	})
 }
@@ -189,3 +196,16 @@ func asTypeMap(accts []authtypes.AccountI) map[reflect.Type]struct{} {
 	}
 	return m
 }
+
+// split into pre and post VM operations
+func splitOpts(opts []Option) ([]Option, []Option) {
+	pre, post := make([]Option, 0), make([]Option, 0)
+	for _, o := range opts {
+		if _, ok := o.(postOptsFn); ok {
+			post = append(post, o)
+		} else {
+			pre = append(pre, o)
+		}
+	}
+	return pre, post
+}

x/wasm/keeper/options_test.go

@@ -20,8 +20,9 @@ import (
 
 func TestConstructorOptions(t *testing.T) {
 	specs := map[string]struct {
-		srcOpt Option
-		verify func(*testing.T, Keeper)
+		srcOpt    Option
+		verify    func(*testing.T, Keeper)
+		isPostOpt bool
 	}{
 		"wasm engine": {
 			srcOpt: WithWasmEngine(&wasmtesting.MockWasmer{}),
@@ -37,6 +38,7 @@ func TestConstructorOptions(t *testing.T) {
 			verify: func(t *testing.T, k Keeper) {
 				assert.IsType(t, &wasmtesting.MockWasmer{}, k.wasmVM)
 			},
+			isPostOpt: true,
 		},
 		"message handler": {
 			srcOpt: WithMessageHandler(&wasmtesting.MockMessageHandler{}),
@@ -58,6 +60,7 @@ func TestConstructorOptions(t *testing.T) {
 			verify: func(t *testing.T, k Keeper) {
 				assert.IsType(t, &wasmtesting.MockMessageHandler{}, k.messenger)
 			},
+			isPostOpt: true,
 		},
 		"query plugins decorator": {
 			srcOpt: WithQueryHandlerDecorator(func(old WasmVMQueryHandler) WasmVMQueryHandler {
@@ -67,6 +70,7 @@ func TestConstructorOptions(t *testing.T) {
 			verify: func(t *testing.T, k Keeper) {
 				assert.IsType(t, &wasmtesting.MockQueryHandler{}, k.wasmVMQueryHandler)
 			},
+			isPostOpt: true,
 		},
 		"coin transferrer": {
 			srcOpt: WithCoinTransferrer(&wasmtesting.MockCoinTransferrer{}),
@@ -123,7 +127,10 @@ func TestConstructorOptions(t *testing.T) {
 	}
 	for name, spec := range specs {
 		t.Run(name, func(t *testing.T) {
-			k := NewKeeper(nil, nil, authkeeper.AccountKeeper{}, &bankkeeper.BaseKeeper{}, stakingkeeper.Keeper{}, nil, nil, nil, nil, nil, nil, nil, nil, "tempDir", types.DefaultWasmConfig(), AvailableCapabilities, "", spec.srcOpt)
+			opt := spec.srcOpt
+			_, gotPostOptMarker := opt.(postOptsFn)
+			require.Equal(t, spec.isPostOpt, gotPostOptMarker)
+			k := NewKeeper(nil, nil, authkeeper.AccountKeeper{}, &bankkeeper.BaseKeeper{}, stakingkeeper.Keeper{}, nil, nil, nil, nil, nil, nil, nil, nil, "tempDir", types.DefaultWasmConfig(), AvailableCapabilities, "", opt)
 			spec.verify(t, k)
 		})
 	}
@@ -133,3 +140,32 @@ func setAPIDefaults() {
 	costHumanize = DefaultGasCostHumanAddress * DefaultGasMultiplier
 	costCanonical = DefaultGasCostCanonicalAddress * DefaultGasMultiplier
 }
+
+func TestSplitOpts(t *testing.T) {
+	a := optsFn(nil)
+	b := optsFn(nil)
+	c := postOptsFn(nil)
+	d := postOptsFn(nil)
+	specs := map[string]struct {
+		src             []Option
+		expPre, expPost []Option
+	}{
+		"by type": {
+			src:     []Option{a, c},
+			expPre:  []Option{a},
+			expPost: []Option{c},
+		},
+		"ordered": {
+			src:     []Option{a, b, c, d},
+			expPre:  []Option{a, b},
+			expPost: []Option{c, d},
+		},
+	}
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			gotPre, gotPost := splitOpts(spec.src)
+			assert.Equal(t, spec.expPre, gotPre)
+			assert.Equal(t, spec.expPost, gotPost)
+		})
+	}
+}