From 97b12630b0c0aaa8c7246c0e17bf730e7ad5b7f5 Mon Sep 17 00:00:00 2001 From: zhiyu1998 Date: Tue, 11 Apr 2023 13:40:35 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20=E5=8A=A0=E5=85=A5=E9=99=90?= =?UTF-8?q?=E6=B5=81=E7=AE=97=E6=B3=95=EF=BC=8C=E9=99=90=E5=88=B6=E6=90=9C?= =?UTF-8?q?=E4=B9=A6=E6=97=B6=E9=97=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/query.js | 108 +++++++++++++++++++++++++++--------------- utils/token-bucket.js | 57 ++++++++++++++++++++++ 2 files changed, 126 insertions(+), 39 deletions(-) create mode 100644 utils/token-bucket.js diff --git a/apps/query.js b/apps/query.js index bdce166..b57eca3 100644 --- a/apps/query.js +++ b/apps/query.js @@ -11,6 +11,7 @@ import { CAT_LIMIT } from "../utils/constant.js"; import { getZHelper, getBookDetail, getYiBook, getZBook } from "../utils/books.js"; // 工具类 import _ from "lodash"; +import TokenBucket from '../utils/token-bucket.js' export class query extends plugin { constructor() { @@ -56,6 +57,10 @@ export class query extends plugin { reg: "^#竹白(.*)", fnc: "zhubaiSearch", }, + { + reg: "^#测试", + fnc: "test1", + }, ], }); } @@ -303,40 +308,43 @@ export class query extends plugin { } // 集成易书、zBook - try { - const bookList = await Promise.allSettled([ - getYiBook(e, keyword), - getZBook(e, keyword), - ]); - // 压缩直链结果 - const combineRet = bookList - .filter(item => item.status === "fulfilled" && item.value && item.value.length > 0) - .flatMap(item => { - return item.value.flat(); - }); - await e.reply(await Bot.makeForwardMsg(combineRet)); - // ZHelper 特殊处理 - const zHelper = await getZHelper(e, keyword); - zHelper.unshift({ - message: "回复格式如下\n" + - "#bookid➕id➕来源\n" + - "\n" + - "示例⬇️\n" + - "#bookid 13366067 superlib \n" + - "\n" + - "注意‼️\n" + - "1⃣️数字字母之间空格\n" + - "2⃣️id就是每条介绍最前面那串短数字不是isbn号\n" + - "3⃣️注意看书籍来源，只回复#bookid ➕id 默认来源zlibrary ", - nickname: e.sender.card || e.user_id, - user_id: e.user_id, - }) - zHelper.length > 1 && + const searchBookFunc = async () => { + try { + const bookList = await Promise.allSettled([ + getYiBook(e, keyword), + getZBook(e, keyword), + ]); + // 压缩直链结果 + const combineRet = bookList + .filter(item => item.status === "fulfilled" && item.value && item.value.length > 0) + .flatMap(item => { + return item.value.flat(); + }); + await e.reply(await Bot.makeForwardMsg(combineRet)); + // ZHelper 特殊处理 + const zHelper = await getZHelper(e, keyword); + zHelper.unshift({ + message: "回复格式如下\n" + + "#bookid➕id➕来源\n" + + "\n" + + "示例⬇️\n" + + "#bookid 13366067 superlib \n" + + "\n" + + "注意‼️\n" + + "1⃣️数字字母之间空格\n" + + "2⃣️id就是每条介绍最前面那串短数字不是isbn号\n" + + "3⃣️注意看书籍来源，只回复#bookid ➕id 默认来源zlibrary ", + nickname: e.sender.card || e.user_id, + user_id: e.user_id, + }) + zHelper.length > 1 && e.reply(await Bot.makeForwardMsg(zHelper)); - } catch (err) { - logger.error(err); - e.reply("部分搜书正在施工🚧"); + } catch (err) { + logger.error(err); + e.reply("部分搜书正在施工🚧"); + } } + await this.limitUserUse(e, searchBookFunc); return true; } @@ -355,13 +363,15 @@ export class query extends plugin { id = /\d+/.exec(keyword)[0]; source = ""; } - try { - const res = await getBookDetail(e, id, source); - e.reply(await Bot.makeForwardMsg(res)); - } catch (err) { - logger.error(err); - e.reply("搜书正在施工🚧"); - } + await this.limitUserUse(e, async () => { + try { + const res = await getBookDetail(e, id, source); + e.reply(await Bot.makeForwardMsg(res)); + } catch (err) { + logger.error(err); + e.reply("搜书正在施工🚧"); + } + }) return true; } @@ -404,9 +414,29 @@ export class query extends plugin { return true; } + /** + * 限制用户调用（默认1分钟1次） + * @param e + * @param func + * @return {Promise} + */ + async limitUserUse(e, func) { + if (query.#tokenBucket.consume(e.user_id, 1)) { + await func(); + } else { + e.reply(`🙅‍${e.nickname}你已经被限流，请稍后再试！`, true); + } + } + // 删除标签 removeTag(title) { const titleRex = /<[^>]+>/g; return title.replace(titleRex, ""); } + + /** + * 令牌桶 拿来限流 + * @type {TokenBucket} + */ + static #tokenBucket = new TokenBucket(1, 1); } diff --git a/utils/token-bucket.js b/utils/token-bucket.js new file mode 100644 index 0000000..61c59b7 --- /dev/null +++ b/utils/token-bucket.js @@ -0,0 +1,57 @@ +export default class TokenBucket { + constructor(rate, capacity) { + this.rate = rate / 60; // 修改为每分钟生成的令牌数量 + this.capacity = capacity; + this.tokens = capacity; + this.tokens = new Map(); // 使用 Map 存储每个用户的令牌桶 + this.lastTime = new Date().getTime(); + } + + /** + * 消耗令牌 + * @param id 用户id + * @param count 请求次数 + * @return {boolean} + */ + consume(id, count = 1) { + const now = new Date().getTime(); + const elapsed = now - this.lastTime; + const addedTokens = elapsed * (this.rate / 1000); + this.lastTime = now; + + // 获取用户的令牌桶，如果不存在则创建一个新的令牌桶 + let userTokens = this.tokens.get(id); + if (!userTokens) { + userTokens = { tokens: this.capacity, lastTime: now }; + this.tokens.set(id, userTokens); + } + + // 更新用户的令牌桶中的令牌数量 + userTokens.tokens = Math.min( + userTokens.tokens + addedTokens, + this.capacity + ); + userTokens.lastTime = now; + + // 判断是否有足够的令牌 + if (count <= userTokens.tokens) { + userTokens.tokens -= count; + return true; + } else { + return false; + } + } + + /** + * 重置令牌 + * @param newCapacity + */ + resetCapacity(newCapacity) { + if (newCapacity >= this.tokens) { + this.capacity = newCapacity; + this.tokens = newCapacity; + } else { + throw new Error('New capacity cannot be less than current tokens'); + } + } +} \ No newline at end of file