diff --git a/api/server.go b/api/server.go index aa1f54d..ba47189 100644 --- a/api/server.go +++ b/api/server.go @@ -57,14 +57,16 @@ func NewServer( p plugin.Plugin, ) *Server { logger := logrus.WithField("service", "plugin").Logger - schedulerService := scheduler.NewSchedulerService( + schedulerService, err := scheduler.NewSchedulerService( db, - logger.WithField("service", "scheduler").Logger, client, redisOpts, ) - schedulerService.Start() - logger.Info("Scheduler service started") + if err != nil { + logger.Fatalf("Failed to initialize scheduler service: %v", err) + return nil + } + policyService, err := service.NewPolicyService(db, schedulerService, logger.WithField("service", "policy").Logger) if err != nil { logger.Fatalf("Failed to initialize policy service: %v", err) diff --git a/cmd/payroll/worker/main.go b/cmd/payroll/worker/main.go index c7c0b10..bcd6a53 100644 --- a/cmd/payroll/worker/main.go +++ b/cmd/payroll/worker/main.go @@ -32,6 +32,7 @@ func main() { Password: cfg.Redis.Password, DB: cfg.Redis.DB, } + logger := logrus.StandardLogger() client := asynq.NewClient(redisOptions) srv := asynq.NewServer( diff --git a/go.mod b/go.mod index 55632c9..9ac4b9c 100644 --- a/go.mod +++ b/go.mod @@ -18,7 +18,7 @@ require ( github.com/spf13/viper v1.20.1 github.com/vultisig/commondata v0.0.0-20250430024109-a2492623ef05 github.com/vultisig/mobile-tss-lib v0.0.0-20250316003201-2e7e570a4a74 - github.com/vultisig/recipes v0.0.0-20250531132511-82a0eb621885 + github.com/vultisig/recipes v0.0.0-20250603213257-b4a6b0afe2a0 github.com/vultisig/verifier v0.0.0-20250531103513-c1f5cb38b103 github.com/vultisig/vultiserver v0.0.0-20250515110921-82d56d3d9cc9 ) diff --git a/go.sum b/go.sum index e7249f6..114e750 100644 --- a/go.sum +++ b/go.sum @@ -740,6 +740,8 @@ github.com/vultisig/recipes v0.0.0-20250531074825-c4872dec4ab9 h1:SSGqON8i7paf3/ github.com/vultisig/recipes v0.0.0-20250531074825-c4872dec4ab9/go.mod h1:tGlNiMDC33gPdX3C1DSESfQeCsaiOx+SNzInnF/Oi68= github.com/vultisig/recipes v0.0.0-20250531132511-82a0eb621885 h1:hXju0II1G1CydXchhjTla1JxK8WceDUGJhQhJHJSUIE= github.com/vultisig/recipes v0.0.0-20250531132511-82a0eb621885/go.mod h1:BXXJ25U75xaexJLoiiaLEZ6TQrcy8+8UGJGd4c+hNSE= +github.com/vultisig/recipes v0.0.0-20250603213257-b4a6b0afe2a0 h1:JauwztTOr7EqVhhXIQw+fWX6+VKDc5T6zLVBFIvGuBk= +github.com/vultisig/recipes v0.0.0-20250603213257-b4a6b0afe2a0/go.mod h1:BXXJ25U75xaexJLoiiaLEZ6TQrcy8+8UGJGd4c+hNSE= github.com/vultisig/verifier v0.0.0-20250531103513-c1f5cb38b103 h1:1Qgt3uwzG11UqKoizd5Ft65Ij/8c2Kxpg3VB6N3/Fb0= github.com/vultisig/verifier v0.0.0-20250531103513-c1f5cb38b103/go.mod h1:8b9k8CJtF1KwjMoZuWFYmz0Q+zdikrx3q0ExbZqsLgk= github.com/vultisig/vultiserver v0.0.0-20250515110921-82d56d3d9cc9 h1:pZhGN8q8+gPB1JJjVDC1hDg8qn6Tbj0XBJymgTQ8qQg= diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index ba34298..3b3285c 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -4,13 +4,13 @@ import ( "context" "encoding/json" "fmt" - "strconv" "time" "github.com/hibiken/asynq" "github.com/jackc/pgx/v5" "github.com/robfig/cron/v3" "github.com/sirupsen/logrus" + rtypes "github.com/vultisig/recipes/types" vtypes "github.com/vultisig/verifier/types" "github.com/vultisig/plugin/internal/tasks" @@ -31,21 +31,22 @@ type SchedulerService struct { done chan struct{} } -func NewSchedulerService(db storage.DatabaseStorage, logger *logrus.Logger, client *asynq.Client, redisOpts asynq.RedisClientOpt) *SchedulerService { +func NewSchedulerService(db storage.DatabaseStorage, + client *asynq.Client, + redisOpts asynq.RedisClientOpt) (*SchedulerService, error) { if db == nil { - logger.Fatal("database connection is nil") + return nil, fmt.Errorf("db is nil") } // create inspector using the same Redis configuration as the client inspector := asynq.NewInspector(redisOpts) - return &SchedulerService{ db: db, - logger: logger, + logger: logrus.WithField("component", "scheduler").Logger, client: client, inspector: inspector, done: make(chan struct{}), - } + }, nil } func (s *SchedulerService) Start() { @@ -182,31 +183,33 @@ func (s *SchedulerService) CreateTimeTrigger(ctx context.Context, policy vtypes. } func (s *SchedulerService) GetTriggerFromPolicy(policy vtypes.PluginPolicy) (*types.TimeTrigger, error) { - var policySchedule struct { - Schedule struct { - Frequency string `json:"frequency"` - StartTime time.Time `json:"start_time"` - Interval string `json:"interval"` - EndTime *time.Time `json:"end_time,omitempty"` - } `json:"schedule"` - } - - if err := json.Unmarshal([]byte(policy.Recipe), &policySchedule); err != nil { + var p rtypes.Policy + if err := json.Unmarshal([]byte(policy.Recipe), &p); err != nil { return nil, fmt.Errorf("failed to parse policy schedule: %w", err) } - - interval, err := strconv.Atoi(policySchedule.Schedule.Interval) - if err != nil { - return nil, fmt.Errorf("failed to parse interval: %w", err) + if p.Schedule == nil { + return nil, fmt.Errorf("policy schedule is nil") } - - cronExpr := frequencyToCron(policySchedule.Schedule.Frequency, policySchedule.Schedule.StartTime, interval) + interval := 1 + if p.Schedule.Interval > 0 { + interval = int(p.Schedule.Interval) + } + startTime := time.Now().UTC() + if p.Schedule.GetStartTime() != nil { + startTime = p.Schedule.GetStartTime().AsTime() + } + var endTime *time.Time + if p.Schedule.GetEndTime() != nil { + tmpTime := p.Schedule.GetEndTime().AsTime() + endTime = &tmpTime + } + cronExpr := frequencyToCron(p.Schedule.GetFrequency(), startTime, interval) trigger := types.TimeTrigger{ PolicyID: policy.ID, CronExpression: cronExpr, - StartTime: time.Now().UTC(), - EndTime: policySchedule.Schedule.EndTime, - Frequency: policySchedule.Schedule.Frequency, + StartTime: startTime, + EndTime: endTime, + Frequency: p.Schedule.Frequency, Interval: interval, Status: types.StatusTimeTriggerPending, } @@ -214,9 +217,12 @@ func (s *SchedulerService) GetTriggerFromPolicy(policy vtypes.PluginPolicy) (*ty return &trigger, nil } -func createSchedule(cronExpr, frequency string, startTime time.Time, interval int) (cron.Schedule, error) { +func createSchedule(cronExpr string, frequency rtypes.ScheduleFrequency, startTime time.Time, interval int) (cron.Schedule, error) { // Use our custom schedule implementation for intervals > 1 and when frequency is daily, weekly, monthly - if interval > 1 && (frequency == "daily" || frequency == "weekly" || frequency == "monthly") { + if interval > 1 && + (frequency == rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_DAILY || + frequency == rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_WEEKLY || + frequency == rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_MONTHLY) { return NewIntervalSchedule(frequency, startTime, interval) } @@ -229,28 +235,30 @@ func createSchedule(cronExpr, frequency string, startTime time.Time, interval in return schedule, nil } -func frequencyToCron(frequency string, startTime time.Time, interval int) string { +func frequencyToCron(frequency rtypes.ScheduleFrequency, startTime time.Time, interval int) string { switch frequency { - case "minutely": - return fmt.Sprintf("*/%d * * * *", interval) - case "hourly": + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_UNSPECIFIED: + return "" + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_DAILY: + return fmt.Sprintf("%d %d * * *", startTime.Minute(), startTime.Hour()) + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_HOURLY: if interval == 1 { return fmt.Sprintf("%d * * * *", startTime.Minute()) } return fmt.Sprintf("%d */%d * * *", startTime.Minute(), interval) - case "daily": - return fmt.Sprintf("%d %d * * *", startTime.Minute(), startTime.Hour()) - case "weekly": + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_WEEKLY: return fmt.Sprintf("%d %d * * %d", startTime.Minute(), startTime.Hour(), startTime.Weekday()) - case "monthly": + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_MONTHLY: return fmt.Sprintf("%d %d %d * *", startTime.Minute(), startTime.Hour(), startTime.Day()) + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_BIWEEKLY: + return fmt.Sprintf("%d %d */14 * %d", startTime.Minute(), startTime.Hour(), startTime.Weekday()) default: return "" } } type IntervalSchedule struct { - Frequency string + Frequency rtypes.ScheduleFrequency Interval int StartTime time.Time Minute int @@ -260,7 +268,7 @@ type IntervalSchedule struct { Location *time.Location } -func NewIntervalSchedule(frequency string, startTime time.Time, interval int) (*IntervalSchedule, error) { +func NewIntervalSchedule(frequency rtypes.ScheduleFrequency, startTime time.Time, interval int) (*IntervalSchedule, error) { if interval < 1 { return nil, fmt.Errorf("interval must be at least 1") } @@ -281,11 +289,11 @@ func (s *IntervalSchedule) Next(t time.Time) time.Time { t = t.In(s.Location) switch s.Frequency { - case "daily": + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_DAILY: return s.nextDaily(t) - case "weekly": + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_WEEKLY: return s.nextWeekly(t) - case "monthly": + case rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_MONTHLY: return s.nextMonthly(t) default: return time.Time{} diff --git a/internal/types/time_trigger.go b/internal/types/time_trigger.go index 8452d35..f0e8de6 100644 --- a/internal/types/time_trigger.go +++ b/internal/types/time_trigger.go @@ -4,6 +4,7 @@ import ( "time" "github.com/google/uuid" + rtypes "github.com/vultisig/recipes/types" ) type TimeTriggerStatus string @@ -14,12 +15,12 @@ const ( ) type TimeTrigger struct { - PolicyID uuid.UUID `json:"policy_id"` - CronExpression string `json:"cron_expression"` - StartTime time.Time `json:"start_time"` - EndTime *time.Time `json:"end_time"` - Frequency string `json:"frequency"` - Interval int `json:"interval"` - LastExecution *time.Time `json:"last_execution"` - Status TimeTriggerStatus `json:"status"` + PolicyID uuid.UUID `json:"policy_id"` + CronExpression string `json:"cron_expression"` + StartTime time.Time `json:"start_time"` + EndTime *time.Time `json:"end_time"` + Frequency rtypes.ScheduleFrequency `json:"frequency"` + Interval int `json:"interval"` + LastExecution *time.Time `json:"last_execution"` + Status TimeTriggerStatus `json:"status"` } diff --git a/plugin/payroll/policy.go b/plugin/payroll/policy.go index 4195ed3..fca7b3b 100644 --- a/plugin/payroll/policy.go +++ b/plugin/payroll/policy.go @@ -152,7 +152,25 @@ func (p *PayrollPlugin) validateAmount(pc *rtypes.ParameterConstraint) error { } return nil } -func (p *PayrollPlugin) CheckRule(rule *rtypes.Rule) error { +func (p *PayrollPlugin) validateSchedule(schedule *rtypes.Schedule) error { + if schedule == nil { + return fmt.Errorf("schedule is nil") + } + if schedule.GetFrequency() == rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_UNSPECIFIED { + return fmt.Errorf("schedule frequency is required") + } + + if schedule.GetStartTime() == nil { + return fmt.Errorf("start time is required") + } + + if schedule.GetEndTime() != nil && schedule.GetEndTime().AsTime().Before(schedule.GetStartTime().AsTime()) { + return fmt.Errorf("end time cannot be before start time") + } + + return nil +} +func (p *PayrollPlugin) checkRule(rule *rtypes.Rule) error { if rule.Effect != rtypes.Effect_EFFECT_ALLOW { return fmt.Errorf("rule effect must be ALLOW, got: %s", rule.Effect) } @@ -182,8 +200,8 @@ func (p *PayrollPlugin) ValidatePluginPolicy(policyDoc vtypes.PluginPolicy) erro if policyDoc.PluginID != vtypes.PluginVultisigPayroll_0000 { return fmt.Errorf("policy does not match plugin type, expected: %s, got: %s", vtypes.PluginVultisigPayroll_0000, policyDoc.PluginID) } - var rPolicy rtypes.Policy + policyBytes, err := base64.RawStdEncoding.DecodeString(policyDoc.Recipe) if err != nil { return fmt.Errorf("failed to decode policy recipe: %w", err) @@ -199,9 +217,11 @@ func (p *PayrollPlugin) ValidatePluginPolicy(policyDoc vtypes.PluginPolicy) erro if len(rPolicy.Rules) == 0 { return fmt.Errorf("no rules") } - + if err := p.validateSchedule(rPolicy.Schedule); err != nil { + return fmt.Errorf("schedule validation failed: %w", err) + } for _, rule := range rPolicy.Rules { - if err := p.CheckRule(rule); err != nil { + if err := p.checkRule(rule); err != nil { return fmt.Errorf("rule validation failed: %w", err) } } diff --git a/storage/postgres/db.go b/storage/postgres/db.go index b7a68b0..3abd6ce 100644 --- a/storage/postgres/db.go +++ b/storage/postgres/db.go @@ -21,7 +21,7 @@ type PostgresBackend struct { } type MigrationOptions struct { - RunSystemMigrations bool + RunSystemMigrations bool RunPluginMigrations bool } @@ -39,7 +39,7 @@ func NewPostgresBackend(dsn string, opts *MigrationOptions) (*PostgresBackend, e // Apply default options if not provided if opts == nil { opts = &MigrationOptions{ - RunSystemMigrations: true, + RunSystemMigrations: true, RunPluginMigrations: true, } } @@ -51,18 +51,18 @@ func NewPostgresBackend(dsn string, opts *MigrationOptions) (*PostgresBackend, e return backend, nil } -func (d *PostgresBackend) Close() error { - d.pool.Close() +func (p *PostgresBackend) Close() error { + p.pool.Close() return nil } -func (d *PostgresBackend) Migrate(opts *MigrationOptions) error { +func (p *PostgresBackend) Migrate(opts *MigrationOptions) error { logrus.Info("Starting database migration...") - + // Run system migrations first (plugin_policies table) if opts.RunSystemMigrations { - systemMgr := NewSystemMigrationManager(d.pool) + systemMgr := NewSystemMigrationManager(p.pool) if err := systemMgr.Migrate(); err != nil { return fmt.Errorf("failed to run system migrations: %w", err) } @@ -70,7 +70,7 @@ func (d *PostgresBackend) Migrate(opts *MigrationOptions) error { // Run plugin migrations (all other tables) if opts.RunPluginMigrations { - pluginMgr := NewPluginMigrationManager(d.pool) + pluginMgr := NewPluginMigrationManager(p.pool) if err := pluginMgr.Migrate(); err != nil { return fmt.Errorf("failed to run plugin migrations: %w", err) } diff --git a/storage/postgres/migrations/plugin/20241211182120_add_time_triggers.sql b/storage/postgres/migrations/plugin/20241211182120_add_time_triggers.sql index bbc0177..d0b3667 100644 --- a/storage/postgres/migrations/plugin/20241211182120_add_time_triggers.sql +++ b/storage/postgres/migrations/plugin/20241211182120_add_time_triggers.sql @@ -11,7 +11,7 @@ CREATE TABLE time_triggers ( cron_expression TEXT NOT NULL, start_time TIMESTAMP NOT NULL, end_time TIMESTAMP, - frequency TEXT NOT NULL, + frequency INT NOT NULL, interval INTEGER NOT NULL, last_execution TIMESTAMP, status trigger_status NOT NULL, diff --git a/storage/postgres/time_trigger.go b/storage/postgres/time_trigger.go index 24b5824..404957f 100644 --- a/storage/postgres/time_trigger.go +++ b/storage/postgres/time_trigger.go @@ -12,9 +12,17 @@ import ( "github.com/vultisig/plugin/internal/types" ) +// checkContext returns ctx.Err() if the context is done (timeout or canceled), otherwise nil. +func checkContext(ctx context.Context) error { + if ctx == nil { + return fmt.Errorf("context is nil") + } + return ctx.Err() +} + func (p *PostgresBackend) CreateTimeTriggerTx(ctx context.Context, tx pgx.Tx, trigger types.TimeTrigger) error { - if p.pool == nil { - return fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return err } query := ` @@ -37,8 +45,8 @@ func (p *PostgresBackend) CreateTimeTriggerTx(ctx context.Context, tx pgx.Tx, tr } func (p *PostgresBackend) DeleteTimeTrigger(ctx context.Context, policyID uuid.UUID) error { - if p.pool == nil { - return fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return err } query := `DELETE FROM time_triggers WHERE policy_id = $1` @@ -48,14 +56,14 @@ func (p *PostgresBackend) DeleteTimeTrigger(ctx context.Context, policyID uuid.U } func (p *PostgresBackend) GetPendingTimeTriggers(ctx context.Context) ([]types.TimeTrigger, error) { - if p.pool == nil { - return nil, fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return nil, err } // TODO: add limit and proper index query := ` WITH active_triggers AS ( - SELECT t.policy_id, t.cron_expression, t.start_time, t.end_time, t.frequency, t.interval, t.last_execution, t.status + SELECT t.policy_id, t.cron_expression, t.start_time, t.end_time, t.frequency, t.interval, t.last_execution, t.status FROM time_triggers t INNER JOIN plugin_policies p ON t.policy_id = p.id WHERE t.start_time <= $1 @@ -96,8 +104,8 @@ func (p *PostgresBackend) GetPendingTimeTriggers(ctx context.Context) ([]types.T } func (p *PostgresBackend) UpdateTimeTriggerLastExecution(ctx context.Context, policyID uuid.UUID) error { - if p.pool == nil { - return fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return err } query := ` @@ -111,8 +119,8 @@ func (p *PostgresBackend) UpdateTimeTriggerLastExecution(ctx context.Context, po } func (p *PostgresBackend) UpdateTimeTriggerTx(ctx context.Context, policyID uuid.UUID, trigger types.TimeTrigger, tx pgx.Tx) error { - if p.pool == nil { - return fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return err } query := ` @@ -134,8 +142,8 @@ func (p *PostgresBackend) UpdateTimeTriggerTx(ctx context.Context, policyID uuid } func (p *PostgresBackend) GetTriggerStatus(ctx context.Context, policyID uuid.UUID) (types.TimeTriggerStatus, error) { - if p.pool == nil { - return "", fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return "", err } query := ` @@ -157,8 +165,8 @@ func (p *PostgresBackend) GetTriggerStatus(ctx context.Context, policyID uuid.UU } func (p *PostgresBackend) UpdateTriggerStatus(ctx context.Context, policyID uuid.UUID, status types.TimeTriggerStatus) error { - if p.pool == nil { - return fmt.Errorf("database pool is nil") + if err := checkContext(ctx); err != nil { + return err } query := ` diff --git a/storage/postgres/time_trigger_test.go b/storage/postgres/time_trigger_test.go new file mode 100644 index 0000000..ad4e63b --- /dev/null +++ b/storage/postgres/time_trigger_test.go @@ -0,0 +1,83 @@ +package postgres + +import ( + "context" + "testing" + "time" + + "github.com/google/uuid" + rtypes "github.com/vultisig/recipes/types" + vtypes "github.com/vultisig/verifier/types" + + "github.com/vultisig/plugin/internal/types" +) + +func TestTimeTrigger(t *testing.T) { + t.SkipNow() + // Initialize Postgres backend + backend, err := NewPostgresBackend("postgres://myuser:mypassword@localhost:5432/vultisig-plugin?sslmode=disable", nil) + if err != nil { + t.Fatalf("Failed to create Postgres backend: %v", err) + } + defer backend.Close() + policyID := uuid.New() + + // Create a new time trigger + trigger := types.TimeTrigger{ + PolicyID: policyID, + CronExpression: "0 0 * * *", + StartTime: time.Now().UTC(), + EndTime: nil, + Frequency: rtypes.ScheduleFrequency_SCHEDULE_FREQUENCY_DAILY, + Interval: 1, + Status: "PENDING", + } + + ctx := context.Background() + tx, err := backend.pool.Begin(ctx) + if err != nil { + t.Fatalf("Failed to begin transaction: %v", err) + } + if _, err := backend.InsertPluginPolicyTx(ctx, tx, vtypes.PluginPolicy{ + ID: policyID, + PublicKey: "4a7b9c2f8e1d3a5b6c4f9e2d7a1b3c5f8d4e6a2b9c1f7e3d5a6b4c2f8e1d3a5b", + PluginID: "vultisig-payroll-0000", + PluginVersion: "1", + PolicyVersion: "1", + Signature: "whatever", + Recipe: "k3uaL47TpWtsT54tfhO7K6Gbyxf4H71jWma0zC2e3kQ=", + Active: true, + }); err != nil { + tx.Rollback(ctx) + t.Fatalf("Failed to insert plugin policy: %v", err) + } + if err := backend.CreateTimeTriggerTx(ctx, tx, trigger); err != nil { + tx.Rollback(ctx) + t.Fatalf("Failed to create time trigger: %v", err) + } + + if err := tx.Commit(ctx); err != nil { + t.Fatalf("Failed to commit transaction: %v", err) + } + + t.Logf("Successfully created time trigger with ID: %s", trigger.PolicyID) + + // Test getting pending time triggers + triggers, err := backend.GetPendingTimeTriggers(ctx) + if err != nil { + t.Fatalf("Failed to get pending time triggers: %v", err) + } + + if len(triggers) == 0 { + t.Fatal("Expected at least one pending time trigger") + } + + t.Logf("Found %d pending time triggers", len(triggers)) + + // Test deleting the time trigger + if err := backend.DeleteTimeTrigger(ctx, trigger.PolicyID); err != nil { + t.Fatalf("Failed to delete time trigger: %v", err) + } + + t.Logf("Successfully deleted time trigger with ID: %s", trigger.PolicyID) +}