diff --git a/src/__main__.py b/src/__main__.py index 96e27b4e..a1b439dc 100644 --- a/src/__main__.py +++ b/src/__main__.py @@ -24,5 +24,4 @@ logger.info("Starting bot & webhook server") logger.debug(f"Starting webhook server listening on port: {settings.WEBHOOK_PORT}") bot.loop.create_task(serve()) - logger.debug(f"Starting bot with token: {settings.bot.TOKEN}") - bot.loop.create_task(bot.run(settings.bot.TOKEN)) + bot.run(settings.bot.TOKEN) diff --git a/src/webhooks/server.py b/src/webhooks/server.py index 7dfd78e6..7aca37b0 100644 --- a/src/webhooks/server.py +++ b/src/webhooks/server.py @@ -1,7 +1,9 @@ +import asyncio import hashlib import hmac -import logging import json +import logging +import signal from typing import Any, Dict from fastapi import FastAPI, HTTPException, Request @@ -100,7 +102,18 @@ async def webhook_handler(request: Request) -> Dict[str, Any]: config = HypercornConfig() config.bind = [f"0.0.0.0:{settings.WEBHOOK_PORT}"] +# Upper bound only: an idle server stops at once. Keep it below the pod's termination grace period. +config.graceful_timeout = 50 + +async def serve() -> None: + """Run the webhook server until SIGTERM/SIGINT, drain in-flight webhooks, then close the bot.""" + shutdown = asyncio.Event() + loop = asyncio.get_running_loop() + # Replaces py-cord's handlers, which stop the loop and cancel in-flight webhooks. + # The bot stays connected during the drain, because the handlers need it. + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, shutdown.set) -async def serve(): - await hypercorn_serve(app, config) + await hypercorn_serve(app, config, shutdown_trigger=shutdown.wait) + await bot.close() diff --git a/tests/src/test_no_token_in_logs.py b/tests/src/test_no_token_in_logs.py new file mode 100644 index 00000000..9f41fdbe --- /dev/null +++ b/tests/src/test_no_token_in_logs.py @@ -0,0 +1,24 @@ +import ast +from pathlib import Path + +SRC = Path(__file__).resolve().parents[2] / "src" +LOG_METHODS = {"debug", "info", "warning", "error", "exception", "critical", "log"} + + +def _log_calls_reading_a_token() -> list[str]: + found = [] + for path in sorted(SRC.rglob("*.py")): + for node in ast.walk(ast.parse(path.read_text())): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in LOG_METHODS + ): + continue + if any(isinstance(n, ast.Attribute) and n.attr == "TOKEN" for n in ast.walk(node)): + found.append(f"{path.relative_to(SRC.parent)}:{node.lineno}") + return found + + +def test_log_calls_do_not_include_a_token(): + assert _log_calls_reading_a_token() == [] diff --git a/tests/src/webhooks/test_server_serve.py b/tests/src/webhooks/test_server_serve.py new file mode 100644 index 00000000..c03598bc --- /dev/null +++ b/tests/src/webhooks/test_server_serve.py @@ -0,0 +1,39 @@ +import asyncio +import signal +from unittest.mock import AsyncMock, patch + +import pytest + +from src.webhooks import server + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sig", [signal.SIGTERM, signal.SIGINT]) +async def test_signal_drains_webhooks_then_closes_the_bot(monkeypatch, sig): + signal_handlers = {} + monkeypatch.setattr( + asyncio.get_running_loop(), "add_signal_handler", lambda s, callback: signal_handlers.__setitem__(s, callback) + ) + calls = [] + + async def fake_hypercorn_serve(app, config, shutdown_trigger): + await shutdown_trigger() + calls.append("drained") + + close = AsyncMock(side_effect=lambda: calls.append("bot closed")) + with ( + patch.object(server, "hypercorn_serve", new=fake_hypercorn_serve), + patch.object(server.bot, "close", new=close), + ): + serving = asyncio.create_task(server.serve()) + await asyncio.sleep(0) + assert not serving.done() + + signal_handlers[sig]() + await asyncio.wait_for(serving, timeout=1) + + assert calls == ["drained", "bot closed"] + + +def test_webhooks_get_up_to_50_seconds_to_drain(): + assert server.config.graceful_timeout == 50