Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: bind quota to account instead of token (close #64, #31)

Browse Source
This commit is contained in:
JustSong 2023-05-16 11:26:09 +08:00
parent 7c56a36a1c
commit 01abed0a30
12 changed files with 240 additions and 176 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
common/constants.go
View File

@ -50,10 +50,10 @@ var WeChatAccountQRCodeImageURL = ""
var TurnstileSiteKey = ""
var TurnstileSecretKey = ""
var QuotaForNewUser = 100
var QuotaForNewUser = 0
var ChannelDisableThreshold = 5.0
var AutomaticDisableChannelEnabled = false
var QuotaRemindThreshold = 1000 // TODO: QuotaRemindThreshold
var QuotaRemindThreshold = 1000
var RootUserEmail = ""

65
controller/token.go
View File

@ -100,7 +100,6 @@ func GetTokenStatus(c *gin.Context) {
}
func AddToken(c *gin.Context) {
isAdmin := c.GetInt("role") >= common.RoleAdminUser
token := model.Token{}
err := c.ShouldBindJSON(&token)
if err != nil {
@ -118,27 +117,14 @@ func AddToken(c *gin.Context) {
return
}
cleanToken := model.Token{
UserId: c.GetInt("id"),
Name: token.Name,
Key: common.GetUUID(),
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: token.ExpiredTime,
}
if isAdmin {
cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedQuota = token.UnlimitedQuota
} 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
}
cleanToken.RemainQuota = quota
UserId: c.GetInt("id"),
Name: token.Name,
Key: common.GetUUID(),
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: token.ExpiredTime,
RemainQuota: token.RemainQuota,
UnlimitedQuota: token.UnlimitedQuota,
}
err = cleanToken.Insert()
if err != nil {
@ -148,10 +134,6 @@ func AddToken(c *gin.Context) {
})
return
}
if !isAdmin {
// update user quota
err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainQuota)
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
@ -240,34 +222,3 @@ func UpdateToken(c *gin.Context) {
})
return
}
type topUpRequest struct {
Id int `json:"id"`
Key string `json:"key"`
}
func TopUp(c *gin.Context) {
req := topUpRequest{}
err := c.ShouldBindJSON(&req)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
quota, err := model.Redeem(req.Key, req.Id)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": quota,
})
return
}

31
controller/user.go
View File

@ -654,3 +654,34 @@ func EmailBind(c *gin.Context) {
})
return
}
type topUpRequest struct {
Key string `json:"key"`
}
func TopUp(c *gin.Context) {
req := topUpRequest{}
err := c.ShouldBindJSON(&req)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
id := c.GetInt("id")
quota, err := model.Redeem(req.Key, id)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": quota,
})
return
}

8
model/redemption.go
View File

