austnv/remnawave-minishop
1
0
Fork 0
Code Issues 1 Pull Requests Activity

feat: add Telegram anti-flood action buckets

Browse Source
This commit is contained in:
3252a8
2026-06-05 11:14:34 +03:00
parent 5e1fe59396
commit 6dc43182ab
3 changed files with 181 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/bot/middlewares/update_antiflood.py
+110 -5
View File
@@ -15,6 +15,24 @@ logger = logging.getLogger(__name__)
DEFAULT_WINDOW_SECONDS = 60
DEFAULT_MAX_UPDATES_PER_WINDOW = 180
DEFAULT_MESSAGE_MAX_PER_WINDOW = 120
DEFAULT_CALLBACK_MAX_PER_WINDOW = 240
DEFAULT_INLINE_MAX_PER_WINDOW = 60
DEFAULT_START_MAX_PER_WINDOW = 30
DEFAULT_EXPENSIVE_CALLBACK_MAX_PER_WINDOW = 60
EXPENSIVE_CALLBACK_PREFIXES = (
"pay_",
"trial_action:confirm_activate",
"main_action:request_trial",
"main_action:apply_promo",
"main_action:bot_apply_promo",
"tariff_change:apply:",
"tariff_change:confirm_pay:",
"tariff_change:pay:",
"autorenew:confirm:",
"disconnect_device:",
)
@dataclass(frozen=True)
@@ -31,6 +49,7 @@ class UpdateAntiFloodMiddleware(BaseMiddleware):
settings: Settings,
*,
default_rule: Optional[RateLimitRule] = None,
action_rules: Optional[Dict[str, RateLimitRule]] = None,
) -> None:
super().__init__()
self.settings = settings
@@ -48,6 +67,7 @@ class UpdateAntiFloodMiddleware(BaseMiddleware):
or DEFAULT_MAX_UPDATES_PER_WINDOW
),
)
self.action_rules = action_rules or _default_action_rules(settings)
self._local_buckets: Dict[str, Deque[float]] = defaultdict(deque)
self._local_lock = asyncio.Lock()
@@ -64,18 +84,23 @@ class UpdateAntiFloodMiddleware(BaseMiddleware):
if not actor_key:
return await handler(event, data)
if await self._is_limited(actor_key, self.default_rule):
action_key = _update_action_key(event)
if await self._is_limited("updates", actor_key, self.default_rule) or (
action_key
and action_key in self.action_rules
and await self._is_limited(action_key, actor_key, self.action_rules[action_key])
):
logger.warning(
"Telegram update dropped by anti-flood: actor=%s update_type=%s",
actor_key,
getattr(event, "event_type", "unknown"),
action_key or getattr(event, "event_type", "unknown"),
)
data["antiflood_dropped"] = True
return None
return await handler(event, data)
async def _is_limited(self, actor_key: str, rule: RateLimitRule) -> bool:
async def _is_limited(self, bucket_name: str, actor_key: str, rule: RateLimitRule) -> bool:
if rule.window_seconds <= 0 or rule.max_events <= 0:
return False
@@ -86,7 +111,7 @@ class UpdateAntiFloodMiddleware(BaseMiddleware):
self.settings,
"rate-limit",
"telegram",
"updates",
bucket_name,
actor_key,
)
current = int(await redis.incr(key))
@@ -96,7 +121,7 @@ class UpdateAntiFloodMiddleware(BaseMiddleware):
except Exception as exc:
logger.warning("Redis telegram anti-flood unavailable; using local fallback: %s", exc)
return await self._is_limited_local(actor_key, rule)
return await self._is_limited_local(f"{bucket_name}:{actor_key}", rule)
async def _is_limited_local(self, actor_key: str, rule: RateLimitRule) -> bool:
now = time.monotonic()
@@ -132,3 +157,83 @@ def _update_actor_key(update: Update) -> Optional[str]:
if chat_id is not None:
return f"chat:{int(chat_id)}"
return None
def _update_action_key(update: Update) -> str:
if update.message:
text = update.message.text or ""
if text.startswith("/start"):
return "start"
return "message"
if update.callback_query:
data = update.callback_query.data or ""
if data.startswith(EXPENSIVE_CALLBACK_PREFIXES):
return "expensive_callback"
return "callback"
if update.inline_query:
return "inline"
return "updates"
def _default_action_rules(settings: Settings) -> Dict[str, RateLimitRule]:
window_seconds = int(
getattr(settings, "TELEGRAM_ANTIFLOOD_WINDOW_SECONDS", DEFAULT_WINDOW_SECONDS)
or DEFAULT_WINDOW_SECONDS
)
return {
"message": RateLimitRule(
window_seconds,
int(
getattr(
settings,
"TELEGRAM_ANTIFLOOD_MESSAGE_MAX_PER_WINDOW",
DEFAULT_MESSAGE_MAX_PER_WINDOW,
)
or DEFAULT_MESSAGE_MAX_PER_WINDOW
),
),
"callback": RateLimitRule(
window_seconds,
int(
getattr(
settings,
"TELEGRAM_ANTIFLOOD_CALLBACK_MAX_PER_WINDOW",
DEFAULT_CALLBACK_MAX_PER_WINDOW,
)
or DEFAULT_CALLBACK_MAX_PER_WINDOW
),
),
"inline": RateLimitRule(
window_seconds,
int(
getattr(
settings,
"TELEGRAM_ANTIFLOOD_INLINE_MAX_PER_WINDOW",
DEFAULT_INLINE_MAX_PER_WINDOW,
)
or DEFAULT_INLINE_MAX_PER_WINDOW
),
),
"start": RateLimitRule(
window_seconds,
int(
getattr(
settings,
"TELEGRAM_ANTIFLOOD_START_MAX_PER_WINDOW",
DEFAULT_START_MAX_PER_WINDOW,
)
or DEFAULT_START_MAX_PER_WINDOW
),
),
"expensive_callback": RateLimitRule(
window_seconds,
int(
getattr(
settings,
"TELEGRAM_ANTIFLOOD_EXPENSIVE_CALLBACK_MAX_PER_WINDOW",
DEFAULT_EXPENSIVE_CALLBACK_MAX_PER_WINDOW,
)
or DEFAULT_EXPENSIVE_CALLBACK_MAX_PER_WINDOW
),
),
}
backend/config/settings.py
+5
View File
@@ -223,6 +223,11 @@ class Settings(BaseSettings):
TELEGRAM_ANTIFLOOD_ENABLED: bool = Field(default=True)
TELEGRAM_ANTIFLOOD_WINDOW_SECONDS: int = Field(default=60)
TELEGRAM_ANTIFLOOD_MAX_UPDATES_PER_WINDOW: int = Field(default=180)
TELEGRAM_ANTIFLOOD_MESSAGE_MAX_PER_WINDOW: int = Field(default=120)
TELEGRAM_ANTIFLOOD_CALLBACK_MAX_PER_WINDOW: int = Field(default=240)
TELEGRAM_ANTIFLOOD_INLINE_MAX_PER_WINDOW: int = Field(default=60)
TELEGRAM_ANTIFLOOD_START_MAX_PER_WINDOW: int = Field(default=30)
TELEGRAM_ANTIFLOOD_EXPENSIVE_CALLBACK_MAX_PER_WINDOW: int = Field(default=60)
WEBHOOK_QUEUE_NAME: str = Field(default="webhook-events")
WEBHOOK_QUEUE_CONCURRENCY: int = Field(default=4)
WORKER_PANEL_SYNC_INTERVAL_SECONDS: int = Field(default=900)
tests/test_update_antiflood_middleware.py
+66
View File
@@ -30,6 +30,35 @@ def _message_update(*, user_id=42, chat_id=42, chat_type="private", text="hello"
)
def _callback_update(
*,
user_id=42,
chat_id=42,
chat_type="private",
data="main_action:back_to_main",
):
return SimpleNamespace(
event_type="callback_query",
message=None,
callback_query=SimpleNamespace(
from_user=SimpleNamespace(id=user_id),
message=SimpleNamespace(chat=SimpleNamespace(id=chat_id, type=chat_type)),
data=data,
answer=AsyncMock(),
),
inline_query=None,
)
def _inline_update(*, user_id=42, query="ref"):
return SimpleNamespace(
event_type="inline_query",
message=None,
callback_query=None,
inline_query=SimpleNamespace(from_user=SimpleNamespace(id=user_id), query=query),
)
class UpdateAntiFloodMiddlewareTests(unittest.IsolatedAsyncioTestCase):
async def test_extreme_update_flood_is_dropped_before_handler(self):
middleware = UpdateAntiFloodMiddleware(
@@ -58,6 +87,43 @@ class UpdateAntiFloodMiddlewareTests(unittest.IsolatedAsyncioTestCase):
handler.assert_awaited_once()
async def test_action_specific_limits_are_counted_separately(self):
middleware = UpdateAntiFloodMiddleware(
_settings(),
default_rule=RateLimitRule(window_seconds=60, max_events=100),
action_rules={
"start": RateLimitRule(window_seconds=60, max_events=1),
"callback": RateLimitRule(window_seconds=60, max_events=2),
"expensive_callback": RateLimitRule(window_seconds=60, max_events=1),
"inline": RateLimitRule(window_seconds=60, max_events=1),
},
)
handler = AsyncMock(return_value="ok")
with patch("bot.middlewares.update_antiflood.get_redis", AsyncMock(return_value=None)):
self.assertEqual(
await middleware(handler, _message_update(text="/start"), {}),
"ok",
)
self.assertIsNone(await middleware(handler, _message_update(text="/start abc"), {}))
self.assertEqual(await middleware(handler, _callback_update(), {}), "ok")
self.assertEqual(await middleware(handler, _callback_update(), {}), "ok")
self.assertIsNone(await middleware(handler, _callback_update(), {}))
self.assertEqual(
await middleware(handler, _callback_update(data="pay_fk:1:100:subscription"), {}),
"ok",
)
self.assertIsNone(
await middleware(handler, _callback_update(data="pay_fk:1:100:subscription"), {})
)
self.assertEqual(await middleware(handler, _inline_update(), {}), "ok")
self.assertIsNone(await middleware(handler, _inline_update(), {}))
self.assertEqual(handler.await_count, 5)
if __name__ == "__main__":
unittest.main()