From 6dc43182ab4fcf99ca5f90ff92a275d1a24ca89e Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Fri, 5 Jun 2026 11:14:34 +0300 Subject: [PATCH] feat: add Telegram anti-flood action buckets --- backend/bot/middlewares/update_antiflood.py | 115 +++++++++++++++++++- backend/config/settings.py | 5 + tests/test_update_antiflood_middleware.py | 66 +++++++++++ 3 files changed, 181 insertions(+), 5 deletions(-) diff --git a/backend/bot/middlewares/update_antiflood.py b/backend/bot/middlewares/update_antiflood.py index 08e763d..6f5ae95 100644 --- a/backend/bot/middlewares/update_antiflood.py +++ b/backend/bot/middlewares/update_antiflood.py @@ -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 + ), + ), + } diff --git a/backend/config/settings.py b/backend/config/settings.py index b31c441..94de9d5 100644 --- a/backend/config/settings.py +++ b/backend/config/settings.py @@ -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) diff --git a/tests/test_update_antiflood_middleware.py b/tests/test_update_antiflood_middleware.py index cc6bec3..5121cb3 100644 --- a/tests/test_update_antiflood_middleware.py +++ b/tests/test_update_antiflood_middleware.py @@ -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()