refactor impl message

demo/app/app.go

@@ -606,7 +606,7 @@ func NewMeshApp(
 	})
 	meshProvMessageHandler := wasmkeeper.WithMessageHandlerDecorator(func(nested wasmkeeper.Messenger) wasmkeeper.Messenger {
 		return wasmkeeper.NewMessageHandlerChain(
-			meshsecprovkeeper.NewDefaultCustomMsgHandler(&app.MeshSecProvKeeper),
+			meshsecprovkeeper.NewCustomMsgHandler(app.MeshSecProvKeeper),
 		)
 	})
 	wasmOpts = append(wasmOpts, meshMessageHandler, meshProvMessageHandler,

x/meshsecurityprovider/keeper/handler_plugin.go

@@ -14,39 +14,13 @@ import (
 	"github.com/osmosis-labs/mesh-security-sdk/x/meshsecurityprovider/types"
 )
 
-// AuthSource abstract type that provides contract authorization.
-// This is an extension point for custom implementations.
-type AuthSource interface {
-	// IsAuthorized returns if the contract authorized to execute a virtual stake message
-	IsAuthorized(ctx sdk.Context, contractAddr sdk.AccAddress) bool
-}
-
-// abstract keeper
-type msKeeper interface {
-	Bond(ctx sdk.Context, actor sdk.AccAddress, delegator sdk.AccAddress, coin sdk.Coin) error
-	Unbond(ctx sdk.Context, actor sdk.AccAddress, delegator sdk.AccAddress, coin sdk.Coin) error
-}
-
 type CustomMsgHandler struct {
-	k    msKeeper
-	auth AuthSource
+	k Keeper
 }
 
-// NewDefaultCustomMsgHandler constructor to set up the CustomMsgHandler with vault contract authorization
-func NewDefaultCustomMsgHandler(k *Keeper) *CustomMsgHandler {
-	return &CustomMsgHandler{k: k, auth: isVaultContract(k)}
-}
-
-func isVaultContract(k *Keeper) AuthSourceFn {
-	return func(ctx sdk.Context, contractAddr sdk.AccAddress) bool {
-		return k.VaultAddress(ctx) == contractAddr.String()
-	}
-}
-
-// NewCustomMsgHandler constructor to set up CustomMsgHandler with an individual auth source.
-// This is an extension point for non default contract authorization logic.
-func NewCustomMsgHandler(k msKeeper, auth AuthSource) *CustomMsgHandler {
-	return &CustomMsgHandler{k: k, auth: auth}
+// NewCustomMsgHandler constructor to set up CustomMsgHandler.
+func NewCustomMsgHandler(k Keeper) *CustomMsgHandler {
+	return &CustomMsgHandler{k: k}
 }
 
 // DispatchMsg handle contract message of type Custom in the mesh-security namespace
@@ -63,10 +37,6 @@ func (h CustomMsgHandler) DispatchMsg(ctx sdk.Context, contractAddr sdk.AccAddre
 		return nil, nil, wasmtypes.ErrUnknownMsg
 	}
 
-	if !h.auth.IsAuthorized(ctx, contractAddr) {
-		return nil, nil, sdkerrors.ErrUnauthorized.Wrapf("contract has no permission for mesh security operations")
-	}
-
 	switch {
 	case customMsg.ProviderMsg.Bond != nil:
 		return h.handleBondMsg(ctx, contractAddr, customMsg.ProviderMsg.Bond)
@@ -78,6 +48,10 @@ func (h CustomMsgHandler) DispatchMsg(ctx sdk.Context, contractAddr sdk.AccAddre
 }
 
 func (h CustomMsgHandler) handleBondMsg(ctx sdk.Context, actor sdk.AccAddress, bondMsg *contract.BondMsg) ([]sdk.Event, [][]byte, error) {
+	if actor.String() != h.k.VaultAddress(ctx) {
+		return nil, nil, sdkerrors.ErrUnauthorized.Wrapf("contract has no permission for mesh security operations")
+	}
+
 	coin, err := wasmkeeper.ConvertWasmCoinToSdkCoin(bondMsg.Amount)
 	if err != nil {
 		return nil, nil, err
@@ -102,6 +76,10 @@ func (h CustomMsgHandler) handleBondMsg(ctx sdk.Context, actor sdk.AccAddress, b
 }
 
 func (h CustomMsgHandler) handleUnbondMsg(ctx sdk.Context, actor sdk.AccAddress, unbondMsg *contract.UnbondMsg) ([]sdk.Event, [][]byte, error) {
+	if actor.String() != h.k.VaultAddress(ctx) {
+		return nil, nil, sdkerrors.ErrUnauthorized.Wrapf("contract has no permission for mesh security operations")
+	}
+
 	coin, err := wasmkeeper.ConvertWasmCoinToSdkCoin(unbondMsg.Amount)
 	if err != nil {
 		return nil, nil, err
@@ -124,11 +102,3 @@ func (h CustomMsgHandler) handleUnbondMsg(ctx sdk.Context, actor sdk.AccAddress,
 		sdk.NewAttribute(types.AttributeKeyDelegator, delAddr.String()),
 	)}, nil, nil
 }
-
-// AuthSourceFn is helper for simple AuthSource types
-type AuthSourceFn func(ctx sdk.Context, contractAddr sdk.AccAddress) bool
-
-// IsAuthorized returns if the contract authorized to execute a stake message
-func (a AuthSourceFn) IsAuthorized(ctx sdk.Context, contractAddr sdk.AccAddress) bool {
-	return a(ctx, contractAddr)
-}