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("a).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 = () => {
/>