feat: added execution stage for schedules #NTRN-339

app/proposals_allowlisting.go

@@ -75,6 +75,8 @@ func isSdkMessageWhitelisted(msg sdk.Msg) bool {
 		*feeburnertypes.MsgUpdateParams,
 		*feerefundertypes.MsgUpdateParams,
 		*crontypes.MsgUpdateParams,
+		*crontypes.MsgAddSchedule,
+		*crontypes.MsgRemoveSchedule,
 		*contractmanagertypes.MsgUpdateParams,
 		*dextypes.MsgUpdateParams,
 		*banktypes.MsgUpdateParams,

proto/neutron/cron/genesis.proto

@@ -8,7 +8,7 @@ import "neutron/cron/schedule.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 
-// GenesisState defines the cron module's genesis state.
+// Defines the cron module's genesis state.
 message GenesisState {
   repeated Schedule scheduleList = 2 [(gogoproto.nullable) = false];
   Params params = 1 [(gogoproto.nullable) = false];

proto/neutron/cron/params.proto

@@ -5,7 +5,7 @@ import "gogoproto/gogo.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 
-// Params defines the parameters for the module.
+// Defines the parameters for the module.
 message Params {
   option (gogoproto.goproto_stringer) = false;
   // Security address that can remove schedules

proto/neutron/cron/query.proto

@@ -10,7 +10,7 @@ import "neutron/cron/schedule.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 
-// Query defines the gRPC querier service.
+// Defines the gRPC querier service.
 service Query {
   // Queries the parameters of the module.
   rpc Params(QueryParamsRequest) returns (QueryParamsResponse) {
@@ -30,25 +30,31 @@ service Query {
   // this line is used by starport scaffolding # 2
 }
 
+// The request type for the Query/Params RPC method.
 message QueryParamsRequest {}
 
+// The response type for the Query/Params RPC method.
 message QueryParamsResponse {
   // params holds all the parameters of this module.
   Params params = 1 [(gogoproto.nullable) = false];
 }
 
+// The request type for the Query/Schedule RPC method.
 message QueryGetScheduleRequest {
   string name = 1;
 }
 
+// The response type for the Query/Params RPC method.
 message QueryGetScheduleResponse {
   Schedule schedule = 1 [(gogoproto.nullable) = false];
 }
 
+// The request type for the Query/Schedules RPC method.
 message QuerySchedulesRequest {
   cosmos.base.query.v1beta1.PageRequest pagination = 1;
 }
 
+// The response type for the Query/Params RPC method.
 message QuerySchedulesResponse {
   repeated Schedule schedules = 1 [(gogoproto.nullable) = false];
   cosmos.base.query.v1beta1.PageResponse pagination = 2;

proto/neutron/cron/schedule.proto

@@ -5,25 +5,38 @@ import "gogoproto/gogo.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 
+// Defines when messages will be executed in the block
+enum ExecutionStage {
+  // Execution at the end of the block
+  EXECUTION_STAGE_END_BLOCKER = 0;
+  // Execution at the beginning of the block
+  EXECUTION_STAGE_BEGIN_BLOCKER = 1;
+}
+
+// Defines the schedule for execution
 message Schedule {
   // Name of schedule
   string name = 1;
   // Period in blocks
   uint64 period = 2;
-  // Msgs that will be executed every period amount of time
+  // Msgs that will be executed every certain number of blocks, specified in the `period` field
   repeated MsgExecuteContract msgs = 3 [(gogoproto.nullable) = false];
   // Last execution's block height
   uint64 last_execute_height = 4;
+  // Stage when messages will be executed
+  ExecutionStage execution_stage = 5;
 }
 
+// Defines the contract and the message to pass
 message MsgExecuteContract {
-  // Contract is the address of the smart contract
+  // The address of the smart contract
   string contract = 1;
-  // Msg is json encoded message to be passed to the contract
+  // JSON encoded message to be passed to the contract
   string msg = 2;
 }
 
+// Defines the number of current schedules
 message ScheduleCount {
-  // Count is the number of current schedules
+  // The number of current schedules
   int32 count = 1;
 }

proto/neutron/cron/tx.proto

@@ -6,32 +6,72 @@ import "cosmos/msg/v1/msg.proto";
 import "cosmos_proto/cosmos.proto";
 import "gogoproto/gogo.proto";
 import "neutron/cron/params.proto";
+import "neutron/cron/schedule.proto";
 
 // this line is used by starport scaffolding # proto/tx/import
 
 option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 
-// Msg defines the Msg service.
+// Defines the Msg service.
 service Msg {
   option (cosmos.msg.v1.service) = true;
 
+  // Adds new schedule.
+  rpc AddSchedule(MsgAddSchedule) returns (MsgAddScheduleResponse);
+  // Removes schedule.
+  rpc RemoveSchedule(MsgRemoveSchedule) returns (MsgRemoveScheduleResponse);
+  // Updates the module parameters.
   rpc UpdateParams(MsgUpdateParams) returns (MsgUpdateParamsResponse);
   // this line is used by starport scaffolding # proto/tx/rpc
 }
 
+// The MsgAddSchedule request type.
+message MsgAddSchedule {
+  option (amino.name) = "cron/MsgAddSchedule";
+  option (cosmos.msg.v1.signer) = "authority";
+
+  // The address of the governance account.
+  string authority = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+  // Name of the schedule
+  string name = 2;
+  // Period in blocks
+  uint64 period = 3;
+  // Msgs that will be executed every certain number of blocks, specified in the `period` field
+  repeated MsgExecuteContract msgs = 4 [(gogoproto.nullable) = false];
+  // Stage when messages will be executed
+  ExecutionStage execution_stage = 5;
+}
+
+// Defines the response structure for executing a MsgAddSchedule message.
+message MsgAddScheduleResponse {}
+
+// The MsgRemoveSchedule request type.
+message MsgRemoveSchedule {
+  option (amino.name) = "cron/MsgRemoveSchedule";
+  option (cosmos.msg.v1.signer) = "authority";
+
+  // The address of the governance account.
+  string authority = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+  // Name of the schedule
+  string name = 2;
+}
+
+// Defines the response structure for executing a MsgRemoveSchedule message.
+message MsgRemoveScheduleResponse {}
+
 // this line is used by starport scaffolding # proto/tx/message
 
-// MsgUpdateParams is the MsgUpdateParams request type.
+// The MsgUpdateParams request type.
 //
 // Since: 0.47
 message MsgUpdateParams {
   option (amino.name) = "cron/MsgUpdateParams";
   option (cosmos.msg.v1.signer) = "authority";
 
-  // Authority is the address of the governance account.
+  // The address of the governance account.
   string authority = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
 
-  // params defines the x/cron parameters to update.
+  // Defines the x/cron parameters to update.
   //
   // NOTE: All parameters must be supplied.
   Params params = 2 [
@@ -40,8 +80,7 @@ message MsgUpdateParams {
   ];
 }
 
-// MsgUpdateParamsResponse defines the response structure for executing a
-// MsgUpdateParams message.
+// Defines the response structure for executing a MsgUpdateParams message.
 //
 // Since: 0.47
 message MsgUpdateParamsResponse {}

proto/neutron/cron/v1/schedule.proto

@@ -0,0 +1,32 @@
+syntax = "proto3";
+package neutron.cron.v1;
+
+import "gogoproto/gogo.proto";
+
+option go_package = "github.com/neutron-org/neutron/v4/x/cron/types/v1";
+
+// Defines the schedule for execution
+message Schedule {
+  // Name of schedule
+  string name = 1;
+  // Period in blocks
+  uint64 period = 2;
+  // Msgs that will be executed every certain number of blocks, specified in the `period` field
+  repeated MsgExecuteContract msgs = 3 [(gogoproto.nullable) = false];
+  // Last execution's block height
+  uint64 last_execute_height = 4;
+}
+
+// Defines the contract and the message to pass
+message MsgExecuteContract {
+  // The address of the smart contract
+  string contract = 1;
+  // JSON encoded message to be passed to the contract
+  string msg = 2;
+}
+
+// Defines the number of current schedules
+message ScheduleCount {
+  // The number of current schedules
+  int32 count = 1;
+}

wasmbinding/bindings/msg.go

@@ -196,9 +196,10 @@ type ForceTransfer struct {
 
 // AddSchedule adds new schedule to the cron module
 type AddSchedule struct {
-	Name   string               `json:"name"`
-	Period uint64               `json:"period"`
-	Msgs   []MsgExecuteContract `json:"msgs"`
+	Name           string               `json:"name"`
+	Period         uint64               `json:"period"`
+	Msgs           []MsgExecuteContract `json:"msgs"`
+	ExecutionStage string               `json:"execution_stage"`
 }
 
 // AddScheduleResponse holds response AddSchedule

wasmbinding/message_plugin.go

@@ -73,7 +73,8 @@ func CustomMessageDecorator(
 			Adminserver:           adminmodulekeeper.NewMsgServerImpl(*adminKeeper),
 			Bank:                  bankKeeper,
 			TokenFactory:          tokenFactoryKeeper,
-			CronKeeper:            cronKeeper,
+			CronMsgServer:         cronkeeper.NewMsgServerImpl(*cronKeeper),
+			CronQueryServer:       cronKeeper,
 			AdminKeeper:           adminKeeper,
 			ContractmanagerKeeper: contractmanagerKeeper,
 			DexMsgServer:          dexkeeper.NewMsgServerImpl(*dexKeeper),
@@ -90,7 +91,8 @@ type CustomMessenger struct {
 	Adminserver           admintypes.MsgServer
 	Bank                  *bankkeeper.BaseKeeper
 	TokenFactory          *tokenfactorykeeper.Keeper
-	CronKeeper            *cronkeeper.Keeper
+	CronMsgServer         crontypes.MsgServer
+	CronQueryServer       crontypes.QueryServer
 	AdminKeeper           *adminmodulekeeper.Keeper
 	ContractmanagerKeeper *contractmanagerkeeper.Keeper
 	DexMsgServer          dextypes.MsgServer
@@ -989,6 +991,8 @@ func (m *CustomMessenger) addSchedule(ctx sdk.Context, contractAddr sdk.AccAddre
 		return nil, nil, nil, errors.Wrap(sdkerrors.ErrUnauthorized, "only admin can add schedule")
 	}
 
+	authority := authtypes.NewModuleAddress(admintypes.ModuleName)
+
 	msgs := make([]crontypes.MsgExecuteContract, 0, len(addSchedule.Msgs))
 	for _, msg := range addSchedule.Msgs {
 		msgs = append(msgs, crontypes.MsgExecuteContract{
@@ -997,13 +1001,20 @@ func (m *CustomMessenger) addSchedule(ctx sdk.Context, contractAddr sdk.AccAddre
 		})
 	}
 
-	err := m.CronKeeper.AddSchedule(ctx, addSchedule.Name, addSchedule.Period, msgs)
+	_, err := m.CronMsgServer.AddSchedule(ctx, &crontypes.MsgAddSchedule{
+		Authority:      authority.String(),
+		Name:           addSchedule.Name,
+		Period:         addSchedule.Period,
+		Msgs:           msgs,
+		ExecutionStage: crontypes.ExecutionStage(crontypes.ExecutionStage_value[addSchedule.ExecutionStage]),
+	})
 	if err != nil {
 		ctx.Logger().Error("failed to addSchedule",
 			"from_address", contractAddr.String(),
+			"name", addSchedule.Name,
 			"error", err,
 		)
-		return nil, nil, nil, errors.Wrap(err, "marshal json failed")
+		return nil, nil, nil, errors.Wrapf(err, "failed to add %s schedule", addSchedule.Name)
 	}
 
 	ctx.Logger().Debug("schedule added",
@@ -1016,12 +1027,30 @@ func (m *CustomMessenger) addSchedule(ctx sdk.Context, contractAddr sdk.AccAddre
 }
 
 func (m *CustomMessenger) removeSchedule(ctx sdk.Context, contractAddr sdk.AccAddress, removeSchedule *bindings.RemoveSchedule) ([]sdk.Event, [][]byte, [][]*types.Any, error) {
-	params := m.CronKeeper.GetParams(ctx)
-	if !m.isAdmin(ctx, contractAddr) && contractAddr.String() != params.SecurityAddress {
+	params, err := m.CronQueryServer.Params(ctx, &crontypes.QueryParamsRequest{})
+	if err != nil {
+		ctx.Logger().Error("failed to removeSchedule", "error", err)
+		return nil, nil, nil, errors.Wrap(err, "failed to removeSchedule")
+	}
+
+	if !m.isAdmin(ctx, contractAddr) && contractAddr.String() != params.Params.SecurityAddress {
 		return nil, nil, nil, errors.Wrap(sdkerrors.ErrUnauthorized, "only admin or security dao can remove schedule")
 	}
 
-	m.CronKeeper.RemoveSchedule(ctx, removeSchedule.Name)
+	authority := authtypes.NewModuleAddress(admintypes.ModuleName)
+
+	_, err = m.CronMsgServer.RemoveSchedule(ctx, &crontypes.MsgRemoveSchedule{
+		Authority: authority.String(),
+		Name:      removeSchedule.Name,
+	})
+	if err != nil {
+		ctx.Logger().Error("failed to removeSchedule",
+			"from_address", contractAddr.String(),
+			"name", removeSchedule.Name,
+			"error", err,
+		)
+		return nil, nil, nil, errors.Wrapf(err, "failed to remove %s schedule", removeSchedule.Name)
+	}
 
 	ctx.Logger().Debug("schedule removed",
 		"from_address", contractAddr.String(),

wasmbinding/test/custom_message_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	contractmanagertypes "github.com/neutron-org/neutron/v4/x/contractmanager/types"
+	types2 "github.com/neutron-org/neutron/v4/x/cron/types"
 
 	"cosmossdk.io/math"
 	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
@@ -23,6 +24,8 @@ import (
 
 	adminkeeper "github.com/cosmos/admin-module/v2/x/adminmodule/keeper"
 
+	cronkeeper "github.com/neutron-org/neutron/v4/x/cron/keeper"
+
 	"github.com/neutron-org/neutron/v4/app/params"
 
 	"github.com/CosmWasm/wasmd/x/wasm/keeper"
@@ -67,7 +70,8 @@ func (suite *CustomMessengerTestSuite) SetupTest() {
 	suite.messenger.Adminserver = adminkeeper.NewMsgServerImpl(suite.neutron.AdminmoduleKeeper)
 	suite.messenger.Bank = &suite.neutron.BankKeeper
 	suite.messenger.TokenFactory = suite.neutron.TokenFactoryKeeper
-	suite.messenger.CronKeeper = &suite.neutron.CronKeeper
+	suite.messenger.CronMsgServer = cronkeeper.NewMsgServerImpl(suite.neutron.CronKeeper)
+	suite.messenger.CronQueryServer = suite.neutron.CronKeeper
 	suite.messenger.AdminKeeper = &suite.neutron.AdminmoduleKeeper
 	suite.messenger.ContractmanagerKeeper = &suite.neutron.ContractManagerKeeper
 	suite.contractOwner = keeper.RandomAccountAddress(suite.T())
@@ -718,6 +722,10 @@ func (suite *CustomMessengerTestSuite) TestAddRemoveSchedule() {
 		},
 	}
 
+	schedule, ok := suite.neutron.CronKeeper.GetSchedule(suite.ctx, "schedule1")
+	suite.True(ok)
+	suite.Equal(types2.ExecutionStage_EXECUTION_STAGE_END_BLOCKER, schedule.ExecutionStage)
+
 	// Dispatch AddSchedule message
 	_, err = suite.executeNeutronMsg(suite.contractAddress, msg)
 	suite.NoError(err)

x/cron/genesis.go

@@ -11,7 +11,7 @@ import (
 func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState) {
 	// Set all the schedules
 	for _, elem := range genState.ScheduleList {
-		err := k.AddSchedule(ctx, elem.Name, elem.Period, elem.Msgs)
+		err := k.AddSchedule(ctx, elem.Name, elem.Period, elem.Msgs, elem.ExecutionStage)
 		if err != nil {
 			panic(err)
 		}

x/cron/keeper/grpc_query_schedule_test.go

@@ -133,8 +133,9 @@ func createNSchedule(t *testing.T, ctx sdk.Context, k *cronkeeper.Keeper, n int3
 		item.Period = 1000
 		item.Msgs = nil
 		item.LastExecuteHeight = uint64(ctx.BlockHeight())
+		item.ExecutionStage = types.ExecutionStage_EXECUTION_STAGE_END_BLOCKER
 
-		err := k.AddSchedule(ctx, item.Name, item.Period, item.Msgs)
+		err := k.AddSchedule(ctx, item.Name, item.Period, item.Msgs, item.ExecutionStage)
 		require.NoError(t, err)
 
 		res[idx] = item