From b9cc5dfa3f9073e27a812102e115b9dbf3efed59 Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Wed, 26 Apr 2023 21:40:56 +0800
Subject: [PATCH] feat: able to set initial quota for new user (close #22)

---
 common/constants.go                 |  2 ++
 controller/token.go                 | 17 +++++++++++++++++
 model/main.go                       |  1 +
 model/option.go                     |  3 +++
 model/user.go                       | 14 ++++++++++++++
 web/src/components/SystemSetting.js | 23 ++++++++++++++++++++++-
 6 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/common/constants.go b/common/constants.go
index 523bf07d1b..e4985241ac 100644
--- a/common/constants.go
+++ b/common/constants.go
@@ -46,6 +46,8 @@ var WeChatAccountQRCodeImageURL = ""
 var TurnstileSiteKey = ""
 var TurnstileSecretKey = ""
 
+var QuotaForNewUser = 100
+
 const (
 	RoleGuestUser  = 0
 	RoleCommonUser = 1
diff --git a/controller/token.go b/controller/token.go
index b12d2e2d39..d5d1474c7d 100644
--- a/controller/token.go
+++ b/controller/token.go
@@ -104,6 +104,19 @@ func AddToken(c *gin.Context) {
 	if isAdmin {
 		cleanToken.RemainTimes = token.RemainTimes
 		cleanToken.UnlimitedTimes = token.UnlimitedTimes
+	} else {
+		userId := c.GetInt("id")
+		quota, err := model.GetUserQuota(userId)
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": err.Error(),
+			})
+			return
+		}
+		if quota > 0 {
+			cleanToken.RemainTimes = quota
+		}
 	}
 	err = cleanToken.Insert()
 	if err != nil {
@@ -113,6 +126,10 @@ func AddToken(c *gin.Context) {
 		})
 		return
 	}
+	if !isAdmin {
+		// update user quota
+		err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainTimes)
+	}
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
diff --git a/model/main.go b/model/main.go
index 85b72acc54..0bc09230a0 100644
--- a/model/main.go
+++ b/model/main.go
@@ -25,6 +25,7 @@ func createRootAccountIfNeed() error {
 			Role:        common.RoleRootUser,
 			Status:      common.UserStatusEnabled,
 			DisplayName: "Root User",
+			AccessToken: common.GetUUID(),
 		}
 		DB.Create(&rootUser)
 	}
diff --git a/model/option.go b/model/option.go
index a30a8b1d16..f06d155108 100644
--- a/model/option.go
+++ b/model/option.go
@@ -46,6 +46,7 @@ func InitOptionMap() {
 	common.OptionMap["WeChatAccountQRCodeImageURL"] = ""
 	common.OptionMap["TurnstileSiteKey"] = ""
 	common.OptionMap["TurnstileSecretKey"] = ""
+	common.OptionMap["QuotaForNewUser"] = strconv.Itoa(common.QuotaForNewUser)
 	common.OptionMapRWMutex.Unlock()
 	options, _ := AllOption()
 	for _, option := range options {
@@ -131,5 +132,7 @@ func updateOptionMap(key string, value string) {
 		common.TurnstileSiteKey = value
 	case "TurnstileSecretKey":
 		common.TurnstileSecretKey = value
+	case "QuotaForNewUser":
+		common.QuotaForNewUser, _ = strconv.Atoi(value)
 	}
 }
diff --git a/model/user.go b/model/user.go
index 2bd0e0ebb6..84f979607c 100644
--- a/model/user.go
+++ b/model/user.go
@@ -2,6 +2,7 @@ package model
 
 import (
 	"errors"
+	"gorm.io/gorm"
 	"one-api/common"
 	"strings"
 )
@@ -21,6 +22,7 @@ type User struct {
 	VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
 	Balance          int    `json:"balance" gorm:"type:int;default:0"`
 	AccessToken      string `json:"access_token" gorm:"column:access_token;uniqueIndex"` // this token is for system management
+	Quota            int    `json:"quota" gorm:"type:int;default:0"`
 }
 
 func GetMaxUserId() int {
@@ -69,6 +71,8 @@ func (user *User) Insert() error {
 			return err
 		}
 	}
+	user.Quota = common.QuotaForNewUser
+	user.AccessToken = common.GetUUID()
 	err = DB.Create(user).Error
 	return err
 }
@@ -202,3 +206,13 @@ func ValidateAccessToken(token string) (user *User) {
 	}
 	return nil
 }
+
+func GetUserQuota(id int) (quota int, err error) {
+	err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
+	return quota, err
+}
+
+func DecreaseUserQuota(id int, quota int) (err error) {
+	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
+	return err
+}
diff --git a/web/src/components/SystemSetting.js b/web/src/components/SystemSetting.js
index f681061501..e0f1b105ee 100644
--- a/web/src/components/SystemSetting.js
+++ b/web/src/components/SystemSetting.js
@@ -24,6 +24,7 @@ const SystemSetting = () => {
     TurnstileSiteKey: '',
     TurnstileSecretKey: '',
     RegisterEnabled: '',
+    QuotaForNewUser: 0,
   });
   let originInputs = {};
   let [loading, setLoading] = useState(false);
@@ -86,7 +87,8 @@ const SystemSetting = () => {
       name === 'WeChatServerToken' ||
       name === 'WeChatAccountQRCodeImageURL' ||
       name === 'TurnstileSiteKey' ||
-      name === 'TurnstileSecretKey'
+      name === 'TurnstileSecretKey' ||
+      name === 'QuotaForNewUser'
     ) {
       setInputs((inputs) => ({ ...inputs, [name]: value }));
     } else {
@@ -228,6 +230,25 @@ const SystemSetting = () => {
             />
           </Form.Group>
           <Divider />
+          <Header as='h3'>
+            运营设置
+          </Header>
+          <Form.Group widths={3}>
+            <Form.Input
+              label='新用户初始配额'
+              name='QuotaForNewUser'
+              onChange={handleInputChange}
+              autoComplete='off'
+              value={inputs.QuotaForNewUser}
+              type='number'
+              min='0'
+              placeholder='例如：100'
+            />
+          </Form.Group>
+          <Form.Button onClick={()=>{
+            updateOption('QuotaForNewUser', inputs.QuotaForNewUser).then();
+          }}>保存运营设置</Form.Button>
+          <Divider />
           <Header as='h3'>
             配置 SMTP
             <Header.Subheader>用以支持系统的邮件发送</Header.Subheader>