feat: able to set initial quota for new user (close #22)

common/constants.go

@@ -46,6 +46,8 @@ var WeChatAccountQRCodeImageURL = ""
 var TurnstileSiteKey = ""
 var TurnstileSecretKey = ""
 
+var QuotaForNewUser = 100
+
 const (
 	RoleGuestUser  = 0
 	RoleCommonUser = 1

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": "",

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)
 	}

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)
 	}
 }

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
+}

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 {
@@ -229,6 +231,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>
           </Header>