@ -40,12 +40,12 @@ func GetRedemptionById(id int) (*Redemption, error) {
return &redemption, err
}
func Redeem(key string, tokenId int) (quota int, err error) {
func Redeem(key string, userId int) (quota int, err error) {
if key == "" {
return 0, errors.New("未提供兑换码")
}
if tokenId == 0 {
return 0, errors.New("未提供 token id")
if userId == 0 {
return 0, errors.New("无效的 user id")
}
redemption := &Redemption{}
err = DB.Where("`key` = ?", key).First(redemption).Error
@ -55,7 +55,7 @@ func Redeem(key string, tokenId int) (quota int, err error) {
if redemption.Status != common.RedemptionCodeStatusEnabled {
return 0, errors.New("该兑换码已被使用")
}
err = IncreaseTokenQuota(tokenId, redemption.Quota)
err = IncreaseUserQuota(userId, redemption.Quota)
if err != nil {
return 0, err
}

74
model/token.go
View File

@ -2,6 +2,7 @@ package model
import (
"errors"
"fmt"
_ "gorm.io/driver/sqlite"
"gorm.io/gorm"
"one-api/common"
@ -82,6 +83,16 @@ func GetTokenByIds(id int, userId int) (*Token, error) {
return &token, err
}
func GetTokenById(id int) (*Token, error) {
if id == 0 {
return nil, errors.New("id 为空！")
}
token := Token{Id: id}
var err error = nil
err = DB.First(&token, "id = ?", id).Error
return &token, err
}
func (token *Token) Insert() error {
var err error
err = DB.Create(token).Error
@ -116,26 +127,53 @@ func DeleteTokenById(id int, userId int) (err error) {
if err != nil {
return err
}
quota := token.RemainQuota
if quota != 0 {
if quota > 0 {
err = IncreaseUserQuota(userId, quota)
} else {
err = DecreaseUserQuota(userId, -quota)
}
}
if err != nil {
return err
}
return token.Delete()
}
func IncreaseTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error
return err
}
func DecreaseTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
func DecreaseTokenQuota(tokenId int, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数！")
}
token, err := GetTokenById(tokenId)
if err != nil {
return err
}
if token.RemainQuota < quota {
return errors.New("令牌额度不足")
}
userQuota, err := GetUserQuota(token.UserId)
if err != nil {
return err
}
if userQuota < quota {
return errors.New("用户额度不足")
}
quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-quota < common.QuotaRemindThreshold
noMoreQuota := userQuota-quota <= 0
if quotaTooLow || noMoreQuota {
go func() {
email, err := GetUserEmail(token.UserId)
if err != nil {
common.SysError("获取用户邮箱失败：" + err.Error())
}
prompt := "您的额度即将用尽"
if noMoreQuota {
prompt = "您的额度已用尽"
}
if email != "" {
topUpLink := fmt.Sprintf("%s/topup", common.ServerAddress)
err = common.SendEmail(prompt, email,
fmt.Sprintf("%s，剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota-quota, topUpLink, topUpLink))
if err != nil {
common.SysError("发送邮件失败：" + err.Error())
}
}
}()
}
err = DB.Model(&Token{}).Where("id = ?", tokenId).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
if err != nil {
return err
}
err = DecreaseUserQuota(token.UserId, quota)
return err
}

11
model/user.go
View File

@ -225,12 +225,23 @@ func GetUserQuota(id int) (quota int, err error) {
return quota, err
}
func GetUserEmail(id int) (email string, err error) {
err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
return email, err
}
func IncreaseUserQuota(id int, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数！")
}
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
return err
}
func DecreaseUserQuota(id int, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数！")
}
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
return err
}

2
router/api-router.go
View File

@ -37,6 +37,7 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.PUT("/self", controller.UpdateSelf)
selfRoute.DELETE("/self", controller.DeleteSelf)
selfRoute.GET("/token", controller.GenerateAccessToken)
selfRoute.POST("/topup", controller.TopUp)
}
adminRoute := userRoute.Group("/")
@ -74,7 +75,6 @@ func SetApiRouter(router *gin.Engine) {
{
tokenRoute.GET("/", controller.GetAllTokens)
tokenRoute.GET("/search", controller.SearchTokens)
tokenRoute.POST("/topup", controller.TopUp)
tokenRoute.GET("/:id", controller.GetToken)
tokenRoute.POST("/", controller.AddToken)
tokenRoute.PUT("/", controller.UpdateToken)

11
web/src/App.js
View File

@ -21,6 +21,7 @@ import EditToken from './pages/Token/EditToken';
import EditChannel from './pages/Channel/EditChannel';
import Redemption from './pages/Redemption';
import EditRedemption from './pages/Redemption/EditRedemption';
import TopUp from './pages/TopUp';
const Home = lazy(() => import('./pages/Home'));
const About = lazy(() => import('./pages/About'));
@ -239,6 +240,16 @@ function App() {
</PrivateRoute>
}
/>
<Route
path='/topup'
element={
<PrivateRoute>
<Suspense fallback={<Loading></Loading>}>
<TopUp />
</Suspense>
</PrivateRoute>
}
/>
<Route
path='/about'
element={

6
web/src/components/Header.js
View File

@ -30,6 +30,12 @@ const headerButtons = [
icon: 'dollar sign',
admin: true,
},
{
name: '充值',
to: '/topup',
icon: 'cart',
admin: true,
},
{
name: '用户',
to: '/user',

73
web/src/components/TokensTable.js
View File

@ -36,8 +36,6 @@ const TokensTable = () => {
const [searching, setSearching] = useState(false);
const [showTopUpModal, setShowTopUpModal] = useState(false);
const [targetTokenIdx, setTargetTokenIdx] = useState(0);
const [redemptionCode, setRedemptionCode] = useState('');
const [topUpLink, setTopUpLink] = useState('');
const loadTokens = async (startIdx) => {
const res = await API.get(`/api/token/?p=${startIdx}`);
@ -77,13 +75,6 @@ const TokensTable = () => {
.catch((reason) => {
showError(reason);
});
let status = localStorage.getItem('status');
if (status) {
status = JSON.parse(status);
if (status.top_up_link) {
setTopUpLink(status.top_up_link);
}
}
}, []);
const manageToken = async (id, action, idx) => {
@ -156,28 +147,6 @@ const TokensTable = () => {
setLoading(false);
};
const topUp = async () => {
if (redemptionCode === '') {
return;
}
const res = await API.post('/api/token/topup/', {
id: tokens[targetTokenIdx].id,
key: redemptionCode
});
const { success, message, data } = res.data;
if (success) {
showSuccess('充值成功！');
let newTokens = [...tokens];
let realIdx = (activePage - 1) * ITEMS_PER_PAGE + targetTokenIdx;
newTokens[realIdx].remain_quota += data;
setTokens(newTokens);
setRedemptionCode('');
setShowTopUpModal(false);
} else {
showError(message);
}
}
return (
<>
<Form onSubmit={searchTokens}>
@ -279,15 +248,6 @@ const TokensTable = () => {
>
复制
</Button>
<Button
size={'small'}
color={'yellow'}
onClick={() => {
setTargetTokenIdx(idx);
setShowTopUpModal(true);
}}>
充值
</Button>
<Popup
trigger={
<Button size='small' negative>
@ -355,39 +315,6 @@ const TokensTable = () => {
</Table.Row>
</Table.Footer>
</Table>
<Modal
onClose={() => setShowTopUpModal(false)}
onOpen={() => setShowTopUpModal(true)}
open={showTopUpModal}
size={'mini'}
>
<Modal.Header>通过兑换码为令牌「{tokens[targetTokenIdx]?.name}」充值</Modal.Header>
<Modal.Content>
<Modal.Description>
{/*<Image src={status.wechat_qrcode} fluid />*/}
{
topUpLink && <p>
<a target='_blank' href={topUpLink}>点击此处获取兑换码</a>
</p>
}
<Form size='large'>
<Form.Input
fluid
placeholder='兑换码'
name='redemptionCode'
value={redemptionCode}
onChange={(e) => {
setRedemptionCode(e.target.value);
}}
/>
<Button color='' fluid size='large' onClick={topUp}>
充值
</Button>
</Form>
</Modal.Description>
</Modal.Content>
</Modal>
</>
);
};

37
web/src/pages/Token/EditToken.js
View File

@ -1,7 +1,7 @@
import React, { useEffect, useState } from 'react';
import { Button, Form, Header, Segment } from 'semantic-ui-react';
import { useParams } from 'react-router-dom';
import { API, isAdmin, showError, showSuccess, timestamp2string } from '../../helpers';
import { API, showError, showSuccess, timestamp2string } from '../../helpers';
const EditToken = () => {
const params = useParams();
@ -14,7 +14,6 @@ const EditToken = () => {
expired_time: -1,
unlimited_quota: false
};
const isAdminUser = isAdmin();
const [inputs, setInputs] = useState(originInputs);
const { name, remain_quota, expired_time, unlimited_quota } = inputs;
@ -107,25 +106,21 @@ const EditToken = () => {
required={!isEdit}
/>
</Form.Field>
{
isAdminUser && <>
<Form.Field>
<Form.Input
label='额度'
name='remain_quota'
placeholder={'请输入额度'}
onChange={handleInputChange}
value={remain_quota}
autoComplete='new-password'
type='number'
disabled={unlimited_quota}
/>
</Form.Field>
<Button type={'button'} style={{marginBottom: '14px'}} onClick={() => {
setUnlimitedQuota();
}}>{unlimited_quota ? '取消无限额度' : '设置为无限额度'}</Button>
</>
}
<Form.Field>
<Form.Input
label='额度'
name='remain_quota'
placeholder={'请输入额度'}
onChange={handleInputChange}
value={remain_quota}
autoComplete='new-password'
type='number'
disabled={unlimited_quota}
/>
</Form.Field>
<Button type={'button'} style={{ marginBottom: '14px' }} onClick={() => {
setUnlimitedQuota();
}}>{unlimited_quota ? '取消无限额度' : '设置为无限额度'}</Button>
<Form.Field>
<Form.Input
label='过期时间'

94
web/src/pages/TopUp/index.js Normal file
View File

@ -0,0 +1,94 @@
import React, { useEffect, useState } from 'react';
import { Button, Form, Grid, Header, Segment, Statistic } from 'semantic-ui-react';
import { API, showError, showSuccess } from '../../helpers';
const TopUp = () => {
const [redemptionCode, setRedemptionCode] = useState('');
const [topUpLink, setTopUpLink] = useState('');
const [userQuota, setUserQuota] = useState(0);
const topUp = async () => {
if (redemptionCode === '') {
return;
}
const res = await API.post('/api/user/topup', {
key: redemptionCode
});
const { success, message, data } = res.data;
if (success) {
showSuccess('充值成功！');
setUserQuota((quota) => {
return quota + data;
});
setRedemptionCode('');
} else {
showError(message);
}
};
const openTopUpLink = () => {
if (!topUpLink) {
showError('超级管理员未设置充值链接！');
return;
}
window.open(topUpLink, '_blank');
};
const getUserQuota = async ()=>{
let res = await API.get(`/api/user/self`);
const {success, message, data} = res.data;
if (success) {
setUserQuota(data.quota);
} else {
showError(message);
}
}
useEffect(() => {
let status = localStorage.getItem('status');
if (status) {
status = JSON.parse(status);
if (status.top_up_link) {
setTopUpLink(status.top_up_link);
}
}
getUserQuota().then();
}, []);
return (
<Segment>
<Header as='h3'>充值额度</Header>
<Grid columns={2} stackable>
<Grid.Column>
<Form>
<Form.Input
placeholder='兑换码'
name='redemptionCode'
value={redemptionCode}
onChange={(e) => {
setRedemptionCode(e.target.value);
}}
/>
<Button color='green' onClick={openTopUpLink}>
获取兑换码
</Button>
<Button color='yellow' onClick={topUp}>
充值
</Button>
</Form>
</Grid.Column>
<Grid.Column>
<Statistic.Group widths='one'>
<Statistic>
<Statistic.Value>{userQuota}</Statistic.Value>
<Statistic.Label>剩余额度</Statistic.Label>
</Statistic>
</Statistic.Group>
</Grid.Column>
</Grid>
</Segment>
);
};
export default TopUp;