Skip to content

Commit

Permalink
Merge pull request #2 from j471782517/main
Browse files Browse the repository at this point in the history
增加环境变量GENERATE_DEFAULT_TOKEN 设置之后将生成初始令牌,默认关闭。
  • Loading branch information
guoruqiang authored Aug 24, 2024
2 parents 2422eb2 + 716bf6f commit 1379d7f
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 18 deletions.
3 changes: 3 additions & 0 deletions constant/env.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,6 @@ func InitEnv() {
}
}
}

// 是否生成初始令牌,默认关闭。
var GenerateDefaultToken = common.GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false)
47 changes: 29 additions & 18 deletions controller/user.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (

"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
"one-api/constant"
)

type LoginRequest struct {
Expand Down Expand Up @@ -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{
Expand Down

0 comments on commit 1379d7f

Please sign in to comment.