fix(rpc): authorize agent task results

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
2026-09-20 18:20:12 +00:00 · 2026-05-21 02:02:02 +00:00
parent 8add559ec8
commit 02129f16fb
4 changed files with 1030 additions and 7 deletions
@@ -8,6 +8,11 @@ import (
 	"time"

 	"github.com/gin-gonic/gin"
+	"github.com/patrickmn/go-cache"
+	"github.com/robfig/cron/v3"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+
 	"github.com/nezhahq/nezha/model"
 	pb "github.com/nezhahq/nezha/proto"
 	"google.golang.org/grpc/metadata"
@@ -335,6 +340,241 @@ func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) {
 	assertNoTask(t, stream)
 }

+func TestAlertTriggerCronResultAuthorizationConsumesOneDispatch(t *testing.T) {
+	cronClass := &CronClass{}
+	cronClass.reserveAlertTriggerCronResult(42, 7)
+	cronClass.reserveAlertTriggerCronResult(42, 7)
+
+	if !cronClass.consumeAlertTriggerCronResult(42, 7) {
+		t.Fatal("expected first alert-trigger authorization to be consumed")
+	}
+	if !cronClass.consumeAlertTriggerCronResult(42, 7) {
+		t.Fatal("expected second alert-trigger authorization to be consumed")
+	}
+	if cronClass.consumeAlertTriggerCronResult(42, 7) {
+		t.Fatal("expected alert-trigger authorization to be consumed only once per dispatch")
+	}
+}
+
+func TestAlertTriggerCronResultAuthorizationExpires(t *testing.T) {
+	cronClass := &CronClass{
+		pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
+			42: {7: {time.Now().Add(-time.Second)}},
+		},
+	}
+
+	if cronClass.consumeAlertTriggerCronResult(42, 7) {
+		t.Fatal("expired alert-trigger authorization must not be accepted")
+	}
+	if len(cronClass.pendingAlertTriggerTasks) != 0 {
+		t.Fatal("expired alert-trigger authorization must be pruned")
+	}
+}
+
+func TestAlertTriggerCronResultAuthorizationRevokeRemovesLatestDispatch(t *testing.T) {
+	existingAuthorizationExpiresAt := time.Now().Add(time.Hour)
+	cronClass := &CronClass{
+		pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
+			42: {7: {existingAuthorizationExpiresAt}},
+		},
+	}
+	cronClass.reserveAlertTriggerCronResult(42, 7)
+
+	cronClass.revokeAlertTriggerCronResult(42, 7)
+
+	authorizations := cronClass.pendingAlertTriggerTasks[42][7]
+	if len(authorizations) != 1 {
+		t.Fatalf("expected one previous alert-trigger authorization to remain, got %d", len(authorizations))
+	}
+	if !authorizations[0].Equal(existingAuthorizationExpiresAt) {
+		t.Fatal("send failure rollback must remove the newest reserved authorization")
+	}
+}
+
+func TestCronClassUpdatePrunesAlertTriggerCronResultAuthorization(t *testing.T) {
+	cronClass := &CronClass{
+		Cron: cron.New(cron.WithSeconds()),
+		class: class[uint64, *model.Cron]{
+			list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}},
+		},
+		pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
+			42: {7: {time.Now().Add(time.Hour)}},
+		},
+	}
+
+	cronClass.Update(&model.Cron{Common: model.Common{ID: 42}})
+
+	if len(cronClass.pendingAlertTriggerTasks) != 0 {
+		t.Fatal("cron update must prune old alert-trigger result authorizations")
+	}
+}
+
+func TestCronClassDeletePrunesAlertTriggerCronResultAuthorization(t *testing.T) {
+	cronClass := &CronClass{
+		Cron: cron.New(cron.WithSeconds()),
+		class: class[uint64, *model.Cron]{
+			list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}},
+		},
+		pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{
+			42: {7: {time.Now().Add(time.Hour)}},
+		},
+	}
+
+	cronClass.Delete([]uint64{42})
+
+	if len(cronClass.pendingAlertTriggerTasks) != 0 {
+		t.Fatal("cron delete must prune alert-trigger result authorizations")
+	}
+}
+
+// CanReportCronResult is the cron-side dual of canReportServiceResult: it gates
+// agent-reported TaskTypeCommand results to only the cron/server pairs the
+// dashboard actually fanned the task out to. Without these inbound checks any
+// authenticated agent could fabricate a TaskResult for an arbitrary cron ID and
+// poison LastResult / fire success/failure notifications belonging to another
+// tenant. The tests below pin each Cover branch end-to-end against the dispatch
+// logic in CronTrigger so the two sides stay symmetric.
+
+func TestCanReportCronResultRejectsNilCronOrReporter(t *testing.T) {
+	cr := &model.Cron{Common: model.Common{ID: 7, UserID: 100}, Cover: model.CronCoverAll}
+	reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
+
+	if CanReportCronResult(nil, reporter) {
+		t.Fatal("nil cron must be rejected — would dereference inside cover branches")
+	}
+	if CanReportCronResult(cr, nil) {
+		t.Fatal("nil reporter must be rejected")
+	}
+}
+
+func TestCanReportCronResultRejectsForeignReporter(t *testing.T) {
+	replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
+		100: {Role: model.RoleMember},
+		200: {Role: model.RoleMember},
+	})
+
+	cr := &model.Cron{
+		Common: model.Common{ID: 7, UserID: 100},
+		Cover:  model.CronCoverAll,
+	}
+	foreign := &model.Server{Common: model.Common{ID: 1, UserID: 200}}
+
+	if CanReportCronResult(cr, foreign) {
+		t.Fatal("foreign-user reporter must be rejected: CronTrigger never dispatched to it")
+	}
+}
+
+func TestCanReportCronResultCronCoverAllRejectsReporterInDenyList(t *testing.T) {
+	cr := &model.Cron{
+		Common:  model.Common{ID: 7, UserID: 100},
+		Cover:   model.CronCoverAll,
+		Servers: []uint64{1},
+	}
+	reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
+
+	if CanReportCronResult(cr, reporter) {
+		t.Fatal("CronCoverAll treats Servers as deny-list; reporter in the list must be rejected")
+	}
+}
+
+func TestCanReportCronResultCronCoverAllAcceptsReporterNotInDenyList(t *testing.T) {
+	cr := &model.Cron{
+		Common:  model.Common{ID: 7, UserID: 100},
+		Cover:   model.CronCoverAll,
+		Servers: []uint64{99},
+	}
+	reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
+
+	if !CanReportCronResult(cr, reporter) {
+		t.Fatal("CronCoverAll with reporter NOT in Servers must accept — CronTrigger dispatches to it")
+	}
+}
+
+func TestCanReportCronResultCronCoverIgnoreAllAcceptsReporterInAllowList(t *testing.T) {
+	cr := &model.Cron{
+		Common:  model.Common{ID: 7, UserID: 100},
+		Cover:   model.CronCoverIgnoreAll,
+		Servers: []uint64{1},
+	}
+	reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
+
+	if !CanReportCronResult(cr, reporter) {
+		t.Fatal("CronCoverIgnoreAll treats Servers as allow-list; reporter in the list must be accepted")
+	}
+}
+
+func TestCanReportCronResultCronCoverIgnoreAllRejectsReporterOutsideAllowList(t *testing.T) {
+	cr := &model.Cron{
+		Common:  model.Common{ID: 7, UserID: 100},
+		Cover:   model.CronCoverIgnoreAll,
+		Servers: []uint64{99},
+	}
+	reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}}
+
+	if CanReportCronResult(cr, reporter) {
+		t.Fatal("CronCoverIgnoreAll with reporter NOT in Servers must reject — CronTrigger never dispatched to it")
+	}
+}
+
+// failingTaskStream simulates a TaskStream whose Send always errors. CronTrigger
+// uses this signal to revoke a reserved alert-trigger authorization, so the
+// agent can't later attach to the cron via CanReportCronResult based on a
+// dispatch that never actually reached the wire.
+type failingTaskStream struct {
+	capturedTaskStream
+	sendErr error
+}
+
+func newFailingTaskStream(err error) *failingTaskStream {
+	return &failingTaskStream{
+		capturedTaskStream: capturedTaskStream{tasks: make(chan *pb.Task, 4)},
+		sendErr:            err,
+	}
+}
+
+func (s *failingTaskStream) Send(task *pb.Task) error {
+	s.tasks <- task
+	return s.sendErr
+}
+
+func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) {
+	failing := newFailingTaskStream(context.Canceled)
+	replaceServerSharedForSecurityTest(t,
+		&model.Server{Common: model.Common{ID: 7, UserID: 100}, Name: "broken-server", TaskStream: failing},
+	)
+	replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
+		100: {Role: model.RoleMember},
+	})
+
+	originalCronShared := CronShared
+	t.Cleanup(func() { CronShared = originalCronShared })
+	CronShared = &CronClass{
+		class:                    class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
+		pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{},
+	}
+
+	cr := &model.Cron{
+		Common: model.Common{ID: 42, UserID: 100},
+		Cover:  model.CronCoverAlertTrigger,
+	}
+
+	CronTrigger(cr, 7)()
+
+	// drain the dispatched task — Send error is what we care about, not the payload
+	select {
+	case <-failing.tasks:
+	case <-time.After(time.Second):
+		t.Fatal("expected CronTrigger to call Send before reacting to the error")
+	}
+
+	if CronShared.consumeAlertTriggerCronResult(42, 7) {
+		t.Fatal("Send failure must revoke the reserved alert-trigger authorization; otherwise a foreign agent could later report a result for a dispatch that never reached the wire")
+	}
+	if len(CronShared.pendingAlertTriggerTasks) != 0 {
+		t.Fatalf("expected pendingAlertTriggerTasks to be empty after revoke, got %d entries", len(CronShared.pendingAlertTriggerTasks))
+	}
+}
+
 func TestClassCheckPermission(t *testing.T) {
 	replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
 		1:   {Role: model.RoleAdmin},
@@ -377,3 +617,308 @@ func TestClassCheckPermission(t *testing.T) {
 		t.Fatal("expected admin to access any resource")
 	}
 }
