97 lines
2.8 KiB
Python
97 lines
2.8 KiB
Python
"""多渠道网关 Redis 滑动窗口限流。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
import uuid
|
|
from typing import TYPE_CHECKING
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
from yuxi.channel.constants import InboundRejectionReason, channel_rate_limit_key
|
|
|
|
from .policy import SecurityCheckResult
|
|
|
|
if TYPE_CHECKING:
|
|
from .registry import SecurityContext
|
|
|
|
|
|
class RateLimiter:
|
|
"""基于 Redis 滑动窗口的限流器。
|
|
|
|
使用 Lua 脚本保证 "检查窗口 -> 计数 -> 写入" 的原子性，
|
|
避免超限请求仍被计入窗口。
|
|
"""
|
|
|
|
name = "rate_limit"
|
|
default_priority = 400
|
|
|
|
_ALLOW_SCRIPT = """
|
|
local key = KEYS[1]
|
|
local window_start = tonumber(ARGV[1])
|
|
local max_requests = tonumber(ARGV[2])
|
|
local member_score = tonumber(ARGV[3])
|
|
local member = ARGV[4]
|
|
local window_seconds = tonumber(ARGV[5])
|
|
|
|
redis.call('zremrangebyscore', key, 0, window_start)
|
|
local current = redis.call('zcard', key)
|
|
if current >= max_requests then
|
|
return 0
|
|
end
|
|
|
|
redis.call('zadd', key, member_score, member)
|
|
redis.call('expire', key, window_seconds)
|
|
return 1
|
|
"""
|
|
|
|
def __init__(self, redis: Redis):
|
|
self.redis = redis
|
|
|
|
async def check(self, ctx: SecurityContext) -> SecurityCheckResult | None:
|
|
"""SecurityChecker 协议入口，保持原 is_allowed 方法可用。"""
|
|
max_requests = self._resolve_max_requests(ctx)
|
|
actor_id = ctx.resolved_sender_id or ctx.inbound.sender_id or ctx.inbound.peer_id or ""
|
|
rate_key = channel_rate_limit_key(
|
|
ctx.inbound.channel_type,
|
|
ctx.inbound.account_id or "",
|
|
actor_id,
|
|
)
|
|
if await self.is_allowed(rate_key, max_requests, window_seconds=60):
|
|
return None
|
|
return SecurityCheckResult(
|
|
allowed=False,
|
|
reason=InboundRejectionReason.RATE_LIMITED,
|
|
)
|
|
|
|
def _resolve_max_requests(self, ctx: SecurityContext) -> int:
|
|
resolve = getattr(ctx.plugin, "resolve_rate_limit_policy", None)
|
|
if resolve is not None:
|
|
rate_policy = resolve(ctx.config, ctx.inbound.account_id)
|
|
if rate_policy is not None:
|
|
return rate_policy.max_requests_per_minute
|
|
return int(ctx.config.get("rate_limit", {}).get("max_requests_per_minute", 60))
|
|
|
|
async def is_allowed(
|
|
self,
|
|
key: str,
|
|
max_requests: int,
|
|
window_seconds: int,
|
|
) -> bool:
|
|
"""原有公开方法，继续可用。"""
|
|
now = time.time()
|
|
window_start = now - window_seconds
|
|
member = f"{now}:{uuid.uuid4().hex}"
|
|
|
|
allowed = await self.redis.eval(
|
|
self._ALLOW_SCRIPT,
|
|
1,
|
|
key,
|
|
window_start,
|
|
max_requests,
|
|
now,
|
|
member,
|
|
window_seconds,
|
|
)
|
|
return bool(allowed)
|