mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 09:40:12 +00:00
fix(rpc): authorize agent task results
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
naiba/CloudCode
@@ -5,6 +5,8 @@ import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/copier"
|
||||
|
||||
@@ -15,9 +17,13 @@ import (
|
||||
pb "github.com/nezhahq/nezha/proto"
|
||||
)
|
||||
|
||||
const alertTriggerCronResultAuthorizationTTL = 24 * time.Hour
|
||||
|
||||
type CronClass struct {
|
||||
class[uint64, *model.Cron]
|
||||
*cron.Cron
|
||||
pendingAlertTriggerTasksMu sync.Mutex
|
||||
pendingAlertTriggerTasks map[uint64]map[uint64][]time.Time
|
||||
}
|
||||
|
||||
func NewCronClass() *CronClass {
|
||||
@@ -64,7 +70,8 @@ func NewCronClass() *CronClass {
|
||||
list: list,
|
||||
sortedList: sortedList,
|
||||
},
|
||||
Cron: cronx,
|
||||
Cron: cronx,
|
||||
pendingAlertTriggerTasks: make(map[uint64]map[uint64][]time.Time),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,6 +85,7 @@ func (c *CronClass) Update(cr *model.Cron) {
|
||||
delete(c.list, cr.ID)
|
||||
c.list[cr.ID] = cr
|
||||
c.listMu.Unlock()
|
||||
c.deleteAlertTriggerCronResultAuthorizations([]uint64{cr.ID})
|
||||
|
||||
c.sortList()
|
||||
}
|
||||
@@ -92,6 +100,7 @@ func (c *CronClass) Delete(idList []uint64) {
|
||||
delete(c.list, id)
|
||||
}
|
||||
c.listMu.Unlock()
|
||||
c.deleteAlertTriggerCronResultAuthorizations(idList)
|
||||
|
||||
c.sortList()
|
||||
}
|
||||
@@ -130,6 +139,113 @@ func cronCanBeTriggeredByOwner(cr *model.Cron, triggerOwner uint64) bool {
|
||||
return cr.UserID == triggerOwner || userIsAdmin(triggerOwner)
|
||||
}
|
||||
|
||||
func CanReportCronResult(cr *model.Cron, reporter *model.Server) bool {
|
||||
if cr == nil || reporter == nil || !cronCanSendToServer(cr, reporter) {
|
||||
return false
|
||||
}
|
||||
if cr.Cover == model.CronCoverAll {
|
||||
return !slices.Contains(cr.Servers, reporter.ID)
|
||||
}
|
||||
if cr.Cover == model.CronCoverIgnoreAll {
|
||||
return slices.Contains(cr.Servers, reporter.ID)
|
||||
}
|
||||
if cr.Cover == model.CronCoverAlertTrigger {
|
||||
return CronShared != nil && CronShared.consumeAlertTriggerCronResult(cr.ID, reporter.ID)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *CronClass) reserveAlertTriggerCronResult(cronID uint64, serverID uint64) {
|
||||
c.pendingAlertTriggerTasksMu.Lock()
|
||||
defer c.pendingAlertTriggerTasksMu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
c.pruneExpiredAlertTriggerCronResultsLocked(now)
|
||||
if c.pendingAlertTriggerTasks == nil {
|
||||
c.pendingAlertTriggerTasks = make(map[uint64]map[uint64][]time.Time)
|
||||
}
|
||||
if c.pendingAlertTriggerTasks[cronID] == nil {
|
||||
c.pendingAlertTriggerTasks[cronID] = make(map[uint64][]time.Time)
|
||||
}
|
||||
c.pendingAlertTriggerTasks[cronID][serverID] = append(c.pendingAlertTriggerTasks[cronID][serverID], now.Add(alertTriggerCronResultAuthorizationTTL))
|
||||
}
|
||||
|
||||
func (c *CronClass) revokeAlertTriggerCronResult(cronID uint64, serverID uint64) {
|
||||
c.pendingAlertTriggerTasksMu.Lock()
|
||||
defer c.pendingAlertTriggerTasksMu.Unlock()
|
||||
|
||||
serverTasks := c.pendingAlertTriggerTasks[cronID]
|
||||
expiresAtList := serverTasks[serverID]
|
||||
if len(expiresAtList) == 0 {
|
||||
return
|
||||
}
|
||||
expiresAtList = expiresAtList[:len(expiresAtList)-1]
|
||||
if len(expiresAtList) == 0 {
|
||||
delete(serverTasks, serverID)
|
||||
} else {
|
||||
serverTasks[serverID] = expiresAtList
|
||||
}
|
||||
if len(serverTasks) == 0 {
|
||||
delete(c.pendingAlertTriggerTasks, cronID)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *CronClass) consumeAlertTriggerCronResult(cronID uint64, serverID uint64) bool {
|
||||
c.pendingAlertTriggerTasksMu.Lock()
|
||||
defer c.pendingAlertTriggerTasksMu.Unlock()
|
||||
|
||||
c.pruneExpiredAlertTriggerCronResultsLocked(time.Now())
|
||||
return c.consumeAlertTriggerCronResultLocked(cronID, serverID)
|
||||
}
|
||||
|
||||
func (c *CronClass) consumeAlertTriggerCronResultLocked(cronID uint64, serverID uint64) bool {
|
||||
serverTasks := c.pendingAlertTriggerTasks[cronID]
|
||||
expiresAtList := serverTasks[serverID]
|
||||
if len(expiresAtList) == 0 {
|
||||
return false
|
||||
}
|
||||
expiresAtList = expiresAtList[1:]
|
||||
if len(expiresAtList) == 0 {
|
||||
delete(serverTasks, serverID)
|
||||
} else {
|
||||
serverTasks[serverID] = expiresAtList
|
||||
}
|
||||
if len(serverTasks) == 0 {
|
||||
delete(c.pendingAlertTriggerTasks, cronID)
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *CronClass) pruneExpiredAlertTriggerCronResultsLocked(now time.Time) {
|
||||
for cronID, serverTasks := range c.pendingAlertTriggerTasks {
|
||||
for serverID, expiresAtList := range serverTasks {
|
||||
validExpiresAtList := expiresAtList[:0]
|
||||
for _, expiresAt := range expiresAtList {
|
||||
if expiresAt.After(now) {
|
||||
validExpiresAtList = append(validExpiresAtList, expiresAt)
|
||||
}
|
||||
}
|
||||
if len(validExpiresAtList) == 0 {
|
||||
delete(serverTasks, serverID)
|
||||
} else {
|
||||
serverTasks[serverID] = validExpiresAtList
|
||||
}
|
||||
}
|
||||
if len(serverTasks) == 0 {
|
||||
delete(c.pendingAlertTriggerTasks, cronID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *CronClass) deleteAlertTriggerCronResultAuthorizations(cronIDs []uint64) {
|
||||
c.pendingAlertTriggerTasksMu.Lock()
|
||||
defer c.pendingAlertTriggerTasksMu.Unlock()
|
||||
|
||||
for _, cronID := range cronIDs {
|
||||
delete(c.pendingAlertTriggerTasks, cronID)
|
||||
}
|
||||
}
|
||||
|
||||
func ManualTrigger(cr *model.Cron) {
|
||||
CronTrigger(cr)()
|
||||
}
|
||||
@@ -149,11 +265,17 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() {
|
||||
return
|
||||
}
|
||||
if s.TaskStream != nil {
|
||||
s.TaskStream.Send(&pb.Task{
|
||||
cronShared := CronShared
|
||||
if cronShared != nil {
|
||||
cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID)
|
||||
}
|
||||
if err := s.TaskStream.Send(&pb.Task{
|
||||
Id: cr.ID,
|
||||
Data: cr.Command,
|
||||
Type: model.TaskTypeCommand,
|
||||
})
|
||||
}); err != nil && cronShared != nil {
|
||||
cronShared.revokeAlertTriggerCronResult(cr.ID, s.ID)
|
||||
}
|
||||
} else {
|
||||
// 保存当前服务器状态信息
|
||||
curServer := model.Server{}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -472,16 +472,37 @@ func (ss *ServiceSentinel) CheckPermission(c *gin.Context, idList iter.Seq[uint6
|
||||
return true
|
||||
}
|
||||
|
||||
func canReportServiceResult(service *model.Service, reporter *model.Server, taskType uint64) bool {
|
||||
if service == nil || reporter == nil || uint64(service.Type) != taskType {
|
||||
return false
|
||||
}
|
||||
switch service.Cover {
|
||||
case model.ServiceCoverAll:
|
||||
if service.SkipServers[reporter.ID] {
|
||||
return false
|
||||
}
|
||||
case model.ServiceCoverIgnoreAll:
|
||||
if !service.SkipServers[reporter.ID] {
|
||||
return false
|
||||
}
|
||||
default:
|
||||
return false
|
||||
}
|
||||
|
||||
return service.UserID == reporter.UserID || userIsAdmin(service.UserID)
|
||||
}
|
||||
|
||||
// worker 服务监控的实际工作流程
|
||||
func (ss *ServiceSentinel) worker() {
|
||||
// 从服务状态汇报管道获取汇报的服务数据
|
||||
for r := range ss.serviceReportChannel {
|
||||
css, _ := ss.Get(r.Data.GetId())
|
||||
if css == nil || css.ID == 0 {
|
||||
cs, _ := ss.Get(r.Data.GetId())
|
||||
reporter, _ := ServerShared.Get(r.Reporter)
|
||||
// 入站结果必须匹配出站任务派发边界，避免 agent 伪造其他服务 ID 写入监控状态。
|
||||
if !canReportServiceResult(cs, reporter, r.Data.GetType()) {
|
||||
log.Printf("NEZHA>> Incorrect service monitor report %+v", r)
|
||||
continue
|
||||
}
|
||||
css = nil
|
||||
|
||||
mh := r.Data
|
||||
if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing {
|
||||
@@ -607,7 +628,7 @@ func (ss *ServiceSentinel) worker() {
|
||||
ss.serviceCurrentStatusData[mh.GetId()].result = ss.serviceCurrentStatusData[mh.GetId()].result[:0]
|
||||
}
|
||||
|
||||
cs, _ := ss.Get(mh.GetId())
|
||||
cs, _ = ss.Get(mh.GetId())
|
||||
m := ServerShared.GetList()
|
||||
// 延迟报警
|
||||
if mh.Delay > 0 {
|
||||
|
||||
Reference in New Issue
Block a user