Jelajahi Sumber

增加环境变量GENERATE_DEFAULT_TOKEN 设置之后将生成初始令牌,默认关闭。

Jin Weihan 1 tahun lalu
induk
melakukan
716bf6f48a
2 mengubah file dengan 32 tambahan dan 18 penghapusan
  1. 3 0
      constant/env.go
  2. 29 18
      controller/user.go

+ 3 - 0
constant/env.go

@@ -44,3 +44,6 @@ func InitEnv() {
 		}
 	}
 }
+
+// 是否生成初始令牌,默认关闭。
+var GenerateDefaultToken = common.GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false)

+ 29 - 18
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{