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

songquanpeng · Apr 26, 2023 · b9cc5df · b9cc5df
1 parent 8c305dc
commit b9cc5df
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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>