Refactor wasmbinding to use message decorator

app/app.go

@@ -474,7 +474,7 @@ func New(
 	// The last arguments can contain custom message handlers, and custom query handlers,
 	// if we want to allow any custom callbacks
 	supportedFeatures := "iterator,staking,stargate,sei"
-	wasmOpts = append(wasmbinding.RegisterCustomPlugins(&app.OracleKeeper, &app.DexKeeper, &app.EpochKeeper, &app.TokenFactoryKeeper), wasmOpts...)
+	wasmOpts = append(wasmbinding.RegisterCustomPlugins(&app.OracleKeeper, &app.DexKeeper, &app.EpochKeeper, &app.TokenFactoryKeeper, &app.AccountKeeper, app.MsgServiceRouter()), wasmOpts...)
 	app.WasmKeeper = wasm.NewKeeper(
 		appCodec,
 		keys[wasm.StoreKey],

cmd/seid/cmd/root.go

@@ -7,7 +7,6 @@ import (
 	"path/filepath"
 
 	"github.com/CosmWasm/wasmd/x/wasm"
-	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/config"
@@ -29,7 +28,6 @@ import (
 	genutilcli "github.com/cosmos/cosmos-sdk/x/genutil/client/cli"
 	"github.com/sei-protocol/sei-chain/app"
 	"github.com/sei-protocol/sei-chain/app/params"
-	"github.com/sei-protocol/sei-chain/wasmbinding"
 	"github.com/spf13/cast"
 	"github.com/spf13/cobra"
 	tmcli "github.com/tendermint/tendermint/libs/cli"
@@ -233,10 +231,6 @@ func newApp(
 		panic(err)
 	}
 
-	wasmopts := []wasm.Option{wasmkeeper.WithMessageEncoders(&wasmkeeper.MessageEncoders{
-		Custom: wasmbinding.CustomEncoder,
-	})}
-
 	return app.New(
 		logger,
 		db,
@@ -248,7 +242,7 @@ func newApp(
 		app.MakeEncodingConfig(),
 		wasm.EnableAllProposals,
 		appOpts,
-		wasmopts,
+		[]wasm.Option{},
 		baseapp.SetPruning(pruningOpts),
 		baseapp.SetMinGasPrices(cast.ToString(appOpts.Get(server.FlagMinGasPrices))),
 		baseapp.SetMinRetainBlocks(cast.ToUint64(appOpts.Get(server.FlagMinRetainBlocks))),

wasmbinding/message_plugin.go

@@ -0,0 +1,128 @@
+package wasmbinding
+
+import (
+	"encoding/json"
+
+	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
+	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
+	dexwasm "github.com/sei-protocol/sei-chain/x/dex/client/wasm"
+	tokenfactorywasm "github.com/sei-protocol/sei-chain/x/tokenfactory/client/wasm"
+)
+
+// CustomMessageDecorator returns decorator for custom CosmWasm bindings messages
+func CustomMessageDecorator(
+	router wasmkeeper.MessageRouter,
+	accountKeeper *authkeeper.AccountKeeper,
+) func(wasmkeeper.Messenger) wasmkeeper.Messenger {
+	return func(old wasmkeeper.Messenger) wasmkeeper.Messenger {
+		return &CustomMessenger{
+			router:        router,
+			wrapped:       old,
+			accountKeeper: accountKeeper,
+		}
+	}
+}
+
+type CustomMessenger struct {
+	router        wasmkeeper.MessageRouter
+	wrapped       wasmkeeper.Messenger
+	accountKeeper *authkeeper.AccountKeeper
+}
+
+type SeiWasmMessage struct {
+	PlaceOrders  json.RawMessage `json:"place_orders,omitempty"`
+	CancelOrders json.RawMessage `json:"cancel_orders,omitempty"`
+	CreateDenom  json.RawMessage `json:"create_denom,omitempty"`
+	MintTokens   json.RawMessage `json:"mint_tokens,omitempty"`
+	BurnTokens   json.RawMessage `json:"burn_tokens,omitempty"`
+	ChangeAdmin  json.RawMessage `json:"change_admin,omitempty"`
+}
+
+var _ wasmkeeper.Messenger = &CustomMessenger{}
+
+// DispatchMsg executes on the bindingMsgs
+func (m *CustomMessenger) DispatchMsg(ctx sdk.Context, contractAddr sdk.AccAddress, contractIBCPortID string, msg wasmvmtypes.CosmosMsg) ([]sdk.Event, [][]byte, error) {
+	if msg.Custom != nil {
+		return m.DispatchCustomMsg(ctx, contractAddr, contractIBCPortID, msg)
+	}
+	return m.wrapped.DispatchMsg(ctx, contractAddr, contractIBCPortID, msg)
+}
+
+// DispatchCustomMsg function is forked from wasmd. sdk.Msg will be validated and routed to the corresponding module msg server in this function.
+func (m *CustomMessenger) DispatchCustomMsg(
+	ctx sdk.Context,
+	contractAddr sdk.AccAddress,
+	contractIBCPortID string,
+	msg wasmvmtypes.CosmosMsg,
+) (events []sdk.Event, data [][]byte, err error) {
+	var parsedMessage SeiWasmMessage
+	if err := json.Unmarshal(msg.Custom, &parsedMessage); err != nil {
+		return nil, nil, sdkerrors.Wrap(err, "Error parsing Sei Wasm Message")
+	}
+
+	var sdkMsgs []sdk.Msg
+	switch {
+	case parsedMessage.PlaceOrders != nil:
+		sdkMsgs, err = dexwasm.EncodeDexPlaceOrders(parsedMessage.PlaceOrders, contractAddr)
+	case parsedMessage.CancelOrders != nil:
+		sdkMsgs, err = dexwasm.EncodeDexCancelOrders(parsedMessage.CancelOrders, contractAddr)
+	case parsedMessage.CreateDenom != nil:
+		sdkMsgs, err = tokenfactorywasm.EncodeTokenFactoryCreateDenom(parsedMessage.CreateDenom, contractAddr)
+	case parsedMessage.MintTokens != nil:
+		sdkMsgs, err = tokenfactorywasm.EncodeTokenFactoryMint(parsedMessage.MintTokens, contractAddr)
+	case parsedMessage.BurnTokens != nil:
+		sdkMsgs, err = tokenfactorywasm.EncodeTokenFactoryBurn(parsedMessage.BurnTokens, contractAddr)
+	case parsedMessage.ChangeAdmin != nil:
+		sdkMsgs, err = tokenfactorywasm.EncodeTokenFactoryChangeAdmin(parsedMessage.ChangeAdmin, contractAddr)
+	default:
+		sdkMsgs, err = []sdk.Msg{}, wasmvmtypes.UnsupportedRequest{Kind: "Unknown Sei Wasm Message"}
+	}
+	if err != nil {
+		return nil, nil, err
+	}
+
+	for _, sdkMsg := range sdkMsgs {
+		res, err := m.handleSdkMessage(ctx, contractAddr, sdkMsg)
+		if err != nil {
+			return nil, nil, err
+		}
+		// append data
+		data = append(data, res.Data)
+		// append events
+		sdkEvents := make([]sdk.Event, len(res.Events))
+		for i := range res.Events {
+			sdkEvents[i] = sdk.Event(res.Events[i])
+		}
+		events = append(events, sdkEvents...)
+	}
+	return events, data, nil
+}
+
+// This function is forked from wasmd. sdk.Msg will be validated and routed to the corresponding module msg server in this function.
+func (m *CustomMessenger) handleSdkMessage(ctx sdk.Context, contractAddr sdk.Address, msg sdk.Msg) (*sdk.Result, error) {
+	if err := msg.ValidateBasic(); err != nil {
+		return nil, err
+	}
+	// make sure this account can send it
+	for _, acct := range msg.GetSigners() {
+		if !acct.Equals(contractAddr) {
+			return nil, sdkerrors.Wrap(sdkerrors.ErrUnauthorized, "contract doesn't have permission")
+		}
+	}
+
+	// find the handler and execute it
+	if handler := m.router.Handler(msg); handler != nil {
+		// ADR 031 request type routing
+		msgResult, err := handler(ctx, msg)
+		return msgResult, err
+	}
+	// legacy sdk.Msg routing
+	// Assuming that the app developer has migrated all their Msgs to
+	// proto messages and has registered all `Msg services`, then this
+	// path should never be called, because all those Msgs should be
+	// registered within the `msgServiceRouter` already.
+	return nil, sdkerrors.Wrapf(sdkerrors.ErrUnknownRequest, "can't route message %+v", msg)
+}

wasmbinding/test/encoder_test.go

@@ -5,10 +5,11 @@ import (
 	"testing"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/sei-protocol/sei-chain/wasmbinding"
 	"github.com/sei-protocol/sei-chain/wasmbinding/bindings"
+	dexwasm "github.com/sei-protocol/sei-chain/x/dex/client/wasm"
 	"github.com/sei-protocol/sei-chain/x/dex/types"
 	dextypes "github.com/sei-protocol/sei-chain/x/dex/types"
+	tokenfactorywasm "github.com/sei-protocol/sei-chain/x/tokenfactory/client/wasm"
 	tokenfactorytypes "github.com/sei-protocol/sei-chain/x/tokenfactory/types"
 	"github.com/stretchr/testify/require"
 )
@@ -36,13 +37,9 @@ func TestEncodePlaceOrder(t *testing.T) {
 		Funds:        []sdk.Coin{fund},
 		ContractAddr: TEST_TARGET_CONTRACT,
 	}
-	serialized, _ := json.Marshal(msg)
-	msgData := wasmbinding.SeiWasmMessage{
-		PlaceOrders: serialized,
-	}
-	serializedMsg, _ := json.Marshal(msgData)
+	serializedMsg, _ := json.Marshal(msg)
 
-	decodedMsgs, err := wasmbinding.CustomEncoder(contractAddr, serializedMsg)
+	decodedMsgs, err := dexwasm.EncodeDexPlaceOrders(serializedMsg, contractAddr)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(decodedMsgs))
 	typedDecodedMsg, ok := decodedMsgs[0].(*dextypes.MsgPlaceOrders)
@@ -65,13 +62,9 @@ func TestDecodeOrderCancellation(t *testing.T) {
 		},
 		ContractAddr: TEST_TARGET_CONTRACT,
 	}
-	serialized, _ := json.Marshal(msg)
-	msgData := wasmbinding.SeiWasmMessage{
-		CancelOrders: serialized,
-	}
-	serializedMsg, _ := json.Marshal(msgData)
+	serializedMsg, _ := json.Marshal(msg)
 
-	decodedMsgs, err := wasmbinding.CustomEncoder(contractAddr, serializedMsg)
+	decodedMsgs, err := dexwasm.EncodeDexCancelOrders(serializedMsg, contractAddr)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(decodedMsgs))
 	typedDecodedMsg, ok := decodedMsgs[0].(*dextypes.MsgCancelOrders)
@@ -94,13 +87,9 @@ func TestEncodeCreateDenom(t *testing.T) {
 	msg := bindings.CreateDenom{
 		Subdenom: "subdenom",
 	}
-	serialized, _ := json.Marshal(msg)
-	msgData := wasmbinding.SeiWasmMessage{
-		CreateDenom: serialized,
-	}
-	serializedMsg, _ := json.Marshal(msgData)
+	serializedMsg, _ := json.Marshal(msg)
 
-	decodedMsgs, err := wasmbinding.CustomEncoder(contractAddr, serializedMsg)
+	decodedMsgs, err := tokenfactorywasm.EncodeTokenFactoryCreateDenom(serializedMsg, contractAddr)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(decodedMsgs))
 	typedDecodedMsg, ok := decodedMsgs[0].(*tokenfactorytypes.MsgCreateDenom)
@@ -118,13 +107,9 @@ func TestEncodeMint(t *testing.T) {
 	msg := bindings.MintTokens{
 		Amount: sdk.Coin{Amount: sdk.NewInt(100), Denom: "subdenom"},
 	}
-	serialized, _ := json.Marshal(msg)
-	msgData := wasmbinding.SeiWasmMessage{
-		MintTokens: serialized,
-	}
-	serializedMsg, _ := json.Marshal(msgData)
+	serializedMsg, _ := json.Marshal(msg)
 
-	decodedMsgs, err := wasmbinding.CustomEncoder(contractAddr, serializedMsg)
+	decodedMsgs, err := tokenfactorywasm.EncodeTokenFactoryMint(serializedMsg, contractAddr)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(decodedMsgs))
 	typedDecodedMsg, ok := decodedMsgs[0].(*tokenfactorytypes.MsgMint)
@@ -142,13 +127,9 @@ func TestEncodeBurn(t *testing.T) {
 	msg := bindings.BurnTokens{
 		Amount: sdk.Coin{Amount: sdk.NewInt(10), Denom: "subdenom"},
 	}
-	serialized, _ := json.Marshal(msg)
-	msgData := wasmbinding.SeiWasmMessage{
-		BurnTokens: serialized,
-	}
-	serializedMsg, _ := json.Marshal(msgData)
+	serializedMsg, _ := json.Marshal(msg)
 
-	decodedMsgs, err := wasmbinding.CustomEncoder(contractAddr, serializedMsg)
+	decodedMsgs, err := tokenfactorywasm.EncodeTokenFactoryBurn(serializedMsg, contractAddr)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(decodedMsgs))
 	typedDecodedMsg, ok := decodedMsgs[0].(*tokenfactorytypes.MsgBurn)
@@ -167,13 +148,9 @@ func TestEncodeChangeAdmin(t *testing.T) {
 		Denom:           "factory/sei1y3pxq5dp900czh0mkudhjdqjq5m8cpmmps8yjw/subdenom",
 		NewAdminAddress: "sei1hjfwcza3e3uzeznf3qthhakdr9juetl7g6esl4",
 	}
-	serialized, _ := json.Marshal(msg)
-	msgData := wasmbinding.SeiWasmMessage{
-		ChangeAdmin: serialized,
-	}
-	serializedMsg, _ := json.Marshal(msgData)
+	serializedMsg, _ := json.Marshal(msg)
 
-	decodedMsgs, err := wasmbinding.CustomEncoder(contractAddr, serializedMsg)
+	decodedMsgs, err := tokenfactorywasm.EncodeTokenFactoryChangeAdmin(serializedMsg, contractAddr)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(decodedMsgs))
 	typedDecodedMsg, ok := decodedMsgs[0].(*tokenfactorytypes.MsgChangeAdmin)

wasmbinding/wasm.go

@@ -3,6 +3,7 @@ package wasmbinding
 import (
 	"github.com/CosmWasm/wasmd/x/wasm"
 	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
+	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	dexwasm "github.com/sei-protocol/sei-chain/x/dex/client/wasm"
 	dexkeeper "github.com/sei-protocol/sei-chain/x/dex/keeper"
 	epochwasm "github.com/sei-protocol/sei-chain/x/epoch/client/wasm"
@@ -18,6 +19,8 @@ func RegisterCustomPlugins(
 	dex *dexkeeper.Keeper,
 	epoch *epochkeeper.Keeper,
 	tokenfactory *tokenfactorykeeper.Keeper,
+	accountKeeper *authkeeper.AccountKeeper,
+	router wasmkeeper.MessageRouter,
 ) []wasmkeeper.Option {
 	dexHandler := dexwasm.NewDexWasmQueryHandler(dex)
 	oracleHandler := oraclewasm.NewOracleWasmQueryHandler(oracle)
@@ -28,8 +31,12 @@ func RegisterCustomPlugins(
 	queryPluginOpt := wasmkeeper.WithQueryPlugins(&wasmkeeper.QueryPlugins{
 		Custom: CustomQuerier(wasmQueryPlugin),
 	})
+	messengerDecoratorOpt := wasmkeeper.WithMessageHandlerDecorator(
+		CustomMessageDecorator(router, accountKeeper),
+	)
 
 	return []wasm.Option{
 		queryPluginOpt,
+		messengerDecoratorOpt,
 	}
 }