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 = () => { /> +
+ 运营设置 +
+ + + + { + updateOption('QuotaForNewUser', inputs.QuotaForNewUser).then(); + }}>保存运营设置 +
配置 SMTP 用以支持系统的邮件发送