tcjlj/nezha
1
0
Fork 0
mirror of https://github.com/nezhahq/nezha.git synced 2025-02-02 01:28:13 -05:00
Code Issues Actions Packages Projects Releases Wiki Activity

1

Browse Source
This commit is contained in:
uubulb 2024-12-21 23:39:27 +08:00
parent 18513110cf
commit 4085a928a8
7 changed files with 89 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
cmd/dashboard/controller/alertrule.go
View File

@ -62,7 +62,7 @@ func createAlertRule(c *gin.Context) (uint64, error) {
r.TriggerMode = arf.TriggerMode
r.Enable = &enable
if err := validateRule(&r); err != nil {
if err := validateRule(c, &r); err != nil {
return 0, err
}
@ -116,7 +116,7 @@ func updateAlertRule(c *gin.Context) (any, error) {
r.TriggerMode = arf.TriggerMode
r.Enable = &enable
if err := validateRule(&r); err != nil {
if err := validateRule(c, &r); err != nil {
return 0, err
}
@ -164,9 +164,34 @@ func batchDeleteAlertRule(c *gin.Context) (any, error) {
return nil, nil
}
func validateRule(r *model.AlertRule) error {
func validateRule(c *gin.Context, r *model.AlertRule) error {
if len(r.Rules) > 0 {
for _, rule := range r.Rules {
singleton.ServerLock.RLock()
isCoverAll := rule.Cover == model.RuleCoverAll
isCoverIgnoreAll := rule.Cover == model.RuleCoverIgnoreAll
for s, enabled := range rule.Ignore {
if isCoverAll {
for id, server := range singleton.ServerList {
if enabled && id == s {
continue
}
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return singleton.Localizer.ErrorT("permission denied")
}
}
} else if isCoverIgnoreAll && enabled {
if server, ok := singleton.ServerList[s]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return singleton.Localizer.ErrorT("permission denied")
}
}
}
}
singleton.ServerLock.RUnlock()
if !rule.IsTransferDurationRule() {
if rule.Duration < 3 {
return singleton.Localizer.ErrorT("duration need to be at least 3")

37
cmd/dashboard/controller/service.go
View File

@ -210,6 +210,10 @@ func createService(c *gin.Context) (uint64, error) {
m.RecoverTriggerTasks = mf.RecoverTriggerTasks
m.FailTriggerTasks = mf.FailTriggerTasks
if err := validateServers(c, &m); err != nil {
return 0, err
}
if err := singleton.DB.Create(&m).Error; err != nil {
return 0, newGormError("%v", err)
}
@ -284,6 +288,10 @@ func updateService(c *gin.Context) (any, error) {
m.RecoverTriggerTasks = mf.RecoverTriggerTasks
m.FailTriggerTasks = mf.FailTriggerTasks
if err := validateServers(c, &m); err != nil {
return 0, err
}
if err := singleton.DB.Save(&m).Error; err != nil {
return nil, newGormError("%v", err)
}
@ -351,3 +359,32 @@ func batchDeleteService(c *gin.Context) (any, error) {
singleton.ServiceSentinelShared.UpdateServiceList()
return nil, nil
}
func validateServers(c *gin.Context, ss *model.Service) error {
singleton.ServerLock.RLock()
defer singleton.ServerLock.RUnlock()
isCoverAll := ss.Cover == model.ServiceCoverAll
isCoverIgnoreAll := ss.Cover == model.ServiceCoverIgnoreAll
for s, enabled := range ss.SkipServers {
if isCoverAll {
for id, server := range singleton.ServerList {
if enabled && id == s {
continue
}
if !server.HasPermission(c) {
return singleton.Localizer.ErrorT("permission denied")
}
}
} else if isCoverIgnoreAll && enabled {
if server, ok := singleton.ServerList[s]; ok {
if !server.HasPermission(c) {
return singleton.Localizer.ErrorT("permission denied")
}
}
}
}
return nil
}

18
cmd/dashboard/rpc/rpc.go
View File

@ -100,12 +100,13 @@ func DispatchTask(serviceSentinelDispatchBus <-chan model.Service) {
continue
}
if task.Cover == model.ServiceCoverIgnoreAll && task.SkipServers[singleton.SortedServerList[workedServerIndex].ID] {
var role uint8 = model.RoleMember
server := singleton.SortedServerList[workedServerIndex]
if err := singleton.DB.Model(&model.User{}).Select("role").Where("id = ?", task.UserID).Limit(1).Scan(&role).Error; err != nil {
workedServerIndex++
continue
singleton.UserLock.RLock()
role, ok := singleton.UserRoleMap[server.UserID]
if !ok {
role = model.RoleMember
}
singleton.UserLock.RUnlock()
if task.UserID == server.UserID || role == model.RoleAdmin {
singleton.SortedServerList[workedServerIndex].TaskStream.Send(task.PB())
}
@ -113,12 +114,13 @@ func DispatchTask(serviceSentinelDispatchBus <-chan model.Service) {
continue
}
if task.Cover == model.ServiceCoverAll && !task.SkipServers[singleton.SortedServerList[workedServerIndex].ID] {
var role uint8 = model.RoleMember
server := singleton.SortedServerList[workedServerIndex]
if err := singleton.DB.Model(&model.User{}).Select("role").Where("id = ?", task.UserID).Limit(1).Scan(&role).Error; err != nil {
workedServerIndex++
continue
singleton.UserLock.RLock()
role, ok := singleton.UserRoleMap[server.UserID]
if !ok {
role = model.RoleMember
}
singleton.UserLock.RUnlock()
if task.UserID == server.UserID || role == model.RoleAdmin {
singleton.SortedServerList[workedServerIndex].TaskStream.Send(task.PB())
}

6
model/alertrule.go
View File

@ -62,13 +62,9 @@ func (r *AlertRule) Enabled() bool {
}
// Snapshot 对传入的Server进行该报警规则下所有type的检查 返回每项检查结果
func (r *AlertRule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server, db *gorm.DB) []bool {
func (r *AlertRule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server, db *gorm.DB, role uint8) []bool {
point := make([]bool, len(r.Rules))
var role uint8 = RoleMember
if err := db.Model(&User{}).Select("role").Where("id = ?", r.UserID).Limit(1).Scan(&role).Error; err != nil {
return point
}
if r.UserID != server.UserID && role != RoleAdmin {
return point
}

3
service/rpc/auth.go
View File

@ -2,7 +2,6 @@ package rpc
import (
"context"
"crypto/subtle"
"strings"
petname "github.com/dustinkirkland/golang-petname"
@ -39,7 +38,7 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
singleton.UserLock.RLock()
userId, ok := singleton.AgentSecretToUserId[clientSecret]
if !ok && subtle.ConstantTimeCompare([]byte(clientSecret), []byte(singleton.Conf.AgentSecretKey)) != 1 {
if !ok && clientSecret != singleton.Conf.AgentSecretKey {
singleton.UserLock.RUnlock()
model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail, model.BlockIDgRPC)
return 0, status.Error(codes.Unauthenticated, "客户端认证失败")

8
service/singleton/alertsentinel.go
View File

@ -143,8 +143,14 @@ func checkStatus() {
}
for _, server := range ServerList {
// 监测点
UserLock.RLock()
role, ok := UserRoleMap[alert.UserID]
if !ok {
role = model.RoleMember
}
UserLock.RUnlock()
alertsStore[alert.ID][server.ID] = append(alertsStore[alert.
ID][server.ID], alert.Snapshot(AlertsCycleTransferStatsStore[alert.ID], server, DB))
ID][server.ID], alert.Snapshot(AlertsCycleTransferStatsStore[alert.ID], server, DB, role))
// 发送通知，分为触发报警和恢复通知
max, passed := alert.Check(alertsStore[alert.ID][server.ID])
// 保存当前服务器状态信息

5
service/singleton/user.go
View File

@ -11,12 +11,15 @@ var (
UserIdToAgentSecret map[uint64]string
AgentSecretToUserId map[string]uint64
UserRoleMap map[uint64]uint8
UserLock sync.RWMutex
)
func initUser() {
UserIdToAgentSecret = make(map[uint64]string)
AgentSecretToUserId = make(map[string]uint64)
UserRoleMap = make(map[uint64]uint8)
var users []model.User
DB.Find(&users)
@ -24,6 +27,7 @@ func initUser() {
for _, u := range users {
UserIdToAgentSecret[u.ID] = u.AgentSecret
AgentSecretToUserId[u.AgentSecret] = u.ID
UserRoleMap[u.ID] = u.Role
}
}
@ -37,6 +41,7 @@ func OnUserUpdate(u *model.User) {
UserIdToAgentSecret[u.ID] = u.AgentSecret
AgentSecretToUserId[u.AgentSecret] = u.ID
UserRoleMap[u.ID] = u.Role
}
func OnUserDelete(id []uint64, errorFunc func(string, ...interface{}) error) error {