Compare commits

..

5 Commits

Author SHA1 Message Date
JustSong
195e94a75d fix: fix MySQL syntax error (#54) 2023-04-27 11:10:10 +08:00
JustSong
5bfc224669 fix: specify type for token (close #23) 2023-04-27 09:32:20 +08:00
JustSong
fd149c242f fix: remove rate limit for relay api 2023-04-26 21:50:09 +08:00
JustSong
b9cc5dfa3f feat: able to set initial quota for new user (close #22) 2023-04-26 21:40:56 +08:00
JustSong
8c305dc1bc feat: able to manage system vai access token (close #12) 2023-04-26 20:54:39 +08:00
13 changed files with 158 additions and 15 deletions

View File

@@ -46,6 +46,8 @@ var WeChatAccountQRCodeImageURL = ""
var TurnstileSiteKey = ""
var TurnstileSecretKey = ""
var QuotaForNewUser = 100
const (
RoleGuestUser = 0
RoleCommonUser = 1
@@ -63,7 +65,7 @@ var (
// All duration's unit is seconds
// Shouldn't larger then RateLimitKeyExpirationDuration
var (
GlobalApiRateLimitNum = 60000 // TODO: temporary set to 60000
GlobalApiRateLimitNum = 180
GlobalApiRateLimitDuration int64 = 3 * 60
GlobalWebRateLimitNum = 60

View File

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

View File

@@ -243,6 +243,42 @@ func GetUser(c *gin.Context) {
return
}
func GenerateAccessToken(c *gin.Context) {
id := c.GetInt("id")
user, err := model.GetUserById(id, true)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
user.AccessToken = common.GetUUID()
if model.DB.Where("token = ?", user.AccessToken).First(user).RowsAffected != 0 {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "请重试,系统生成的 UUID 竟然重复了!",
})
return
}
if err := user.Update(false); err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": user.AccessToken,
})
return
}
func GetSelf(c *gin.Context) {
id := c.GetInt("id")
user, err := model.GetUserById(id, false)

View File

@@ -16,12 +16,31 @@ func authHelper(c *gin.Context, minRole int) {
id := session.Get("id")
status := session.Get("status")
if username == nil {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": "无权进行此操作,未登录",
})
c.Abort()
return
// Check access token
accessToken := c.Request.Header.Get("Authorization")
if accessToken == "" {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": "无权进行此操作,未登录且未提供 access token",
})
c.Abort()
return
}
user := model.ValidateAccessToken(accessToken)
if user != nil && user.Username != "" {
// Token is valid
username = user.Username
role = user.Role
id = user.Id
status = user.Status
} else {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "无权进行此操作access token 无效",
})
c.Abort()
return
}
}
if status.(int) == common.UserStatusDisabled {
c.JSON(http.StatusOK, gin.H{

View File

@@ -25,6 +25,7 @@ func createRootAccountIfNeed() error {
Role: common.RoleRootUser,
Status: common.UserStatusEnabled,
DisplayName: "Root User",
AccessToken: common.GetUUID(),
}
DB.Create(&rootUser)
}

View File

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

View File

@@ -9,7 +9,7 @@ import (
type Redemption struct {
Id int `json:"id"`
UserId int `json:"user_id"`
Key string `json:"key" gorm:"uniqueIndex"`
Key string `json:"key" gorm:"type:char(32);uniqueIndex"`
Status int `json:"status" gorm:"default:1"`
Name string `json:"name" gorm:"index"`
Quota int `json:"quota" gorm:"default:100"`
@@ -48,7 +48,7 @@ func Redeem(key string, tokenId int) (quota int, err error) {
return 0, errors.New("未提供 token id")
}
redemption := &Redemption{}
err = DB.Where("key = ?", key).First(redemption).Error
err = DB.Where("`key` = ?", key).First(redemption).Error
if err != nil {
return 0, errors.New("无效的兑换码")
}

View File

@@ -11,7 +11,7 @@ import (
type Token struct {
Id int `json:"id"`
UserId int `json:"user_id"`
Key string `json:"key" gorm:"uniqueIndex"`
Key string `json:"key" gorm:"type:char(32);uniqueIndex"`
Status int `json:"status" gorm:"default:1"`
Name string `json:"name" gorm:"index" `
CreatedTime int64 `json:"created_time" gorm:"bigint"`
@@ -39,7 +39,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
}
key = strings.Replace(key, "Bearer ", "", 1)
token = &Token{}
err = DB.Where("key = ?", key).First(token).Error
err = DB.Where("`key` = ?", key).First(token).Error
if err == nil {
if token.Status != common.TokenStatusEnabled {
return nil, errors.New("该 token 状态不可用")

View File

@@ -2,7 +2,9 @@ package model
import (
"errors"
"gorm.io/gorm"
"one-api/common"
"strings"
)
// User if you add sensitive fields, don't forget to clean them in setupLogin function.
@@ -19,6 +21,8 @@ type User struct {
WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
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:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
Quota int `json:"quota" gorm:"type:int;default:0"`
}
func GetMaxUserId() int {
@@ -67,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
}
@@ -188,3 +194,25 @@ func IsAdmin(userId int) bool {
}
return user.Role >= common.RoleAdminUser
}
func ValidateAccessToken(token string) (user *User) {
if token == "" {
return nil
}
token = strings.Replace(token, "Bearer ", "", 1)
user = &User{}
if DB.Where("access_token = ?", token).First(user).RowsAffected == 1 {
return 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
}

View File

@@ -35,6 +35,7 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.GET("/self", controller.GetSelf)
selfRoute.PUT("/self", controller.UpdateSelf)
selfRoute.DELETE("/self", controller.DeleteSelf)
selfRoute.GET("/token", controller.GenerateAccessToken)
}
adminRoute := userRoute.Group("/")

View File

@@ -8,12 +8,12 @@ import (
func SetRelayRouter(router *gin.Engine) {
relayV1Router := router.Group("/v1")
relayV1Router.Use(middleware.GlobalAPIRateLimit(), middleware.TokenAuth(), middleware.Distribute())
relayV1Router.Use(middleware.TokenAuth(), middleware.Distribute())
{
relayV1Router.Any("/*path", controller.Relay)
}
relayDashboardRouter := router.Group("/dashboard")
relayDashboardRouter.Use(middleware.GlobalAPIRateLimit(), middleware.TokenAuth(), middleware.Distribute())
relayDashboardRouter.Use(middleware.TokenAuth(), middleware.Distribute())
{
relayDashboardRouter.Any("/*path", controller.Relay)
}

View File

@@ -1,5 +1,5 @@
import React, { useEffect, useState } from 'react';
import { Button, Divider, Form, Header, Image, Modal } from 'semantic-ui-react';
import { Button, Divider, Form, Header, Image, Message, Modal } from 'semantic-ui-react';
import { Link } from 'react-router-dom';
import { API, copy, showError, showInfo, showSuccess } from '../helpers';
import Turnstile from 'react-turnstile';
@@ -34,6 +34,17 @@ const PersonalSetting = () => {
setInputs((inputs) => ({ ...inputs, [name]: value }));
};
const generateAccessToken = async () => {
const res = await API.get('/api/user/token');
const { success, message, data } = res.data;
if (success) {
await copy(data);
showSuccess(`令牌已重置并已复制到剪贴板:${data}`);
} else {
showError(message);
}
};
const bindWeChat = async () => {
if (inputs.wechat_verification_code === '') return;
const res = await API.get(
@@ -92,9 +103,13 @@ const PersonalSetting = () => {
return (
<div style={{ lineHeight: '40px' }}>
<Header as='h3'>通用设置</Header>
<Message>
注意此处生成的令牌用于系统管理而非用于请求 OpenAI 相关的服务请知悉
</Message>
<Button as={Link} to={`/user/edit/`}>
更新个人信息
</Button>
<Button onClick={generateAccessToken}>生成系统访问令牌</Button>
<Divider />
<Header as='h3'>账号绑定</Header>
<Button

View File

@@ -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>