feat: able to set the token's expiration time and number of uses

2025-11-04 13:23:42 +08:00 · 2023-04-24 20:52:40 +08:00
parent c0bb2338de
commit 918ba60802
7 changed files with 228 additions and 42 deletions
--- a/README.md
+++ b/README.md
@@ -52,12 +52,13 @@ _✨ All in one 的 OpenAI 接口，整合各种 API 访问方式，开箱即用
   + [x] 自定义渠道
 2. 支持通过负载均衡的方式访问多个渠道。
 3. 支持单个访问渠道设置多个 API Key，利用起来你的多个 API Key。
-4. 支持 HTTP SSE。
-5. 多种用户登录注册方式：
+4. 支持设置令牌的过期时间和使用次数。
+5. 支持 HTTP SSE。
+6. 多种用户登录注册方式：
   + 邮箱登录注册以及通过邮箱进行密码重置。
   + [GitHub 开放授权](https://github.com/settings/applications/new)。
   + 微信公众号授权（需要额外部署 [WeChat Server](https://github.com/songquanpeng/wechat-server)）。
-6. 支持用户管理。
+7. 支持用户管理。

 ## 部署
 ### 基于 Docker 进行部署
--- a/common/constants.go
+++ b/common/constants.go
@@ -87,8 +87,10 @@ const (
 )

 const (
-	TokenStatusEnabled  = 1 // don't use 0, 0 is the default value!
-	TokenStatusDisabled = 2 // also don't use 0
+	TokenStatusEnabled   = 1 // don't use 0, 0 is the default value!
+	TokenStatusDisabled  = 2 // also don't use 0
+	TokenStatusExpired   = 3
+	TokenStatusExhausted = 4
 )

 const (
--- a/controller/token.go
+++ b/controller/token.go
@@ -98,6 +98,8 @@ func AddToken(c *gin.Context) {
 		Key:          common.GetUUID(),
 		CreatedTime:  common.GetTimestamp(),
 		AccessedTime: common.GetTimestamp(),
+		ExpiredTime:  token.ExpiredTime,
+		RemainTimes:  token.RemainTimes,
 	}
 	err = cleanToken.Insert()
 	if err != nil {
@@ -151,8 +153,27 @@ func UpdateToken(c *gin.Context) {
 		})
 		return
 	}
+	if token.Status == common.TokenStatusEnabled {
+		if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "令牌已过期，无法启用，请先修改令牌过期时间",
+			})
+			return
+		}
+		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes == 0 {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数",
+			})
+			return
+		}
+	}
+
 	cleanToken.Name = token.Name
 	cleanToken.Status = token.Status
