feat: user roles (#852)

* [WIP] feat: user roles

* update

* update

* admin handler

* update

* feat: user-specific connection secret

* simplify some logics

* cleanup

* update waf

* update user api error handling

* update waf api

* fix codeql

* update waf table

* fix several problems

* add pagination for waf api

* update permission checks

* switch to runtime check

* 1

* cover?

* some changes
This commit is contained in:
UUBulb 2024-12-22 00:05:41 +08:00 committed by GitHub
parent 50ee62172f
commit 653d0cf2e9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
35 changed files with 841 additions and 180 deletions

View File

@ -50,6 +50,9 @@ func createAlertRule(c *gin.Context) (uint64, error) {
return 0, err return 0, err
} }
uid := getUid(c)
r.UserID = uid
r.Name = arf.Name r.Name = arf.Name
r.Rules = arf.Rules r.Rules = arf.Rules
r.FailTriggerTasks = arf.FailTriggerTasks r.FailTriggerTasks = arf.FailTriggerTasks
@ -59,7 +62,7 @@ func createAlertRule(c *gin.Context) (uint64, error) {
r.TriggerMode = arf.TriggerMode r.TriggerMode = arf.TriggerMode
r.Enable = &enable r.Enable = &enable
if err := validateRule(&r); err != nil { if err := validateRule(c, &r); err != nil {
return 0, err return 0, err
} }
@ -100,6 +103,10 @@ func updateAlertRule(c *gin.Context) (any, error) {
return nil, singleton.Localizer.ErrorT("alert id %d does not exist", id) return nil, singleton.Localizer.ErrorT("alert id %d does not exist", id)
} }
if !r.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
r.Name = arf.Name r.Name = arf.Name
r.Rules = arf.Rules r.Rules = arf.Rules
r.FailTriggerTasks = arf.FailTriggerTasks r.FailTriggerTasks = arf.FailTriggerTasks
@ -109,7 +116,7 @@ func updateAlertRule(c *gin.Context) (any, error) {
r.TriggerMode = arf.TriggerMode r.TriggerMode = arf.TriggerMode
r.Enable = &enable r.Enable = &enable
if err := validateRule(&r); err != nil { if err := validateRule(c, &r); err != nil {
return 0, err return 0, err
} }
@ -134,11 +141,21 @@ func updateAlertRule(c *gin.Context) (any, error) {
// @Router /batch-delete/alert-rule [post] // @Router /batch-delete/alert-rule [post]
func batchDeleteAlertRule(c *gin.Context) (any, error) { func batchDeleteAlertRule(c *gin.Context) (any, error) {
var ar []uint64 var ar []uint64
if err := c.ShouldBindJSON(&ar); err != nil { if err := c.ShouldBindJSON(&ar); err != nil {
return nil, err return nil, err
} }
var ars []model.AlertRule
if err := singleton.DB.Where("id in (?)", ar).Find(&ars).Error; err != nil {
return nil, err
}
for _, a := range ars {
if !a.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
if err := singleton.DB.Unscoped().Delete(&model.AlertRule{}, "id in (?)", ar).Error; err != nil { if err := singleton.DB.Unscoped().Delete(&model.AlertRule{}, "id in (?)", ar).Error; err != nil {
return nil, newGormError("%v", err) return nil, newGormError("%v", err)
} }
@ -147,9 +164,20 @@ func batchDeleteAlertRule(c *gin.Context) (any, error) {
return nil, nil return nil, nil
} }
func validateRule(r *model.AlertRule) error { func validateRule(c *gin.Context, r *model.AlertRule) error {
if len(r.Rules) > 0 { if len(r.Rules) > 0 {
for _, rule := range r.Rules { for _, rule := range r.Rules {
singleton.ServerLock.RLock()
for s := range rule.Ignore {
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.IsTransferDurationRule() {
if rule.Duration < 3 { if rule.Duration < 3 {
return singleton.Localizer.ErrorT("duration need to be at least 3") return singleton.Localizer.ErrorT("duration need to be at least 3")

View File

@ -9,6 +9,7 @@ import (
"net/http" "net/http"
"os" "os"
"path" "path"
"slices"
"strings" "strings"
jwt "github.com/appleboy/gin-jwt/v2" jwt "github.com/appleboy/gin-jwt/v2"
@ -78,11 +79,11 @@ func routers(r *gin.Engine, frontendDist fs.FS) {
auth.GET("/profile", commonHandler(getProfile)) auth.GET("/profile", commonHandler(getProfile))
auth.POST("/profile", commonHandler(updateProfile)) auth.POST("/profile", commonHandler(updateProfile))
auth.GET("/user", commonHandler(listUser)) auth.GET("/user", adminHandler(listUser))
auth.POST("/user", commonHandler(createUser)) auth.POST("/user", adminHandler(createUser))
auth.POST("/batch-delete/user", commonHandler(batchDeleteUser)) auth.POST("/batch-delete/user", adminHandler(batchDeleteUser))
auth.GET("/service/list", commonHandler(listService)) auth.GET("/service/list", listHandler(listService))
auth.POST("/service", commonHandler(createService)) auth.POST("/service", commonHandler(createService))
auth.PATCH("/service/:id", commonHandler(updateService)) auth.PATCH("/service/:id", commonHandler(updateService))
auth.POST("/batch-delete/service", commonHandler(batchDeleteService)) auth.POST("/batch-delete/service", commonHandler(batchDeleteService))
@ -96,42 +97,42 @@ func routers(r *gin.Engine, frontendDist fs.FS) {
auth.PATCH("/notification-group/:id", commonHandler(updateNotificationGroup)) auth.PATCH("/notification-group/:id", commonHandler(updateNotificationGroup))
auth.POST("/batch-delete/notification-group", commonHandler(batchDeleteNotificationGroup)) auth.POST("/batch-delete/notification-group", commonHandler(batchDeleteNotificationGroup))
auth.GET("/server", commonHandler(listServer)) auth.GET("/server", listHandler(listServer))
auth.PATCH("/server/:id", commonHandler(updateServer)) auth.PATCH("/server/:id", commonHandler(updateServer))
auth.POST("/batch-delete/server", commonHandler(batchDeleteServer)) auth.POST("/batch-delete/server", commonHandler(batchDeleteServer))
auth.POST("/force-update/server", commonHandler(forceUpdateServer)) auth.POST("/force-update/server", commonHandler(forceUpdateServer))
auth.GET("/notification", commonHandler(listNotification)) auth.GET("/notification", listHandler(listNotification))
auth.POST("/notification", commonHandler(createNotification)) auth.POST("/notification", commonHandler(createNotification))
auth.PATCH("/notification/:id", commonHandler(updateNotification)) auth.PATCH("/notification/:id", commonHandler(updateNotification))
auth.POST("/batch-delete/notification", commonHandler(batchDeleteNotification)) auth.POST("/batch-delete/notification", commonHandler(batchDeleteNotification))
auth.GET("/alert-rule", commonHandler(listAlertRule)) auth.GET("/alert-rule", listHandler(listAlertRule))
auth.POST("/alert-rule", commonHandler(createAlertRule)) auth.POST("/alert-rule", commonHandler(createAlertRule))
auth.PATCH("/alert-rule/:id", commonHandler(updateAlertRule)) auth.PATCH("/alert-rule/:id", commonHandler(updateAlertRule))
auth.POST("/batch-delete/alert-rule", commonHandler(batchDeleteAlertRule)) auth.POST("/batch-delete/alert-rule", commonHandler(batchDeleteAlertRule))
auth.GET("/cron", commonHandler(listCron)) auth.GET("/cron", listHandler(listCron))
auth.POST("/cron", commonHandler(createCron)) auth.POST("/cron", commonHandler(createCron))
auth.PATCH("/cron/:id", commonHandler(updateCron)) auth.PATCH("/cron/:id", commonHandler(updateCron))
auth.GET("/cron/:id/manual", commonHandler(manualTriggerCron)) auth.GET("/cron/:id/manual", commonHandler(manualTriggerCron))
auth.POST("/batch-delete/cron", commonHandler(batchDeleteCron)) auth.POST("/batch-delete/cron", commonHandler(batchDeleteCron))
auth.GET("/ddns", commonHandler(listDDNS)) auth.GET("/ddns", listHandler(listDDNS))
auth.GET("/ddns/providers", commonHandler(listProviders)) auth.GET("/ddns/providers", commonHandler(listProviders))
auth.POST("/ddns", commonHandler(createDDNS)) auth.POST("/ddns", commonHandler(createDDNS))
auth.PATCH("/ddns/:id", commonHandler(updateDDNS)) auth.PATCH("/ddns/:id", commonHandler(updateDDNS))
auth.POST("/batch-delete/ddns", commonHandler(batchDeleteDDNS)) auth.POST("/batch-delete/ddns", commonHandler(batchDeleteDDNS))
auth.GET("/nat", commonHandler(listNAT)) auth.GET("/nat", listHandler(listNAT))
auth.POST("/nat", commonHandler(createNAT)) auth.POST("/nat", commonHandler(createNAT))
auth.PATCH("/nat/:id", commonHandler(updateNAT)) auth.PATCH("/nat/:id", commonHandler(updateNAT))
auth.POST("/batch-delete/nat", commonHandler(batchDeleteNAT)) auth.POST("/batch-delete/nat", commonHandler(batchDeleteNAT))
auth.GET("/waf", commonHandler(listBlockedAddress)) auth.GET("/waf", pCommonHandler(listBlockedAddress))
auth.POST("/batch-delete/waf", commonHandler(batchDeleteBlockedAddress)) auth.POST("/batch-delete/waf", adminHandler(batchDeleteBlockedAddress))
auth.PATCH("/setting", commonHandler(updateConfig)) auth.PATCH("/setting", adminHandler(updateConfig))
r.NoRoute(fallbackToFrontend(frontendDist)) r.NoRoute(fallbackToFrontend(frontendDist))
} }
@ -152,6 +153,7 @@ func newErrorResponse(err error) model.CommonResponse[any] {
} }
type handlerFunc[T any] func(c *gin.Context) (T, error) type handlerFunc[T any] func(c *gin.Context) (T, error)
type pHandlerFunc[S ~[]E, E any] func(c *gin.Context) (*model.Value[S], error)
// There are many error types in gorm, so create a custom type to represent all // There are many error types in gorm, so create a custom type to represent all
// gorm errors here instead // gorm errors here instead
@ -189,29 +191,86 @@ func (we *wsError) Error() string {
func commonHandler[T any](handler handlerFunc[T]) func(*gin.Context) { func commonHandler[T any](handler handlerFunc[T]) func(*gin.Context) {
return func(c *gin.Context) { return func(c *gin.Context) {
data, err := handler(c) handle(c, handler)
if err == nil { }
c.JSON(http.StatusOK, model.CommonResponse[T]{Success: true, Data: data}) }
func adminHandler[T any](handler handlerFunc[T]) func(*gin.Context) {
return func(c *gin.Context) {
auth, ok := c.Get(model.CtxKeyAuthorizedUser)
if !ok {
c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("unauthorized")))
return return
} }
switch err.(type) {
case *gormError: user := *auth.(*model.User)
log.Printf("NEZHA>> gorm error: %v", err) if user.Role != model.RoleAdmin {
c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("database error"))) c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("permission denied")))
return return
case *wsError: }
// Connection is upgraded to WebSocket, so c.Writer is no longer usable
if msg := err.Error(); msg != "" { handle(c, handler)
log.Printf("NEZHA>> websocket error: %v", err) }
} }
return
default: func handle[T any](c *gin.Context, handler handlerFunc[T]) {
data, err := handler(c)
if err == nil {
c.JSON(http.StatusOK, model.CommonResponse[T]{Success: true, Data: data})
return
}
switch err.(type) {
case *gormError:
log.Printf("NEZHA>> gorm error: %v", err)
c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("database error")))
return
case *wsError:
// Connection is upgraded to WebSocket, so c.Writer is no longer usable
if msg := err.Error(); msg != "" {
log.Printf("NEZHA>> websocket error: %v", err)
}
return
default:
c.JSON(http.StatusOK, newErrorResponse(err))
return
}
}
func listHandler[S ~[]E, E model.CommonInterface](handler handlerFunc[S]) func(*gin.Context) {
return func(c *gin.Context) {
data, err := handler(c)
if err != nil {
c.JSON(http.StatusOK, newErrorResponse(err)) c.JSON(http.StatusOK, newErrorResponse(err))
return return
} }
c.JSON(http.StatusOK, model.CommonResponse[S]{Success: true, Data: filter(c, data)})
} }
} }
func pCommonHandler[S ~[]E, E any](handler pHandlerFunc[S, E]) func(*gin.Context) {
return func(c *gin.Context) {
data, err := handler(c)
if err != nil {
c.JSON(http.StatusOK, newErrorResponse(err))
return
}
c.JSON(http.StatusOK, model.PaginatedResponse[S, E]{Success: true, Data: data})
}
}
func filter[S ~[]E, E model.CommonInterface](ctx *gin.Context, s S) S {
return slices.DeleteFunc(s, func(e E) bool {
return !e.HasPermission(ctx)
})
}
func getUid(c *gin.Context) uint64 {
user, _ := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User)
return user.ID
}
func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) { func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) {
checkLocalFileOrFs := func(c *gin.Context, fs fs.FS, path string) bool { checkLocalFileOrFs := func(c *gin.Context, fs fs.FS, path string) bool {
if _, err := os.Stat(path); err == nil { if _, err := os.Stat(path); err == nil {

View File

@ -1,7 +1,6 @@
package controller package controller
import ( import (
"fmt"
"strconv" "strconv"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
@ -50,6 +49,18 @@ func createCron(c *gin.Context) (uint64, error) {
return 0, err return 0, err
} }
singleton.ServerLock.RLock()
for _, sid := range cf.Servers {
if server, ok := singleton.ServerList[sid]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return 0, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.ServerLock.RUnlock()
cr.UserID = getUid(c)
cr.TaskType = cf.TaskType cr.TaskType = cf.TaskType
cr.Name = cf.Name cr.Name = cf.Name
cr.Scheduler = cf.Scheduler cr.Scheduler = cf.Scheduler
@ -104,9 +115,24 @@ func updateCron(c *gin.Context) (any, error) {
return 0, err return 0, err
} }
singleton.ServerLock.RLock()
for _, sid := range cf.Servers {
if server, ok := singleton.ServerList[sid]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.ServerLock.RUnlock()
var cr model.Cron var cr model.Cron
if err := singleton.DB.First(&cr, id).Error; err != nil { if err := singleton.DB.First(&cr, id).Error; err != nil {
return nil, fmt.Errorf("task id %d does not exist", id) return nil, singleton.Localizer.ErrorT("task id %d does not exist", id)
}
if !cr.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
} }
cr.TaskType = cf.TaskType cr.TaskType = cf.TaskType
@ -156,12 +182,19 @@ func manualTriggerCron(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
var cr model.Cron singleton.CronLock.RLock()
if err := singleton.DB.First(&cr, id).Error; err != nil { cr, ok := singleton.Crons[id]
if !ok {
singleton.CronLock.RUnlock()
return nil, singleton.Localizer.ErrorT("task id %d does not exist", id) return nil, singleton.Localizer.ErrorT("task id %d does not exist", id)
} }
singleton.CronLock.RUnlock()
singleton.ManualTrigger(&cr) if !cr.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
singleton.ManualTrigger(cr)
return nil, nil return nil, nil
} }
@ -178,11 +211,21 @@ func manualTriggerCron(c *gin.Context) (any, error) {
// @Router /batch-delete/cron [post] // @Router /batch-delete/cron [post]
func batchDeleteCron(c *gin.Context) (any, error) { func batchDeleteCron(c *gin.Context) (any, error) {
var cr []uint64 var cr []uint64
if err := c.ShouldBindJSON(&cr); err != nil { if err := c.ShouldBindJSON(&cr); err != nil {
return nil, err return nil, err
} }
singleton.CronLock.RLock()
for _, crID := range cr {
if crn, ok := singleton.Crons[crID]; ok {
if !crn.HasPermission(c) {
singleton.CronLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.CronLock.RUnlock()
if err := singleton.DB.Unscoped().Delete(&model.Cron{}, "id in (?)", cr).Error; err != nil { if err := singleton.DB.Unscoped().Delete(&model.Cron{}, "id in (?)", cr).Error; err != nil {
return nil, newGormError("%v", err) return nil, newGormError("%v", err)
} }

View File

@ -56,6 +56,7 @@ func createDDNS(c *gin.Context) (uint64, error) {
return 0, singleton.Localizer.ErrorT("the retry count must be an integer between 1 and 10") return 0, singleton.Localizer.ErrorT("the retry count must be an integer between 1 and 10")
} }
p.UserID = getUid(c)
p.Name = df.Name p.Name = df.Name
enableIPv4 := df.EnableIPv4 enableIPv4 := df.EnableIPv4
enableIPv6 := df.EnableIPv6 enableIPv6 := df.EnableIPv6
@ -125,6 +126,10 @@ func updateDDNS(c *gin.Context) (any, error) {
return nil, singleton.Localizer.ErrorT("profile id %d does not exist", id) return nil, singleton.Localizer.ErrorT("profile id %d does not exist", id)
} }
if !p.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
p.Name = df.Name p.Name = df.Name
enableIPv4 := df.EnableIPv4 enableIPv4 := df.EnableIPv4
enableIPv6 := df.EnableIPv6 enableIPv6 := df.EnableIPv6
@ -178,6 +183,17 @@ func batchDeleteDDNS(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
singleton.DDNSCacheLock.RLock()
for _, pid := range ddnsConfigs {
if p, ok := singleton.DDNSCache[pid]; ok {
if !p.HasPermission(c) {
singleton.DDNSCacheLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.DDNSCacheLock.RUnlock()
if err := singleton.DB.Unscoped().Delete(&model.DDNSProfile{}, "id in (?)", ddnsConfigs).Error; err != nil { if err := singleton.DB.Unscoped().Delete(&model.DDNSProfile{}, "id in (?)", ddnsConfigs).Error; err != nil {
return nil, newGormError("%v", err) return nil, newGormError("%v", err)
} }

View File

@ -7,6 +7,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
"github.com/hashicorp/go-uuid" "github.com/hashicorp/go-uuid"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/pkg/utils"
"github.com/nezhahq/nezha/pkg/websocketx" "github.com/nezhahq/nezha/pkg/websocketx"
@ -31,13 +32,6 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
return nil, err return nil, err
} }
streamId, err := uuid.GenerateUUID()
if err != nil {
return nil, err
}
rpc.NezhaHandlerSingleton.CreateStream(streamId)
singleton.ServerLock.RLock() singleton.ServerLock.RLock()
server := singleton.ServerList[id] server := singleton.ServerList[id]
singleton.ServerLock.RUnlock() singleton.ServerLock.RUnlock()
@ -45,6 +39,17 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
return nil, singleton.Localizer.ErrorT("server not found or not connected") return nil, singleton.Localizer.ErrorT("server not found or not connected")
} }
if !server.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
streamId, err := uuid.GenerateUUID()
if err != nil {
return nil, err
}
rpc.NezhaHandlerSingleton.CreateStream(streamId)
fmData, _ := utils.Json.Marshal(&model.TaskFM{ fmData, _ := utils.Json.Marshal(&model.TaskFM{
StreamID: streamId, StreamID: streamId,
}) })

View File

@ -88,18 +88,21 @@ func authenticator() func(c *gin.Context) (interface{}, error) {
} }
var user model.User var user model.User
realip := c.GetString(model.CtxKeyRealIPStr)
if err := singleton.DB.Select("id", "password").Where("username = ?", loginVals.Username).First(&user).Error; err != nil { if err := singleton.DB.Select("id", "password").Where("username = ?", loginVals.Username).First(&user).Error; err != nil {
if err == gorm.ErrRecordNotFound { if err == gorm.ErrRecordNotFound {
model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeLoginFail) model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeLoginFail, model.BlockIDUnknownUser)
} }
return nil, jwt.ErrFailedAuthentication return nil, jwt.ErrFailedAuthentication
} }
if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(loginVals.Password)); err != nil { if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(loginVals.Password)); err != nil {
model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeLoginFail) model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeLoginFail, int64(user.ID))
return nil, jwt.ErrFailedAuthentication return nil, jwt.ErrFailedAuthentication
} }
model.ClearIP(singleton.DB, realip, model.BlockIDUnknownUser)
model.ClearIP(singleton.DB, realip, int64(user.ID))
return utils.Itoa(user.ID), nil return utils.Itoa(user.ID), nil
} }
} }
@ -169,10 +172,10 @@ func optionalAuthMiddleware(mw *jwt.GinJWTMiddleware) func(c *gin.Context) {
identity := mw.IdentityHandler(c) identity := mw.IdentityHandler(c)
if identity != nil { if identity != nil {
model.ClearIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr)) model.ClearIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.BlockIDToken)
c.Set(mw.IdentityKey, identity) c.Set(mw.IdentityKey, identity)
} else { } else {
if err := model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeBruteForceToken); err != nil { if err := model.BlockIP(singleton.DB, c.GetString(model.CtxKeyRealIPStr), model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken); err != nil {
waf.ShowBlockPage(c, err) waf.ShowBlockPage(c, err)
return return
} }

View File

@ -51,6 +51,18 @@ func createNAT(c *gin.Context) (uint64, error) {
return 0, err return 0, err
} }
singleton.ServerLock.RLock()
if server, ok := singleton.ServerList[nf.ServerID]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return 0, singleton.Localizer.ErrorT("permission denied")
}
}
singleton.ServerLock.RUnlock()
uid := getUid(c)
n.UserID = uid
n.Name = nf.Name n.Name = nf.Name
n.Domain = nf.Domain n.Domain = nf.Domain
n.Host = nf.Host n.Host = nf.Host
@ -90,11 +102,24 @@ func updateNAT(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
singleton.ServerLock.RLock()
if server, ok := singleton.ServerList[nf.ServerID]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
singleton.ServerLock.RUnlock()
var n model.NAT var n model.NAT
if err = singleton.DB.First(&n, id).Error; err != nil { if err = singleton.DB.First(&n, id).Error; err != nil {
return nil, singleton.Localizer.ErrorT("profile id %d does not exist", id) return nil, singleton.Localizer.ErrorT("profile id %d does not exist", id)
} }
if !n.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
n.Name = nf.Name n.Name = nf.Name
n.Domain = nf.Domain n.Domain = nf.Domain
n.Host = nf.Host n.Host = nf.Host
@ -122,11 +147,21 @@ func updateNAT(c *gin.Context) (any, error) {
// @Router /batch-delete/nat [post] // @Router /batch-delete/nat [post]
func batchDeleteNAT(c *gin.Context) (any, error) { func batchDeleteNAT(c *gin.Context) (any, error) {
var n []uint64 var n []uint64
if err := c.ShouldBindJSON(&n); err != nil { if err := c.ShouldBindJSON(&n); err != nil {
return nil, err return nil, err
} }
singleton.NATCacheRwLock.RLock()
for _, id := range n {
if p, ok := singleton.NATCache[singleton.NATIDToDomain[id]]; ok {
if !p.HasPermission(c) {
singleton.NATCacheRwLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.NATCacheRwLock.RUnlock()
if err := singleton.DB.Unscoped().Delete(&model.NAT{}, "id in (?)", n).Error; err != nil { if err := singleton.DB.Unscoped().Delete(&model.NAT{}, "id in (?)", n).Error; err != nil {
return nil, newGormError("%v", err) return nil, newGormError("%v", err)
} }

View File

@ -48,6 +48,7 @@ func createNotification(c *gin.Context) (uint64, error) {
} }
var n model.Notification var n model.Notification
n.UserID = getUid(c)
n.Name = nf.Name n.Name = nf.Name
n.RequestMethod = nf.RequestMethod n.RequestMethod = nf.RequestMethod
n.RequestType = nf.RequestType n.RequestType = nf.RequestType
@ -106,6 +107,10 @@ func updateNotification(c *gin.Context) (any, error) {
return nil, singleton.Localizer.ErrorT("notification id %d does not exist", id) return nil, singleton.Localizer.ErrorT("notification id %d does not exist", id)
} }
if !n.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
n.Name = nf.Name n.Name = nf.Name
n.RequestMethod = nf.RequestMethod n.RequestMethod = nf.RequestMethod
n.RequestType = nf.RequestType n.RequestType = nf.RequestType
@ -149,11 +154,20 @@ func updateNotification(c *gin.Context) (any, error) {
// @Router /batch-delete/notification [post] // @Router /batch-delete/notification [post]
func batchDeleteNotification(c *gin.Context) (any, error) { func batchDeleteNotification(c *gin.Context) (any, error) {
var n []uint64 var n []uint64
if err := c.ShouldBindJSON(&n); err != nil { if err := c.ShouldBindJSON(&n); err != nil {
return nil, err return nil, err
} }
singleton.NotificationsLock.RLock()
for _, nid := range n {
if ns, ok := singleton.NotificationMap[nid]; ok {
if !ns.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.NotificationsLock.RUnlock()
err := singleton.DB.Transaction(func(tx *gorm.DB) error { err := singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Unscoped().Delete(&model.Notification{}, "id in (?)", n).Error; err != nil { if err := tx.Unscoped().Delete(&model.Notification{}, "id in (?)", n).Error; err != nil {
return err return err

View File

@ -20,7 +20,7 @@ import (
// @Produce json // @Produce json
// @Success 200 {object} model.CommonResponse[[]model.NotificationGroupResponseItem] // @Success 200 {object} model.CommonResponse[[]model.NotificationGroupResponseItem]
// @Router /notification-group [get] // @Router /notification-group [get]
func listNotificationGroup(c *gin.Context) ([]model.NotificationGroupResponseItem, error) { func listNotificationGroup(c *gin.Context) ([]*model.NotificationGroupResponseItem, error) {
var ng []model.NotificationGroup var ng []model.NotificationGroup
if err := singleton.DB.Find(&ng).Error; err != nil { if err := singleton.DB.Find(&ng).Error; err != nil {
return nil, err return nil, err
@ -39,9 +39,9 @@ func listNotificationGroup(c *gin.Context) ([]model.NotificationGroupResponseIte
groupNotifications[n.NotificationGroupID] = append(groupNotifications[n.NotificationGroupID], n.NotificationID) groupNotifications[n.NotificationGroupID] = append(groupNotifications[n.NotificationGroupID], n.NotificationID)
} }
ngRes := make([]model.NotificationGroupResponseItem, 0, len(ng)) ngRes := make([]*model.NotificationGroupResponseItem, 0, len(ng))
for _, n := range ng { for _, n := range ng {
ngRes = append(ngRes, model.NotificationGroupResponseItem{ ngRes = append(ngRes, &model.NotificationGroupResponseItem{
Group: n, Group: n,
Notifications: groupNotifications[n.ID], Notifications: groupNotifications[n.ID],
}) })
@ -68,8 +68,22 @@ func createNotificationGroup(c *gin.Context) (uint64, error) {
} }
ngf.Notifications = slices.Compact(ngf.Notifications) ngf.Notifications = slices.Compact(ngf.Notifications)
singleton.NotificationsLock.RLock()
for _, nid := range ngf.Notifications {
if n, ok := singleton.NotificationMap[nid]; ok {
if !n.HasPermission(c) {
singleton.NotificationsLock.RUnlock()
return 0, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.NotificationsLock.RUnlock()
uid := getUid(c)
var ng model.NotificationGroup var ng model.NotificationGroup
ng.Name = ngf.Name ng.Name = ngf.Name
ng.UserID = uid
var count int64 var count int64
if err := singleton.DB.Model(&model.Notification{}).Where("id in (?)", ngf.Notifications).Count(&count).Error; err != nil { if err := singleton.DB.Model(&model.Notification{}).Where("id in (?)", ngf.Notifications).Count(&count).Error; err != nil {
@ -86,6 +100,9 @@ func createNotificationGroup(c *gin.Context) (uint64, error) {
} }
for _, n := range ngf.Notifications { for _, n := range ngf.Notifications {
if err := tx.Create(&model.NotificationGroupNotification{ if err := tx.Create(&model.NotificationGroupNotification{
Common: model.Common{
UserID: uid,
},
NotificationGroupID: ng.ID, NotificationGroupID: ng.ID,
NotificationID: n, NotificationID: n,
}).Error; err != nil { }).Error; err != nil {
@ -126,11 +143,27 @@ func updateNotificationGroup(c *gin.Context) (any, error) {
if err := c.ShouldBindJSON(&ngf); err != nil { if err := c.ShouldBindJSON(&ngf); err != nil {
return nil, err return nil, err
} }
singleton.NotificationsLock.RLock()
for _, nid := range ngf.Notifications {
if n, ok := singleton.NotificationMap[nid]; ok {
if !n.HasPermission(c) {
singleton.NotificationsLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.NotificationsLock.RUnlock()
var ngDB model.NotificationGroup var ngDB model.NotificationGroup
if err := singleton.DB.First(&ngDB, id).Error; err != nil { if err := singleton.DB.First(&ngDB, id).Error; err != nil {
return nil, singleton.Localizer.ErrorT("group id %d does not exist", id) return nil, singleton.Localizer.ErrorT("group id %d does not exist", id)
} }
if !ngDB.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
ngDB.Name = ngf.Name ngDB.Name = ngf.Name
ngf.Notifications = slices.Compact(ngf.Notifications) ngf.Notifications = slices.Compact(ngf.Notifications)
@ -142,6 +175,8 @@ func updateNotificationGroup(c *gin.Context) (any, error) {
return nil, singleton.Localizer.ErrorT("have invalid notification id") return nil, singleton.Localizer.ErrorT("have invalid notification id")
} }
uid := getUid(c)
err = singleton.DB.Transaction(func(tx *gorm.DB) error { err = singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Save(&ngDB).Error; err != nil { if err := tx.Save(&ngDB).Error; err != nil {
return err return err
@ -152,6 +187,9 @@ func updateNotificationGroup(c *gin.Context) (any, error) {
for _, n := range ngf.Notifications { for _, n := range ngf.Notifications {
if err := tx.Create(&model.NotificationGroupNotification{ if err := tx.Create(&model.NotificationGroupNotification{
Common: model.Common{
UserID: uid,
},
NotificationGroupID: ngDB.ID, NotificationGroupID: ngDB.ID,
NotificationID: n, NotificationID: n,
}).Error; err != nil { }).Error; err != nil {
@ -185,6 +223,17 @@ func batchDeleteNotificationGroup(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
var ng []model.NotificationGroup
if err := singleton.DB.Where("id in (?)", ngn).Find(&ng).Error; err != nil {
return nil, err
}
for _, n := range ng {
if !n.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
err := singleton.DB.Transaction(func(tx *gorm.DB) error { err := singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Unscoped().Delete(&model.NotificationGroup{}, "id in (?)", ngn).Error; err != nil { if err := tx.Unscoped().Delete(&model.NotificationGroup{}, "id in (?)", ngn).Error; err != nil {
return err return err

View File

@ -56,11 +56,26 @@ func updateServer(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
singleton.DDNSCacheLock.RLock()
for _, pid := range sf.DDNSProfiles {
if p, ok := singleton.DDNSCache[pid]; ok {
if !p.HasPermission(c) {
singleton.DDNSCacheLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.DDNSCacheLock.RUnlock()
var s model.Server var s model.Server
if err := singleton.DB.First(&s, id).Error; err != nil { if err := singleton.DB.First(&s, id).Error; err != nil {
return nil, singleton.Localizer.ErrorT("server id %d does not exist", id) return nil, singleton.Localizer.ErrorT("server id %d does not exist", id)
} }
if !s.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
s.Name = sf.Name s.Name = sf.Name
s.DisplayIndex = sf.DisplayIndex s.DisplayIndex = sf.DisplayIndex
s.Note = sf.Note s.Note = sf.Note
@ -104,6 +119,17 @@ func batchDeleteServer(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
singleton.ServerLock.RLock()
for _, sid := range servers {
if s, ok := singleton.ServerList[sid]; ok {
if !s.HasPermission(c) {
singleton.ServerLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.ServerLock.RUnlock()
err := singleton.DB.Transaction(func(tx *gorm.DB) error { err := singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Unscoped().Delete(&model.Server{}, "id in (?)", servers).Error; err != nil { if err := tx.Unscoped().Delete(&model.Server{}, "id in (?)", servers).Error; err != nil {
return err return err
@ -161,6 +187,9 @@ func forceUpdateServer(c *gin.Context) (*model.ForceUpdateResponse, error) {
server := singleton.ServerList[sid] server := singleton.ServerList[sid]
singleton.ServerLock.RUnlock() singleton.ServerLock.RUnlock()
if server != nil && server.TaskStream != nil { if server != nil && server.TaskStream != nil {
if !server.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
if err := server.TaskStream.Send(&pb.Task{ if err := server.TaskStream.Send(&pb.Task{
Type: model.TaskTypeUpgrade, Type: model.TaskTypeUpgrade,
}); err != nil { }); err != nil {

View File

@ -20,7 +20,7 @@ import (
// @Produce json // @Produce json
// @Success 200 {object} model.CommonResponse[[]model.ServerGroupResponseItem] // @Success 200 {object} model.CommonResponse[[]model.ServerGroupResponseItem]
// @Router /server-group [get] // @Router /server-group [get]
func listServerGroup(c *gin.Context) ([]model.ServerGroupResponseItem, error) { func listServerGroup(c *gin.Context) ([]*model.ServerGroupResponseItem, error) {
var sg []model.ServerGroup var sg []model.ServerGroup
if err := singleton.DB.Find(&sg).Error; err != nil { if err := singleton.DB.Find(&sg).Error; err != nil {
return nil, err return nil, err
@ -38,9 +38,9 @@ func listServerGroup(c *gin.Context) ([]model.ServerGroupResponseItem, error) {
groupServers[s.ServerGroupId] = append(groupServers[s.ServerGroupId], s.ServerId) groupServers[s.ServerGroupId] = append(groupServers[s.ServerGroupId], s.ServerId)
} }
var sgRes []model.ServerGroupResponseItem var sgRes []*model.ServerGroupResponseItem
for _, s := range sg { for _, s := range sg {
sgRes = append(sgRes, model.ServerGroupResponseItem{ sgRes = append(sgRes, &model.ServerGroupResponseItem{
Group: s, Group: s,
Servers: groupServers[s.ID], Servers: groupServers[s.ID],
}) })
@ -67,8 +67,22 @@ func createServerGroup(c *gin.Context) (uint64, error) {
} }
sgf.Servers = slices.Compact(sgf.Servers) sgf.Servers = slices.Compact(sgf.Servers)
singleton.ServerLock.RLock()
for _, sid := range sgf.Servers {
if server, ok := singleton.ServerList[sid]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return 0, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.ServerLock.RUnlock()
uid := getUid(c)
var sg model.ServerGroup var sg model.ServerGroup
sg.Name = sgf.Name sg.Name = sgf.Name
sg.UserID = uid
var count int64 var count int64
if err := singleton.DB.Model(&model.Server{}).Where("id in (?)", sgf.Servers).Count(&count).Error; err != nil { if err := singleton.DB.Model(&model.Server{}).Where("id in (?)", sgf.Servers).Count(&count).Error; err != nil {
@ -84,6 +98,9 @@ func createServerGroup(c *gin.Context) (uint64, error) {
} }
for _, s := range sgf.Servers { for _, s := range sgf.Servers {
if err := tx.Create(&model.ServerGroupServer{ if err := tx.Create(&model.ServerGroupServer{
Common: model.Common{
UserID: uid,
},
ServerGroupId: sg.ID, ServerGroupId: sg.ID,
ServerId: s, ServerId: s,
}).Error; err != nil { }).Error; err != nil {
@ -125,10 +142,26 @@ func updateServerGroup(c *gin.Context) (any, error) {
} }
sg.Servers = slices.Compact(sg.Servers) sg.Servers = slices.Compact(sg.Servers)
singleton.ServerLock.RLock()
for _, sid := range sg.Servers {
if server, ok := singleton.ServerList[sid]; ok {
if !server.HasPermission(c) {
singleton.ServerLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.ServerLock.RUnlock()
var sgDB model.ServerGroup var sgDB model.ServerGroup
if err := singleton.DB.First(&sgDB, id).Error; err != nil { if err := singleton.DB.First(&sgDB, id).Error; err != nil {
return nil, singleton.Localizer.ErrorT("group id %d does not exist", id) return nil, singleton.Localizer.ErrorT("group id %d does not exist", id)
} }
if !sgDB.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("unauthorized")
}
sgDB.Name = sg.Name sgDB.Name = sg.Name
var count int64 var count int64
@ -139,6 +172,8 @@ func updateServerGroup(c *gin.Context) (any, error) {
return nil, singleton.Localizer.ErrorT("have invalid server id") return nil, singleton.Localizer.ErrorT("have invalid server id")
} }
uid := getUid(c)
err = singleton.DB.Transaction(func(tx *gorm.DB) error { err = singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Save(&sgDB).Error; err != nil { if err := tx.Save(&sgDB).Error; err != nil {
return err return err
@ -149,6 +184,9 @@ func updateServerGroup(c *gin.Context) (any, error) {
for _, s := range sg.Servers { for _, s := range sg.Servers {
if err := tx.Create(&model.ServerGroupServer{ if err := tx.Create(&model.ServerGroupServer{
Common: model.Common{
UserID: uid,
},
ServerGroupId: sgDB.ID, ServerGroupId: sgDB.ID,
ServerId: s, ServerId: s,
}).Error; err != nil { }).Error; err != nil {
@ -181,6 +219,17 @@ func batchDeleteServerGroup(c *gin.Context) (any, error) {
return nil, err return nil, err
} }
var sg []model.ServerGroup
if err := singleton.DB.Where("id in (?)", sgs).Find(&sg).Error; err != nil {
return nil, err
}
for _, s := range sg {
if !s.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
err := singleton.DB.Transaction(func(tx *gorm.DB) error { err := singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Unscoped().Delete(&model.ServerGroup{}, "id in (?)", sgs).Error; err != nil { if err := tx.Unscoped().Delete(&model.ServerGroup{}, "id in (?)", sgs).Error; err != nil {
return err return err

View File

@ -190,7 +190,10 @@ func createService(c *gin.Context) (uint64, error) {
return 0, err return 0, err
} }
uid := getUid(c)
var m model.Service var m model.Service
m.UserID = uid
m.Name = mf.Name m.Name = mf.Name
m.Target = strings.TrimSpace(mf.Target) m.Target = strings.TrimSpace(mf.Target)
m.Type = mf.Type m.Type = mf.Type
@ -207,6 +210,10 @@ func createService(c *gin.Context) (uint64, error) {
m.RecoverTriggerTasks = mf.RecoverTriggerTasks m.RecoverTriggerTasks = mf.RecoverTriggerTasks
m.FailTriggerTasks = mf.FailTriggerTasks m.FailTriggerTasks = mf.FailTriggerTasks
if err := validateServers(c, &m); err != nil {
return 0, err
}
if err := singleton.DB.Create(&m).Error; err != nil { if err := singleton.DB.Create(&m).Error; err != nil {
return 0, newGormError("%v", err) return 0, newGormError("%v", err)
} }
@ -260,6 +267,11 @@ func updateService(c *gin.Context) (any, error) {
if err := singleton.DB.First(&m, id).Error; err != nil { if err := singleton.DB.First(&m, id).Error; err != nil {
return nil, singleton.Localizer.ErrorT("service id %d does not exist", id) return nil, singleton.Localizer.ErrorT("service id %d does not exist", id)
} }
if !m.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
m.Name = mf.Name m.Name = mf.Name
m.Target = strings.TrimSpace(mf.Target) m.Target = strings.TrimSpace(mf.Target)
m.Type = mf.Type m.Type = mf.Type
@ -276,6 +288,10 @@ func updateService(c *gin.Context) (any, error) {
m.RecoverTriggerTasks = mf.RecoverTriggerTasks m.RecoverTriggerTasks = mf.RecoverTriggerTasks
m.FailTriggerTasks = mf.FailTriggerTasks m.FailTriggerTasks = mf.FailTriggerTasks
if err := validateServers(c, &m); err != nil {
return 0, err
}
if err := singleton.DB.Save(&m).Error; err != nil { if err := singleton.DB.Save(&m).Error; err != nil {
return nil, newGormError("%v", err) return nil, newGormError("%v", err)
} }
@ -318,6 +334,18 @@ func batchDeleteService(c *gin.Context) (any, error) {
if err := c.ShouldBindJSON(&ids); err != nil { if err := c.ShouldBindJSON(&ids); err != nil {
return nil, err return nil, err
} }
singleton.ServiceSentinelShared.ServicesLock.RLock()
for _, id := range ids {
if ss, ok := singleton.ServiceSentinelShared.Services[id]; ok {
if !ss.HasPermission(c) {
singleton.ServiceSentinelShared.ServicesLock.RUnlock()
return nil, singleton.Localizer.ErrorT("permission denied")
}
}
}
singleton.ServiceSentinelShared.ServicesLock.RUnlock()
err := singleton.DB.Transaction(func(tx *gorm.DB) error { err := singleton.DB.Transaction(func(tx *gorm.DB) error {
if err := tx.Unscoped().Delete(&model.Service{}, "id in (?)", ids).Error; err != nil { if err := tx.Unscoped().Delete(&model.Service{}, "id in (?)", ids).Error; err != nil {
return err return err
@ -331,3 +359,18 @@ func batchDeleteService(c *gin.Context) (any, error) {
singleton.ServiceSentinelShared.UpdateServiceList() singleton.ServiceSentinelShared.UpdateServiceList()
return nil, nil return nil, nil
} }
func validateServers(c *gin.Context, ss *model.Service) error {
singleton.ServerLock.RLock()
defer singleton.ServerLock.RUnlock()
for s := range ss.SkipServers {
if server, ok := singleton.ServerList[s]; ok {
if !server.HasPermission(c) {
return singleton.Localizer.ErrorT("permission denied")
}
}
}
return nil
}

View File

@ -6,6 +6,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
"github.com/hashicorp/go-uuid" "github.com/hashicorp/go-uuid"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/pkg/utils"
"github.com/nezhahq/nezha/pkg/websocketx" "github.com/nezhahq/nezha/pkg/websocketx"
@ -29,13 +30,6 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
return nil, err return nil, err
} }
streamId, err := uuid.GenerateUUID()
if err != nil {
return nil, err
}
rpc.NezhaHandlerSingleton.CreateStream(streamId)
singleton.ServerLock.RLock() singleton.ServerLock.RLock()
server := singleton.ServerList[createTerminalReq.ServerID] server := singleton.ServerList[createTerminalReq.ServerID]
singleton.ServerLock.RUnlock() singleton.ServerLock.RUnlock()
@ -43,6 +37,17 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
return nil, singleton.Localizer.ErrorT("server not found or not connected") return nil, singleton.Localizer.ErrorT("server not found or not connected")
} }
if !server.HasPermission(c) {
return nil, singleton.Localizer.ErrorT("permission denied")
}
streamId, err := uuid.GenerateUUID()
if err != nil {
return nil, err
}
rpc.NezhaHandlerSingleton.CreateStream(streamId)
terminalData, _ := utils.Json.Marshal(&model.TerminalTask{ terminalData, _ := utils.Json.Marshal(&model.TerminalTask{
StreamID: streamId, StreamID: streamId,
}) })

View File

@ -114,6 +114,7 @@ func createUser(c *gin.Context) (uint64, error) {
var u model.User var u model.User
u.Username = uf.Username u.Username = uf.Username
u.Role = model.RoleMember
hash, err := bcrypt.GenerateFromPassword([]byte(uf.Password), bcrypt.DefaultCost) hash, err := bcrypt.GenerateFromPassword([]byte(uf.Password), bcrypt.DefaultCost)
if err != nil { if err != nil {
@ -125,6 +126,7 @@ func createUser(c *gin.Context) (uint64, error) {
return 0, err return 0, err
} }
singleton.OnUserUpdate(&u)
return u.ID, nil return u.ID, nil
} }
@ -149,5 +151,6 @@ func batchDeleteUser(c *gin.Context) (any, error) {
return nil, singleton.Localizer.ErrorT("can't delete yourself") return nil, singleton.Localizer.ErrorT("can't delete yourself")
} }
return nil, singleton.DB.Where("id IN (?)", ids).Delete(&model.User{}).Error err := singleton.OnUserDelete(ids, newGormError)
return nil, err
} }

View File

@ -1,6 +1,8 @@
package controller package controller
import ( import (
"strconv"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
@ -13,16 +15,40 @@ import (
// @Schemes // @Schemes
// @Description List server // @Description List server
// @Tags auth required // @Tags auth required
// @Param limit query uint false "Page limit"
// @Param offset query uint false "Page offset"
// @Produce json // @Produce json
// @Success 200 {object} model.CommonResponse[[]model.WAFApiMock] // @Success 200 {object} model.PaginatedResponse[[]model.WAFApiMock, model.WAFApiMock]
// @Router /waf [get] // @Router /waf [get]
func listBlockedAddress(c *gin.Context) ([]*model.WAF, error) { func listBlockedAddress(c *gin.Context) (*model.Value[[]*model.WAF], error) {
limit, err := strconv.Atoi(c.Query("limit"))
if err != nil || limit < 1 {
limit = 25
}
offset, err := strconv.Atoi(c.Query("offset"))
if err != nil || offset < 0 {
offset = 0
}
var waf []*model.WAF var waf []*model.WAF
if err := singleton.DB.Find(&waf).Error; err != nil { if err := singleton.DB.Limit(limit).Offset(offset).Find(&waf).Error; err != nil {
return nil, err return nil, err
} }
return waf, nil var total int64
if err := singleton.DB.Model(&model.WAF{}).Count(&total).Error; err != nil {
return nil, err
}
return &model.Value[[]*model.WAF]{
Value: waf,
Pagination: model.Pagination{
Offset: offset,
Limit: limit,
Total: total,
},
}, nil
} }
// Batch delete blocked addresses // Batch delete blocked addresses

View File

@ -100,12 +100,34 @@ func DispatchTask(serviceSentinelDispatchBus <-chan model.Service) {
continue continue
} }
if task.Cover == model.ServiceCoverIgnoreAll && task.SkipServers[singleton.SortedServerList[workedServerIndex].ID] { if task.Cover == model.ServiceCoverIgnoreAll && task.SkipServers[singleton.SortedServerList[workedServerIndex].ID] {
singleton.SortedServerList[workedServerIndex].TaskStream.Send(task.PB()) server := singleton.SortedServerList[workedServerIndex]
singleton.UserLock.RLock()
var role uint8
if u, ok := singleton.UserInfoMap[server.UserID]; !ok {
role = model.RoleMember
} else {
role = u.Role
}
singleton.UserLock.RUnlock()
if task.UserID == server.UserID || role == model.RoleAdmin {
singleton.SortedServerList[workedServerIndex].TaskStream.Send(task.PB())
}
workedServerIndex++ workedServerIndex++
continue continue
} }
if task.Cover == model.ServiceCoverAll && !task.SkipServers[singleton.SortedServerList[workedServerIndex].ID] { if task.Cover == model.ServiceCoverAll && !task.SkipServers[singleton.SortedServerList[workedServerIndex].ID] {
singleton.SortedServerList[workedServerIndex].TaskStream.Send(task.PB()) server := singleton.SortedServerList[workedServerIndex]
singleton.UserLock.RLock()
var role uint8
if u, ok := singleton.UserInfoMap[server.UserID]; !ok {
role = model.RoleMember
} else {
role = u.Role
}
singleton.UserLock.RUnlock()
if task.UserID == server.UserID || role == model.RoleAdmin {
singleton.SortedServerList[workedServerIndex].TaskStream.Send(task.PB())
}
workedServerIndex++ workedServerIndex++
continue continue
} }

4
go.mod
View File

@ -1,8 +1,8 @@
module github.com/nezhahq/nezha module github.com/nezhahq/nezha
go 1.22.7 go 1.23.0
toolchain go1.23.1 toolchain go1.23.2
require ( require (
github.com/appleboy/gin-jwt/v2 v2.10.0 github.com/appleboy/gin-jwt/v2 v2.10.0

View File

@ -62,10 +62,15 @@ func (r *AlertRule) Enabled() bool {
} }
// Snapshot 对传入的Server进行该报警规则下所有type的检查 返回每项检查结果 // 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, 0, len(r.Rules)) point := make([]bool, len(r.Rules))
for _, rule := range r.Rules {
point = append(point, rule.Snapshot(cycleTransferStats, server, db)) if r.UserID != server.UserID && role != RoleAdmin {
return point
}
for i, rule := range r.Rules {
point[i] = rule.Snapshot(cycleTransferStats, server, db)
} }
return point return point
} }

View File

@ -15,6 +15,23 @@ type CommonResponse[T any] struct {
Error string `json:"error,omitempty"` Error string `json:"error,omitempty"`
} }
type PaginatedResponse[S ~[]E, E any] struct {
Success bool `json:"success,omitempty"`
Data *Value[S] `json:"data,omitempty"`
Error string `json:"error,omitempty"`
}
type Value[T any] struct {
Value T `json:"value,omitempty"`
Pagination Pagination `json:"pagination,omitempty"`
}
type Pagination struct {
Offset int `json:"offset,omitempty"`
Limit int `json:"limit,omitempty"`
Total int64 `json:"total,omitempty"`
}
type LoginResponse struct { type LoginResponse struct {
Token string `json:"token,omitempty"` Token string `json:"token,omitempty"`
Expire string `json:"expire,omitempty"` Expire string `json:"expire,omitempty"`

View File

@ -2,6 +2,8 @@ package model
import ( import (
"time" "time"
"github.com/gin-gonic/gin"
) )
const ( const (
@ -18,6 +20,47 @@ type Common struct {
UpdatedAt time.Time `gorm:"autoUpdateTime" json:"updated_at,omitempty"` UpdatedAt time.Time `gorm:"autoUpdateTime" json:"updated_at,omitempty"`
// Do not use soft deletion // Do not use soft deletion
// DeletedAt gorm.DeletedAt `gorm:"index" json:"deleted_at,omitempty"` // DeletedAt gorm.DeletedAt `gorm:"index" json:"deleted_at,omitempty"`
UserID uint64 `json:"-"`
}
func (c *Common) GetID() uint64 {
return c.ID
}
func (c *Common) GetUserID() uint64 {
return c.UserID
}
func (c *Common) HasPermission(ctx *gin.Context) bool {
auth, ok := ctx.Get(CtxKeyAuthorizedUser)
if !ok {
return false
}
user := *auth.(*User)
if user.Role == RoleAdmin {
return true
}
return user.ID == c.UserID
}
type CommonInterface interface {
GetID() uint64
GetUserID() uint64
HasPermission(*gin.Context) bool
}
func FindByUserID[S ~[]E, E CommonInterface](s S, uid uint64) []uint64 {
var list []uint64
for _, v := range s {
if v.GetUserID() == uid {
list = append(list, v.GetID())
}
}
return list
} }
type Response struct { type Response struct {

View File

@ -1,9 +1,41 @@
package model package model
import (
"github.com/nezhahq/nezha/pkg/utils"
"gorm.io/gorm"
)
const (
RoleAdmin uint8 = iota
RoleMember
)
type User struct { type User struct {
Common Common
Username string `json:"username,omitempty" gorm:"uniqueIndex"` Username string `json:"username,omitempty" gorm:"uniqueIndex"`
Password string `json:"password,omitempty" gorm:"type:char(72)"` Password string `json:"password,omitempty" gorm:"type:char(72)"`
Role uint8 `json:"role,omitempty"`
AgentSecret string `json:"agent_secret,omitempty" gorm:"type:char(32)"`
}
type UserInfo struct {
Role uint8
_ [3]byte
AgentSecret string
}
func (u *User) BeforeSave(tx *gorm.DB) error {
if u.AgentSecret != "" {
return nil
}
key, err := utils.GenerateRandomString(32)
if err != nil {
return err
}
u.AgentSecret = key
return nil
} }
type Profile struct { type Profile struct {

View File

@ -1,6 +0,0 @@
package model
type UserGroup struct {
Common
Name string `json:"name"`
}

View File

@ -1,7 +0,0 @@
package model
type UserGroupUser struct {
Common
UserGroupId uint64 `json:"user_group_id"`
UserId uint64 `json:"user_id"`
}

View File

@ -16,22 +16,30 @@ const (
WAFBlockReasonTypeAgentAuthFail WAFBlockReasonTypeAgentAuthFail
) )
const (
BlockIDgRPC = -127 + iota
BlockIDToken
BlockIDUnknownUser
)
type WAFApiMock struct { type WAFApiMock struct {
IP string `json:"ip,omitempty"` IP string `json:"ip,omitempty"`
Count uint64 `json:"count,omitempty"` BlockIdentifier int64 `json:"block_identifier,omitempty"`
LastBlockReason uint8 `json:"last_block_reason,omitempty"` BlockReason uint8 `json:"block_reason,omitempty"`
LastBlockTimestamp uint64 `json:"last_block_timestamp,omitempty"` BlockTimestamp uint64 `json:"block_timestamp,omitempty"`
Count uint64 `json:"count,omitempty"`
} }
type WAF struct { type WAF struct {
IP []byte `gorm:"type:binary(16);primaryKey" json:"ip,omitempty"` IP []byte `gorm:"type:binary(16);primaryKey" json:"ip,omitempty"`
Count uint64 `json:"count,omitempty"` BlockIdentifier int64 `gorm:"primaryKey" json:"block_identifier,omitempty"`
LastBlockReason uint8 `json:"last_block_reason,omitempty"` BlockReason uint8 `json:"block_reason,omitempty"`
LastBlockTimestamp uint64 `json:"last_block_timestamp,omitempty"` BlockTimestamp uint64 `gorm:"index" json:"block_timestamp,omitempty"`
Count uint64 `json:"count,omitempty"`
} }
func (w *WAF) TableName() string { func (w *WAF) TableName() string {
return "waf" return "nz_waf"
} }
func CheckIP(db *gorm.DB, ip string) error { func CheckIP(db *gorm.DB, ip string) error {
@ -42,22 +50,31 @@ func CheckIP(db *gorm.DB, ip string) error {
if err != nil { if err != nil {
return err return err
} }
var w WAF
result := db.Limit(1).Find(&w, "ip = ?", ipBinary) var blockTimestamp uint64
result := db.Model(&WAF{}).Order("block_timestamp desc").Select("block_timestamp").Where("ip = ?", ipBinary).Limit(1).Find(&blockTimestamp)
if result.Error != nil { if result.Error != nil {
return result.Error return result.Error
} }
if result.RowsAffected == 0 { // 检查是否未找到记录
// 检查是否未找到记录
if result.RowsAffected < 1 {
return nil return nil
} }
var count uint64
if err := db.Model(&WAF{}).Select("SUM(count)").Where("ip = ?", ipBinary).Scan(&count).Error; err != nil {
return err
}
now := time.Now().Unix() now := time.Now().Unix()
if powAdd(w.Count, 4, w.LastBlockTimestamp) > uint64(now) { if powAdd(count, 4, blockTimestamp) > uint64(now) {
return errors.New("you are blocked by nezha WAF") return errors.New("you are blocked by nezha WAF")
} }
return nil return nil
} }
func ClearIP(db *gorm.DB, ip string) error { func ClearIP(db *gorm.DB, ip string, uid int64) error {
if ip == "" { if ip == "" {
return nil return nil
} }
@ -65,7 +82,7 @@ func ClearIP(db *gorm.DB, ip string) error {
if err != nil { if err != nil {
return err return err
} }
return db.Unscoped().Delete(&WAF{}, "ip = ?", ipBinary).Error return db.Unscoped().Delete(&WAF{}, "ip = ? and block_identifier = ?", ipBinary, uid).Error
} }
func BatchClearIP(db *gorm.DB, ip []string) error { func BatchClearIP(db *gorm.DB, ip []string) error {
@ -83,7 +100,7 @@ func BatchClearIP(db *gorm.DB, ip []string) error {
return db.Unscoped().Delete(&WAF{}, "ip in (?)", ips).Error return db.Unscoped().Delete(&WAF{}, "ip in (?)", ips).Error
} }
func BlockIP(db *gorm.DB, ip string, reason uint8) error { func BlockIP(db *gorm.DB, ip string, reason uint8, uid int64) error {
if ip == "" { if ip == "" {
return nil return nil
} }
@ -91,16 +108,19 @@ func BlockIP(db *gorm.DB, ip string, reason uint8) error {
if err != nil { if err != nil {
return err return err
} }
var w WAF w := WAF{
w.IP = ipBinary IP: ipBinary,
BlockIdentifier: uid,
}
now := uint64(time.Now().Unix())
return db.Transaction(func(tx *gorm.DB) error { return db.Transaction(func(tx *gorm.DB) error {
if err := tx.Where(&w).Attrs(WAF{ if err := tx.Where(&w).Attrs(WAF{
LastBlockReason: reason, BlockReason: reason,
LastBlockTimestamp: uint64(time.Now().Unix()), BlockTimestamp: now,
}).FirstOrCreate(&w).Error; err != nil { }).FirstOrCreate(&w).Error; err != nil {
return err return err
} }
return tx.Exec("UPDATE waf SET count = count + 1, last_block_reason = ?, last_block_timestamp = ? WHERE ip = ?", reason, uint64(time.Now().Unix()), ipBinary).Error return tx.Exec("UPDATE nz_waf SET count = count + 1, block_reason = ?, block_timestamp = ? WHERE ip = ? and block_identifier = ?", reason, now, ipBinary, uid).Error
}) })
} }

View File

@ -3,10 +3,12 @@ package utils
import ( import (
"crypto/rand" "crypto/rand"
"errors" "errors"
"maps"
"math/big" "math/big"
"net/netip" "net/netip"
"os" "os"
"regexp" "regexp"
"slices"
"strconv" "strconv"
"strings" "strings"
@ -145,3 +147,8 @@ func Itoa[T constraints.Integer](i T) string {
return "" return ""
} }
} }
func MapValuesToSlice[Map ~map[K]V, K comparable, V any](m Map) []V {
s := make([]V, 0, len(m))
return slices.AppendSeq(s, maps.Values(m))
}

View File

@ -36,12 +36,16 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string)
if clientSecret != singleton.Conf.AgentSecretKey { singleton.UserLock.RLock()
model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail) userId, ok := singleton.AgentSecretToUserId[clientSecret]
if !ok && clientSecret != singleton.Conf.AgentSecretKey {
singleton.UserLock.RUnlock()
model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail, model.BlockIDgRPC)
return 0, status.Error(codes.Unauthenticated, "客户端认证失败") return 0, status.Error(codes.Unauthenticated, "客户端认证失败")
} }
singleton.UserLock.RUnlock()
model.ClearIP(singleton.DB, ip) model.ClearIP(singleton.DB, ip, model.BlockIDgRPC)
var clientUUID string var clientUUID string
if value, ok := md["client_uuid"]; ok { if value, ok := md["client_uuid"]; ok {
@ -53,21 +57,26 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
} }
singleton.ServerLock.RLock() singleton.ServerLock.RLock()
defer singleton.ServerLock.RUnlock()
clientID, hasID := singleton.ServerUUIDToID[clientUUID] clientID, hasID := singleton.ServerUUIDToID[clientUUID]
singleton.ServerLock.RUnlock()
if !hasID { if !hasID {
s := model.Server{UUID: clientUUID, Name: petname.Generate(2, "-")} s := model.Server{UUID: clientUUID, Name: petname.Generate(2, "-"), Common: model.Common{
UserID: userId,
}}
if err := singleton.DB.Create(&s).Error; err != nil { if err := singleton.DB.Create(&s).Error; err != nil {
return 0, status.Error(codes.Unauthenticated, err.Error()) return 0, status.Error(codes.Unauthenticated, err.Error())
} }
s.Host = &model.Host{} s.Host = &model.Host{}
s.State = &model.HostState{} s.State = &model.HostState{}
s.GeoIP = &model.GeoIP{} s.GeoIP = &model.GeoIP{}
// generate a random silly server name
singleton.ServerLock.Lock()
singleton.ServerList[s.ID] = &s singleton.ServerList[s.ID] = &s
singleton.ServerUUIDToID[clientUUID] = s.ID singleton.ServerUUIDToID[clientUUID] = s.ID
singleton.ServerLock.Unlock()
singleton.ReSortServer() singleton.ReSortServer()
clientID = s.ID clientID = s.ID
} }

View File

@ -143,8 +143,16 @@ func checkStatus() {
} }
for _, server := range ServerList { for _, server := range ServerList {
// 监测点 // 监测点
UserLock.RLock()
var role uint8
if u, ok := UserInfoMap[server.UserID]; !ok {
role = model.RoleMember
} else {
role = u.Role
}
UserLock.RUnlock()
alertsStore[alert.ID][server.ID] = append(alertsStore[alert. 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]) max, passed := alert.Check(alertsStore[alert.ID][server.ID])
// 保存当前服务器状态信息 // 保存当前服务器状态信息

View File

@ -12,6 +12,7 @@ import (
"github.com/robfig/cron/v3" "github.com/robfig/cron/v3"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
pb "github.com/nezhahq/nezha/proto" pb "github.com/nezhahq/nezha/proto"
) )
@ -79,10 +80,7 @@ func UpdateCronList() {
CronLock.RLock() CronLock.RLock()
defer CronLock.RUnlock() defer CronLock.RUnlock()
CronList = make([]*model.Cron, 0, len(Crons)) CronList = utils.MapValuesToSlice(Crons)
for _, c := range Crons {
CronList = append(CronList, c)
}
slices.SortFunc(CronList, func(a, b *model.Cron) int { slices.SortFunc(CronList, func(a, b *model.Cron) int {
return cmp.Compare(a.ID, b.ID) return cmp.Compare(a.ID, b.ID)
}) })

View File

@ -13,6 +13,7 @@ import (
ddns2 "github.com/nezhahq/nezha/pkg/ddns" ddns2 "github.com/nezhahq/nezha/pkg/ddns"
"github.com/nezhahq/nezha/pkg/ddns/dummy" "github.com/nezhahq/nezha/pkg/ddns/dummy"
"github.com/nezhahq/nezha/pkg/ddns/webhook" "github.com/nezhahq/nezha/pkg/ddns/webhook"
"github.com/nezhahq/nezha/pkg/utils"
) )
var ( var (
@ -24,12 +25,10 @@ var (
func initDDNS() { func initDDNS() {
DB.Find(&DDNSList) DB.Find(&DDNSList)
DDNSCacheLock.Lock()
DDNSCache = make(map[uint64]*model.DDNSProfile) DDNSCache = make(map[uint64]*model.DDNSProfile)
for i := 0; i < len(DDNSList); i++ { for i := 0; i < len(DDNSList); i++ {
DDNSCache[DDNSList[i].ID] = DDNSList[i] DDNSCache[DDNSList[i].ID] = DDNSList[i]
} }
DDNSCacheLock.Unlock()
OnNameserverUpdate() OnNameserverUpdate()
} }
@ -56,10 +55,7 @@ func UpdateDDNSList() {
DDNSListLock.Lock() DDNSListLock.Lock()
defer DDNSListLock.Unlock() defer DDNSListLock.Unlock()
DDNSList = make([]*model.DDNSProfile, 0, len(DDNSCache)) DDNSList = utils.MapValuesToSlice(DDNSCache)
for _, p := range DDNSCache {
DDNSList = append(DDNSList, p)
}
slices.SortFunc(DDNSList, func(a, b *model.DDNSProfile) int { slices.SortFunc(DDNSList, func(a, b *model.DDNSProfile) int {
return cmp.Compare(a.ID, b.ID) return cmp.Compare(a.ID, b.ID)
}) })

View File

@ -6,6 +6,7 @@ import (
"sync" "sync"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
) )
var ( var (
@ -19,8 +20,6 @@ var (
func initNAT() { func initNAT() {
DB.Find(&NATList) DB.Find(&NATList)
NATCacheRwLock.Lock()
defer NATCacheRwLock.Unlock()
NATCache = make(map[string]*model.NAT) NATCache = make(map[string]*model.NAT)
for i := 0; i < len(NATList); i++ { for i := 0; i < len(NATList); i++ {
NATCache[NATList[i].Domain] = NATList[i] NATCache[NATList[i].Domain] = NATList[i]
@ -59,10 +58,7 @@ func UpdateNATList() {
NATListLock.Lock() NATListLock.Lock()
defer NATListLock.Unlock() defer NATListLock.Unlock()
NATList = make([]*model.NAT, 0, len(NATCache)) NATList = utils.MapValuesToSlice(NATCache)
for _, n := range NATCache {
NATList = append(NATList, n)
}
slices.SortFunc(NATList, func(a, b *model.NAT) int { slices.SortFunc(NATList, func(a, b *model.NAT) int {
return cmp.Compare(a.ID, b.ID) return cmp.Compare(a.ID, b.ID)
}) })

View File

@ -9,6 +9,7 @@ import (
"time" "time"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
) )
const ( const (
@ -30,7 +31,7 @@ var (
) )
// InitNotification 初始化 GroupID <-> ID <-> Notification 的映射 // InitNotification 初始化 GroupID <-> ID <-> Notification 的映射
func InitNotification() { func initNotification() {
NotificationList = make(map[uint64]map[uint64]*model.Notification) NotificationList = make(map[uint64]map[uint64]*model.Notification)
NotificationIDToGroups = make(map[uint64]map[uint64]struct{}) NotificationIDToGroups = make(map[uint64]map[uint64]struct{})
NotificationGroup = make(map[uint64]string) NotificationGroup = make(map[uint64]string)
@ -38,9 +39,7 @@ func InitNotification() {
// loadNotifications 从 DB 初始化通知方式相关参数 // loadNotifications 从 DB 初始化通知方式相关参数
func loadNotifications() { func loadNotifications() {
InitNotification() initNotification()
NotificationsLock.Lock()
groupNotifications := make(map[uint64][]uint64) groupNotifications := make(map[uint64][]uint64)
var ngn []model.NotificationGroupNotification var ngn []model.NotificationGroupNotification
if err := DB.Find(&ngn).Error; err != nil { if err := DB.Find(&ngn).Error; err != nil {
@ -74,8 +73,6 @@ func loadNotifications() {
} }
} }
} }
NotificationsLock.Unlock()
} }
func UpdateNotificationList() { func UpdateNotificationList() {
@ -85,10 +82,7 @@ func UpdateNotificationList() {
NotificationSortedLock.Lock() NotificationSortedLock.Lock()
defer NotificationSortedLock.Unlock() defer NotificationSortedLock.Unlock()
NotificationListSorted = make([]*model.Notification, 0, len(NotificationMap)) NotificationListSorted = utils.MapValuesToSlice(NotificationMap)
for _, n := range NotificationMap {
NotificationListSorted = append(NotificationListSorted, n)
}
slices.SortFunc(NotificationListSorted, func(a, b *model.Notification) int { slices.SortFunc(NotificationListSorted, func(a, b *model.Notification) int {
return cmp.Compare(a.ID, b.ID) return cmp.Compare(a.ID, b.ID)
}) })

View File

@ -1,10 +1,12 @@
package singleton package singleton
import ( import (
"sort" "cmp"
"slices"
"sync" "sync"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
) )
var ( var (
@ -45,29 +47,21 @@ func ReSortServer() {
SortedServerLock.Lock() SortedServerLock.Lock()
defer SortedServerLock.Unlock() defer SortedServerLock.Unlock()
SortedServerList = make([]*model.Server, 0, len(ServerList)) SortedServerList = utils.MapValuesToSlice(ServerList)
SortedServerListForGuest = make([]*model.Server, 0) // 按照服务器 ID 排序的具体实现ID越大越靠前
for _, s := range ServerList { slices.SortStableFunc(SortedServerList, func(a, b *model.Server) int {
SortedServerList = append(SortedServerList, s) if a.DisplayIndex == b.DisplayIndex {
return cmp.Compare(a.ID, b.ID)
}
return cmp.Compare(b.DisplayIndex, a.DisplayIndex)
})
SortedServerListForGuest = make([]*model.Server, 0, len(SortedServerList))
for _, s := range SortedServerList {
if !s.HideForGuest { if !s.HideForGuest {
SortedServerListForGuest = append(SortedServerListForGuest, s) SortedServerListForGuest = append(SortedServerListForGuest, s)
} }
} }
// 按照服务器 ID 排序的具体实现ID越大越靠前
sort.SliceStable(SortedServerList, func(i, j int) bool {
if SortedServerList[i].DisplayIndex == SortedServerList[j].DisplayIndex {
return SortedServerList[i].ID < SortedServerList[j].ID
}
return SortedServerList[i].DisplayIndex > SortedServerList[j].DisplayIndex
})
sort.SliceStable(SortedServerListForGuest, func(i, j int) bool {
if SortedServerListForGuest[i].DisplayIndex == SortedServerListForGuest[j].DisplayIndex {
return SortedServerListForGuest[i].ID < SortedServerListForGuest[j].ID
}
return SortedServerListForGuest[i].DisplayIndex > SortedServerListForGuest[j].DisplayIndex
})
} }
func OnServerDelete(sid []uint64) { func OnServerDelete(sid []uint64) {

View File

@ -11,6 +11,7 @@ import (
"github.com/jinzhu/copier" "github.com/jinzhu/copier"
"github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/utils"
pb "github.com/nezhahq/nezha/proto" pb "github.com/nezhahq/nezha/proto"
) )
@ -174,11 +175,7 @@ func (ss *ServiceSentinel) UpdateServiceList() {
ss.ServiceListLock.Lock() ss.ServiceListLock.Lock()
defer ss.ServiceListLock.Unlock() defer ss.ServiceListLock.Unlock()
ss.ServiceList = make([]*model.Service, 0, len(ss.Services)) ss.ServiceList = utils.MapValuesToSlice(ss.Services)
for _, v := range ss.Services {
ss.ServiceList = append(ss.ServiceList, v)
}
slices.SortFunc(ss.ServiceList, func(a, b *model.Service) int { slices.SortFunc(ss.ServiceList, func(a, b *model.Service) int {
return cmp.Compare(a.ID, b.ID) return cmp.Compare(a.ID, b.ID)
}) })
@ -192,13 +189,6 @@ func (ss *ServiceSentinel) loadServiceHistory() {
panic(err) panic(err)
} }
ss.serviceResponseDataStoreLock.Lock()
defer ss.serviceResponseDataStoreLock.Unlock()
ss.monthlyStatusLock.Lock()
defer ss.monthlyStatusLock.Unlock()
ss.ServicesLock.Lock()
defer ss.ServicesLock.Unlock()
for i := 0; i < len(services); i++ { for i := 0; i < len(services); i++ {
task := *services[i] task := *services[i]
// 通过cron定时将服务监控任务传递给任务调度管道 // 通过cron定时将服务监控任务传递给任务调度管道

View File

@ -42,6 +42,7 @@ func InitTimezoneAndCache() {
// LoadSingleton 加载子服务并执行 // LoadSingleton 加载子服务并执行
func LoadSingleton() { func LoadSingleton() {
initUser() // 加载用户ID绑定表
initI18n() // 加载本地化服务 initI18n() // 加载本地化服务
loadNotifications() // 加载通知服务 loadNotifications() // 加载通知服务
loadServers() // 加载服务器列表 loadServers() // 加载服务器列表
@ -81,8 +82,8 @@ func InitDBFromPath(path string) {
} }
err = DB.AutoMigrate(model.Server{}, model.User{}, model.ServerGroup{}, model.NotificationGroup{}, err = DB.AutoMigrate(model.Server{}, model.User{}, model.ServerGroup{}, model.NotificationGroup{},
model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{}, model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{},
model.ServiceHistory{}, model.Cron{}, model.Transfer{}, model.ServerGroupServer{}, model.UserGroup{}, model.ServiceHistory{}, model.Cron{}, model.Transfer{}, model.ServerGroupServer{},
model.UserGroupUser{}, model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{}, model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{},
model.WAF{}) model.WAF{})
if err != nil { if err != nil {
panic(err) panic(err)

135
service/singleton/user.go Normal file
View File

@ -0,0 +1,135 @@
package singleton
import (
"sync"
"github.com/nezhahq/nezha/model"
"gorm.io/gorm"
)
var (
UserInfoMap map[uint64]model.UserInfo
AgentSecretToUserId map[string]uint64
UserLock sync.RWMutex
)
func initUser() {
UserInfoMap = make(map[uint64]model.UserInfo)
AgentSecretToUserId = make(map[string]uint64)
var users []model.User
DB.Find(&users)
for _, u := range users {
UserInfoMap[u.ID] = model.UserInfo{
Role: u.Role,
AgentSecret: u.AgentSecret,
}
AgentSecretToUserId[u.AgentSecret] = u.ID
}
}
func OnUserUpdate(u *model.User) {
UserLock.Lock()
defer UserLock.Unlock()
if u == nil {
return
}
UserInfoMap[u.ID] = model.UserInfo{
Role: u.Role,
AgentSecret: u.AgentSecret,
}
AgentSecretToUserId[u.AgentSecret] = u.ID
}
func OnUserDelete(id []uint64, errorFunc func(string, ...interface{}) error) error {
UserLock.Lock()
defer UserLock.Unlock()
if len(id) < 1 {
return Localizer.ErrorT("user id not specified")
}
var (
cron, server bool
crons, servers []uint64
)
for _, uid := range id {
err := DB.Transaction(func(tx *gorm.DB) error {
CronLock.RLock()
crons = model.FindByUserID(CronList, uid)
CronLock.RUnlock()
cron = len(crons) > 0
if cron {
if err := tx.Unscoped().Delete(&model.Cron{}, "id in (?)", crons).Error; err != nil {
return err
}
}
SortedServerLock.RLock()
servers = model.FindByUserID(SortedServerList, uid)
SortedServerLock.RUnlock()
server = len(servers) > 0
if server {
if err := tx.Unscoped().Delete(&model.Server{}, "id in (?)", servers).Error; err != nil {
return err
}
if err := tx.Unscoped().Delete(&model.ServerGroupServer{}, "server_id in (?)", servers).Error; err != nil {
return err
}
}
if err := tx.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers).Error; err != nil {
return err
}
if err := tx.Where("id IN (?)", id).Delete(&model.User{}).Error; err != nil {
return err
}
return nil
})
if err != nil {
return errorFunc("%v", err)
}
if cron {
OnDeleteCron(crons)
}
if server {
AlertsLock.Lock()
for _, sid := range servers {
for _, alert := range Alerts {
if AlertsCycleTransferStatsStore[alert.ID] != nil {
delete(AlertsCycleTransferStatsStore[alert.ID].ServerName, sid)
delete(AlertsCycleTransferStatsStore[alert.ID].Transfer, sid)
delete(AlertsCycleTransferStatsStore[alert.ID].NextUpdate, sid)
}
}
}
AlertsLock.Unlock()
OnServerDelete(servers)
}
secret := UserInfoMap[uid].AgentSecret
delete(AgentSecretToUserId, secret)
delete(UserInfoMap, uid)
}
if cron {
UpdateCronList()
}
if server {
ReSortServer()
}
return nil
}