diff --git a/backend/bot/middlewares/update_antiflood.py b/backend/bot/middlewares/update_antiflood.py index 6f5ae95..9e43478 100644 --- a/backend/bot/middlewares/update_antiflood.py +++ b/backend/bot/middlewares/update_antiflood.py @@ -77,6 +77,17 @@ class UpdateAntiFloodMiddleware(BaseMiddleware): event: Update, data: Dict[str, Any], ) -> Any: + if bool(getattr(self.settings, "TELEGRAM_DROP_NON_PRIVATE_UPDATES", True)): + chat_type = _message_or_callback_chat_type(event) + if chat_type is not None and chat_type != "private": + logger.info( + "Telegram update dropped outside private chat: chat_type=%s update_type=%s", + chat_type, + getattr(event, "event_type", "unknown"), + ) + data["antiflood_dropped"] = True + return None + if not bool(getattr(self.settings, "TELEGRAM_ANTIFLOOD_ENABLED", True)): return await handler(event, data) @@ -159,6 +170,18 @@ def _update_actor_key(update: Update) -> Optional[str]: return None +def _message_or_callback_chat_type(update: Update) -> Optional[str]: + if update.message and update.message.chat: + return str(update.message.chat.type) + if ( + update.callback_query + and update.callback_query.message + and update.callback_query.message.chat + ): + return str(update.callback_query.message.chat.type) + return None + + def _update_action_key(update: Update) -> str: if update.message: text = update.message.text or "" diff --git a/backend/config/settings.py b/backend/config/settings.py index 94de9d5..348c047 100644 --- a/backend/config/settings.py +++ b/backend/config/settings.py @@ -220,6 +220,7 @@ class Settings(BaseSettings): PANEL_SYNC_LIFETIME_TRAFFIC_MIN_DELTA_BYTES: int = Field(default=104857600) WEBAPP_RATE_LIMIT_TTL_SECONDS: int = Field(default=60) WEBAPP_RATE_LIMIT_MAX_REQUESTS: int = Field(default=30) + TELEGRAM_DROP_NON_PRIVATE_UPDATES: bool = Field(default=True) 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) diff --git a/tests/test_update_antiflood_middleware.py b/tests/test_update_antiflood_middleware.py index 5121cb3..45c0c54 100644 --- a/tests/test_update_antiflood_middleware.py +++ b/tests/test_update_antiflood_middleware.py @@ -9,6 +9,7 @@ def _settings(**overrides): base = { "REDIS_URL": None, "REDIS_KEY_PREFIX": "test-shop", + "TELEGRAM_DROP_NON_PRIVATE_UPDATES": True, "TELEGRAM_ANTIFLOOD_ENABLED": True, "TELEGRAM_ANTIFLOOD_WINDOW_SECONDS": 60, "TELEGRAM_ANTIFLOOD_MAX_UPDATES_PER_WINDOW": 180, @@ -124,6 +125,33 @@ class UpdateAntiFloodMiddlewareTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(handler.await_count, 5) + async def test_non_private_message_is_dropped_before_handler(self): + middleware = UpdateAntiFloodMiddleware(_settings()) + handler = AsyncMock(return_value="ok") + + result = await middleware( + handler, + _message_update(chat_id=-100123, chat_type="supergroup"), + {}, + ) + + self.assertIsNone(result) + handler.assert_not_awaited() + + async def test_non_private_drop_can_be_disabled(self): + middleware = UpdateAntiFloodMiddleware(_settings(TELEGRAM_DROP_NON_PRIVATE_UPDATES=False)) + handler = AsyncMock(return_value="ok") + + with patch("bot.middlewares.update_antiflood.get_redis", AsyncMock(return_value=None)): + result = await middleware( + handler, + _callback_update(chat_id=-100123, chat_type="group"), + {}, + ) + + self.assertEqual(result, "ok") + handler.assert_awaited_once() + if __name__ == "__main__": unittest.main()