From 9ee5595da728d32761b094f8766857dfcdd02fc5 Mon Sep 17 00:00:00 2001
From: UUBulb <35923940+uubulb@users.noreply.github.com>
Date: Sat, 8 Mar 2025 18:47:42 +0800
Subject: [PATCH] generate agent_secret for old users (#1021)

---
 cmd/dashboard/main.go     | 17 ++++++++++-------
 model/user.go             |  4 +++-
 pkg/utils/utils.go        |  9 +++++++++
 service/singleton/user.go |  9 +++++++++
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/cmd/dashboard/main.go b/cmd/dashboard/main.go
index 6d1b8f2..3130063 100644
--- a/cmd/dashboard/main.go
+++ b/cmd/dashboard/main.go
@@ -40,23 +40,23 @@ var (
 	frontendDist embed.FS
 )
 
-func initSystem() {
+func initSystem() error {
 	// 初始化管理员账户
 	var usersCount int64
 	if err := singleton.DB.Model(&model.User{}).Count(&usersCount).Error; err != nil {
-		panic(err)
+		return err
 	}
 	if usersCount == 0 {
 		hash, err := bcrypt.GenerateFromPassword([]byte("admin"), bcrypt.DefaultCost)
 		if err != nil {
-			panic(err)
+			return err
 		}
 		admin := model.User{
 			Username: "admin",
 			Password: string(hash),
 		}
 		if err := singleton.DB.Create(&admin).Error; err != nil {
-			panic(err)
+			return err
 		}
 	}
 
@@ -65,13 +65,14 @@ func initSystem() {
 
 	// 每天的3:30 对 监控记录 和 流量记录 进行清理
 	if _, err := singleton.CronShared.AddFunc("0 30 3 * * *", singleton.CleanServiceHistory); err != nil {
-		panic(err)
+		return err
 	}
 
 	// 每小时对流量记录进行打点
 	if _, err := singleton.CronShared.AddFunc("0 0 * * * *", singleton.RecordTransferHourlyUsage); err != nil {
-		panic(err)
+		return err
 	}
+	return nil
 }
 
 // @title           Nezha Monitoring API
@@ -111,7 +112,9 @@ func main() {
 	singleton.InitConfigFromPath(dashboardCliParam.ConfigFile)
 	singleton.InitTimezoneAndCache()
 	singleton.InitDBFromPath(dashboardCliParam.DatabaseLocation)
-	initSystem()
+	if err := initSystem(); err != nil {
+		log.Fatal(err)
+	}
 
 	l, err := net.Listen("tcp", fmt.Sprintf("%s:%d", singleton.Conf.ListenHost, singleton.Conf.ListenPort))
 	if err != nil {
diff --git a/model/user.go b/model/user.go
index 9e214a5..3d07b79 100644
--- a/model/user.go
+++ b/model/user.go
@@ -13,6 +13,8 @@ const (
 	RoleMember
 )
 
+const DefaultAgentSecretLength = 32
+
 type User struct {
 	Common
 	Username       string `json:"username,omitempty" gorm:"uniqueIndex"`
@@ -32,7 +34,7 @@ func (u *User) BeforeSave(tx *gorm.DB) error {
 		return nil
 	}
 
-	key, err := utils.GenerateRandomString(32)
+	key, err := utils.GenerateRandomString(DefaultAgentSecretLength)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
index e1476bd..53178ba 100644
--- a/pkg/utils/utils.go
+++ b/pkg/utils/utils.go
@@ -4,6 +4,7 @@ import (
 	"cmp"
 	"crypto/rand"
 	"errors"
+	"fmt"
 	"iter"
 	"maps"
 	"math/big"
@@ -82,6 +83,14 @@ func GenerateRandomString(n int) (string, error) {
 	return string(ret), nil
 }
 
+func MustGenerateRandomString(n int) string {
+	str, err := GenerateRandomString(n)
+	if err != nil {
+		panic(fmt.Errorf("MustGenerateRandomString: %v", err))
+	}
+	return str
+}
+
 func Uint64SubInt64(a uint64, b int64) uint64 {
 	if b < 0 {
 		return a + uint64(-b)
diff --git a/service/singleton/user.go b/service/singleton/user.go
index 79c90eb..f3368bd 100644
--- a/service/singleton/user.go
+++ b/service/singleton/user.go
@@ -1,9 +1,11 @@
 package singleton
 
 import (
+	"fmt"
 	"sync"
 
 	"github.com/nezhahq/nezha/model"
+	"github.com/nezhahq/nezha/pkg/utils"
 	"gorm.io/gorm"
 )
 
@@ -29,6 +31,13 @@ func initUser() {
 	AgentSecretToUserId[Conf.AgentSecretKey] = 0
 
 	for _, u := range users {
+		if u.AgentSecret == "" {
+			u.AgentSecret = utils.MustGenerateRandomString(model.DefaultAgentSecretLength)
+			if err := DB.Save(&u).Error; err != nil {
+				panic(fmt.Errorf("update of user %d failed: %v", u.ID, err))
+			}
+		}
+
 		UserInfoMap[u.ID] = model.UserInfo{
 			Role:        u.Role,
 			AgentSecret: u.AgentSecret,