From 716bf6f48a24756f32b3d552f838099b49611d47 Mon Sep 17 00:00:00 2001
From: Jin Weihan <155160686+j471782517@users.noreply.github.com>
Date: Sat, 24 Aug 2024 18:44:37 +0000
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E7=8E=AF=E5=A2=83=E5=8F=98?=
 =?UTF-8?q?=E9=87=8FGENERATE=5FDEFAULT=5FTOKEN=20=E8=AE=BE=E7=BD=AE?=
 =?UTF-8?q?=E4=B9=8B=E5=90=8E=E5=B0=86=E7=94=9F=E6=88=90=E5=88=9D=E5=A7=8B?=
 =?UTF-8?q?=E4=BB=A4=E7=89=8C=EF=BC=8C=E9=BB=98=E8=AE=A4=E5=85=B3=E9=97=AD?=
 =?UTF-8?q?=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 constant/env.go    |  3 +++
 controller/user.go | 47 ++++++++++++++++++++++++++++------------------
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/constant/env.go b/constant/env.go
index dd3ae65..fc9fd8b 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 7a3edf4..f33b9de 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{