+	cleanToken.ExpiredTime = token.ExpiredTime
+	cleanToken.RemainTimes = token.RemainTimes
 	err = cleanToken.Update()
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
--- a/model/token.go
+++ b/model/token.go
@@ -15,6 +15,8 @@ type Token struct {
 	Name         string `json:"name" gorm:"index" `
 	CreatedTime  int64  `json:"created_time" gorm:"bigint"`
 	AccessedTime int64  `json:"accessed_time" gorm:"bigint"`
+	ExpiredTime  int64  `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
+	RemainTimes  int    `json:"remain_times" gorm:"default:-1"`        // -1 means infinite times
 }

 func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@@ -38,13 +40,27 @@ func ValidateUserToken(key string) (token *Token, err error) {
 	err = DB.Where("key = ?", key).First(token).Error
 	if err == nil {
 		if token.Status != common.TokenStatusEnabled {
-			return nil, errors.New("该 token 已被禁用")
+			return nil, errors.New("该 token 状态不可用")
+		}
+		if token.ExpiredTime != -1 && token.ExpiredTime < common.GetTimestamp() {
+			token.Status = common.TokenStatusExpired
+			err := token.SelectUpdate()
+			if err != nil {
+				common.SysError("更新 token 状态失败：" + err.Error())
+			}
+			return nil, errors.New("该 token 已过期")
 		}
 		go func() {
 			token.AccessedTime = common.GetTimestamp()
-			err := token.Update()
+			if token.RemainTimes > 0 {
+				token.RemainTimes--
+				if token.RemainTimes == 0 {
+					token.Status = common.TokenStatusExhausted
+				}
+			}
+			err := token.SelectUpdate()
 			if err != nil {
-				common.SysError("更新 token 访问时间失败：" + err.Error())
+				common.SysError("更新 token 失败：" + err.Error())
 			}
 		}()
 		return token, nil
@@ -74,6 +90,11 @@ func (token *Token) Update() error {
 	return err
 }

+func (token *Token) SelectUpdate() error {
+	// This can update zero values
+	return DB.Model(token).Select("accessed_time", "remain_times", "status").Updates(token).Error
+}
+
 func (token *Token) Delete() error {
 	var err error
 	err = DB.Delete(token).Error
--- a/web/src/components/TokensTable.js
+++ b/web/src/components/TokensTable.js
@@ -1,5 +1,5 @@
 import React, { useEffect, useState } from 'react';
-import { Button, Form, Label, Pagination, Table } from 'semantic-ui-react';
+import { Button, Form, Label, Message, Pagination, Table } from 'semantic-ui-react';
 import { Link } from 'react-router-dom';
 import { API, copy, showError, showSuccess, timestamp2string } from '../helpers';

@@ -13,6 +13,21 @@ function renderTimestamp(timestamp) {
  );
 }

+function renderStatus(status) {
+  switch (status) {
+    case 1:
+      return <Label basic color='green'>已启用</Label>;
+    case 2:
+      return <Label basic color='red'> 已禁用 </Label>;
+    case 3:
+      return <Label basic color='yellow'> 已过期 </Label>;
+    case 4:
+      return <Label basic color='grey'> 已耗尽 </Label>;
+    default:
+      return <Label basic color='black'> 未知状态 </Label>;
+  }
+}
+
 const TokensTable = () => {
  const [tokens, setTokens] = useState([]);
  const [loading, setLoading] = useState(true);
@@ -88,25 +103,6 @@ const TokensTable = () => {
    }
  };

-  const renderStatus = (status) => {
-    switch (status) {
-      case 1:
-        return <Label basic color='green'>已启用</Label>;
-      case 2:
-        return (
-          <Label basic color='red'>
-            已禁用
-          </Label>
-        );
-      default:
-        return (
-          <Label basic color='grey'>
-            未知状态
-          </Label>
-        );
-    }
-  };
-
  const searchTokens = async () => {
    if (searchKeyword === '') {
      // if keyword is blank, load files instead.
@@ -185,6 +181,14 @@ const TokensTable = () => {
            >
              状态
            </Table.HeaderCell>
+            <Table.HeaderCell
+              style={{ cursor: 'pointer' }}
+              onClick={() => {
+                sortToken('remain_times');
+              }}
+            >
+              剩余次数
+            </Table.HeaderCell>
            <Table.HeaderCell
              style={{ cursor: 'pointer' }}
              onClick={() => {
@@ -201,6 +205,14 @@ const TokensTable = () => {
            >
              访问时间
            </Table.HeaderCell>
+            <Table.HeaderCell
+              style={{ cursor: 'pointer' }}
+              onClick={() => {
+                sortToken('expired_time');
+              }}
+            >
+              过期时间
+            </Table.HeaderCell>
            <Table.HeaderCell>操作</Table.HeaderCell>
          </Table.Row>
        </Table.Header>
@@ -218,8 +230,10 @@ const TokensTable = () => {
                  <Table.Cell>{token.id}</Table.Cell>
                  <Table.Cell>{token.name ? token.name : '无'}</Table.Cell>
                  <Table.Cell>{renderStatus(token.status)}</Table.Cell>
+                  <Table.Cell>{token.remain_times === -1 ? "无限制" : token.remain_times}</Table.Cell>
                  <Table.Cell>{renderTimestamp(token.created_time)}</Table.Cell>
                  <Table.Cell>{renderTimestamp(token.accessed_time)}</Table.Cell>
+                  <Table.Cell>{token.expired_time === -1 ? "永不过期" : renderTimestamp(token.expired_time)}</Table.Cell>
                  <Table.Cell>
                    <div>
                      <Button
@@ -272,7 +286,7 @@ const TokensTable = () => {

        <Table.Footer>
          <Table.Row>
-            <Table.HeaderCell colSpan='6'>
+            <Table.HeaderCell colSpan='8'>
              <Button size='small' as={Link} to='/token/add' loading={loading}>
                添加新的令牌
              </Button>
--- a/web/src/pages/Token/AddToken.js
+++ b/web/src/pages/Token/AddToken.js
@@ -1,21 +1,48 @@
 import React, { useState } from 'react';
 import { Button, Form, Header, Segment } from 'semantic-ui-react';
-import { API, showError, showSuccess } from '../../helpers';
+import { API, showError, showSuccess, timestamp2string } from '../../helpers';

 const AddToken = () => {
  const originInputs = {
    name: '',
+    remain_times: -1,
+    expired_time: -1
  };
  const [inputs, setInputs] = useState(originInputs);
-  const { name, display_name, password } = inputs;
+  const { name, remain_times, expired_time } = inputs;

  const handleInputChange = (e, { name, value }) => {
    setInputs((inputs) => ({ ...inputs, [name]: value }));
  };

+  const setExpiredTime = (month, day, hour, minute) => {
+    let now = new Date();
+    let timestamp = now.getTime() / 1000;
+    let seconds = month * 30 * 24 * 60 * 60;
+    seconds += day * 24 * 60 * 60;
+    seconds += hour * 60 * 60;
+    seconds += minute * 60;
+    if (seconds !== 0) {
+      timestamp += seconds;
+      setInputs({ ...inputs, expired_time: timestamp2string(timestamp) });
+    } else {
+      setInputs({ ...inputs, expired_time: -1 });
+    }
+  };
+
  const submit = async () => {
    if (inputs.name === '') return;
-    const res = await API.post(`/api/token/`, inputs);
+    let localInputs = inputs;
+    localInputs.remain_times = parseInt(localInputs.remain_times);
+    if (localInputs.expired_time !== -1) {
+      let time = Date.parse(localInputs.expired_time);
+      if (isNaN(time)) {
+        showError('过期时间格式错误！');
+        return;
+      }
+      localInputs.expired_time = Math.ceil(time / 1000);
+    }
+    const res = await API.post(`/api/token/`, localInputs);
    const { success, message } = res.data;
    if (success) {
      showSuccess('令牌创建成功！');
@@ -28,19 +55,54 @@ const AddToken = () => {
  return (
    <>
      <Segment>
-        <Header as="h3">创建新的令牌</Header>
-        <Form autoComplete="off">
+        <Header as='h3'>创建新的令牌</Header>
+        <Form autoComplete='off'>
          <Form.Field>
            <Form.Input
-              label="名称"
-              name="name"
+              label='名称'
+              name='name'
              placeholder={'请输入名称'}
              onChange={handleInputChange}
              value={name}
-              autoComplete="off"
+              autoComplete='off'
              required
            />
          </Form.Field>
+          <Form.Field>
+            <Form.Input
+              label='剩余次数'
+              name='remain_times'
+              placeholder={'请输入剩余次数，-1 表示无限制'}
+              onChange={handleInputChange}
+              value={remain_times}
+              autoComplete='off'
+            />
+          </Form.Field>
+          <Form.Field>
+            <Form.Input
+              label='过期时间'
+              name='expired_time'
+              placeholder={'请输入过期时间，格式为 yyyy-MM-dd HH:mm:ss，-1 表示无限制'}
+              onChange={handleInputChange}
+              value={expired_time}
+              autoComplete='off'
+            />
+          </Form.Field>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 0, 0, 0);
+          }}>永不过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(1, 0, 0, 0);
+          }}>一个月后过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 1, 0, 0);
+          }}>一天后过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 0, 1, 0);
+          }}>一小时后过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 0, 0, 1);
+          }}>一分钟后过期</Button>
          <Button type={'submit'} onClick={submit}>
            提交
          </Button>
--- a/web/src/pages/Token/EditToken.js
+++ b/web/src/pages/Token/EditToken.js
@@ -1,25 +1,45 @@
 import React, { useEffect, useState } from 'react';
 import { Button, Form, Header, Segment } from 'semantic-ui-react';
 import { useParams } from 'react-router-dom';
-import { API, showError, showSuccess } from '../../helpers';
+import { API, showError, showSuccess, timestamp2string } from '../../helpers';

 const EditToken = () => {
  const params = useParams();
  const tokenId = params.id;
  const [loading, setLoading] = useState(true);
  const [inputs, setInputs] = useState({
-    name: ''
+    name: '',
+    remain_times: -1,
+    expired_time: -1
  });
-  const { name } = inputs;
+  const { name, remain_times, expired_time } = inputs;
+
  const handleInputChange = (e, { name, value }) => {
    setInputs((inputs) => ({ ...inputs, [name]: value }));
  };

+  const setExpiredTime = (month, day, hour, minute) => {
+    let now = new Date();
+    let timestamp = now.getTime() / 1000;
+    let seconds = month * 30 * 24 * 60 * 60;
+    seconds += day * 24 * 60 * 60;
+    seconds += hour * 60 * 60;
+    seconds += minute * 60;
+    if (seconds !== 0) {
+      timestamp += seconds;
+      setInputs({ ...inputs, expired_time: timestamp2string(timestamp) });
+    } else {
+      setInputs({ ...inputs, expired_time: -1 });
+    }
+  };
+
  const loadToken = async () => {
    let res = await API.get(`/api/token/${tokenId}`);
    const { success, message, data } = res.data;
    if (success) {
-      data.password = '';
+      if (data.expired_time !== -1) {
+        data.expired_time = timestamp2string(data.expired_time);
+      }
      setInputs(data);
    } else {
      showError(message);
@@ -31,7 +51,17 @@ const EditToken = () => {
  }, []);

  const submit = async () => {
-    let res = await API.put(`/api/token/`, { ...inputs, id: parseInt(tokenId) });
+    let localInputs = inputs;
+    localInputs.remain_times = parseInt(localInputs.remain_times);
+    if (localInputs.expired_time !== -1) {
+      let time = Date.parse(localInputs.expired_time);
+      if (isNaN(time)) {
+        showError('过期时间格式错误！');
+        return;
+      }
+      localInputs.expired_time = Math.ceil(time / 1000);
+    }
+    let res = await API.put(`/api/token/`, { ...localInputs, id: parseInt(tokenId) });
    const { success, message } = res.data;
    if (success) {
      showSuccess('令牌更新成功！');
@@ -55,6 +85,41 @@ const EditToken = () => {
              autoComplete='off'
            />
          </Form.Field>
+          <Form.Field>
+            <Form.Input
+              label='剩余次数'
+              name='remain_times'
+              placeholder={'请输入剩余次数，-1 表示无限制'}
+              onChange={handleInputChange}
+              value={remain_times}
+              autoComplete='off'
+            />
+          </Form.Field>
+          <Form.Field>
+            <Form.Input
+              label='过期时间'
+              name='expired_time'
+              placeholder={'请输入过期时间，格式为 yyyy-MM-dd HH:mm:ss，-1 表示无限制'}
+              onChange={handleInputChange}
+              value={expired_time}
+              autoComplete='off'
+            />
+          </Form.Field>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 0, 0, 0);
+          }}>永不过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(1, 0, 0, 0);
+          }}>一个月后过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 1, 0, 0);
+          }}>一天后过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 0, 1, 0);
+          }}>一小时后过期</Button>
+          <Button type={'button'} onClick={() => {
+            setExpiredTime(0, 0, 0, 1);
+          }}>一分钟后过期</Button>
          <Button onClick={submit}>提交</Button>
        </Form>
      </Segment>