diff --git a/cmd/dashboard/controller/alertrule.go b/cmd/dashboard/controller/alertrule.go index 708f9cf..5ebb56d 100644 --- a/cmd/dashboard/controller/alertrule.go +++ b/cmd/dashboard/controller/alertrule.go @@ -50,6 +50,9 @@ func createAlertRule(c *gin.Context) (uint64, error) { return 0, err } + uid := getUid(c) + + r.UserID = uid r.Name = arf.Name r.Rules = arf.Rules r.FailTriggerTasks = arf.FailTriggerTasks @@ -100,6 +103,10 @@ func updateAlertRule(c *gin.Context) (any, error) { 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.Rules = arf.Rules r.FailTriggerTasks = arf.FailTriggerTasks @@ -133,12 +140,24 @@ func updateAlertRule(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/alert-rule [post] func batchDeleteAlertRule(c *gin.Context) (any, error) { - var ar []uint64 - - if err := c.ShouldBindJSON(&ar); err != nil { + var arr []uint64 + if err := c.ShouldBindJSON(&arr); err != nil { return nil, err } + var ars []model.AlertRule + if err := singleton.DB.Where("id in (?)", arr).Find(&ars).Error; err != nil { + return nil, err + } + + var ar []uint64 + for _, a := range ars { + if !a.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + ar = append(ar, a.ID) + } + if err := singleton.DB.Unscoped().Delete(&model.AlertRule{}, "id in (?)", ar).Error; err != nil { return nil, newGormError("%v", err) } diff --git a/cmd/dashboard/controller/controller.go b/cmd/dashboard/controller/controller.go index 9f0f4af..1251bbe 100644 --- a/cmd/dashboard/controller/controller.go +++ b/cmd/dashboard/controller/controller.go @@ -9,6 +9,7 @@ import ( "net/http" "os" "path" + "slices" "strings" 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.POST("/profile", commonHandler(updateProfile)) - auth.GET("/user", commonHandler(listUser)) - auth.POST("/user", commonHandler(createUser)) - auth.POST("/batch-delete/user", commonHandler(batchDeleteUser)) + auth.GET("/user", adminHandler(listUser)) + auth.POST("/user", adminHandler(createUser)) + 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.PATCH("/service/:id", commonHandler(updateService)) 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.POST("/batch-delete/notification-group", commonHandler(batchDeleteNotificationGroup)) - auth.GET("/server", commonHandler(listServer)) + auth.GET("/server", listHandler(listServer)) auth.PATCH("/server/:id", commonHandler(updateServer)) auth.POST("/batch-delete/server", commonHandler(batchDeleteServer)) auth.POST("/force-update/server", commonHandler(forceUpdateServer)) - auth.GET("/notification", commonHandler(listNotification)) + auth.GET("/notification", listHandler(listNotification)) auth.POST("/notification", commonHandler(createNotification)) auth.PATCH("/notification/:id", commonHandler(updateNotification)) 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.PATCH("/alert-rule/:id", commonHandler(updateAlertRule)) auth.POST("/batch-delete/alert-rule", commonHandler(batchDeleteAlertRule)) - auth.GET("/cron", commonHandler(listCron)) + auth.GET("/cron", listHandler(listCron)) auth.POST("/cron", commonHandler(createCron)) auth.PATCH("/cron/:id", commonHandler(updateCron)) auth.GET("/cron/:id/manual", commonHandler(manualTriggerCron)) auth.POST("/batch-delete/cron", commonHandler(batchDeleteCron)) - auth.GET("/ddns", commonHandler(listDDNS)) + auth.GET("/ddns", listHandler(listDDNS)) auth.GET("/ddns/providers", commonHandler(listProviders)) auth.POST("/ddns", commonHandler(createDDNS)) auth.PATCH("/ddns/:id", commonHandler(updateDDNS)) auth.POST("/batch-delete/ddns", commonHandler(batchDeleteDDNS)) - auth.GET("/nat", commonHandler(listNAT)) + auth.GET("/nat", listHandler(listNAT)) auth.POST("/nat", commonHandler(createNAT)) auth.PATCH("/nat/:id", commonHandler(updateNAT)) auth.POST("/batch-delete/nat", commonHandler(batchDeleteNAT)) auth.GET("/waf", commonHandler(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)) } @@ -189,29 +190,74 @@ func (we *wsError) Error() string { func commonHandler[T any](handler handlerFunc[T]) func(*gin.Context) { return func(c *gin.Context) { - data, err := handler(c) - if err == nil { - c.JSON(http.StatusOK, model.CommonResponse[T]{Success: true, Data: data}) + handle(c, handler) + } +} + +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 } - switch err.(type) { - case *gormError: - log.Printf("NEZHA>> gorm error: %v", err) - c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("database error"))) + + user := *auth.(*model.User) + if user.Role != model.RoleAdmin { + c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("permission denied"))) 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: + } + + handle(c, handler) + } +} + +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)) return } + + c.JSON(http.StatusOK, model.CommonResponse[S]{Success: true, Data: filter(c, 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) { checkLocalFileOrFs := func(c *gin.Context, fs fs.FS, path string) bool { if _, err := os.Stat(path); err == nil { diff --git a/cmd/dashboard/controller/cron.go b/cmd/dashboard/controller/cron.go index d8d95de..08ee7f1 100644 --- a/cmd/dashboard/controller/cron.go +++ b/cmd/dashboard/controller/cron.go @@ -1,7 +1,6 @@ package controller import ( - "fmt" "strconv" "github.com/gin-gonic/gin" @@ -50,6 +49,7 @@ func createCron(c *gin.Context) (uint64, error) { return 0, err } + cr.UserID = getUid(c) cr.TaskType = cf.TaskType cr.Name = cf.Name cr.Scheduler = cf.Scheduler @@ -106,7 +106,11 @@ func updateCron(c *gin.Context) (any, error) { var cr model.Cron 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 @@ -156,12 +160,19 @@ func manualTriggerCron(c *gin.Context) (any, error) { return nil, err } - var cr model.Cron - if err := singleton.DB.First(&cr, id).Error; err != nil { + singleton.CronLock.RLock() + cr, ok := singleton.Crons[id] + if !ok { + singleton.CronLock.RUnlock() 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 } @@ -177,12 +188,24 @@ func manualTriggerCron(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/cron [post] func batchDeleteCron(c *gin.Context) (any, error) { - var cr []uint64 - - if err := c.ShouldBindJSON(&cr); err != nil { + var crr []uint64 + if err := c.ShouldBindJSON(&crr); err != nil { return nil, err } + var cr []uint64 + singleton.CronLock.RLock() + for _, crID := range crr { + if crn, ok := singleton.Crons[crID]; ok { + if !crn.HasPermission(c) { + singleton.CronLock.RUnlock() + return nil, singleton.Localizer.ErrorT("permission denied") + } + cr = append(cr, crn.ID) + } + } + singleton.CronLock.RUnlock() + if err := singleton.DB.Unscoped().Delete(&model.Cron{}, "id in (?)", cr).Error; err != nil { return nil, newGormError("%v", err) } diff --git a/cmd/dashboard/controller/ddns.go b/cmd/dashboard/controller/ddns.go index b58c28e..a0e52b7 100644 --- a/cmd/dashboard/controller/ddns.go +++ b/cmd/dashboard/controller/ddns.go @@ -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") } + p.UserID = getUid(c) p.Name = df.Name enableIPv4 := df.EnableIPv4 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) } + if !p.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + p.Name = df.Name enableIPv4 := df.EnableIPv4 enableIPv6 := df.EnableIPv6 @@ -172,12 +177,25 @@ func updateDDNS(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/ddns [post] func batchDeleteDDNS(c *gin.Context) (any, error) { - var ddnsConfigs []uint64 + var ddnsConfigsr []uint64 - if err := c.ShouldBindJSON(&ddnsConfigs); err != nil { + if err := c.ShouldBindJSON(&ddnsConfigsr); err != nil { return nil, err } + var ddnsConfigs []uint64 + singleton.DDNSCacheLock.RLock() + for _, pid := range ddnsConfigsr { + if p, ok := singleton.DDNSCache[pid]; ok { + if !p.HasPermission(c) { + singleton.DDNSCacheLock.RUnlock() + return nil, singleton.Localizer.ErrorT("permission denied") + } + ddnsConfigs = append(ddnsConfigs, p.ID) + } + } + singleton.DDNSCacheLock.RUnlock() + if err := singleton.DB.Unscoped().Delete(&model.DDNSProfile{}, "id in (?)", ddnsConfigs).Error; err != nil { return nil, newGormError("%v", err) } diff --git a/cmd/dashboard/controller/fm.go b/cmd/dashboard/controller/fm.go index c413d62..955970b 100644 --- a/cmd/dashboard/controller/fm.go +++ b/cmd/dashboard/controller/fm.go @@ -7,6 +7,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/hashicorp/go-uuid" + "github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/pkg/websocketx" @@ -31,13 +32,6 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) { return nil, err } - streamId, err := uuid.GenerateUUID() - if err != nil { - return nil, err - } - - rpc.NezhaHandlerSingleton.CreateStream(streamId) - singleton.ServerLock.RLock() server := singleton.ServerList[id] 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") } + 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{ StreamID: streamId, }) diff --git a/cmd/dashboard/controller/nat.go b/cmd/dashboard/controller/nat.go index 6a2737e..4a75de7 100644 --- a/cmd/dashboard/controller/nat.go +++ b/cmd/dashboard/controller/nat.go @@ -51,6 +51,9 @@ func createNAT(c *gin.Context) (uint64, error) { return 0, err } + uid := getUid(c) + + n.UserID = uid n.Name = nf.Name n.Domain = nf.Domain n.Host = nf.Host @@ -95,6 +98,10 @@ func updateNAT(c *gin.Context) (any, error) { 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.Domain = nf.Domain n.Host = nf.Host @@ -121,12 +128,24 @@ func updateNAT(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/nat [post] func batchDeleteNAT(c *gin.Context) (any, error) { - var n []uint64 - - if err := c.ShouldBindJSON(&n); err != nil { + var nr []uint64 + if err := c.ShouldBindJSON(&nr); err != nil { return nil, err } + var n []uint64 + singleton.NATCacheRwLock.RLock() + for _, id := range nr { + if p, ok := singleton.NATCache[singleton.NATIDToDomain[id]]; ok { + if !p.HasPermission(c) { + singleton.NATCacheRwLock.RUnlock() + return nil, singleton.Localizer.ErrorT("permission denied") + } + n = append(n, p.ID) + } + } + singleton.NATCacheRwLock.RUnlock() + if err := singleton.DB.Unscoped().Delete(&model.NAT{}, "id in (?)", n).Error; err != nil { return nil, newGormError("%v", err) } diff --git a/cmd/dashboard/controller/notification.go b/cmd/dashboard/controller/notification.go index 78ac770..7a03937 100644 --- a/cmd/dashboard/controller/notification.go +++ b/cmd/dashboard/controller/notification.go @@ -48,6 +48,7 @@ func createNotification(c *gin.Context) (uint64, error) { } var n model.Notification + n.UserID = getUid(c) n.Name = nf.Name n.RequestMethod = nf.RequestMethod 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) } + if !n.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + n.Name = nf.Name n.RequestMethod = nf.RequestMethod n.RequestType = nf.RequestType @@ -148,12 +153,23 @@ func updateNotification(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/notification [post] func batchDeleteNotification(c *gin.Context) (any, error) { - var n []uint64 - - if err := c.ShouldBindJSON(&n); err != nil { + var nr []uint64 + if err := c.ShouldBindJSON(&nr); err != nil { return nil, err } + var n []uint64 + singleton.NotificationsLock.RLock() + for _, nid := range nr { + if ns, ok := singleton.NotificationMap[nid]; ok { + if !ns.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + n = append(n, ns.ID) + } + } + singleton.NotificationsLock.RUnlock() + err := singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Delete(&model.Notification{}, "id in (?)", n).Error; err != nil { return err diff --git a/cmd/dashboard/controller/notification_group.go b/cmd/dashboard/controller/notification_group.go index 2e74dba..6b27471 100644 --- a/cmd/dashboard/controller/notification_group.go +++ b/cmd/dashboard/controller/notification_group.go @@ -20,7 +20,7 @@ import ( // @Produce json // @Success 200 {object} model.CommonResponse[[]model.NotificationGroupResponseItem] // @Router /notification-group [get] -func listNotificationGroup(c *gin.Context) ([]model.NotificationGroupResponseItem, error) { +func listNotificationGroup(c *gin.Context) ([]*model.NotificationGroupResponseItem, error) { var ng []model.NotificationGroup if err := singleton.DB.Find(&ng).Error; err != nil { return nil, err @@ -39,9 +39,9 @@ func listNotificationGroup(c *gin.Context) ([]model.NotificationGroupResponseIte 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 { - ngRes = append(ngRes, model.NotificationGroupResponseItem{ + ngRes = append(ngRes, &model.NotificationGroupResponseItem{ Group: n, Notifications: groupNotifications[n.ID], }) @@ -68,8 +68,11 @@ func createNotificationGroup(c *gin.Context) (uint64, error) { } ngf.Notifications = slices.Compact(ngf.Notifications) + uid := getUid(c) + var ng model.NotificationGroup ng.Name = ngf.Name + ng.UserID = uid var count int64 if err := singleton.DB.Model(&model.Notification{}).Where("id in (?)", ngf.Notifications).Count(&count).Error; err != nil { @@ -86,6 +89,9 @@ func createNotificationGroup(c *gin.Context) (uint64, error) { } for _, n := range ngf.Notifications { if err := tx.Create(&model.NotificationGroupNotification{ + Common: model.Common{ + UserID: uid, + }, NotificationGroupID: ng.ID, NotificationID: n, }).Error; err != nil { @@ -131,6 +137,10 @@ func updateNotificationGroup(c *gin.Context) (any, error) { 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 ngf.Notifications = slices.Compact(ngf.Notifications) @@ -142,6 +152,8 @@ func updateNotificationGroup(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("have invalid notification id") } + uid := getUid(c) + err = singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Save(&ngDB).Error; err != nil { return err @@ -152,6 +164,9 @@ func updateNotificationGroup(c *gin.Context) (any, error) { for _, n := range ngf.Notifications { if err := tx.Create(&model.NotificationGroupNotification{ + Common: model.Common{ + UserID: uid, + }, NotificationGroupID: ngDB.ID, NotificationID: n, }).Error; err != nil { @@ -180,11 +195,24 @@ func updateNotificationGroup(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/notification-group [post] func batchDeleteNotificationGroup(c *gin.Context) (any, error) { - var ngn []uint64 - if err := c.ShouldBindJSON(&ngn); err != nil { + var ngnr []uint64 + if err := c.ShouldBindJSON(&ngnr); err != nil { return nil, err } + var ng []model.NotificationGroup + if err := singleton.DB.Where("id in (?)", ngnr).Find(&ng).Error; err != nil { + return nil, err + } + + var ngn []uint64 + for _, n := range ng { + if !n.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + ngn = append(ngn, n.ID) + } + err := singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Delete(&model.NotificationGroup{}, "id in (?)", ngn).Error; err != nil { return err diff --git a/cmd/dashboard/controller/server.go b/cmd/dashboard/controller/server.go index 48bd56e..69bf57a 100644 --- a/cmd/dashboard/controller/server.go +++ b/cmd/dashboard/controller/server.go @@ -61,6 +61,10 @@ func updateServer(c *gin.Context) (any, error) { 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.DisplayIndex = sf.DisplayIndex s.Note = sf.Note @@ -99,11 +103,24 @@ func updateServer(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/server [post] func batchDeleteServer(c *gin.Context) (any, error) { - var servers []uint64 - if err := c.ShouldBindJSON(&servers); err != nil { + var serversRaw []uint64 + if err := c.ShouldBindJSON(&serversRaw); err != nil { return nil, err } + var servers []uint64 + singleton.ServerLock.RLock() + for _, sid := range serversRaw { + if s, ok := singleton.ServerList[sid]; ok { + if !s.HasPermission(c) { + singleton.ServerLock.RUnlock() + return nil, singleton.Localizer.ErrorT("permission denied") + } + servers = append(servers, s.ID) + } + } + singleton.ServerLock.RUnlock() + err := singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Delete(&model.Server{}, "id in (?)", servers).Error; err != nil { return err @@ -161,6 +178,9 @@ func forceUpdateServer(c *gin.Context) (*model.ForceUpdateResponse, error) { server := singleton.ServerList[sid] singleton.ServerLock.RUnlock() if server != nil && server.TaskStream != nil { + if !server.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } if err := server.TaskStream.Send(&pb.Task{ Type: model.TaskTypeUpgrade, }); err != nil { diff --git a/cmd/dashboard/controller/server_group.go b/cmd/dashboard/controller/server_group.go index 98b5c6f..1b4fd4c 100644 --- a/cmd/dashboard/controller/server_group.go +++ b/cmd/dashboard/controller/server_group.go @@ -20,7 +20,7 @@ import ( // @Produce json // @Success 200 {object} model.CommonResponse[[]model.ServerGroupResponseItem] // @Router /server-group [get] -func listServerGroup(c *gin.Context) ([]model.ServerGroupResponseItem, error) { +func listServerGroup(c *gin.Context) ([]*model.ServerGroupResponseItem, error) { var sg []model.ServerGroup if err := singleton.DB.Find(&sg).Error; err != nil { return nil, err @@ -38,9 +38,9 @@ func listServerGroup(c *gin.Context) ([]model.ServerGroupResponseItem, error) { groupServers[s.ServerGroupId] = append(groupServers[s.ServerGroupId], s.ServerId) } - var sgRes []model.ServerGroupResponseItem + var sgRes []*model.ServerGroupResponseItem for _, s := range sg { - sgRes = append(sgRes, model.ServerGroupResponseItem{ + sgRes = append(sgRes, &model.ServerGroupResponseItem{ Group: s, Servers: groupServers[s.ID], }) @@ -67,8 +67,11 @@ func createServerGroup(c *gin.Context) (uint64, error) { } sgf.Servers = slices.Compact(sgf.Servers) + uid := getUid(c) + var sg model.ServerGroup sg.Name = sgf.Name + sg.UserID = uid var count int64 if err := singleton.DB.Model(&model.Server{}).Where("id in (?)", sgf.Servers).Count(&count).Error; err != nil { @@ -84,6 +87,9 @@ func createServerGroup(c *gin.Context) (uint64, error) { } for _, s := range sgf.Servers { if err := tx.Create(&model.ServerGroupServer{ + Common: model.Common{ + UserID: uid, + }, ServerGroupId: sg.ID, ServerId: s, }).Error; err != nil { @@ -129,6 +135,11 @@ func updateServerGroup(c *gin.Context) (any, error) { if err := singleton.DB.First(&sgDB, id).Error; err != nil { 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 var count int64 @@ -139,6 +150,8 @@ func updateServerGroup(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("have invalid server id") } + uid := getUid(c) + err = singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Save(&sgDB).Error; err != nil { return err @@ -149,6 +162,9 @@ func updateServerGroup(c *gin.Context) (any, error) { for _, s := range sg.Servers { if err := tx.Create(&model.ServerGroupServer{ + Common: model.Common{ + UserID: uid, + }, ServerGroupId: sgDB.ID, ServerId: s, }).Error; err != nil { @@ -176,11 +192,24 @@ func updateServerGroup(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/server-group [post] func batchDeleteServerGroup(c *gin.Context) (any, error) { - var sgs []uint64 - if err := c.ShouldBindJSON(&sgs); err != nil { + var sgsr []uint64 + if err := c.ShouldBindJSON(&sgsr); err != nil { return nil, err } + var sg []model.ServerGroup + if err := singleton.DB.Where("id in (?)", sgsr).Find(&sg).Error; err != nil { + return nil, err + } + + var sgs []uint64 + for _, s := range sg { + if !s.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + sgs = append(sgs, s.ID) + } + err := singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Delete(&model.ServerGroup{}, "id in (?)", sgs).Error; err != nil { return err diff --git a/cmd/dashboard/controller/service.go b/cmd/dashboard/controller/service.go index 2956681..78be9a6 100644 --- a/cmd/dashboard/controller/service.go +++ b/cmd/dashboard/controller/service.go @@ -190,7 +190,10 @@ func createService(c *gin.Context) (uint64, error) { return 0, err } + uid := getUid(c) + var m model.Service + m.UserID = uid m.Name = mf.Name m.Target = strings.TrimSpace(mf.Target) m.Type = mf.Type @@ -260,6 +263,11 @@ func updateService(c *gin.Context) (any, error) { if err := singleton.DB.First(&m, id).Error; err != nil { 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.Target = strings.TrimSpace(mf.Target) m.Type = mf.Type @@ -314,10 +322,24 @@ func updateService(c *gin.Context) (any, error) { // @Success 200 {object} model.CommonResponse[any] // @Router /batch-delete/service [post] func batchDeleteService(c *gin.Context) (any, error) { - var ids []uint64 - if err := c.ShouldBindJSON(&ids); err != nil { + var idsr []uint64 + if err := c.ShouldBindJSON(&idsr); err != nil { return nil, err } + + var ids []uint64 + singleton.ServiceSentinelShared.ServicesLock.RLock() + for _, id := range idsr { + if ss, ok := singleton.ServiceSentinelShared.Services[id]; ok { + if !ss.HasPermission(c) { + singleton.ServiceSentinelShared.ServicesLock.RUnlock() + return nil, singleton.Localizer.ErrorT("permission denied") + } + ids = append(ids, ss.ID) + } + } + singleton.ServiceSentinelShared.ServicesLock.RUnlock() + err := singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Delete(&model.Service{}, "id in (?)", ids).Error; err != nil { return err diff --git a/cmd/dashboard/controller/terminal.go b/cmd/dashboard/controller/terminal.go index a5011b0..3c1ef0c 100644 --- a/cmd/dashboard/controller/terminal.go +++ b/cmd/dashboard/controller/terminal.go @@ -6,6 +6,7 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/hashicorp/go-uuid" + "github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/pkg/websocketx" @@ -29,13 +30,6 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) { return nil, err } - streamId, err := uuid.GenerateUUID() - if err != nil { - return nil, err - } - - rpc.NezhaHandlerSingleton.CreateStream(streamId) - singleton.ServerLock.RLock() server := singleton.ServerList[createTerminalReq.ServerID] 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") } + 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{ StreamID: streamId, }) diff --git a/cmd/dashboard/controller/user.go b/cmd/dashboard/controller/user.go index 7c8b9b9..12b236f 100644 --- a/cmd/dashboard/controller/user.go +++ b/cmd/dashboard/controller/user.go @@ -114,6 +114,7 @@ func createUser(c *gin.Context) (uint64, error) { var u model.User u.Username = uf.Username + u.Role = model.RoleMember hash, err := bcrypt.GenerateFromPassword([]byte(uf.Password), bcrypt.DefaultCost) if err != nil { @@ -125,6 +126,7 @@ func createUser(c *gin.Context) (uint64, error) { return 0, err } + singleton.OnUserUpdate(&u) return u.ID, nil } @@ -149,5 +151,6 @@ func batchDeleteUser(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("can't delete yourself") } + singleton.OnUserDelete(ids) return nil, singleton.DB.Where("id IN (?)", ids).Delete(&model.User{}).Error } diff --git a/go.mod b/go.mod index 28ef269..f0c6f0e 100644 --- a/go.mod +++ b/go.mod @@ -1,8 +1,8 @@ module github.com/nezhahq/nezha -go 1.22.7 +go 1.23.0 -toolchain go1.23.1 +toolchain go1.23.2 require ( github.com/appleboy/gin-jwt/v2 v2.10.0 diff --git a/model/common.go b/model/common.go index 294295e..cc2a08f 100644 --- a/model/common.go +++ b/model/common.go @@ -2,6 +2,8 @@ package model import ( "time" + + "github.com/gin-gonic/gin" ) const ( @@ -18,6 +20,47 @@ type Common struct { UpdatedAt time.Time `gorm:"autoUpdateTime" json:"updated_at,omitempty"` // Do not use soft deletion // 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 FindUserID[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 { diff --git a/model/user.go b/model/user.go index e1f297b..3e987d3 100644 --- a/model/user.go +++ b/model/user.go @@ -1,9 +1,31 @@ package model +import ( + "github.com/nezhahq/nezha/pkg/utils" + "gorm.io/gorm" +) + +const ( + RoleAdmin uint8 = iota + RoleMember +) + type User struct { Common - Username string `json:"username,omitempty" gorm:"uniqueIndex"` - Password string `json:"password,omitempty" gorm:"type:char(72)"` + Username string `json:"username,omitempty" gorm:"uniqueIndex"` + Password string `json:"password,omitempty" gorm:"type:char(72)"` + Role uint8 `json:"role,omitempty"` + AgentSecret string `json:"agent_secret,omitempty" gorm:"type:char(32)"` +} + +func (u *User) BeforeSave(tx *gorm.DB) error { + key, err := utils.GenerateRandomString(32) + if err != nil { + return err + } + + u.AgentSecret = key + return nil } type Profile struct { diff --git a/model/user_group.go b/model/user_group.go deleted file mode 100644 index 45b0ecf..0000000 --- a/model/user_group.go +++ /dev/null @@ -1,6 +0,0 @@ -package model - -type UserGroup struct { - Common - Name string `json:"name"` -} diff --git a/model/user_group_user.go b/model/user_group_user.go deleted file mode 100644 index 080f35e..0000000 --- a/model/user_group_user.go +++ /dev/null @@ -1,7 +0,0 @@ -package model - -type UserGroupUser struct { - Common - UserGroupId uint64 `json:"user_group_id"` - UserId uint64 `json:"user_id"` -} diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go index 1ec9985..d3efd5a 100644 --- a/pkg/utils/utils.go +++ b/pkg/utils/utils.go @@ -3,10 +3,12 @@ package utils import ( "crypto/rand" "errors" + "maps" "math/big" "net/netip" "os" "regexp" + "slices" "strconv" "strings" @@ -145,3 +147,8 @@ func Itoa[T constraints.Integer](i T) string { 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)) +} diff --git a/service/rpc/auth.go b/service/rpc/auth.go index 1709168..817f981 100644 --- a/service/rpc/auth.go +++ b/service/rpc/auth.go @@ -2,6 +2,7 @@ package rpc import ( "context" + "crypto/subtle" "strings" petname "github.com/dustinkirkland/golang-petname" @@ -36,10 +37,14 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) { ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) - if clientSecret != singleton.Conf.AgentSecretKey { + singleton.UserLock.RLock() + userId, ok := singleton.AgentSecretToUserId[clientSecret] + if !ok && subtle.ConstantTimeCompare([]byte(clientSecret), []byte(singleton.Conf.AgentSecretKey)) != 1 { + singleton.UserLock.RUnlock() model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail) return 0, status.Error(codes.Unauthenticated, "客户端认证失败") } + singleton.UserLock.RUnlock() model.ClearIP(singleton.DB, ip) @@ -53,21 +58,26 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) { } singleton.ServerLock.RLock() - defer singleton.ServerLock.RUnlock() - clientID, hasID := singleton.ServerUUIDToID[clientUUID] + singleton.ServerLock.RUnlock() + 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 { return 0, status.Error(codes.Unauthenticated, err.Error()) } s.Host = &model.Host{} s.State = &model.HostState{} s.GeoIP = &model.GeoIP{} - // generate a random silly server name + + singleton.ServerLock.Lock() singleton.ServerList[s.ID] = &s singleton.ServerUUIDToID[clientUUID] = s.ID + singleton.ServerLock.Unlock() singleton.ReSortServer() + clientID = s.ID } diff --git a/service/singleton/crontask.go b/service/singleton/crontask.go index 0005ef7..8fc63e4 100644 --- a/service/singleton/crontask.go +++ b/service/singleton/crontask.go @@ -12,6 +12,7 @@ import ( "github.com/robfig/cron/v3" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" pb "github.com/nezhahq/nezha/proto" ) @@ -79,10 +80,7 @@ func UpdateCronList() { CronLock.RLock() defer CronLock.RUnlock() - CronList = make([]*model.Cron, 0, len(Crons)) - for _, c := range Crons { - CronList = append(CronList, c) - } + CronList = utils.MapValuesToSlice(Crons) slices.SortFunc(CronList, func(a, b *model.Cron) int { return cmp.Compare(a.ID, b.ID) }) diff --git a/service/singleton/ddns.go b/service/singleton/ddns.go index 7f35dab..ddbc040 100644 --- a/service/singleton/ddns.go +++ b/service/singleton/ddns.go @@ -13,6 +13,7 @@ import ( ddns2 "github.com/nezhahq/nezha/pkg/ddns" "github.com/nezhahq/nezha/pkg/ddns/dummy" "github.com/nezhahq/nezha/pkg/ddns/webhook" + "github.com/nezhahq/nezha/pkg/utils" ) var ( @@ -24,12 +25,10 @@ var ( func initDDNS() { DB.Find(&DDNSList) - DDNSCacheLock.Lock() DDNSCache = make(map[uint64]*model.DDNSProfile) for i := 0; i < len(DDNSList); i++ { DDNSCache[DDNSList[i].ID] = DDNSList[i] } - DDNSCacheLock.Unlock() OnNameserverUpdate() } @@ -56,10 +55,7 @@ func UpdateDDNSList() { DDNSListLock.Lock() defer DDNSListLock.Unlock() - DDNSList = make([]*model.DDNSProfile, 0, len(DDNSCache)) - for _, p := range DDNSCache { - DDNSList = append(DDNSList, p) - } + DDNSList = utils.MapValuesToSlice(DDNSCache) slices.SortFunc(DDNSList, func(a, b *model.DDNSProfile) int { return cmp.Compare(a.ID, b.ID) }) diff --git a/service/singleton/nat.go b/service/singleton/nat.go index 7ac2897..e6f0323 100644 --- a/service/singleton/nat.go +++ b/service/singleton/nat.go @@ -6,6 +6,7 @@ import ( "sync" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" ) var ( @@ -19,8 +20,6 @@ var ( func initNAT() { DB.Find(&NATList) - NATCacheRwLock.Lock() - defer NATCacheRwLock.Unlock() NATCache = make(map[string]*model.NAT) for i := 0; i < len(NATList); i++ { NATCache[NATList[i].Domain] = NATList[i] @@ -59,10 +58,7 @@ func UpdateNATList() { NATListLock.Lock() defer NATListLock.Unlock() - NATList = make([]*model.NAT, 0, len(NATCache)) - for _, n := range NATCache { - NATList = append(NATList, n) - } + NATList = utils.MapValuesToSlice(NATCache) slices.SortFunc(NATList, func(a, b *model.NAT) int { return cmp.Compare(a.ID, b.ID) }) diff --git a/service/singleton/notification.go b/service/singleton/notification.go index 498f000..cb63087 100644 --- a/service/singleton/notification.go +++ b/service/singleton/notification.go @@ -9,6 +9,7 @@ import ( "time" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" ) const ( @@ -30,7 +31,7 @@ var ( ) // InitNotification 初始化 GroupID <-> ID <-> Notification 的映射 -func InitNotification() { +func initNotification() { NotificationList = make(map[uint64]map[uint64]*model.Notification) NotificationIDToGroups = make(map[uint64]map[uint64]struct{}) NotificationGroup = make(map[uint64]string) @@ -38,9 +39,7 @@ func InitNotification() { // loadNotifications 从 DB 初始化通知方式相关参数 func loadNotifications() { - InitNotification() - NotificationsLock.Lock() - + initNotification() groupNotifications := make(map[uint64][]uint64) var ngn []model.NotificationGroupNotification if err := DB.Find(&ngn).Error; err != nil { @@ -74,8 +73,6 @@ func loadNotifications() { } } } - - NotificationsLock.Unlock() } func UpdateNotificationList() { @@ -85,10 +82,7 @@ func UpdateNotificationList() { NotificationSortedLock.Lock() defer NotificationSortedLock.Unlock() - NotificationListSorted = make([]*model.Notification, 0, len(NotificationMap)) - for _, n := range NotificationMap { - NotificationListSorted = append(NotificationListSorted, n) - } + NotificationListSorted = utils.MapValuesToSlice(NotificationMap) slices.SortFunc(NotificationListSorted, func(a, b *model.Notification) int { return cmp.Compare(a.ID, b.ID) }) diff --git a/service/singleton/server.go b/service/singleton/server.go index 9450d9f..c6afb32 100644 --- a/service/singleton/server.go +++ b/service/singleton/server.go @@ -1,10 +1,12 @@ package singleton import ( - "sort" + "cmp" + "slices" "sync" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" ) var ( @@ -45,29 +47,21 @@ func ReSortServer() { SortedServerLock.Lock() defer SortedServerLock.Unlock() - SortedServerList = make([]*model.Server, 0, len(ServerList)) - SortedServerListForGuest = make([]*model.Server, 0) - for _, s := range ServerList { - SortedServerList = append(SortedServerList, s) + SortedServerList = utils.MapValuesToSlice(ServerList) + // 按照服务器 ID 排序的具体实现(ID越大越靠前) + slices.SortStableFunc(SortedServerList, func(a, b *model.Server) int { + 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 { 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) { diff --git a/service/singleton/servicesentinel.go b/service/singleton/servicesentinel.go index 2ef5c56..8cbfd1b 100644 --- a/service/singleton/servicesentinel.go +++ b/service/singleton/servicesentinel.go @@ -11,6 +11,7 @@ import ( "github.com/jinzhu/copier" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" pb "github.com/nezhahq/nezha/proto" ) @@ -174,11 +175,7 @@ func (ss *ServiceSentinel) UpdateServiceList() { ss.ServiceListLock.Lock() defer ss.ServiceListLock.Unlock() - ss.ServiceList = make([]*model.Service, 0, len(ss.Services)) - for _, v := range ss.Services { - ss.ServiceList = append(ss.ServiceList, v) - } - + ss.ServiceList = utils.MapValuesToSlice(ss.Services) slices.SortFunc(ss.ServiceList, func(a, b *model.Service) int { return cmp.Compare(a.ID, b.ID) }) @@ -192,13 +189,6 @@ func (ss *ServiceSentinel) loadServiceHistory() { 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++ { task := *services[i] // 通过cron定时将服务监控任务传递给任务调度管道 diff --git a/service/singleton/singleton.go b/service/singleton/singleton.go index 0d2bedb..c284e91 100644 --- a/service/singleton/singleton.go +++ b/service/singleton/singleton.go @@ -40,6 +40,7 @@ func InitTimezoneAndCache() { // LoadSingleton 加载子服务并执行 func LoadSingleton() { + initUser() // 加载用户ID绑定表 initI18n() // 加载本地化服务 loadNotifications() // 加载通知服务 loadServers() // 加载服务器列表 @@ -79,8 +80,8 @@ func InitDBFromPath(path string) { } err = DB.AutoMigrate(model.Server{}, model.User{}, model.ServerGroup{}, model.NotificationGroup{}, model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{}, - model.ServiceHistory{}, model.Cron{}, model.Transfer{}, model.ServerGroupServer{}, model.UserGroup{}, - model.UserGroupUser{}, model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{}, + model.ServiceHistory{}, model.Cron{}, model.Transfer{}, model.ServerGroupServer{}, + model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{}, model.WAF{}) if err != nil { panic(err) diff --git a/service/singleton/user.go b/service/singleton/user.go new file mode 100644 index 0000000..88c333d --- /dev/null +++ b/service/singleton/user.go @@ -0,0 +1,109 @@ +package singleton + +import ( + "sync" + + "github.com/nezhahq/nezha/model" + "gorm.io/gorm" +) + +var ( + UserIdToAgentSecret map[uint64]string + AgentSecretToUserId map[string]uint64 + + UserLock sync.RWMutex +) + +func initUser() { + UserIdToAgentSecret = make(map[uint64]string) + AgentSecretToUserId = make(map[string]uint64) + + var users []model.User + DB.Find(&users) + + for _, u := range users { + UserIdToAgentSecret[u.ID] = u.AgentSecret + AgentSecretToUserId[u.AgentSecret] = u.ID + } +} + +func OnUserUpdate(u *model.User) { + UserLock.Lock() + defer UserLock.Unlock() + + if u == nil { + return + } + + UserIdToAgentSecret[u.ID] = u.AgentSecret + AgentSecretToUserId[u.AgentSecret] = u.ID +} + +func OnUserDelete(id []uint64) { + UserLock.Lock() + defer UserLock.Unlock() + + if len(id) < 1 { + return + } + + var ( + cron bool + server bool + ) + + for _, uid := range id { + secret := UserIdToAgentSecret[uid] + delete(AgentSecretToUserId, secret) + delete(UserIdToAgentSecret, uid) + + CronLock.RLock() + crons := model.FindUserID(CronList, uid) + CronLock.RUnlock() + + cron = len(crons) > 0 + if cron { + DB.Unscoped().Delete(&model.Cron{}, "id in (?)", crons) + OnDeleteCron(crons) + } + + SortedServerLock.RLock() + servers := model.FindUserID(SortedServerList, uid) + SortedServerLock.RUnlock() + + server = len(servers) > 0 + if server { + DB.Transaction(func(tx *gorm.DB) error { + 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 + } + return nil + }) + + 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) + } + } + } + DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers) + AlertsLock.Unlock() + OnServerDelete(servers) + } + } + + if cron { + UpdateCronList() + } + + if server { + ReSortServer() + } +}