diff --git a/constant/env.go b/constant/env.go index dd3ae65a9..fc9fd8b71 100644 --- a/constant/env.go +++ b/constant/env.go @@ -44,3 +44,6 @@ func InitEnv() { } } } + +// 是否生成初始令牌,默认关闭。 +var GenerateDefaultToken = common.GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false) diff --git a/controller/user.go b/controller/user.go index 7a3edf4c6..f33b9de1d 100644 --- a/controller/user.go +++ b/controller/user.go @@ -11,6 +11,7 @@ import ( "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" + "one-api/constant" ) type LoginRequest struct { @@ -197,24 +198,34 @@ func Register(c *gin.Context) { return } // 生成默认令牌 - // tokenName := cleanUser.Username + "的初始令牌" - token := model.Token{ - UserId: insertedUser.Id, // 使用插入后的用户ID - Name: cleanUser.Username + "的初始令牌", - Key: common.GenerateKey(), - CreatedTime: common.GetTimestamp(), - AccessedTime: common.GetTimestamp(), - ExpiredTime: -1, // 永不过期 - RemainQuota: 500000, // 示例额度 - UnlimitedQuota: true, - ModelLimitsEnabled: false, - } - if err := token.Insert(); err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "创建默认令牌失败", - }) - return + if constant.GenerateDefaultToken { + var insertedUser model.User + if err := model.DB.Where("username = ?", cleanUser.Username).First(&insertedUser).Error; err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "用户注册失败或用户ID获取失败", + }) + return + } + // 生成默认令牌 + token := model.Token{ + UserId: insertedUser.Id, // 使用插入后的用户ID + Name: cleanUser.Username + "的初始令牌", + Key: common.GenerateKey(), + CreatedTime: common.GetTimestamp(), + AccessedTime: common.GetTimestamp(), + ExpiredTime: -1, // 永不过期 + RemainQuota: 500000, // 示例额度 + UnlimitedQuota: true, + ModelLimitsEnabled: false, + } + if err := token.Insert(); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "创建默认令牌失败", + }) + return + } } c.JSON(http.StatusOK, gin.H{