+
+func TestServiceMonitorResultSkipsReporterOutsideServiceCover(t *testing.T) {
+	ss := newServiceMonitorSecurityHarness(t,
+		&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "covered-server"},
+		&model.Server{Common: model.Common{ID: 2, UserID: 100}, Name: "uncovered-server"},
+	)
+	addServiceMonitorSecurityService(t, ss, &model.Service{
+		Common:      model.Common{ID: 10, UserID: 100},
+		Name:        "selected-only-service",
+		Type:        model.TaskTypeTCPPing,
+		Target:      "example.invalid:443",
+		Duration:    3600,
+		Cover:       model.ServiceCoverIgnoreAll,
+		SkipServers: map[uint64]bool{1: true},
+	})
+
+	ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
+	ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
+
+	waitForServiceHistory(t, 10, 1)
+	assertNoServiceHistory(t, 10, 2)
+}
+
+func TestServiceMonitorResultSkipsCoveredReporterOwnedByAnotherUser(t *testing.T) {
+	ss := newServiceMonitorSecurityHarness(t,
+		&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
+		&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "foreign-server"},
+	)
+	addServiceMonitorSecurityService(t, ss, &model.Service{
+		Common:      model.Common{ID: 10, UserID: 100},
+		Name:        "owner-only-service",
+		Type:        model.TaskTypeTCPPing,
+		Target:      "example.invalid:443",
+		Duration:    3600,
+		Cover:       model.ServiceCoverIgnoreAll,
+		SkipServers: map[uint64]bool{1: true, 2: true},
+	})
+
+	ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
+	ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
+
+	waitForServiceHistory(t, 10, 1)
+	assertNoServiceHistory(t, 10, 2)
+}
+
+func TestServiceMonitorResultSkipsMismatchedTaskType(t *testing.T) {
+	ss := newServiceMonitorSecurityHarness(t,
+		&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
+	)
+	addServiceMonitorSecurityService(t, ss, &model.Service{
+		Common:      model.Common{ID: 10, UserID: 100},
+		Name:        "http-service",
+		Type:        model.TaskTypeHTTPGet,
+		Target:      "https://example.invalid",
+		Duration:    3600,
+		Cover:       model.ServiceCoverIgnoreAll,
+		SkipServers: map[uint64]bool{1: true},
+	})
+
+	ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, false))
+	ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeHTTPGet, true))
+
+	waitForTodayStats(t, ss, 10, 1, 0)
+}
+
+func TestServiceMonitorResultSkipsUnknownReporter(t *testing.T) {
+	ss := newServiceMonitorSecurityHarness(t,
+		&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
+	)
+	addServiceMonitorSecurityService(t, ss, &model.Service{
+		Common:      model.Common{ID: 10, UserID: 100},
+		Name:        "known-reporter-service",
+		Type:        model.TaskTypeTCPPing,
+		Target:      "example.invalid:443",
+		Duration:    3600,
+		Cover:       model.ServiceCoverIgnoreAll,
+		SkipServers: map[uint64]bool{1: true},
+	})
+
+	ss.Dispatch(serviceMonitorResult(999, 10, model.TaskTypeTCPPing, true))
+	ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
+
+	waitForServiceHistory(t, 10, 1)
+	assertNoServiceHistory(t, 10, 999)
+}
+
+func TestServiceMonitorResultAllowsCoveredReporterOwnedByServiceOwner(t *testing.T) {
+	ss := newServiceMonitorSecurityHarness(t,
+		&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"},
+	)
+	addServiceMonitorSecurityService(t, ss, &model.Service{
+		Common:      model.Common{ID: 10, UserID: 100},
+		Name:        "owner-service",
+		Type:        model.TaskTypeTCPPing,
+		Target:      "example.invalid:443",
+		Duration:    3600,
+		Cover:       model.ServiceCoverIgnoreAll,
+		SkipServers: map[uint64]bool{1: true},
+	})
+
+	ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true))
+
+	waitForServiceHistory(t, 10, 1)
+}
+
+func TestServiceMonitorResultAllowsCoveredReporterForAdminOwnedService(t *testing.T) {
+	replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{
+		1:   {Role: model.RoleAdmin},
+		200: {Role: model.RoleMember},
+	})
+	ss := newServiceMonitorSecurityHarness(t,
+		&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "member-server"},
+	)
+	addServiceMonitorSecurityService(t, ss, &model.Service{
+		Common:      model.Common{ID: 10, UserID: 1},
+		Name:        "admin-service",
+		Type:        model.TaskTypeTCPPing,
+		Target:      "example.invalid:443",
+		Duration:    3600,
+		Cover:       model.ServiceCoverIgnoreAll,
+		SkipServers: map[uint64]bool{2: true},
+	})
+
+	ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true))
+
+	waitForServiceHistory(t, 10, 2)
+}
+
+func newServiceMonitorSecurityHarness(t *testing.T, servers ...*model.Server) *ServiceSentinel {
+	t.Helper()
+
+	originalDB := DB
+	originalConf := Conf
+	originalCache := Cache
+	originalCronShared := CronShared
+	originalServerShared := ServerShared
+	originalServiceSentinelShared := ServiceSentinelShared
+	originalNotificationShared := NotificationShared
+	originalTSDBShared := TSDBShared
+	originalLoc := Loc
+	var sqlDBClose func() error
+
+	t.Cleanup(func() {
+		DB = originalDB
+		Conf = originalConf
+		Cache = originalCache
+		CronShared = originalCronShared
+		ServerShared = originalServerShared
+		ServiceSentinelShared = originalServiceSentinelShared
+		NotificationShared = originalNotificationShared
+		TSDBShared = originalTSDBShared
+		Loc = originalLoc
+		if sqlDBClose != nil {
+			_ = sqlDBClose()
+		}
+	})
+
+	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		t.Fatal(err)
+	}
+	sqlDB.SetMaxOpenConns(1)
+	sqlDBClose = sqlDB.Close
+	DB = db
+	if err := DB.AutoMigrate(
+		model.Server{},
+		model.Service{},
+		model.ServiceHistory{},
+		model.Notification{},
+		model.NotificationGroup{},
+		model.NotificationGroupNotification{},
+	); err != nil {
+		t.Fatal(err)
+	}
+
+	Conf = &ConfigClass{Config: &model.Config{AvgPingCount: 1}}
+	Cache = cache.New(time.Minute, time.Minute)
+	CronShared = &CronClass{
+		Cron:  cron.New(cron.WithSeconds()),
+		class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}},
+	}
+	NotificationShared = &NotificationClass{
+		class:         class[uint64, *model.Notification]{list: map[uint64]*model.Notification{}},
+		groupToIDList: map[uint64]map[uint64]*model.Notification{},
+		idToGroupList: map[uint64]map[uint64]struct{}{},
+		groupList:     map[uint64]string{},
+	}
+	TSDBShared = nil
+	Loc = time.UTC
+
+	serverClass := &ServerClass{
+		class: class[uint64, *model.Server]{
+			list: make(map[uint64]*model.Server),
+		},
+		uuidToID: make(map[string]uint64),
+	}
+	for _, server := range servers {
+		serverClass.list[server.ID] = server
+	}
+	ServerShared = serverClass
+
+	bus := make(chan *model.Service, 1)
+	ss, err := NewServiceSentinel(bus)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ServiceSentinelShared = ss
+	return ss
+}
+
+func addServiceMonitorSecurityService(t *testing.T, ss *ServiceSentinel, service *model.Service) {
+	t.Helper()
+
+	if err := DB.Create(service).Error; err != nil {
+		t.Fatal(err)
+	}
+	if err := ss.Update(service); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func serviceMonitorResult(reporter, serviceID uint64, taskType uint8, successful bool) ReportData {
+	return ReportData{
+		Reporter: reporter,
+		Data: &pb.TaskResult{
+			Id:         serviceID,
+			Type:       uint64(taskType),
+			Delay:      12,
+			Data:       "service monitor result",
+			Successful: successful,
+		},
+	}
+}
+
+func waitForServiceHistory(t *testing.T, serviceID, serverID uint64) {
+	t.Helper()
+
+	deadline := time.After(time.Second)
+	for {
+		var count int64
+		if err := DB.Model(&model.ServiceHistory{}).
+			Where("service_id = ? AND server_id = ?", serviceID, serverID).
+			Count(&count).Error; err != nil {
+			t.Fatal(err)
+		}
+		if count > 0 {
+			return
+		}
+
+		select {
+		case <-deadline:
+			t.Fatalf("expected service history for service %d from server %d", serviceID, serverID)
+		default:
+			time.Sleep(10 * time.Millisecond)
+		}
+	}
+}
+
+func assertNoServiceHistory(t *testing.T, serviceID, serverID uint64) {
+	t.Helper()
+
+	var count int64
+	if err := DB.Model(&model.ServiceHistory{}).
+		Where("service_id = ? AND server_id = ?", serviceID, serverID).
+		Count(&count).Error; err != nil {
+		t.Fatal(err)
+	}
+	if count != 0 {
+		t.Fatalf("expected no service history for service %d from server %d, got %d", serviceID, serverID, count)
+	}
+}
+
+func waitForTodayStats(t *testing.T, ss *ServiceSentinel, serviceID uint64, wantUp, wantDown uint64) {
+	t.Helper()
+
+	deadline := time.After(time.Second)
+	for {
+		ss.serviceResponseDataStoreLock.RLock()
+		stats := ss.serviceStatusToday[serviceID]
+		var up, down uint64
+		if stats != nil {
+			up = stats.Up
+			down = stats.Down
+		}
+		ss.serviceResponseDataStoreLock.RUnlock()
+
+		if up == wantUp && down == wantDown {
+			return
+		}
+		if down > wantDown {
+			t.Fatalf("expected service %d down count %d, got %d", serviceID, wantDown, down)
+		}
+
+		select {
+		case <-deadline:
+			t.Fatalf("expected service %d stats up=%d down=%d", serviceID, wantUp, wantDown)
+		default:
+			time.Sleep(10 * time.Millisecond)
+		}
+	}
+}