From 2e4599d06878a72879f410f09c1cbefab6f73b68 Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Fri, 5 Jun 2026 11:19:48 +0300 Subject: [PATCH] feat: add Telegram payment callback cooldowns --- backend/bot/middlewares/update_antiflood.py | 104 ++++++++++++++++++++ backend/config/settings.py | 3 + tests/test_update_antiflood_middleware.py | 37 +++++++ 3 files changed, 144 insertions(+) diff --git a/backend/bot/middlewares/update_antiflood.py b/backend/bot/middlewares/update_antiflood.py index 9e43478..28b31a7 100644 --- a/backend/bot/middlewares/update_antiflood.py +++ b/backend/bot/middlewares/update_antiflood.py @@ -1,4 +1,5 @@ import asyncio +import hashlib import logging import time from collections import defaultdict, deque @@ -20,6 +21,8 @@ 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 +DEFAULT_PAYMENT_CALLBACK_COOLDOWN_SECONDS = 20 +DEFAULT_TRIAL_CALLBACK_COOLDOWN_SECONDS = 30 EXPENSIVE_CALLBACK_PREFIXES = ( "pay_", @@ -34,6 +37,11 @@ EXPENSIVE_CALLBACK_PREFIXES = ( "disconnect_device:", ) +TRIAL_CALLBACK_PREFIXES = ( + "trial_action:confirm_activate", + "main_action:request_trial", +) + @dataclass(frozen=True) class RateLimitRule: @@ -69,6 +77,7 @@ class UpdateAntiFloodMiddleware(BaseMiddleware): ) self.action_rules = action_rules or _default_action_rules(settings) self._local_buckets: Dict[str, Deque[float]] = defaultdict(deque) + self._local_cooldowns: Dict[str, float] = {} self._local_lock = asyncio.Lock() async def __call__( @@ -96,6 +105,17 @@ class UpdateAntiFloodMiddleware(BaseMiddleware): return await handler(event, data) action_key = _update_action_key(event) + cooldown = _update_action_cooldown(event, self.settings) + if cooldown and await self._is_cooldown_active(cooldown[0], cooldown[1]): + logger.info( + "Telegram callback dropped by action cooldown: actor=%s cooldown=%s", + actor_key, + cooldown[0], + ) + data["antiflood_dropped"] = True + await _quietly_answer_callback(event) + return None + if await self._is_limited("updates", actor_key, self.default_rule) or ( action_key and action_key in self.action_rules @@ -134,6 +154,42 @@ class UpdateAntiFloodMiddleware(BaseMiddleware): return await self._is_limited_local(f"{bucket_name}:{actor_key}", rule) + async def _is_cooldown_active(self, cooldown_key: str, ttl_seconds: int) -> bool: + if ttl_seconds <= 0: + return False + + try: + redis = await get_redis(self.settings) + if redis is not None: + key = redis_key( + self.settings, + "cooldown", + "telegram", + cooldown_key, + ) + acquired = await redis.set(key, "1", nx=True, ex=ttl_seconds) + return not bool(acquired) + except Exception as exc: + logger.warning("Redis telegram cooldown unavailable; using local fallback: %s", exc) + + return await self._is_cooldown_active_local(cooldown_key, ttl_seconds) + + async def _is_cooldown_active_local(self, cooldown_key: str, ttl_seconds: int) -> bool: + now = time.monotonic() + async with self._local_lock: + expired = [ + key for key, expires_at in self._local_cooldowns.items() if expires_at <= now + ] + for key in expired: + self._local_cooldowns.pop(key, None) + + expires_at = self._local_cooldowns.get(cooldown_key) + if expires_at and expires_at > now: + return True + + self._local_cooldowns[cooldown_key] = now + ttl_seconds + return False + async def _is_limited_local(self, actor_key: str, rule: RateLimitRule) -> bool: now = time.monotonic() cutoff = now - rule.window_seconds @@ -198,6 +254,54 @@ def _update_action_key(update: Update) -> str: return "updates" +def _update_action_cooldown(update: Update, settings: Settings) -> Optional[tuple[str, int]]: + if not bool(getattr(settings, "TELEGRAM_ACTION_COOLDOWN_ENABLED", True)): + return None + if not update.callback_query or not update.callback_query.from_user: + return None + + callback_data = update.callback_query.data or "" + if not callback_data: + return None + + user_id = int(update.callback_query.from_user.id) + data_digest = hashlib.sha256(callback_data.encode("utf-8")).hexdigest()[:24] + + if callback_data.startswith("pay_"): + ttl = int( + getattr( + settings, + "TELEGRAM_PAYMENT_CALLBACK_COOLDOWN_SECONDS", + DEFAULT_PAYMENT_CALLBACK_COOLDOWN_SECONDS, + ) + or DEFAULT_PAYMENT_CALLBACK_COOLDOWN_SECONDS + ) + return f"payment:user:{user_id}:data:{data_digest}", ttl + + if callback_data.startswith(TRIAL_CALLBACK_PREFIXES): + ttl = int( + getattr( + settings, + "TELEGRAM_TRIAL_CALLBACK_COOLDOWN_SECONDS", + DEFAULT_TRIAL_CALLBACK_COOLDOWN_SECONDS, + ) + or DEFAULT_TRIAL_CALLBACK_COOLDOWN_SECONDS + ) + return f"trial:user:{user_id}:data:{data_digest}", ttl + + return None + + +async def _quietly_answer_callback(update: Update) -> None: + callback = update.callback_query + if not callback: + return + try: + await callback.answer() + except Exception: + pass + + def _default_action_rules(settings: Settings) -> Dict[str, RateLimitRule]: window_seconds = int( getattr(settings, "TELEGRAM_ANTIFLOOD_WINDOW_SECONDS", DEFAULT_WINDOW_SECONDS) diff --git a/backend/config/settings.py b/backend/config/settings.py index 348c047..7fc6ec6 100644 --- a/backend/config/settings.py +++ b/backend/config/settings.py @@ -229,6 +229,9 @@ class Settings(BaseSettings): 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) + TELEGRAM_ACTION_COOLDOWN_ENABLED: bool = Field(default=True) + TELEGRAM_PAYMENT_CALLBACK_COOLDOWN_SECONDS: int = Field(default=20) + TELEGRAM_TRIAL_CALLBACK_COOLDOWN_SECONDS: int = Field(default=30) 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 45c0c54..57d6102 100644 --- a/tests/test_update_antiflood_middleware.py +++ b/tests/test_update_antiflood_middleware.py @@ -13,6 +13,9 @@ def _settings(**overrides): "TELEGRAM_ANTIFLOOD_ENABLED": True, "TELEGRAM_ANTIFLOOD_WINDOW_SECONDS": 60, "TELEGRAM_ANTIFLOOD_MAX_UPDATES_PER_WINDOW": 180, + "TELEGRAM_ACTION_COOLDOWN_ENABLED": True, + "TELEGRAM_PAYMENT_CALLBACK_COOLDOWN_SECONDS": 20, + "TELEGRAM_TRIAL_CALLBACK_COOLDOWN_SECONDS": 30, } base.update(overrides) return SimpleNamespace(**base) @@ -152,6 +155,40 @@ class UpdateAntiFloodMiddlewareTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(result, "ok") handler.assert_awaited_once() + async def test_duplicate_payment_callback_is_cooled_down_by_exact_payload(self): + middleware = UpdateAntiFloodMiddleware( + _settings(), + default_rule=RateLimitRule(window_seconds=60, max_events=100), + ) + handler = AsyncMock(return_value="ok") + first = _callback_update(data="pay_fk:1:100:subscription") + duplicate = _callback_update(data="pay_fk:1:100:subscription") + different_payment = _callback_update(data="pay_fk:3:250:subscription") + + with patch("bot.middlewares.update_antiflood.get_redis", AsyncMock(return_value=None)): + self.assertEqual(await middleware(handler, first, {}), "ok") + self.assertIsNone(await middleware(handler, duplicate, {})) + self.assertEqual(await middleware(handler, different_payment, {}), "ok") + + self.assertEqual(handler.await_count, 2) + duplicate.callback_query.answer.assert_awaited_once() + + async def test_duplicate_trial_callback_is_cooled_down(self): + middleware = UpdateAntiFloodMiddleware( + _settings(), + default_rule=RateLimitRule(window_seconds=60, max_events=100), + ) + handler = AsyncMock(return_value="ok") + first = _callback_update(data="main_action:request_trial") + duplicate = _callback_update(data="main_action:request_trial") + + with patch("bot.middlewares.update_antiflood.get_redis", AsyncMock(return_value=None)): + self.assertEqual(await middleware(handler, first, {}), "ok") + self.assertIsNone(await middleware(handler, duplicate, {})) + + handler.assert_awaited_once() + duplicate.callback_query.answer.assert_awaited_once() + if __name__ == "__main__": unittest.main()