feat: user-specific connection secret

cmd/dashboard/controller/controller.go

@@ -79,7 +79,7 @@ func routers(r *gin.Engine, frontendDist fs.FS) {
 
 	auth.GET("/profile", commonHandler(getProfile))
 	auth.POST("/profile", commonHandler(updateProfile))
-	auth.GET("/user", commonHandler(listUser))
+	auth.GET("/user", adminHandler(listUser))
 	auth.POST("/user", adminHandler(createUser))
 	auth.POST("/batch-delete/user", adminHandler(batchDeleteUser))
 

cmd/dashboard/controller/user.go

@@ -126,6 +126,7 @@ func createUser(c *gin.Context) (uint64, error) {
 		return 0, err
 	}
 
+	singleton.OnUserUpdate(&u)
 	return u.ID, nil
 }
 
@@ -150,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
 }

model/common.go

@@ -20,7 +20,7 @@ type Common struct {
 	// Do not use soft deletion
 	// DeletedAt gorm.DeletedAt `gorm:"index" json:"deleted_at,omitempty"`
 
-	UserID uint64 `json:"user_id,omitempty"`
+	UserID uint64 `json:"-"`
 }
 
 func (c *Common) GetID() uint64 {

model/user.go

@@ -1,5 +1,10 @@
 package model
 
+import (
+	"github.com/nezhahq/nezha/pkg/utils"
+	"gorm.io/gorm"
+)
+
 const (
 	RoleAdmin uint8 = iota
 	RoleMember
@@ -7,9 +12,20 @@ const (
 
 type User struct {
 	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"`
+	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 {

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
 	}
 

service/singleton/ddns.go

@@ -24,12 +24,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()
 }

service/singleton/nat.go

@@ -19,8 +19,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]

service/singleton/notification.go

@@ -30,7 +30,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 +38,7 @@ func InitNotification() {
 
 // loadNotifications 从 DB 初始化通知方式相关参数
 func loadNotifications() {
-	InitNotification()
+	initNotification()
-	NotificationsLock.Lock()
-
 	groupNotifications := make(map[uint64][]uint64)
 	var ngn []model.NotificationGroupNotification
 	if err := DB.Find(&ngn).Error; err != nil {
@@ -74,8 +72,6 @@ func loadNotifications() {
 			}
 		}
 	}
-
-	NotificationsLock.Unlock()
 }
 
 func UpdateNotificationList() {

service/singleton/servicesentinel.go

@@ -192,13 +192,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定时将服务监控任务传递给任务调度管道

service/singleton/singleton.go

@@ -40,6 +40,7 @@ func InitTimezoneAndCache() {
 
 // LoadSingleton 加载子服务并执行
 func LoadSingleton() {
+	initUser()          // 加载用户ID绑定表
 	initI18n()          // 加载本地化服务
 	loadNotifications() // 加载通知服务
 	loadServers()       // 加载服务器列表

service/singleton/user.go

@@ -0,0 +1,54 @@
+package singleton
+
+import (
+	"sync"
+
+	"github.com/nezhahq/nezha/model"
+)
+
+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
+	}
+
+	for _, uid := range id {
+		secret := UserIdToAgentSecret[uid]
+		delete(AgentSecretToUserId, secret)
+		delete(UserIdToAgentSecret, uid)
+	}
+}