Diuse/zhenxun_bot
1
0
Fork 0
mirror of https://github.com/zhenxun-org/zhenxun_bot.git synced 2026-10-15 01:00:00 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

🎉🍾♻️refactor(core): 重构核心鉴权与缓存机制，引入异步发送队列以提升性能 🚀 (#2088)

Browse Source 
* 添加测试插件

* ✨ feat(auth): 添加缓存机制以优化用户和插件数据查询性能

* 🗑️ chore(help): 删除笨蛋检测插件代码

* ✨ feat(bot): 添加对扩展插件的加载支持

* ✨ feat(auth): 优化权限检查和缓存机制，增加用户和插件数据的并行查询

* ✨ feat(cache): 引入运行时缓存机制，优化用户和插件的ban记录管理

* ✨ feat(auth): 更新is_ban函数文档，添加参数和返回值说明

* ```
feat(auth): 使用内存缓存优化权限验证性能

- 移除数据库查询超时控制，改用 LevelUserMemoryCache、BotMemoryCache
  和 PluginLimitMemoryCache 进行缓存查询
- 优化 auth_admin、auth_bot、auth_limit 权限验证逻辑，提升响应速度
- 添加 background 参数支持异步发送权限不足提示消息
- 移除 asyncio 依赖，简化代码结构

fix(auth): 修复限制通知频率控制问题

- 实现限制通知冷却机制，避免重复发送相同限制消息
- 添加 AUTH_LIMIT_NOTICE_CD 配置项，默认值为 2 秒
- 使用 FreqLimiter 控制限制通知发送频率

refactor(models): 增强模型数据变更时的缓存同步

- 在 BotConsole、GroupConsole、LevelUser、PluginLimit 模型的
  create、update_or_create、save、delete 方法中自动更新对应缓存
- 确保数据库和内存缓存数据一致性

docs(ban_console): 修正文档注释并优化日志信息

- 修正 BanConsole 类中方法的文档字符串，使用标准参数和返回值格式
- 优化调试日志信息，使描述更加清晰准确
```

* ✨ feat(auth): 更新Limit类以支持PluginLimitSnapshot，优化限制信息处理

* ✨ feat(bot_manage): 优化Bot控制台初始化逻辑，处理IntegrityError异常

* ✨ feat(mmm1): 新增消息推送功能，支持私聊和群聊事件处理

* ✨ feat(group_member_update): 优化群组成员更新逻辑，增加活动跟踪和消息记录功能

* ✨ feat(mmm1): 删除冗余的消息推送功能代码

* ✨ feat(auth): 优化权限检查逻辑，增加模块阻止功能和缓存处理

* ✨ feat(auth): 优化权限检查逻辑，增加快速ban检测和前置检查功能

* ✨ feat(chat_history): 增强消息处理规则，添加时间间隔限制以防止重复消息
✨ feat(data_source): 引入异步获取群成员信息的功能，优化用户信息更新逻辑

* ✨ feat(send_queue): 添加异步发送队列以优化API调用和速率限制

* ✨ feat(auth): 添加缓存就绪检查以优化权限处理逻辑

* ✨ feat(plugins): 移除不必要的插件加载以简化插件管理

* ✨ feat(group_console): 优化群组获取逻辑，添加缓存检查以提升性能
✨ 只接收缓存完成之后时间的消息

* ✨ feat(auth): 添加异步任务管理和超载检测，优化权限处理逻辑
✨ feat(chat_history): 修改规则函数为异步，提升消息处理效率
✨ feat(group_handle): 增加安全获取群组信息的异步方法，添加超时处理
✨ feat(record_request): 引入安全获取群组信息的异步方法，优化群邀请处理
✨ feat(ban_memory_cache): 增强禁言内存缓存，添加负缓存机制
✨ feat(message_load): 新增消息负载检测功能，优化任务调度
✨ feat(scheduler): 在调度器中集成消息压力检测，优化任务执行
✨ feat(send_queue): 引入异步任务管理，优化发送队列处理

* ✨ feat(auth): 添加对 LevelUserSnapshot 和 BotSnapshot 的支持，优化权限检查逻辑
✨ 格式化

* ✨ feat(db_context): 增强 get_or_create 方法，处理并发创建冲突并回退查询已存在记录

* ✨ feat(bot_manage): 增强 init_bot_console 方法，处理并发创建冲突并回退查询已存在的 bot 数据

* 🚨 auto fix by pre-commit hooks

* ✨ feat(runtime_cache): 优化消息处理逻辑，支持 bytes 和 bytearray 类型的联合判断

* ✨ style(runtime_cache): 格式化代码，优化多行表达式的可读性

* 🚨 auto fix by pre-commit hooks

* ✨ refactor: 优化类型注解，改进群组成员更新逻辑和平台处理
✨ 格式化

* 🚨 auto fix by pre-commit hooks

* ✨ refactor: 优化代码格式，增强可读性并修复类型注解

* 🚨 auto fix by pre-commit hooks

* ✨ refactor: 更新文档注释，增强is_ban函数的可读性

* ✨ refactor: 优化代码结构，移除冗余函数，增强可读性并改进任务调度逻辑

* ✨ refactor: 调整定时任务时间，优化渲染服务的初始化逻辑，增强代码可读性

---------

Co-authored-by: ATTomatoo <1126160939@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
ManyManyTomato
2026-01-28 10:15:24 +08:00
committed by GitHub
co-authored by ATTomatoo pre-commit-ci[bot]
parent c9f0a8b9d9
commit 837330e30a
41 changed files with 3656 additions and 763 deletions
Download Patch File Download Diff File Expand all files Collapse all files
zhenxun/builtin_plugins/hooks/auth/auth_admin.py
+18 -31
View File
@@ -1,4 +1,3 @@
import asyncio
import time
from nonebot_plugin_alconna import At
@@ -6,8 +5,7 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.models.level_user import LevelUser
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.data_access import DataAccess
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.cache.runtime_cache import LevelUserMemoryCache, LevelUserSnapshot
from zhenxun.services.log import logger
from zhenxun.utils.utils import get_entity_ids
@@ -16,7 +14,14 @@ from .exception import SkipPluginException
from .utils import send_message
async def auth_admin(plugin: PluginInfo, session: Uninfo):
async def auth_admin(
plugin: PluginInfo,
session: Uninfo,
cached_levels: tuple[
LevelUser | LevelUserSnapshot | None, LevelUser | LevelUserSnapshot | None
]
| None = None,
):
"""管理员命令 个人权限
参数:
@@ -30,37 +35,17 @@ async def auth_admin(plugin: PluginInfo, session: Uninfo):
try:
entity = get_entity_ids(session)
level_dao = DataAccess(LevelUser)
# 并行查询用户权限数据
global_user: LevelUser | None = None
group_users: LevelUser | None = None
global_user: LevelUser | LevelUserSnapshot | None = None
group_users: LevelUser | LevelUserSnapshot | None = None
# 查询全局权限
global_user_task = level_dao.safe_get_or_none(
user_id=session.user.id, group_id__isnull=True
)
# 如果在群组中，查询群组权限
group_users_task = None
if entity.group_id:
group_users_task = level_dao.safe_get_or_none(
user_id=session.user.id, group_id=entity.group_id
if cached_levels is not None:
global_user, group_users = cached_levels
else:
global_user, group_users = await LevelUserMemoryCache.get_levels(
session.user.id, entity.group_id
)
# 等待查询完成，添加超时控制
try:
results = await asyncio.wait_for(
asyncio.gather(global_user_task, group_users_task or asyncio.sleep(0)),
timeout=DB_TIMEOUT_SECONDS,
)
global_user = results[0]
group_users = results[1] if group_users_task else None
except asyncio.TimeoutError:
logger.error(f"查询用户权限超时: user_id={session.user.id}", LOGGER_COMMAND)
# 超时时不阻塞，继续执行
return
user_level = global_user.user_level if global_user else 0
if entity.group_id and group_users:
user_level = max(user_level, group_users.user_level)
@@ -73,6 +58,7 @@ async def auth_admin(plugin: PluginInfo, session: Uninfo):
f"你的权限不足喔，该功能需要的权限等级: {plugin.admin_level}",
],
entity.user_id,
background=True,
)
raise SkipPluginException(
@@ -83,6 +69,7 @@ async def auth_admin(plugin: PluginInfo, session: Uninfo):
await send_message(
session,
f"你的权限不足喔，该功能需要的权限等级: {plugin.admin_level}",
background=True,
)
raise SkipPluginException(
zhenxun/builtin_plugins/hooks/auth/auth_ban.py
+76 -71
View File
@@ -9,7 +9,8 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.configs.config import Config
from zhenxun.models.ban_console import BanConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.data_access import DataAccess
from zhenxun.services.cache.cache_containers import CacheDict
from zhenxun.services.cache.runtime_cache import BanMemoryCache
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.log import logger
from zhenxun.utils.enum import PluginType
@@ -25,6 +26,76 @@ Config.add_plugin_config(
"才不会给你发消息.",
help="对被ban用户发送的消息",
)
Config.add_plugin_config(
"hook",
"BAN_CACHE_TTL",
2,
help="ban cache ttl seconds",
)
Config.add_plugin_config(
"hook",
"BAN_CACHE_TTL_POSITIVE",
30,
help="ban cache ttl seconds for banned users",
)
Config.add_plugin_config(
"hook",
"BAN_CACHE_TTL_NEGATIVE",
5,
help="ban cache ttl seconds for non-banned users",
)
def _coerce_ttl(value, default):
try:
value_int = int(value)
except (TypeError, ValueError):
return default
return value_int if value_int >= 0 else default
_ban_cache_ttl_value = Config.get_config("hook", "BAN_CACHE_TTL", 2)
try:
_ban_cache_ttl_value = int(_ban_cache_ttl_value)
except (TypeError, ValueError):
_ban_cache_ttl_value = 2
_ban_cache_ttl_positive = _coerce_ttl(
Config.get_config("hook", "BAN_CACHE_TTL_POSITIVE", _ban_cache_ttl_value),
_ban_cache_ttl_value,
)
_ban_cache_ttl_negative = _coerce_ttl(
Config.get_config("hook", "BAN_CACHE_TTL_NEGATIVE", _ban_cache_ttl_value),
_ban_cache_ttl_value,
)
BAN_CACHE = (
CacheDict("AUTH_BAN_CACHE", expire=0)
if max(_ban_cache_ttl_positive, _ban_cache_ttl_negative) > 0
else None
)
def _ban_cache_key(user_id: str | None, group_id: str | None) -> str:
return f"{user_id or ''}:{group_id or ''}"
def _ban_cache_get(key: str) -> int | None:
if not BAN_CACHE:
return None
try:
return BAN_CACHE[key]
except KeyError:
return None
def _ban_cache_set(key: str, value: int) -> None:
if not BAN_CACHE:
return
ttl = _ban_cache_ttl_positive if value else _ban_cache_ttl_negative
if ttl <= 0:
return
BAN_CACHE.set(key, value, expire=ttl)
async def calculate_ban_time(ban_record: BanConsole | None) -> int:
@@ -57,79 +128,13 @@ async def is_ban(user_id: str | None, group_id: str | None) -> int:
group_id: 群组ID
返回:
int: ban的剩余时间，0表示未被ban
int: ban剩余时长，-1时为永久ban，0表示未被ban
"""
if not user_id and not group_id:
return 0
start_time = time.time()
ban_dao = DataAccess(BanConsole)
# 分别获取用户在群组中的ban记录和全局ban记录
group_user = None
user = None
try:
# 并行查询用户和群组的 ban 记录
tasks = []
if user_id and group_id:
tasks.append(ban_dao.safe_get_or_none(user_id=user_id, group_id=group_id))
if user_id:
tasks.append(
ban_dao.safe_get_or_none(user_id=user_id, group_id__isnull=True)
)
# 等待所有查询完成，添加超时控制
if tasks:
try:
ban_records = await asyncio.wait_for(
asyncio.gather(*tasks), timeout=DB_TIMEOUT_SECONDS
)
if len(tasks) == 2:
group_user, user = ban_records
elif user_id and group_id:
group_user = ban_records[0]
else:
user = ban_records[0]
except asyncio.TimeoutError:
logger.error(
f"查询ban记录超时: user_id={user_id}, group_id={group_id}",
LOGGER_COMMAND,
)
return 0
# 检查记录并计算ban时间
results = []
if group_user:
results.append(group_user)
if user:
results.append(user)
# 如果没有找到记录，返回0
if not results:
return 0
logger.debug(f"查询到的ban记录: {results}", LOGGER_COMMAND)
# 检查所有记录，找出最严格的ban（时间最长的）
max_ban_time: int = 0
for result in results:
if result.duration > 0 or result.duration == -1:
# 直接计算ban时间，避免再次查询数据库
ban_time = await calculate_ban_time(result)
if ban_time == -1 or ban_time > max_ban_time:
max_ban_time = ban_time
return max_ban_time
finally:
# 记录执行时间
elapsed = time.time() - start_time
if elapsed > WARNING_THRESHOLD: # 记录耗时超过500ms的检查
logger.warning(
f"is_ban 耗时: {elapsed:.3f}s",
LOGGER_COMMAND,
session=user_id,
group_id=group_id,
)
if not BanMemoryCache.is_loaded():
return 0
return BanMemoryCache.remaining_time(user_id, group_id)
def check_plugin_type(matcher: Matcher) -> bool:
zhenxun/builtin_plugins/hooks/auth/auth_bot.py
+10 -15
View File
@@ -1,10 +1,8 @@
import asyncio
import time
from zhenxun.models.bot_console import BotConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.data_access import DataAccess
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.cache.runtime_cache import BotMemoryCache, BotSnapshot
from zhenxun.services.log import logger
from zhenxun.utils.common_utils import CommonUtils
@@ -12,7 +10,12 @@ from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .exception import SkipPluginException
async def auth_bot(plugin: PluginInfo, bot_id: str):
async def auth_bot(
plugin: PluginInfo,
bot_id: str,
bot_data: BotConsole | BotSnapshot | None = None,
skip_fetch: bool = False,
):
"""bot层面的权限检查
参数:
@@ -26,17 +29,9 @@ async def auth_bot(plugin: PluginInfo, bot_id: str):
start_time = time.time()
try:
# 从数据库或缓存中获取 bot 信息
bot_dao = DataAccess(BotConsole)
try:
bot: BotConsole | None = await asyncio.wait_for(
bot_dao.safe_get_or_none(bot_id=bot_id), timeout=DB_TIMEOUT_SECONDS
)
except asyncio.TimeoutError:
logger.error(f"查询Bot信息超时: bot_id={bot_id}", LOGGER_COMMAND)
# 超时时不阻塞，继续执行
return
bot: BotConsole | BotSnapshot | None = bot_data
if bot is None and not skip_fetch:
bot = await BotMemoryCache.get(bot_id)
if not bot or not bot.status:
raise SkipPluginException("Bot不存在或休眠中阻断权限检测...")
zhenxun/builtin_plugins/hooks/auth/auth_cost.py
+8 -3
View File
@@ -10,12 +10,16 @@ from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .exception import SkipPluginException
from .utils import send_message
DEFAULT_GOLD = 100
async def auth_cost(user: UserConsole, plugin: PluginInfo, session: Uninfo) -> int:
async def auth_cost(
user: UserConsole | None, plugin: PluginInfo, session: Uninfo
) -> int:
"""检测是否满足金币条件
参数:
user: UserConsole
user: UserConsole | None
plugin: PluginInfo
session: Uninfo
@@ -25,7 +29,8 @@ async def auth_cost(user: UserConsole, plugin: PluginInfo, session: Uninfo) -> i
start_time = time.time()
try:
if user.gold < plugin.cost_gold:
user_gold = user.gold if user else DEFAULT_GOLD
if user_gold < plugin.cost_gold:
"""插件消耗金币不足"""
await send_message(session, f"金币不足..该功能需要{plugin.cost_gold}金币..")
raise SkipPluginException(f"{plugin.name}({plugin.module}) 金币限制...")
zhenxun/builtin_plugins/hooks/auth/auth_group.py
+4 -5
View File
@@ -1,9 +1,8 @@
import time
from nonebot_plugin_alconna import UniMsg
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.cache.runtime_cache import GroupSnapshot
from zhenxun.services.log import logger
from .config import LOGGER_COMMAND, WARNING_THRESHOLD, SwitchEnum
@@ -12,8 +11,8 @@ from .exception import SkipPluginException
async def auth_group(
plugin: PluginInfo,
group: GroupConsole | None,
message: UniMsg,
group: GroupConsole | GroupSnapshot | None,
text: str | None,
group_id: str | None,
):
"""群黑名单检测 群总开关检测
@@ -29,7 +28,7 @@ async def auth_group(
start_time = time.time()
try:
text = message.extract_plain_text()
text = text or ""
if not group:
raise SkipPluginException("群组信息不存在...")
zhenxun/builtin_plugins/hooks/auth/auth_limit.py
+71 -44
View File
@@ -1,13 +1,18 @@
import asyncio
import time
from typing import ClassVar
from typing import Any, ClassVar
import nonebot
from nonebot_plugin_uninfo import Uninfo
from pydantic import BaseModel
from zhenxun.configs.config import Config
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.plugin_limit import PluginLimit
from zhenxun.services.cache.runtime_cache import (
PluginLimitMemoryCache,
PluginLimitSnapshot,
)
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.log import logger
from zhenxun.utils.enum import LimitWatchType, PluginLimitType
@@ -22,6 +27,16 @@ from .exception import SkipPluginException
driver = nonebot.get_driver()
Config.add_plugin_config(
"hook",
"AUTH_LIMIT_NOTICE_CD",
2,
help="auth limit notice cooldown seconds",
)
_LIMIT_NOTICE_CD = int(Config.get_config("hook", "AUTH_LIMIT_NOTICE_CD", 2) or 2)
_LIMIT_NOTICE_LIMITER = FreqLimiter(_LIMIT_NOTICE_CD)
_LIMIT_NOTICE_TASKS: set[asyncio.Task] = set()
@PriorityLifecycle.on_startup(priority=5)
async def _():
@@ -30,13 +45,41 @@ async def _():
class Limit(BaseModel):
limit: PluginLimit
limit: PluginLimit | PluginLimitSnapshot
limiter: FreqLimiter | UserBlockLimiter | CountLimiter
class Config:
arbitrary_types_allowed = True
def _limit_notice_key(
limit: PluginLimit | PluginLimitSnapshot,
user_id: str,
group_id: str | None,
channel_id: str | None,
) -> str:
key = user_id
if group_id and limit.watch_type == LimitWatchType.GROUP:
key = channel_id or group_id
return f"{limit.module}:{limit.limit_type}:{key}"
def _send_limit_notice(message: str, format_kwargs: dict[str, Any], key: str) -> None:
if not _LIMIT_NOTICE_LIMITER.check(key):
return
_LIMIT_NOTICE_LIMITER.start_cd(key)
async def _send():
try:
await MessageUtils.build_message(message, format_args=format_kwargs).send()
except Exception as exc:
logger.error("limit notice send failed", LOGGER_COMMAND, e=exc)
task = asyncio.create_task(_send())
_LIMIT_NOTICE_TASKS.add(task)
task.add_done_callback(_LIMIT_NOTICE_TASKS.discard)
class LimitManager:
add_module: ClassVar[list] = []
last_update_time: ClassVar[float] = 0
@@ -48,8 +91,11 @@ class LimitManager:
count_limit: ClassVar[dict[str, Limit]] = {}
# 模块限制缓存，避免频繁查询数据库
module_limit_cache: ClassVar[dict[str, tuple[float, list[PluginLimit]]]] = {}
module_limit_cache: ClassVar[
dict[str, tuple[float, list[PluginLimitSnapshot], bool]]
] = {}
module_cache_ttl: ClassVar[float] = 60 # 模块缓存有效期（秒）
module_cache_error_ttl: ClassVar[float] = 5 # 超时缓存有效期（秒）
@classmethod
async def init_limit(cls):
@@ -70,14 +116,8 @@ class LimitManager:
cls.is_updating = True
try:
start_time = time.time()
try:
limit_list = await asyncio.wait_for(
PluginLimit.filter(status=True).all(), timeout=DB_TIMEOUT_SECONDS
)
except asyncio.TimeoutError:
logger.error("查询限制信息超时", LOGGER_COMMAND)
cls.is_updating = False
return
await PluginLimitMemoryCache.ensure_loaded()
limit_list = PluginLimitMemoryCache.get_all_limits()
# 清空旧数据
cls.add_module = []
@@ -96,7 +136,7 @@ class LimitManager:
cls.is_updating = False
@classmethod
def add_limit(cls, limit: PluginLimit):
def add_limit(cls, limit: PluginLimit | PluginLimitSnapshot):
"""添加限制
参数:
@@ -109,12 +149,16 @@ class LimitManager:
limit=limit, limiter=UserBlockLimiter()
)
elif limit.limit_type == PluginLimitType.CD:
cd_value = int(limit.cd or 0)
cls.cd_limit[limit.module] = Limit(
limit=limit, limiter=FreqLimiter(limit.cd)
limit=limit, limiter=FreqLimiter(cd_value)
)
elif limit.limit_type == PluginLimitType.COUNT:
max_count = int(limit.max_count or 0)
if max_count <= 0:
return
cls.count_limit[limit.module] = Limit(
limit=limit, limiter=CountLimiter(limit.max_count)
limit=limit, limiter=CountLimiter(max_count)
)
@classmethod
@@ -144,7 +188,7 @@ class LimitManager:
limiter.set_false(key_type)
@classmethod
async def get_module_limits(cls, module: str) -> list[PluginLimit]:
async def get_module_limits(cls, module: str) -> list[PluginLimitSnapshot]:
"""获取模块的限制信息，使用缓存减少数据库查询
参数:
@@ -157,30 +201,20 @@ class LimitManager:
# 检查缓存
if module in cls.module_limit_cache:
cache_time, limits = cls.module_limit_cache[module]
if current_time - cache_time < cls.module_cache_ttl:
cache_time, limits, is_error = cls.module_limit_cache[module]
ttl = cls.module_cache_error_ttl if is_error else cls.module_cache_ttl
if current_time - cache_time < ttl:
return limits
# 缓存不存在或已过期，从数据库查询
# 缓存不存在或已过期，从内存缓存获取
try:
start_time = time.time()
limits = await asyncio.wait_for(
PluginLimit.filter(module=module, status=True).all(),
timeout=DB_TIMEOUT_SECONDS,
)
elapsed = time.time() - start_time
if elapsed > WARNING_THRESHOLD: # 记录耗时超过500ms的查询
logger.warning(
f"查询模块限制信息耗时: {elapsed:.3f}s, 模块: {module}",
LOGGER_COMMAND,
)
# 更新缓存
cls.module_limit_cache[module] = (current_time, limits)
await PluginLimitMemoryCache.ensure_loaded()
limits = await PluginLimitMemoryCache.get_limits(module)
cls.module_limit_cache[module] = (current_time, limits, False)
return limits
except asyncio.TimeoutError:
logger.error(f"查询模块限制信息超时: {module}", LOGGER_COMMAND)
# 超时时返回空列表，避免阻塞
except Exception as exc:
logger.error(f"get module limits failed: {module}", LOGGER_COMMAND, e=exc)
cls.module_limit_cache[module] = (current_time, [], True)
return []
@classmethod
@@ -275,15 +309,8 @@ class LimitManager:
left_time = limiter.left_time(key_type)
cd_str = TimeUtils.format_duration(left_time)
format_kwargs = {"cd": cd_str}
try:
await asyncio.wait_for(
MessageUtils.build_message(
limit.result, format_args=format_kwargs
).send(),
timeout=DB_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
logger.error(f"发送限制消息超时: {limit.module}", LOGGER_COMMAND)
notice_key = _limit_notice_key(limit, user_id, group_id, channel_id)
_send_limit_notice(limit.result, format_kwargs, notice_key)
raise SkipPluginException(
f"{limit.module}({limit.limit_type}) 正在限制中..."
)
zhenxun/builtin_plugins/hooks/auth/auth_plugin.py
+98 -79
View File
@@ -1,4 +1,3 @@
import asyncio
import time
from nonebot.adapters import Event
@@ -6,9 +5,8 @@ from nonebot_plugin_uninfo import Uninfo
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.cache.runtime_cache import GroupSnapshot, _parse_block_modules
from zhenxun.services.log import logger
from zhenxun.utils.common_utils import CommonUtils
from zhenxun.utils.enum import BlockType
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
@@ -16,74 +14,89 @@ from .exception import IsSuperuserException, SkipPluginException
from .utils import freq, is_poke, send_message
def _get_group_block_sets(
group: GroupConsole | GroupSnapshot,
) -> tuple[frozenset[str], frozenset[str]]:
block_set = getattr(group, "block_plugin_set", None)
super_block_set = getattr(group, "superuser_block_plugin_set", None)
if block_set is None:
block_set = _parse_block_modules(getattr(group, "block_plugin", "") or "")
setattr(group, "block_plugin_set", block_set)
if super_block_set is None:
super_block_set = _parse_block_modules(
getattr(group, "superuser_block_plugin", "") or ""
)
setattr(group, "superuser_block_plugin_set", super_block_set)
return block_set, super_block_set
class GroupCheck:
def __init__(
self, plugin: PluginInfo, group: GroupConsole, session: Uninfo, is_poke: bool
self,
plugin: PluginInfo,
group: GroupConsole | GroupSnapshot,
session: Uninfo,
is_poke: bool,
skip_group_block: bool,
) -> None:
self.session = session
self.is_poke = is_poke
self.plugin = plugin
self.group_data = group
self.group_id = group.group_id
self.skip_group_block = skip_group_block
(
self.block_plugin_set,
self.superuser_block_plugin_set,
) = _get_group_block_sets(group)
async def check(self):
start_time = time.time()
try:
# 检查超级用户禁用
if (
self.group_data
and CommonUtils.format(self.plugin.module)
in self.group_data.superuser_block_plugin
):
if freq.is_send_limit_message(self.plugin, self.group_id, self.is_poke):
try:
await asyncio.wait_for(
send_message(
self.session,
"超级管理员禁用了该群此功能...",
self.group_id,
),
timeout=DB_TIMEOUT_SECONDS,
if not self.skip_group_block:
# 检查超级用户禁用
if (
self.group_data
and self.plugin.module in self.superuser_block_plugin_set
):
if freq.is_send_limit_message(
self.plugin, self.group_id, self.is_poke
):
await send_message(
self.session,
"超级管理员禁用了该群此功能...",
self.group_id,
background=True,
)
except asyncio.TimeoutError:
logger.error(f"发送消息超时: {self.group_id}", LOGGER_COMMAND)
raise SkipPluginException(
f"{self.plugin.name}({self.plugin.module})"
f" 超级管理员禁用了该群此功能..."
)
raise SkipPluginException(
f"{self.plugin.name}({self.plugin.module})"
f" 超级管理员禁用了该群此功能..."
)
# 检查普通禁用
if (
self.group_data
and CommonUtils.format(self.plugin.module)
in self.group_data.block_plugin
):
if freq.is_send_limit_message(self.plugin, self.group_id, self.is_poke):
try:
await asyncio.wait_for(
send_message(
self.session, "该群未开启此功能...", self.group_id
),
timeout=DB_TIMEOUT_SECONDS,
# 检查普通禁用
if self.group_data and self.plugin.module in self.block_plugin_set:
if freq.is_send_limit_message(
self.plugin, self.group_id, self.is_poke
):
await send_message(
self.session,
"该群未开启此功能...",
self.group_id,
background=True,
)
except asyncio.TimeoutError:
logger.error(f"发送消息超时: {self.group_id}", LOGGER_COMMAND)
raise SkipPluginException(
f"{self.plugin.name}({self.plugin.module}) 未开启此功能..."
)
raise SkipPluginException(
f"{self.plugin.name}({self.plugin.module}) 未开启此功能..."
)
# 检查全局禁用
if self.plugin.block_type == BlockType.GROUP:
if freq.is_send_limit_message(self.plugin, self.group_id, self.is_poke):
try:
await asyncio.wait_for(
send_message(
self.session, "该功能在群组中已被禁用...", self.group_id
),
timeout=DB_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
logger.error(f"发送消息超时: {self.group_id}", LOGGER_COMMAND)
await send_message(
self.session,
"该功能在群组中已被禁用...",
self.group_id,
background=True,
)
raise SkipPluginException(
f"{self.plugin.name}({self.plugin.module})该插件在群组中已被禁用..."
)
@@ -98,7 +111,9 @@ class GroupCheck:
class PluginCheck:
def __init__(self, group: GroupConsole | None, session: Uninfo, is_poke: bool):
def __init__(
self, group: GroupConsole | GroupSnapshot | None, session: Uninfo, is_poke: bool
):
self.session = session
self.is_poke = is_poke
self.group_data = group
@@ -117,13 +132,11 @@ class PluginCheck:
"""
if plugin.block_type == BlockType.PRIVATE:
if freq.is_send_limit_message(plugin, self.session.user.id, self.is_poke):
try:
await asyncio.wait_for(
send_message(self.session, "该功能在私聊中已被禁用..."),
timeout=DB_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
logger.error("发送消息超时", LOGGER_COMMAND)
await send_message(
self.session,
"该功能在私聊中已被禁用...",
background=True,
)
raise SkipPluginException(
f"{plugin.name}({plugin.module}) 该插件在私聊中已被禁用..."
)
@@ -147,13 +160,12 @@ class PluginCheck:
sid = self.group_id or self.session.user.id
if freq.is_send_limit_message(plugin, sid, self.is_poke):
try:
await asyncio.wait_for(
send_message(self.session, "全局未开启此功能...", sid),
timeout=DB_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
logger.error(f"发送消息超时: {sid}", LOGGER_COMMAND)
await send_message(
self.session,
"全局未开启此功能...",
sid,
background=True,
)
raise SkipPluginException(
f"{plugin.name}({plugin.module}) 全局未开启此功能..."
)
@@ -167,7 +179,12 @@ class PluginCheck:
async def auth_plugin(
plugin: PluginInfo, group: GroupConsole | None, session: Uninfo, event: Event
plugin: PluginInfo,
group: GroupConsole | GroupSnapshot | None,
session: Uninfo,
event: Event,
*,
skip_group_block: bool = False,
):
"""插件状态
@@ -181,19 +198,21 @@ async def auth_plugin(
is_poke_event = is_poke(event)
user_check = PluginCheck(group, session, is_poke_event)
tasks = []
if group:
tasks.append(GroupCheck(plugin, group, session, is_poke_event).check())
block_set, super_block_set = _get_group_block_sets(group)
if (
plugin.status
and plugin.block_type != BlockType.GROUP
and not block_set
and not super_block_set
):
return
await GroupCheck(
plugin, group, session, is_poke_event, skip_group_block
).check()
else:
tasks.append(user_check.check_user(plugin))
tasks.append(user_check.check_global(plugin))
try:
await asyncio.wait_for(
asyncio.gather(*tasks), timeout=DB_TIMEOUT_SECONDS * 2
)
except asyncio.TimeoutError:
logger.error("插件用户/群组/全局检查超时...", LOGGER_COMMAND)
await user_check.check_user(plugin)
await user_check.check_global(plugin)
finally:
# 记录总执行时间
zhenxun/builtin_plugins/hooks/auth/utils.py
+28 -14
View File
@@ -1,3 +1,4 @@
import asyncio
import contextlib
from nonebot.adapters import Event
@@ -13,6 +14,7 @@ from zhenxun.utils.utils import FreqLimiter
from .config import LOGGER_COMMAND
base_config = Config.get("hook")
_SEND_TASKS: set[asyncio.Task] = set()
def is_poke(event: Event) -> bool:
@@ -32,7 +34,10 @@ def is_poke(event: Event) -> bool:
async def send_message(
session: Uninfo, message: list | str, check_tag: str | None = None
session: Uninfo,
message: list | str,
check_tag: str | None = None,
background: bool = False,
):
"""发送消息
@@ -41,19 +46,28 @@ async def send_message(
message: 消息
check_tag: cd flag
"""
try:
if not check_tag:
await MessageUtils.build_message(message).send(reply_to=True)
elif freq._flmt.check(check_tag):
freq._flmt.start_cd(check_tag)
await MessageUtils.build_message(message).send(reply_to=True)
except Exception as e:
logger.error(
"发送消息失败",
LOGGER_COMMAND,
session=session,
e=e,
)
async def _send():
try:
if not check_tag:
await MessageUtils.build_message(message).send(reply_to=True)
elif freq._flmt.check(check_tag):
freq._flmt.start_cd(check_tag)
await MessageUtils.build_message(message).send(reply_to=True)
except Exception as e:
logger.error(
"发送消息失败",
LOGGER_COMMAND,
session=session,
e=e,
)
if background:
task = asyncio.create_task(_send())
_SEND_TASKS.add(task)
task.add_done_callback(_SEND_TASKS.discard)
return
await _send()
class FreqUtils:
zhenxun/builtin_plugins/hooks/auth_checker.py
+721 -102
View File
@@ -1,25 +1,39 @@
import asyncio
import contextlib
import time
from typing import cast
from nonebot import get_loaded_plugins
from nonebot.adapters import Bot, Event
from nonebot.exception import IgnoredException
from nonebot.matcher import Matcher
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo
from tortoise.exceptions import IntegrityError
from zhenxun.models.group_console import GroupConsole
from zhenxun.configs.config import Config
from zhenxun.configs.utils import PluginExtraData
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.user_console import UserConsole
from zhenxun.services.cache.cache_containers import CacheDict
from zhenxun.services.cache.runtime_cache import (
BotMemoryCache,
BotSnapshot,
GroupMemoryCache,
GroupSnapshot,
LevelUserMemoryCache,
LevelUserSnapshot,
PluginInfoMemoryCache,
)
from zhenxun.services.data_access import DataAccess
from zhenxun.services.log import logger
from zhenxun.utils.enum import GoldHandle, PluginType
from zhenxun.services.message_load import is_overloaded
from zhenxun.utils.enum import BlockType, GoldHandle, PluginType
from zhenxun.utils.exception import InsufficientGold
from zhenxun.utils.platform import PlatformUtils
from zhenxun.utils.utils import get_entity_ids
from .auth.auth_admin import auth_admin
from .auth.auth_ban import auth_ban
from .auth.auth_ban import auth_ban, is_ban
from .auth.auth_bot import auth_bot
from .auth.auth_cost import auth_cost
from .auth.auth_group import auth_group
@@ -34,6 +48,54 @@ from .auth.exception import (
)
from .auth.utils import base_config
Config.add_plugin_config(
"hook",
"AUTH_HOOKS_CONCURRENCY_LIMIT",
6,
help="auth hooks concurrency limit",
)
Config.add_plugin_config(
"hook",
"AUTH_DB_CONCURRENCY_LIMIT",
6,
help="auth db concurrency limit",
)
Config.add_plugin_config(
"hook",
"AUTH_PLUGIN_CACHE_TTL",
30,
help="plugin info cache ttl seconds",
)
Config.add_plugin_config(
"hook",
"AUTH_USER_CACHE_TTL",
5,
help="user cache ttl seconds",
)
Config.add_plugin_config(
"hook",
"AUTH_EVENT_CACHE_TTL",
2,
help="event auth cache ttl seconds",
)
def _coerce_positive_int(value, default):
try:
value_int = int(value)
except (TypeError, ValueError):
return default
return value_int if value_int > 0 else default
def _coerce_cache_ttl(value, default):
try:
value_int = int(value)
except (TypeError, ValueError):
return default
return value_int if value_int >= 0 else default
# 超时设置（秒）
TIMEOUT_SECONDS = 5.0
# 熔断计数器
@@ -51,13 +113,302 @@ CIRCUIT_RESET_TIME = 300 # 5分钟
# 并发控制：限制同时进入 hooks 并行检查的协程数
# 默认为 6，可通过环境变量 AUTH_HOOKS_CONCURRENCY_LIMIT 调整
HOOKS_CONCURRENCY_LIMIT = base_config.get("AUTH_HOOKS_CONCURRENCY_LIMIT")
HOOKS_CONCURRENCY_LIMIT = _coerce_positive_int(
base_config.get("AUTH_HOOKS_CONCURRENCY_LIMIT", 6), 6
)
DB_CONCURRENCY_LIMIT = _coerce_positive_int(
base_config.get("AUTH_DB_CONCURRENCY_LIMIT", HOOKS_CONCURRENCY_LIMIT),
HOOKS_CONCURRENCY_LIMIT,
)
PLUGIN_CACHE_TTL = _coerce_cache_ttl(base_config.get("AUTH_PLUGIN_CACHE_TTL", 30), 30)
USER_CACHE_TTL = _coerce_cache_ttl(base_config.get("AUTH_USER_CACHE_TTL", 5), 5)
PLUGIN_CACHE = (
CacheDict("AUTH_PLUGIN_CACHE", expire=PLUGIN_CACHE_TTL)
if PLUGIN_CACHE_TTL > 0
else None
)
USER_CACHE = (
CacheDict("AUTH_USER_CACHE", expire=USER_CACHE_TTL) if USER_CACHE_TTL > 0 else None
)
EVENT_CACHE_TTL = _coerce_cache_ttl(base_config.get("AUTH_EVENT_CACHE_TTL", 2), 2)
EVENT_CACHE = (
CacheDict("AUTH_EVENT_CACHE", expire=EVENT_CACHE_TTL)
if EVENT_CACHE_TTL > 0
else None
)
# 路由索引缓存
_ROUTE_INDEX_LOCK = asyncio.Lock()
_ROUTE_INDEX_READY = False
_ROUTE_COMMAND_MAP: dict[str, set[str]] = {}
_ROUTE_PREFIX_MAP: dict[str, set[str]] = {}
_ROUTE_MODULES_WITH_COMMANDS: set[str] = set()
# 全局信号量与计数器
HOOKS_SEMAPHORE = asyncio.Semaphore(HOOKS_CONCURRENCY_LIMIT)
HOOKS_ACTIVE_COUNT = 0
HOOKS_ACTIVE_LOCK = asyncio.Lock()
DB_SEMAPHORE = asyncio.Semaphore(DB_CONCURRENCY_LIMIT)
DB_ACTIVE_COUNT = 0
DB_ACTIVE_LOCK = asyncio.Lock()
def _cache_get(cache: CacheDict | None, key: str):
if not cache:
return None
try:
return cache[key]
except KeyError:
return None
def _cache_set(cache: CacheDict | None, key: str, value):
if cache:
cache[key] = value
def _debug_log(message: str, *args, **kwargs) -> None:
if is_overloaded():
return
logger.debug(message, *args, **kwargs)
def _event_cache_key(event: Event, session: Uninfo, entity) -> str:
msg_id = getattr(event, "message_id", None)
if msg_id is None:
msg_id = getattr(event, "id", None)
if msg_id is None:
msg_id = id(event)
platform = PlatformUtils.get_platform(session)
group_id = entity.group_id or ""
channel_id = entity.channel_id or ""
return (
f"{platform}:{session.self_id}:{entity.user_id}:"
f"{group_id}:{channel_id}:{msg_id}"
)
def _get_event_cache(event: Event, session: Uninfo, entity):
if not EVENT_CACHE:
return None
key = _event_cache_key(event, session, entity)
try:
return EVENT_CACHE[key]
except KeyError:
cache = {}
EVENT_CACHE[key] = cache
return cache
def _normalize_command(command: str) -> str:
return command.strip()
def _extract_commands(extra: PluginExtraData | None) -> set[str]:
if not extra:
return set()
commands = {c.command for c in extra.commands if c.command}
commands.update(extra.aliases or set())
return {cmd.strip() for cmd in commands if cmd and cmd.strip()}
async def _ensure_route_index():
global _ROUTE_INDEX_READY
if _ROUTE_INDEX_READY:
return
async with _ROUTE_INDEX_LOCK:
if _ROUTE_INDEX_READY:
return
_ROUTE_COMMAND_MAP.clear()
_ROUTE_PREFIX_MAP.clear()
_ROUTE_MODULES_WITH_COMMANDS.clear()
for plugin in get_loaded_plugins():
if not plugin.metadata:
continue
extra = plugin.metadata.extra or {}
try:
extra_data = PluginExtraData(**extra)
except Exception:
continue
command_set = _extract_commands(extra_data)
if not command_set:
continue
module = plugin.name
_ROUTE_MODULES_WITH_COMMANDS.add(module)
for command in command_set:
normalized = _normalize_command(command)
if not normalized:
continue
_ROUTE_COMMAND_MAP.setdefault(normalized, set()).add(module)
_ROUTE_PREFIX_MAP.setdefault(normalized[0], set()).add(normalized)
_ROUTE_INDEX_READY = True
def _command_matches(text: str, command: str) -> bool:
if not text or not command:
return False
if text == command:
return True
if text.startswith(command):
if len(text) == len(command):
return True
next_char = text[len(command)]
return next_char.isspace()
return False
def _match_route_modules(text: str) -> set[str]:
text = text.strip()
if not text:
return set()
commands = _ROUTE_PREFIX_MAP.get(text[0])
if not commands:
return set()
matched_modules: set[str] = set()
for command in commands:
if _command_matches(text, command):
modules = _ROUTE_COMMAND_MAP.get(command)
if modules:
matched_modules.update(modules)
return matched_modules
def _get_message_text(message: UniMsg, event_cache: dict | None) -> str:
if event_cache is None:
return message.extract_plain_text()
cached = event_cache.get("plain_text")
if cached is None:
cached = message.extract_plain_text()
event_cache["plain_text"] = cached
return cached
async def _get_route_context(text: str, event_cache: dict | None) -> set[str]:
if not text:
return set()
if event_cache is not None and "route_modules" in event_cache:
return event_cache["route_modules"]
await _ensure_route_index()
matched = _match_route_modules(text)
if event_cache is not None:
event_cache["route_modules"] = matched
return matched
async def _has_limits_cached(module: str, event_cache: dict | None) -> bool:
module_limit_cache: dict[str, bool] = {}
if event_cache is not None:
module_limit_cache = event_cache.setdefault("module_limits", {})
if module in module_limit_cache:
return module_limit_cache[module]
limits = await LimitManager.get_module_limits(module)
has_limits = bool(limits)
module_limit_cache[module] = has_limits
return has_limits
@contextlib.asynccontextmanager
async def _db_section():
global DB_ACTIVE_COUNT
await DB_SEMAPHORE.acquire()
async with DB_ACTIVE_LOCK:
DB_ACTIVE_COUNT += 1
_debug_log(f"current db auth concurrency: {DB_ACTIVE_COUNT}", LOGGER_COMMAND)
try:
yield
finally:
with contextlib.suppress(Exception):
DB_SEMAPHORE.release()
async with DB_ACTIVE_LOCK:
DB_ACTIVE_COUNT = max(DB_ACTIVE_COUNT - 1, 0)
_debug_log(
f"current db auth concurrency: {DB_ACTIVE_COUNT}", LOGGER_COMMAND
)
async def _get_group_cached(entity, event_cache) -> GroupSnapshot | None:
if not entity.group_id:
return None
if event_cache is not None and "group" in event_cache:
return event_cache["group"]
group = GroupMemoryCache.get_if_ready(entity.group_id, entity.channel_id)
if event_cache is not None:
event_cache["group"] = group
return group
def _module_in_block_string(module: str, value: str | None) -> bool:
if not value:
return False
return f"<{module}," in value
def _group_has_plugin_block(group, module: str) -> bool:
if not group:
return False
block_set = getattr(group, "block_plugin_set", None)
super_block_set = getattr(group, "superuser_block_plugin_set", None)
if block_set is not None or super_block_set is not None:
if block_set and module in block_set:
return True
if super_block_set and module in super_block_set:
return True
return False
block_plugin = getattr(group, "block_plugin", "") or ""
super_block_plugin = getattr(group, "superuser_block_plugin", "") or ""
return _module_in_block_string(module, block_plugin) or _module_in_block_string(
module, super_block_plugin
)
def _needs_auth_plugin(plugin: PluginInfo, group, entity) -> bool:
if plugin.block_type == BlockType.ALL and not plugin.status:
if group and getattr(group, "is_super", False):
return False
return True
if entity.group_id:
if plugin.block_type == BlockType.GROUP:
return True
return _group_has_plugin_block(group, plugin.module)
return plugin.block_type == BlockType.PRIVATE
def _needs_admin_check(plugin: PluginInfo) -> bool:
if plugin.admin_level and plugin.admin_level > 0:
return True
return plugin.plugin_type in {
PluginType.ADMIN,
PluginType.SUPERUSER,
PluginType.SUPER_AND_ADMIN,
}
async def _get_bot_data_cached(
bot_id: str, event_cache
) -> tuple[BotSnapshot | None, bool]:
if event_cache is not None and "bot_data" in event_cache:
return event_cache.get("bot_data"), event_cache.get("bot_timeout", False)
bot = await BotMemoryCache.get(bot_id)
if event_cache is not None:
event_cache["bot_data"] = bot
event_cache["bot_timeout"] = False
return bot, False
async def _get_admin_levels_cached(
session: Uninfo, entity, event_cache
) -> tuple[tuple[LevelUserSnapshot | None, LevelUserSnapshot | None] | None, bool]:
if event_cache is not None and "admin_levels" in event_cache:
return event_cache.get("admin_levels"), event_cache.get("admin_timeout", False)
levels = await LevelUserMemoryCache.get_levels(session.user.id, entity.group_id)
if event_cache is not None:
event_cache["admin_levels"] = levels
event_cache["admin_timeout"] = False
return levels, False
# 超时装饰器
async def with_timeout(coro, timeout=TIMEOUT_SECONDS, name=None):
@@ -120,76 +471,73 @@ def check_circuit_breaker(name):
return CIRCUIT_BREAKERS[name]["active"]
async def get_plugin_and_user(
module: str, user_id: str
) -> tuple[PluginInfo, UserConsole]:
"""获取用户数据和插件信息
def _is_hidden_plugin(matcher: Matcher) -> bool:
plugin = matcher.plugin
if not plugin or not plugin.metadata:
return False
extra = plugin.metadata.extra or {}
return extra.get("plugin_type") == PluginType.HIDDEN
参数:
module: 模块名
user_id: 用户id
异常:
PermissionExemption: 插件数据不存在
PermissionExemption: 插件类型为HIDDEN
PermissionExemption: 重复创建用户
PermissionExemption: 用户数据不存在
返回:
tuple[PluginInfo, UserConsole]: 插件信息，用户信息
"""
user_dao = DataAccess(UserConsole)
plugin_dao = DataAccess(PluginInfo)
# 并行查询插件和用户数据
plugin_task = plugin_dao.safe_get_or_none(module=module)
user_task = user_dao.get_by_func_or_none(
UserConsole.get_user, False, user_id=user_id
async def _fetch_user_readonly(
user_dao: DataAccess, user_id: str
) -> UserConsole | None:
return await with_timeout(
user_dao.safe_get_or_none(user_id=user_id), name="get_user"
)
try:
plugin, user = await with_timeout(
asyncio.gather(plugin_task, user_task), name="get_plugin_and_user"
)
except asyncio.TimeoutError:
# 如果并行查询超时，尝试串行查询
logger.warning("并行查询超时，尝试串行查询", LOGGER_COMMAND)
plugin = await with_timeout(
plugin_dao.safe_get_or_none(module=module), name="get_plugin"
)
user = await with_timeout(
user_dao.safe_get_or_none(user_id=user_id), name="get_user"
)
except IntegrityError:
await asyncio.sleep(0.5)
plugin_task = plugin_dao.safe_get_or_none(module=module)
user_task = user_dao.get_by_func_or_none(
UserConsole.get_user, False, user_id=user_id
)
plugin, user = await with_timeout(
asyncio.gather(plugin_task, user_task), name="get_plugin_and_user"
)
async def _fetch_plugin(plugin_dao: DataAccess, module: str) -> PluginInfo | None:
return await with_timeout(
plugin_dao.safe_get_or_none(module=module), name="get_plugin"
)
async def get_plugin_and_user(
module: str,
user_id: str,
platform: str | None = None,
event_cache: dict | None = None,
need_user: bool = True,
) -> tuple[PluginInfo, UserConsole | None]:
"""Fetch plugin info and read user only when cost is required."""
user_dao = DataAccess(UserConsole)
plugin = None
if event_cache is not None:
plugin_cache = event_cache.setdefault("plugin_cache", {})
if module in plugin_cache:
plugin = plugin_cache[module]
if plugin is None:
plugin = await PluginInfoMemoryCache.get_by_module(module)
if event_cache is not None:
event_cache.setdefault("plugin_cache", {})[module] = plugin
plugin = cast(PluginInfo | None, plugin)
if not plugin:
raise PermissionExemption(f"插件:{module} 数据不存在，已跳过权限检查...")
raise PermissionExemption(f"plugin:{module} not found, skip permission check")
if plugin.plugin_type == PluginType.HIDDEN:
raise PermissionExemption(
f"插件: {plugin.name}:{plugin.module} 为HIDDEN，已跳过权限检查..."
)
raise PermissionExemption(f"plugin {plugin.name}:{plugin.module} hidden, skip")
user = None
try:
user = await user_dao.get_by_func_or_none(
UserConsole.get_user, False, user_id=user_id
)
except IntegrityError as e:
raise PermissionExemption("重复创建用户，已跳过该次权限检查...") from e
if not user:
raise PermissionExemption("用户数据不存在，已跳过权限检查...")
if need_user and plugin.cost_gold > 0:
if event_cache is not None:
user_cache = event_cache.setdefault("user_cache", {})
if user_id in user_cache:
user = user_cache[user_id]
else:
async with _db_section():
user = await _fetch_user_readonly(user_dao, user_id)
user_cache[user_id] = user
else:
async with _db_section():
user = await _fetch_user_readonly(user_dao, user_id)
return plugin, user
async def get_plugin_cost(
bot: Bot, user: UserConsole, plugin: PluginInfo, session: Uninfo
bot: Bot, user: UserConsole | None, plugin: PluginInfo, session: Uninfo
) -> int:
"""获取插件费用
@@ -278,7 +626,7 @@ async def _enter_hooks_section():
await HOOKS_SEMAPHORE.acquire()
async with HOOKS_ACTIVE_LOCK:
HOOKS_ACTIVE_COUNT += 1
logger.debug(f"当前并发权限检查数量: {HOOKS_ACTIVE_COUNT}", LOGGER_COMMAND)
_debug_log(f"当前并发权限检查数量: {HOOKS_ACTIVE_COUNT}", LOGGER_COMMAND)
async def _leave_hooks_section():
@@ -292,7 +640,85 @@ async def _leave_hooks_section():
HOOKS_ACTIVE_COUNT -= 1
# 保证计数不为负
HOOKS_ACTIVE_COUNT = max(HOOKS_ACTIVE_COUNT, 0)
logger.debug(f"当前并发权限检查数量: {HOOKS_ACTIVE_COUNT}", LOGGER_COMMAND)
_debug_log(f"当前并发权限检查数量: {HOOKS_ACTIVE_COUNT}", LOGGER_COMMAND)
async def auth_ban_fast(
matcher: Matcher, event: Event, bot: Bot, session: Uninfo
) -> None:
"""快速 ban 检测（仅使用内存缓存），用于前置快速裁决。"""
entity = get_entity_ids(session)
event_cache = _get_event_cache(event, session, entity)
if event_cache is not None and event_cache.get("ban_state") is True:
raise SkipPluginException("user or group banned (cached)")
if entity.user_id in bot.config.superusers:
if event_cache is not None:
event_cache["ban_state"] = False
return
if entity.group_id and await is_ban(None, entity.group_id):
if event_cache is not None:
event_cache["ban_state"] = True
raise SkipPluginException("group banned (fast)")
if entity.user_id and await is_ban(entity.user_id, entity.group_id):
if event_cache is not None:
event_cache["ban_state"] = True
raise SkipPluginException("user banned (fast)")
if event_cache is not None:
event_cache["ban_state"] = False
async def route_precheck(
matcher: Matcher,
event: Event,
session: Uninfo,
message: UniMsg,
) -> bool:
module = matcher.plugin_name or ""
if not module:
return False
if _is_hidden_plugin(matcher):
return False
entity = get_entity_ids(session)
event_cache = _get_event_cache(event, session, entity)
text = _get_message_text(message, event_cache)
route_modules = await _get_route_context(text, event_cache)
await _ensure_route_index()
if module in _ROUTE_MODULES_WITH_COMMANDS and module not in route_modules:
if event_cache is not None:
event_cache["route_skip"] = True
return True
return False
async def auth_precheck(
matcher: Matcher,
event: Event,
bot: Bot,
session: Uninfo,
message: UniMsg,
) -> None:
"""轻量前置检查：命令路由 + 必要管理员权限。"""
module = matcher.plugin_name or ""
if not module:
return
if _is_hidden_plugin(matcher):
return
entity = get_entity_ids(session)
if session.user.id in bot.config.superusers:
return
plugin = cast(PluginInfo | None, await PluginInfoMemoryCache.get_by_module(module))
if not plugin:
return
if plugin.plugin_type == PluginType.SUPERUSER:
raise SkipPluginException("超级管理员权限不足...")
if _needs_admin_check(plugin):
await LevelUserMemoryCache.ensure_fresh()
levels = await LevelUserMemoryCache.get_levels(session.user.id, entity.group_id)
await auth_admin(plugin, session, cached_levels=levels)
async def auth(
@@ -301,6 +727,8 @@ async def auth(
bot: Bot,
session: Uninfo,
message: UniMsg,
*,
skip_ban: bool = False,
):
"""权限检查
@@ -316,6 +744,10 @@ async def auth(
ignore_flag = False
entity = get_entity_ids(session)
module = matcher.plugin_name or ""
event_cache = _get_event_cache(event, session, entity)
auth_allowed = None
auth_result_cache = None
admin_checked_pre = False
# 用于记录各个 hook 的执行时间
hook_times = {}
@@ -328,11 +760,46 @@ async def auth(
if not module:
raise PermissionExemption("Matcher插件名称不存在...")
if event_cache is not None:
auth_result_cache = event_cache.setdefault("auth_result", {})
cached_result = auth_result_cache.get(module)
if cached_result is not None:
allowed, reason = cached_result
if not allowed:
raise SkipPluginException(reason or "auth cached skip")
return
if _is_hidden_plugin(matcher):
raise PermissionExemption(f"plugin {module} hidden, skip")
if event_cache is not None and event_cache.get("ban_state") is True:
raise SkipPluginException("user or group banned (cached)")
text = _get_message_text(message, event_cache)
route_modules = await _get_route_context(text, event_cache)
await _ensure_route_index()
route_skip_checks = (
module in _ROUTE_MODULES_WITH_COMMANDS and module not in route_modules
)
if route_skip_checks:
if event_cache is not None:
event_cache["route_skip"] = True
hook_times["route"] = "miss"
auth_allowed = True
return
platform = PlatformUtils.get_platform(session)
# 获取插件和用户数据
plugin_user_start = time.time()
try:
plugin, user = await with_timeout(
get_plugin_and_user(module, entity.user_id), name="get_plugin_and_user"
get_plugin_and_user(
module,
entity.user_id,
platform,
event_cache=event_cache,
need_user=not route_skip_checks,
),
name="get_plugin_and_user",
)
hook_times["get_plugin_user"] = f"{time.time() - plugin_user_start:.3f}s"
except asyncio.TimeoutError:
@@ -343,54 +810,200 @@ async def auth(
)
raise PermissionExemption("获取插件和用户数据超时，请稍后再试...")
# 进入 hooks 并行检查区域（会在高并发时排队）
await _enter_hooks_section()
entered_hooks = True
if not route_skip_checks and _needs_admin_check(plugin):
if plugin.plugin_type in {
PluginType.SUPERUSER,
PluginType.SUPER_AND_ADMIN,
}:
if session.user.id in bot.config.superusers:
hook_times["auth_admin"] = "superuser"
admin_checked_pre = True
elif plugin.plugin_type == PluginType.SUPERUSER:
raise SkipPluginException("超级管理员权限不足...")
if not admin_checked_pre:
await LevelUserMemoryCache.ensure_fresh()
admin_levels = None
admin_timeout = False
if event_cache is not None:
admin_levels, admin_timeout = await _get_admin_levels_cached(
session, entity, event_cache
)
if admin_timeout:
hook_times["auth_admin"] = "timeout"
else:
admin_start = time.time()
await auth_admin(plugin, session, cached_levels=admin_levels)
hook_times["auth_admin"] = f"{time.time() - admin_start:.3f}s(pre)"
admin_checked_pre = True
ban_cache_state = None
if event_cache is not None:
ban_cache_state = event_cache.get("ban_state")
if skip_ban:
if ban_cache_state is True:
hook_times["auth_ban"] = "cached"
raise SkipPluginException("user or group banned (cached)")
if ban_cache_state is None:
ban_start = time.time()
try:
await auth_ban(matcher, bot, session, plugin)
hook_times["auth_ban"] = f"{time.time() - ban_start:.3f}s"
if event_cache is not None:
event_cache["ban_state"] = False
except SkipPluginException:
hook_times["auth_ban"] = f"{time.time() - ban_start:.3f}s"
if event_cache is not None:
event_cache["ban_state"] = True
raise
else:
hook_times["auth_ban"] = "skipped"
else:
if ban_cache_state is True:
hook_times["auth_ban"] = "cached"
raise SkipPluginException("user or group banned (cached)")
if ban_cache_state is None:
ban_start = time.time()
try:
await auth_ban(matcher, bot, session, plugin)
hook_times["auth_ban"] = f"{time.time() - ban_start:.3f}s"
if event_cache is not None:
event_cache["ban_state"] = False
except SkipPluginException:
hook_times["auth_ban"] = f"{time.time() - ban_start:.3f}s"
if event_cache is not None:
event_cache["ban_state"] = True
raise
else:
hook_times["auth_ban"] = "cached"
# 获取插件费用
cost_start = time.time()
try:
cost_gold = await with_timeout(
get_plugin_cost(bot, user, plugin, session), name="get_plugin_cost"
)
hook_times["cost_gold"] = f"{time.time() - cost_start:.3f}s"
except asyncio.TimeoutError:
logger.error(
f"获取插件费用超时，模块: {module}", LOGGER_COMMAND, session=session
)
# 继续执行，不阻止权限检查
if not route_skip_checks and plugin.cost_gold > 0:
cost_start = time.time()
try:
cost_gold = await with_timeout(
get_plugin_cost(bot, user, plugin, session), name="get_plugin_cost"
)
hook_times["cost_gold"] = f"{time.time() - cost_start:.3f}s"
except asyncio.TimeoutError:
logger.error(
f"获取插件费用超时，模块: {module}", LOGGER_COMMAND, session=session
)
# 继续执行，不阻止权限检查
else:
hook_times["cost_gold"] = "skipped"
# 执行 bot_filter
bot_filter(session)
group = None
if entity.group_id:
group_dao = DataAccess(GroupConsole)
group = await with_timeout(
group_dao.safe_get_or_none(
group_id=entity.group_id, channel_id__isnull=True
),
name="get_group",
group = await _get_group_cached(entity, event_cache)
bot_data = None
bot_timeout = False
if event_cache is not None:
bot_data, bot_timeout = await _get_bot_data_cached(bot.self_id, event_cache)
admin_levels = None
admin_timeout = False
if (
not admin_checked_pre
and plugin.admin_level
and event_cache is not None
and not route_skip_checks
):
admin_levels, admin_timeout = await _get_admin_levels_cached(
session, entity, event_cache
)
# 并行执行所有 hook 检查，并记录执行时间
hooks_start = time.time()
# 创建所有 hook 任务
hook_tasks = [
time_hook(auth_ban(matcher, bot, session, plugin), "auth_ban", hook_times),
time_hook(auth_bot(plugin, bot.self_id), "auth_bot", hook_times),
time_hook(
auth_group(plugin, group, message, entity.group_id),
"auth_group",
hook_times,
),
time_hook(auth_admin(plugin, session), "auth_admin", hook_times),
time_hook(
auth_plugin(plugin, group, session, event), "auth_plugin", hook_times
),
time_hook(auth_limit(plugin, session), "auth_limit", hook_times),
]
hook_tasks = []
if event_cache is None:
hook_tasks.append(
time_hook(auth_bot(plugin, bot.self_id), "auth_bot", hook_times)
)
else:
if bot_timeout:
hook_times["auth_bot"] = "timeout"
else:
hook_tasks.append(
time_hook(
auth_bot(
plugin,
bot.self_id,
bot_data=bot_data,
skip_fetch=True,
),
"auth_bot",
hook_times,
)
)
if session.user.id in bot.config.superusers:
hook_times["auth_group"] = "superuser"
else:
hook_tasks.append(
time_hook(
auth_group(plugin, group, text, entity.group_id),
"auth_group",
hook_times,
)
)
if not route_skip_checks and plugin.admin_level and not admin_checked_pre:
if event_cache is None:
hook_tasks.append(
time_hook(auth_admin(plugin, session), "auth_admin", hook_times)
)
else:
if admin_timeout:
hook_times["auth_admin"] = "timeout"
else:
hook_tasks.append(
time_hook(
auth_admin(plugin, session, cached_levels=admin_levels),
"auth_admin",
hook_times,
)
)
else:
hook_times.setdefault("auth_admin", "skipped")
if session.user.id in bot.config.superusers:
hook_times["auth_plugin"] = "superuser"
elif not route_skip_checks and _needs_auth_plugin(plugin, group, entity):
hook_tasks.append(
time_hook(
auth_plugin(
plugin,
group,
session,
event,
skip_group_block=session.user.id in bot.config.superusers,
),
"auth_plugin",
hook_times,
)
)
else:
hook_times["auth_plugin"] = "skipped"
if not route_skip_checks:
has_limits = await _has_limits_cached(module, event_cache)
if has_limits:
hook_tasks.append(
time_hook(auth_limit(plugin, session), "auth_limit", hook_times)
)
else:
hook_times["auth_limit"] = "skipped"
else:
hook_times["auth_limit"] = "skipped"
if hook_tasks:
# 进入 hooks 并行检查区域（会在高并发时排队）
await _enter_hooks_section()
entered_hooks = True
# 使用 gather 并行执行所有 hook，但添加总体超时控制
try:
@@ -408,15 +1021,19 @@ async def auth(
# 不抛出异常，允许继续执行
hooks_time = time.time() - hooks_start
auth_allowed = True
except SkipPluginException as e:
LimitManager.unblock(module, entity.user_id, entity.group_id, entity.channel_id)
logger.info(str(e), LOGGER_COMMAND, session=session)
ignore_flag = True
auth_allowed = False
except IsSuperuserException:
logger.debug("超级用户跳过权限检测...", LOGGER_COMMAND, session=session)
auth_allowed = True
except PermissionExemption as e:
logger.info(str(e), LOGGER_COMMAND, session=session)
auth_allowed = True
finally:
# 如果进入过 hooks 区域，确保释放信号量（即使上层处理抛出了异常）
if entered_hooks:
@@ -428,6 +1045,8 @@ async def auth(
LOGGER_COMMAND,
session=session,
)
if auth_result_cache is not None and auth_allowed is not None:
auth_result_cache[module] = (auth_allowed, None)
# 扣除金币
if not ignore_flag and cost_gold > 0:
gold_start = time.time()
zhenxun/builtin_plugins/hooks/auth_hook.py
+139 -14
View File
@@ -1,34 +1,159 @@
import asyncio
import time
from nonebot import get_driver
from nonebot.adapters import Bot, Event
from nonebot.exception import IgnoredException
from nonebot.matcher import Matcher
from nonebot.message import run_postprocessor, run_preprocessor
from nonebot.message import event_preprocessor, run_postprocessor, run_preprocessor
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo
from zhenxun.services.cache.runtime_cache import is_cache_ready
from zhenxun.services.log import logger
from zhenxun.services.message_load import is_overloaded, signal_overload
from zhenxun.utils.utils import get_entity_ids
from .auth.config import LOGGER_COMMAND
from .auth_checker import LimitManager, auth
from .auth.exception import SkipPluginException
from .auth_checker import (
LimitManager,
_get_event_cache,
auth,
auth_ban_fast,
auth_precheck,
route_precheck,
)
_SKIP_AUTH_PLUGINS = {"chat_history", "chat_message"}
_BOT_CONNECT_TS: float | None = None
_AUTH_QUEUE_MAXSIZE = 200
_AUTH_QUEUE_HIGH_WATER = 160
_AUTH_OVERLOAD_WINDOW = 5.0
_AUTH_QUEUE: asyncio.Queue[tuple[Matcher, Event, Bot, Uninfo, UniMsg]] = asyncio.Queue(
maxsize=_AUTH_QUEUE_MAXSIZE
)
_AUTH_QUEUE_STARTED = False
_AUTH_WORKERS: list[asyncio.Task] = []
_LAST_DROP_LOG = 0.0
driver = get_driver()
@driver.on_bot_connect
async def _mark_bot_connected(bot: Bot):
del bot
global _BOT_CONNECT_TS
_BOT_CONNECT_TS = time.time()
async def _auth_worker(worker_id: int) -> None:
while True:
matcher, event, bot, session, message = await _AUTH_QUEUE.get()
try:
await auth(
matcher,
event,
bot,
session,
message,
skip_ban=True,
)
except IgnoredException:
pass
except Exception as exc:
if not is_overloaded():
logger.error("async auth failed", LOGGER_COMMAND, e=exc)
finally:
_AUTH_QUEUE.task_done()
@driver.on_startup
async def _start_auth_queue():
global _AUTH_QUEUE_STARTED
if _AUTH_QUEUE_STARTED:
return
_AUTH_QUEUE_STARTED = True
worker_count = max(1, min(6, _AUTH_QUEUE_MAXSIZE // 50))
for idx in range(worker_count):
_AUTH_WORKERS.append(asyncio.create_task(_auth_worker(idx)))
def _skip_auth_for_plugin(matcher: Matcher) -> bool:
if not matcher.plugin:
return False
name = (matcher.plugin.name or "").lower()
if name in _SKIP_AUTH_PLUGINS:
return True
module_name = getattr(matcher.plugin, "module_name", "") or ""
return "chat_history" in module_name
@event_preprocessor
async def _drop_message_before_cache_ready(event: Event):
if event.get_type() != "message":
return
if not is_cache_ready():
raise IgnoredException("cache not ready ignore")
if _BOT_CONNECT_TS is not None:
event_ts = getattr(event, "time", None)
if event_ts is not None and event_ts < _BOT_CONNECT_TS:
raise IgnoredException("drop backlog message")
# # 权限检测
@run_preprocessor
async def _(matcher: Matcher, event: Event, bot: Bot, session: Uninfo, message: UniMsg):
async def _auth_preprocessor(
matcher: Matcher, event: Event, bot: Bot, session: Uninfo, message: UniMsg
):
if event.get_type() == "message" and not is_cache_ready():
raise IgnoredException("cache not ready ignore")
if _skip_auth_for_plugin(matcher):
return
start_time = time.time()
await auth(
matcher,
event,
bot,
session,
message,
)
logger.debug(f"权限检测耗时：{time.time() - start_time}秒", LOGGER_COMMAND)
entity = get_entity_ids(session)
event_cache = _get_event_cache(event, session, entity)
if await route_precheck(matcher, event, session, message):
return
try:
await auth_ban_fast(matcher, event, bot, session)
except SkipPluginException as exc:
logger.info(str(exc), LOGGER_COMMAND, session=session)
raise IgnoredException("ban fast ignore") from exc
try:
await auth_precheck(matcher, event, bot, session, message)
except SkipPluginException as exc:
logger.info(str(exc), LOGGER_COMMAND, session=session)
raise IgnoredException("precheck ignore") from exc
if event_cache is not None and event_cache.get("route_skip") is True:
if not is_overloaded():
logger.debug("route miss skip auth task", LOGGER_COMMAND)
return
try:
_AUTH_QUEUE.put_nowait((matcher, event, bot, session, message))
except asyncio.QueueFull:
signal_overload(_AUTH_OVERLOAD_WINDOW)
now = time.monotonic()
global _LAST_DROP_LOG
if now - _LAST_DROP_LOG > 1.0:
_LAST_DROP_LOG = now
logger.warning("auth queue full, skip auth task", LOGGER_COMMAND)
return
if _AUTH_QUEUE.qsize() >= _AUTH_QUEUE_HIGH_WATER:
signal_overload(_AUTH_OVERLOAD_WINDOW)
now = time.monotonic()
last_log = getattr(_auth_preprocessor, "_last_log", 0.0)
if now - last_log > 1.0 and not is_overloaded():
setattr(_auth_preprocessor, "_last_log", now)
logger.debug(
f"auth check cost: {time.time() - start_time:.3f}s",
LOGGER_COMMAND,
)
# 解除命令block阻塞
@run_postprocessor
async def _(matcher: Matcher, session: Uninfo):
async def _unblock_after_matcher(matcher: Matcher, session: Uninfo):
user_id = session.user.id
group_id = None
channel_id = None