diff --git a/drb-c2-core/app/config.py b/drb-c2-core/app/config.py index a770d95..2ed7c1e 100644 --- a/drb-c2-core/app/config.py +++ b/drb-c2-core/app/config.py @@ -256,6 +256,44 @@ class Settings(BaseSettings): # never falls back to this org either way. discord_unlinked_guild_org_id: str = "founding" + # --- Discord permission checks, REST client, OAuth linking (server-26#225) + # Every setting the #225 slices use lives here, added once (S2). + # Kill switch (S3): False restores today's token picking exactly — no + # REST permission check before a bot is handed to a node. + discord_perm_check_enabled: bool = True + # Cache lifetimes for per-(bot, server) results (app/internal/discord_rest.py): + # a result that lets a bot in; denied / not-a-member / invalid token; and + # "couldn't tell" — short, but long enough that the reconciler's 30 s loop + # never becomes a REST loop against Discord. + discord_access_ttl_seconds: int = 120 + discord_access_negative_ttl_seconds: int = 600 + discord_access_unknown_ttl_seconds: int = 30 + # S3: how many bots start_stream checks at once, and the total time it + # waits before failing open (unknown = allowed). + discord_perm_check_concurrency: int = 4 + discord_perm_check_budget_seconds: float = 3.0 + # Process-wide breaker: Discord bans the egress IP (shared with the main + # bot) at 10,000 401/403/429 responses per 10 min. Stop well short. + discord_rest_invalid_budget: int = 4000 + discord_rest_invalid_window_seconds: int = 600 + # A rate-limited request waits at most this long (one retry); longer = unknown. + discord_rest_max_retry_after_seconds: float = 3.0 + discord_rest_timeout_seconds: float = 10.0 + # S5: refresh a server's voice-channel list from a radio bot's REST view + # when the stored directory is older than this. + discord_directory_refresh_seconds: int = 300 + # S6: per-bot OAuth "Connect a Discord server". Empty = OAuth off (routes + # 404, UI hides the button). Must match the redirect URI registered on + # each bot application exactly, e.g. https://drb.cusano.net/discord/connected + discord_oauth_redirect_uri: str = "" + discord_oauth_state_ttl_seconds: int = 600 + # S6: time allowed for the permission check run right after a link. + discord_oauth_link_check_seconds: float = 5.0 + # Pool bots: unverified Discord bots can't join more than 100 servers. + # Connect picks a pool bot under this cap; add bots when all are full + # (owner decision Q3, 2026-10-05). + discord_pool_bot_guild_cap: int = 90 + class Config: env_file = ".env" diff --git a/drb-c2-core/app/internal/discord_leases.py b/drb-c2-core/app/internal/discord_leases.py index 0b3c2ba..869be06 100644 --- a/drb-c2-core/app/internal/discord_leases.py +++ b/drb-c2-core/app/internal/discord_leases.py @@ -429,107 +429,7 @@ async def force_release_token(token_id: str) -> Optional[dict]: # --------------------------------------------------------------------------- -# Token validation / guild membership (server-26#204 item 5) +# Token validation / guild membership / leaving a server moved to +# app/internal/discord_rest.py (server-26#225 S2): inspect_token, +# refresh_membership, leave_guild, TokenRejected. # --------------------------------------------------------------------------- - -DISCORD_API = "https://discord.com/api/v10" - - -class TokenRejected(Exception): - """Discord answered 401: the string is not a working bot token.""" - - -async def inspect_token(token: str) -> dict: - """Ask Discord who this token is and which servers it is in. - - Returns ``{bot_user_id, bot_name, application_id, guild_ids, valid, - validated_at, check_error}``. Raises TokenRejected on a 401. Any other - failure (Discord down, rate limit) is reported in ``check_error`` with - ``valid``/``guild_ids`` = None — unknown, not bad: a token is never - rejected because Discord had a bad minute. - """ - import httpx - headers = {"Authorization": f"Bot {token}"} - out: dict = {"valid": None, "guild_ids": None, "check_error": None, - "validated_at": _now().isoformat()} - try: - async with httpx.AsyncClient(timeout=10, base_url=DISCORD_API, headers=headers) as client: - me = await client.get("/users/@me") - if me.status_code == 401: - raise TokenRejected("Discord rejected this bot token (401).") - me.raise_for_status() - user = me.json() - out["bot_user_id"] = str(user.get("id")) - out["bot_name"] = user.get("username") - app = await client.get("/oauth2/applications/@me") - if app.status_code == 200: - out["application_id"] = str(app.json().get("id")) - guilds = await client.get("/users/@me/guilds", params={"limit": 200}) - guilds.raise_for_status() - out["guild_ids"] = [str(g["id"]) for g in guilds.json()] - out["valid"] = True - except TokenRejected: - raise - except Exception as e: # network / 5xx / 429 — unknown, not invalid - out["check_error"] = str(e)[:200] - return out - - -async def refresh_membership(token_ids: list[str]) -> int: - """Re-read guild membership for these tokens. Returns how many changed.""" - changed = 0 - for tok_id in token_ids: - tok = await fstore.doc_get(TOKENS, tok_id) - if not tok or not tok.get("token"): - continue - try: - info = await inspect_token(tok["token"]) - except TokenRejected as e: - await fstore.doc_update(TOKENS, tok_id, {"valid": False, "check_error": str(e), - "validated_at": _now().isoformat()}) - changed += 1 - continue - updates = {k: v for k, v in info.items() if v is not None or k == "check_error"} - if updates.get("guild_ids") != tok.get("guild_ids") or updates.get("valid") != tok.get("valid"): - changed += 1 - await fstore.doc_update(TOKENS, tok_id, updates) - return changed - - -async def leave_guild(guild_id: str, org_id: Optional[str]) -> list[dict]: - """Make the radio bots DRB holds tokens for leave one Discord server - (Discord REST ``DELETE /users/@me/guilds/{id}``, the bot's own token). - - Only platform-pool bots and ``org_id``'s own BYO bots — never another - org's branded bot. Bots whose cached membership says they are not in the - server are skipped; unknown membership is tried (404 = not a member). - The DRB *server* bot (the discord-bot service) is not in this pool, so it - stays in the server until someone kicks it. Returns one row per bot tried: - ``{token_id, name, left, error}``. Never raises for a Discord failure.""" - import httpx - guild_id = str(guild_id) - out: list[dict] = [] - for tok in await fstore.collection_list(TOKENS): - tok_id = tok.get("token_id") - if not tok_id or not tok.get("token"): - continue - if tok.get("org_id") is not None and tok.get("org_id") != org_id: - continue - members = tok.get("guild_ids") - if members is not None and guild_id not in members: - continue - row = {"token_id": tok_id, "name": tok.get("name") or tok_id, "left": False, "error": None} - try: - async with httpx.AsyncClient(timeout=10, base_url=DISCORD_API, - headers={"Authorization": f"Bot {tok['token']}"}) as client: - r = await client.delete(f"/users/@me/guilds/{guild_id}") - if r.status_code in (200, 204, 404): # 404: Unknown Guild — not a member - row["left"] = True - else: - row["error"] = f"Discord answered {r.status_code}" - except Exception as e: # network — report, do not fail the removal - row["error"] = str(e)[:200] - if row["left"] and members is not None: - await fstore.doc_update(TOKENS, tok_id, {"guild_ids": [g for g in members if g != guild_id]}) - out.append(row) - return out diff --git a/drb-c2-core/app/internal/discord_rest.py b/drb-c2-core/app/internal/discord_rest.py new file mode 100644 index 0000000..5129fe7 --- /dev/null +++ b/drb-c2-core/app/internal/discord_rest.py @@ -0,0 +1,657 @@ +""" +Discord REST client for C2 (server-26#225 slice S2). + +C2 calls Discord's REST API with the radio-bot tokens it holds: to check a +token when an admin adds it, to re-read which servers a bot is in, to make +bots leave a server, and (from S3 on) to read a server's roles, the bot's own +member record and the channel list, so C2 can tell whether a bot may join and +speak in a voice channel before it hands the bot to a node. + +Everything here is built so that a Discord problem can never become a C2 +problem: + +* **One shared ``httpx.AsyncClient``** with the User-Agent Discord asks bots + to send (``DiscordBot (url, version)``). The token travels per request in + the ``Authorization`` header and is never logged, never put in an error + string and never used as a dict key (``_token_key`` hashes it). +* **Rate limits per token.** Buckets are keyed by the ``X-RateLimit-Bucket`` + Discord returns plus the route's major parameter (the guild or channel id). + A request that would have to wait more than + ``discord_rest_max_retry_after_seconds`` is not sent: its result is + *unknown*. A 429 is retried once, only when ``retry_after`` is within that + limit. A global 429 pauses that token. +* **A process-wide breaker on invalid requests.** Discord bans the egress IP + for an hour after 10,000 responses of 401, 403 or 429 (except 429s with + ``X-RateLimit-Scope: shared``) in 10 minutes. The main DRB bot shares that + IP. Once ``discord_rest_invalid_budget`` (default 4000) is reached inside + ``discord_rest_invalid_window_seconds``, nothing more is sent and every + result is unknown until the window drains. +* **Unknown is not bad.** Callers fail open on an unknown result: a Discord + outage must never block a stream start or starve the reconciler. + +Status mapping for the per-server snapshot (``guild_snapshot``): + + 401 token invalid: ``bot_tokens.valid = False``; never retried + 404 10004 / 10007 not a member: the server is dropped from ``guild_ids`` + 403 50001 unknown, then confirmed with ONE read of the bot's + server list; ``guild_ids`` is never changed from the + 403 alone + anything else unknown + +Snapshots are cached in-process per (token_id, guild_id), one c2-core +process being the same assumption ``_RateLimiter`` and ``_publish_locks`` +already make. Results that let a bot in live ``discord_access_ttl_seconds``; +not-member and invalid results live ``discord_access_negative_ttl_seconds``; +unknown results live ``discord_access_unknown_ttl_seconds``, so the +reconciler's 30 s loop can't turn into a REST loop against Discord. +Concurrent loads of one key share one request (single flight). +``invalidate(token_id=, guild_id=)`` drops entries. +""" +from __future__ import annotations + +import asyncio +import hashlib +import re +import time +from collections import deque +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Awaitable, Callable, Optional + +import httpx + +from app.config import settings +from app.internal import firestore as fstore +from app.internal.logger import logger + +DISCORD_API = "https://discord.com/api/v10" +USER_AGENT = "DiscordBot (https://drb.cusano.net, 1)" +TOKENS = "bot_tokens" + +# GET /users/@me/guilds returns at most 200 per page. 50 pages is 10,000 +# servers, well past what an unsharded bot can be in; the cap only stops a +# misbehaving response from looping forever. +GUILDS_PAGE = 200 +MAX_GUILD_PAGES = 50 + +# Discord JSON error codes this module acts on. +UNKNOWN_GUILD = 10004 +UNKNOWN_MEMBER = 10007 +MISSING_ACCESS = 50001 + +# Snapshot / access outcomes. ``denied`` is not produced here; the permission +# check built on these snapshots (S3) uses it, with the same negative TTL. +OK = "ok" +DENIED = "denied" +NOT_MEMBER = "not_member" +INVALID = "invalid" +UNKNOWN = "unknown" + + +class TokenRejected(Exception): + """Discord answered 401: the string is not a working bot token.""" + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def _token_key(token: str) -> str: + """A stable, non-reversible handle for a token, for rate-limit state.""" + return hashlib.sha256(token.encode()).hexdigest()[:16] + + +# --------------------------------------------------------------------------- +# Results +# --------------------------------------------------------------------------- + +@dataclass +class RestResult: + """One Discord REST answer. ``status`` is None when nothing came back + (not sent, network error): the result is unknown.""" + status: Optional[int] + data: Any = None + code: Optional[int] = None # Discord's JSON error code, on an error + error: Optional[str] = None # short reason; never holds the token + + @property + def ok(self) -> bool: + return self.status is not None and 200 <= self.status < 300 + + @property + def unknown(self) -> bool: + """No usable answer: not sent, network error, rate limited or 5xx.""" + return self.status is None or self.status == 429 or self.status >= 500 + + +# --------------------------------------------------------------------------- +# Rate limits and the invalid-request breaker +# --------------------------------------------------------------------------- + +_ID_RE = re.compile(r"\d{5,}") +_MAJOR_RE = re.compile(r"^/(?:guilds|channels|webhooks)/(\d+)") + + +def _route(method: str, path: str) -> tuple[str, str]: + """(route key, major parameter) for a path. Snowflakes in the path are + collapsed so ``/guilds/1/channels`` and ``/guilds/2/channels`` share a + route; the major parameter keeps their buckets apart.""" + m = _MAJOR_RE.match(path) + return f"{method.upper()} {_ID_RE.sub('{id}', path)}", (m.group(1) if m else "") + + +class _RateLimits: + def __init__(self) -> None: + self.route_bucket: dict[tuple[str, str], str] = {} + self.buckets: dict[tuple[str, str, str], tuple[int, float]] = {} + self.paused_until: dict[str, float] = {} + + def wait_for(self, tok: str, method: str, path: str, now: float) -> float: + """Seconds to wait before this request may be sent (0 = go now).""" + wait = max(0.0, self.paused_until.get(tok, 0.0) - now) + route, major = _route(method, path) + bucket = self.route_bucket.get((tok, route)) + if bucket: + remaining, reset_at = self.buckets.get((tok, bucket, major), (1, 0.0)) + if remaining <= 0 and reset_at > now: + wait = max(wait, reset_at - now) + return wait + + def update(self, tok: str, method: str, path: str, headers: httpx.Headers, now: float) -> None: + bucket = headers.get("x-ratelimit-bucket") + if not bucket: + return + route, major = _route(method, path) + self.route_bucket[(tok, route)] = bucket + try: + remaining = int(headers.get("x-ratelimit-remaining", "1")) + reset_after = float(headers.get("x-ratelimit-reset-after", "0")) + except ValueError: + return + self.buckets[(tok, bucket, major)] = (remaining, now + reset_after) + + def pause(self, tok: str, seconds: float, now: float) -> None: + self.paused_until[tok] = max(self.paused_until.get(tok, 0.0), now + seconds) + + +class _Breaker: + """Counts 401 / 403 / non-shared 429 responses in a sliding window.""" + + def __init__(self) -> None: + self.hits: deque[float] = deque() + self.open_logged = False + + def _prune(self, now: float) -> None: + horizon = now - settings.discord_rest_invalid_window_seconds + while self.hits and self.hits[0] <= horizon: + self.hits.popleft() + + def record(self, now: float) -> None: + self.hits.append(now) + + def tripped(self, now: float) -> bool: + self._prune(now) + tripped = len(self.hits) >= settings.discord_rest_invalid_budget + if tripped and not self.open_logged: + logger.warning( + f"Discord REST breaker open: {len(self.hits)} invalid responses in " + f"{settings.discord_rest_invalid_window_seconds}s; REST checks report unknown " + "until the window drains (protects the shared egress IP from a Discord ban)") + self.open_logged = True + elif not tripped: + self.open_logged = False + return tripped + + +_limits = _RateLimits() +_breaker = _Breaker() +_sleep = asyncio.sleep # tests replace this, not asyncio.sleep itself + +# --------------------------------------------------------------------------- +# The shared HTTP client +# --------------------------------------------------------------------------- + +_client_obj: Optional[httpx.AsyncClient] = None +_client_loop: Optional[asyncio.AbstractEventLoop] = None +_transport: Optional[httpx.AsyncBaseTransport] = None + + +def _client() -> httpx.AsyncClient: + """The one shared client, recreated only if the event loop changed (the + connection pool belongs to a loop; tests run one loop per test).""" + global _client_obj, _client_loop + loop = asyncio.get_running_loop() + if _client_obj is None or _client_loop is not loop or _client_obj.is_closed: + _client_obj = httpx.AsyncClient( + base_url=DISCORD_API, + timeout=settings.discord_rest_timeout_seconds, + headers={"User-Agent": USER_AGENT}, + transport=_transport, + ) + _client_loop = loop + return _client_obj + + +def set_transport(transport: Optional[httpx.AsyncBaseTransport]) -> None: + """Tests: route every request through ``transport`` (e.g. httpx.MockTransport).""" + global _transport, _client_obj + _transport = transport + _client_obj = None + + +def reset() -> None: + """Forget all in-process state: client, rate limits, breaker, cache.""" + global _limits, _breaker, _client_obj + _limits = _RateLimits() + _breaker = _Breaker() + _client_obj = None + access_cache.clear() + + +def _retry_after(r: httpx.Response) -> tuple[float, bool]: + """(seconds, is_global) from a 429's body or headers.""" + body: Any = None + try: + body = r.json() + except Exception: + body = None + seconds: Optional[float] = None + is_global = r.headers.get("x-ratelimit-global", "").lower() == "true" + if isinstance(body, dict): + if isinstance(body.get("retry_after"), (int, float)): + seconds = float(body["retry_after"]) + is_global = is_global or bool(body.get("global")) + if seconds is None: + try: + seconds = float(r.headers.get("retry-after", "1")) + except ValueError: + seconds = 1.0 + return max(0.0, seconds), is_global + + +async def request(token: str, method: str, path: str, *, + params: Optional[dict] = None) -> RestResult: + """Send one request as the bot ``token``. Never raises for a Discord or + network failure: those come back as an unknown RestResult.""" + tok = _token_key(token) + max_wait = settings.discord_rest_max_retry_after_seconds + for attempt in (0, 1): + now = time.monotonic() + if _breaker.tripped(now): + return RestResult(None, error="Discord REST paused (invalid-request budget reached)") + wait = _limits.wait_for(tok, method, path, now) + if wait > max_wait: + return RestResult(None, error=f"rate limited for {wait:.0f}s") + if wait > 0: + await _sleep(wait) + try: + r = await _client().request(method, path, params=params, + headers={"Authorization": f"Bot {token}"}) + except httpx.HTTPError as e: + # The exception text holds at most the URL, never the header. + return RestResult(None, error=f"{type(e).__name__}: {str(e)[:120]}") + now = time.monotonic() + _limits.update(tok, method, path, r.headers, now) + status = r.status_code + if status in (401, 403) or (status == 429 and r.headers.get("x-ratelimit-scope") != "shared"): + _breaker.record(now) + if status == 429: + seconds, is_global = _retry_after(r) + if is_global: + _limits.pause(tok, seconds, now) + if attempt == 0 and seconds <= max_wait: + await _sleep(seconds) + continue + return RestResult(429, error=f"rate limited (retry after {seconds:.1f}s)") + data: Any = None + if r.content: + try: + data = r.json() + except ValueError: + data = None + code = data.get("code") if (status >= 400 and isinstance(data, dict)) else None + err = None if 200 <= status < 300 else f"Discord answered {status}" + (f" ({code})" if code else "") + return RestResult(status, data, code if isinstance(code, int) else None, err) + return RestResult(429, error="rate limited") # pragma: no cover - loop always returns + + +# --------------------------------------------------------------------------- +# Cache (in-process, single flight) +# --------------------------------------------------------------------------- + +def ttl_for(outcome: str) -> float: + """How long a result with this outcome stays cached.""" + if outcome == OK: + return float(settings.discord_access_ttl_seconds) + if outcome in (DENIED, NOT_MEMBER, INVALID): + return float(settings.discord_access_negative_ttl_seconds) + return float(settings.discord_access_unknown_ttl_seconds) + + +class AccessCache: + """TTL cache keyed by tuples that start ``(token_id, guild_id, ...)``. + + ``get_or_load`` runs one load per key at a time; concurrent callers await + the same load. A load that overlaps an ``invalidate`` touching its key is + returned to its callers but not stored, so an invalidation is never undone + by a load that started before it.""" + + def __init__(self) -> None: + self._entries: dict[tuple, tuple[float, Any]] = {} + self._inflight: dict[tuple, asyncio.Task] = {} + self._gen: dict[tuple, int] = {} + + def clear(self) -> None: + self._entries.clear() + self._inflight.clear() + self._gen.clear() + + def peek(self, key: tuple) -> Any: + hit = self._entries.get(key) + if hit and hit[0] > time.monotonic(): + return hit[1] + return None + + async def get_or_load(self, key: tuple, loader: Callable[[], Awaitable[Any]], + ttl_of: Callable[[Any], float], *, refresh: bool = False) -> Any: + if not refresh: + hit = self._entries.get(key) + if hit and hit[0] > time.monotonic(): + return hit[1] + task = self._inflight.get(key) + loop = asyncio.get_running_loop() + if task is None or task.done() or task.get_loop() is not loop: + gen = self._gen.get(key, 0) + + async def _load() -> Any: + try: + value = await loader() + if self._gen.get(key, 0) == gen: + self._entries[key] = (time.monotonic() + ttl_of(value), value) + return value + finally: + if self._inflight.get(key) is task_ref[0]: + self._inflight.pop(key, None) + + task_ref: list = [None] + task = loop.create_task(_load()) + task_ref[0] = task + self._inflight[key] = task + return await asyncio.shield(task) + + def invalidate(self, token_id: Optional[str] = None, guild_id: Optional[str] = None) -> int: + """Drop entries for a token, a guild, both, or (no args) everything. + Returns how many cached entries were dropped.""" + def match(k: tuple) -> bool: + return ((token_id is None or k[0] == token_id) + and (guild_id is None or (len(k) > 1 and k[1] == str(guild_id)))) + dropped = [k for k in self._entries if match(k)] + for k in dropped: + del self._entries[k] + for k in set(self._inflight) | set(self._entries) | set(dropped): + if match(k): + self._gen[k] = self._gen.get(k, 0) + 1 + for k in [k for k in self._inflight if match(k)]: + self._inflight.pop(k, None) + return len(dropped) + + +access_cache = AccessCache() + + +def invalidate(token_id: Optional[str] = None, guild_id: Optional[str] = None) -> int: + return access_cache.invalidate(token_id=token_id, guild_id=guild_id) + + +# --------------------------------------------------------------------------- +# Token inspection / membership (moved from discord_leases, server-26#204) +# --------------------------------------------------------------------------- + +async def list_guild_ids(token: str) -> RestResult: + """Every server this bot is in, following ``after=`` pages (the old code + read one page of 200 and silently dropped the rest). On success ``data`` + is the list of guild id strings; any failed page makes the whole result + that page's failure, never a partial list.""" + out: list[str] = [] + after: Optional[str] = None + for _ in range(MAX_GUILD_PAGES): + params: dict = {"limit": GUILDS_PAGE} + if after: + params["after"] = after + r = await request(token, "GET", "/users/@me/guilds", params=params) + if not r.ok: + return r + page = r.data if isinstance(r.data, list) else [] + ids = [str(g["id"]) for g in page if isinstance(g, dict) and g.get("id") is not None] + out.extend(ids) + if len(page) < GUILDS_PAGE or not ids: + return RestResult(200, out) + after = max(ids, key=int) + logger.warning(f"Discord /users/@me/guilds: stopped after {MAX_GUILD_PAGES} pages") + return RestResult(None, error=f"more than {MAX_GUILD_PAGES * GUILDS_PAGE} servers") + + +async def inspect_token(token: str) -> dict: + """Ask Discord who this token is and which servers it is in. + + Returns ``{bot_user_id, bot_name, application_id, guild_ids, valid, + validated_at, check_error}``. Raises TokenRejected on a 401. Any other + failure (Discord down, rate limit) is reported in ``check_error`` with + ``valid``/``guild_ids`` = None — unknown, not bad: a token is never + rejected because Discord had a bad minute. + """ + out: dict = {"valid": None, "guild_ids": None, "check_error": None, + "validated_at": _now().isoformat()} + me = await request(token, "GET", "/users/@me") + if me.status == 401: + raise TokenRejected("Discord rejected this bot token (401).") + if not me.ok or not isinstance(me.data, dict): + out["check_error"] = (me.error or "unexpected answer from Discord")[:200] + return out + out["bot_user_id"] = str(me.data.get("id")) + out["bot_name"] = me.data.get("username") + app = await request(token, "GET", "/oauth2/applications/@me") + if app.ok and isinstance(app.data, dict) and app.data.get("id") is not None: + out["application_id"] = str(app.data["id"]) + guilds = await list_guild_ids(token) + if guilds.status == 401: + raise TokenRejected("Discord rejected this bot token (401).") + if not guilds.ok: + out["check_error"] = (guilds.error or "could not list servers")[:200] + return out + out["guild_ids"] = guilds.data + out["valid"] = True + return out + + +async def refresh_membership(token_ids: list[str]) -> int: + """Re-read guild membership for these tokens. Returns how many changed. + Every token re-read has its cached snapshots dropped.""" + changed = 0 + for tok_id in token_ids: + tok = await fstore.doc_get(TOKENS, tok_id) + if not tok or not tok.get("token"): + continue + try: + info = await inspect_token(tok["token"]) + except TokenRejected as e: + await fstore.doc_update(TOKENS, tok_id, {"valid": False, "check_error": str(e), + "validated_at": _now().isoformat()}) + invalidate(token_id=tok_id) + changed += 1 + continue + updates = {k: v for k, v in info.items() if v is not None or k == "check_error"} + if updates.get("guild_ids") != tok.get("guild_ids") or updates.get("valid") != tok.get("valid"): + changed += 1 + await fstore.doc_update(TOKENS, tok_id, updates) + invalidate(token_id=tok_id) + return changed + + +async def leave_guild(guild_id: str, org_id: Optional[str]) -> list[dict]: + """Make the radio bots DRB holds tokens for leave one Discord server + (Discord REST ``DELETE /users/@me/guilds/{id}``, the bot's own token). + + Only platform-pool bots and ``org_id``'s own BYO bots — never another + org's branded bot. Bots whose cached membership says they are not in the + server are skipped; unknown membership is tried (404 = not a member). + The DRB *server* bot (the discord-bot service) is not in this pool, so it + stays in the server until someone kicks it. Returns one row per bot tried: + ``{token_id, name, left, error}``. Never raises for a Discord failure.""" + guild_id = str(guild_id) + out: list[dict] = [] + for tok in await fstore.collection_list(TOKENS): + tok_id = tok.get("token_id") + if not tok_id or not tok.get("token"): + continue + if tok.get("org_id") is not None and tok.get("org_id") != org_id: + continue + members = tok.get("guild_ids") + if members is not None and guild_id not in members: + continue + row = {"token_id": tok_id, "name": tok.get("name") or tok_id, "left": False, "error": None} + r = await request(tok["token"], "DELETE", f"/users/@me/guilds/{guild_id}") + if r.status in (200, 204, 404): # 404: Unknown Guild — not a member + row["left"] = True + else: + row["error"] = r.error or f"Discord answered {r.status}" + if row["left"]: + invalidate(token_id=tok_id, guild_id=guild_id) + if members is not None: + await fstore.doc_update(TOKENS, tok_id, {"guild_ids": [g for g in members if g != guild_id]}) + out.append(row) + return out + + +# --------------------------------------------------------------------------- +# Per-(token, guild) snapshot — what S3's permission check reads +# --------------------------------------------------------------------------- + +@dataclass +class GuildSnapshot: + """What one bot can see of one server. ``status`` is OK, NOT_MEMBER, + INVALID or UNKNOWN; ``guild``/``member``/``channels`` are set only on OK. + ``guild`` is Discord's guild object (``owner_id``, ``roles``), ``member`` + the bot's own member object (``roles``, ``communication_disabled_until``), + ``channels`` the server's channel list.""" + token_id: str + guild_id: str + status: str + bot_user_id: Optional[str] = None + guild: Optional[dict] = None + member: Optional[dict] = None + channels: Optional[list] = None + error: Optional[str] = None + fetched_at: float = field(default_factory=time.time) + + +async def _mark_invalid(token_id: str) -> None: + await fstore.doc_update(TOKENS, token_id, { + "valid": False, "check_error": "Discord rejected this bot token (401).", + "validated_at": _now().isoformat()}) + + +async def _drop_membership(token_id: str, guild_id: str) -> None: + tok = await fstore.doc_get(TOKENS, token_id) or {} + members = tok.get("guild_ids") + if isinstance(members, list) and guild_id in members: + await fstore.doc_update(TOKENS, token_id, {"guild_ids": [g for g in members if g != guild_id]}) + + +async def _confirm_membership(token_id: str, token: str, guild_id: str) -> Optional[bool]: + """After a 403 Missing Access: one read of the bot's server list decides. + True/False = member or not (``guild_ids`` written from the full list); + None = still unknown.""" + r = await list_guild_ids(token) + if r.status == 401: + await _mark_invalid(token_id) + return None + if not r.ok: + return None + await fstore.doc_update(TOKENS, token_id, {"guild_ids": r.data, "validated_at": _now().isoformat()}) + return guild_id in r.data + + +async def _classify(snap: GuildSnapshot, r: RestResult, token: str) -> Optional[GuildSnapshot]: + """None when ``r`` is a success; otherwise the finished snapshot.""" + if r.ok: + return None + snap.error = r.error + if r.status == 401: + await _mark_invalid(snap.token_id) + snap.status = INVALID + elif r.status == 404 and r.code in (UNKNOWN_GUILD, UNKNOWN_MEMBER): + await _drop_membership(snap.token_id, snap.guild_id) + snap.status = NOT_MEMBER + elif r.status == 403 and r.code == MISSING_ACCESS: + member = await _confirm_membership(snap.token_id, token, snap.guild_id) + snap.status = NOT_MEMBER if member is False else UNKNOWN + else: + snap.status = UNKNOWN + return snap + + +async def _load_snapshot(token_doc: dict, guild_id: str) -> GuildSnapshot: + token_id = str(token_doc.get("token_id")) + snap = GuildSnapshot(token_id=token_id, guild_id=guild_id, status=UNKNOWN, + bot_user_id=token_doc.get("bot_user_id")) + token = token_doc.get("token") + if not token: + snap.error = "no token stored" + return snap + if token_doc.get("valid") is False: + snap.status, snap.error = INVALID, "token marked invalid" + return snap + if not snap.bot_user_id: + me = await request(token, "GET", "/users/@me") + done = await _classify(snap, me, token) + if done: + return done + snap.bot_user_id = str((me.data or {}).get("id") or "") or None + if not snap.bot_user_id: + snap.error = "Discord did not say who this bot is" + return snap + await fstore.doc_update(TOKENS, token_id, {"bot_user_id": snap.bot_user_id}) + + g = await request(token, "GET", f"/guilds/{guild_id}") + done = await _classify(snap, g, token) + if done: + return done + m, c = await asyncio.gather( + request(token, "GET", f"/guilds/{guild_id}/members/{snap.bot_user_id}"), + request(token, "GET", f"/guilds/{guild_id}/channels"), + ) + for r in (m, c): + done = await _classify(snap, r, token) + if done: + return done + if not isinstance(g.data, dict) or not isinstance(m.data, dict) or not isinstance(c.data, list): + snap.error = "unexpected answer from Discord" + return snap + snap.guild, snap.member, snap.channels = g.data, m.data, c.data + snap.status, snap.error = OK, None + return snap + + +async def guild_snapshot(token_doc: dict, guild_id: str, *, refresh: bool = False) -> GuildSnapshot: + """Cached snapshot of one server as one bot sees it: three GETs + (``/guilds/{g}``, ``/guilds/{g}/members/{bot}`` — no privileged intent + needed — and ``/guilds/{g}/channels``). ``refresh=True`` skips the cache + (a concurrent load is still shared). Never raises for a Discord failure.""" + guild_id = str(guild_id) + token_id = str(token_doc.get("token_id")) + key = (token_id, guild_id, "snapshot") + + async def _load() -> GuildSnapshot: + # Fail open on anything else too (a Firestore write inside the + # classification, a malformed body): the caller gets UNKNOWN, never + # an exception, and a shielded load never dies unretrieved. + try: + return await _load_snapshot(token_doc, guild_id) + except Exception as e: + logger.warning(f"Discord snapshot for bot {token_id} in {guild_id} failed: {type(e).__name__}") + return GuildSnapshot(token_id=token_id, guild_id=guild_id, status=UNKNOWN, + bot_user_id=token_doc.get("bot_user_id"), + error=f"check failed ({type(e).__name__})") + + return await access_cache.get_or_load(key, _load, lambda s: ttl_for(s.status), refresh=refresh) diff --git a/drb-c2-core/app/internal/discord_streams.py b/drb-c2-core/app/internal/discord_streams.py index 0c1ff65..13fa8d3 100644 --- a/drb-c2-core/app/internal/discord_streams.py +++ b/drb-c2-core/app/internal/discord_streams.py @@ -49,6 +49,7 @@ from typing import Optional from app.config import settings from app.internal import firestore as fstore from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal.logger import logger STREAMS = leases.STREAMS @@ -416,7 +417,7 @@ async def _start_locked(*, org_id: str, guild_id: str, channel_id: str, system_i tokens = await fstore.collection_list(leases.TOKENS) stale = [t["token_id"] for t in tokens if t.get("token_id") and t.get("guild_ids") is not None and guild_id not in t["guild_ids"]] - if await leases.refresh_membership(stale): + if await discord_rest.refresh_membership(stale): res = await leases.lease_for_stream(stream_id, node_id=node["node_id"], org_id=org_id, preferred_token_id=preferred_token_id, stream_fields=fields) if res.ok: @@ -604,7 +605,7 @@ async def handle_actual(node_id: str, payload: dict) -> None: "actual": {"state": state, "error": error, "code": code, "since": since, "at": now}, }) if recheck: - await leases.refresh_membership(sorted(recheck)) + await discord_rest.refresh_membership(sorted(recheck)) if "sessions" in payload: sessions = {str(t) for t in (payload.get("sessions") or [])} diff --git a/drb-c2-core/app/routers/discord.py b/drb-c2-core/app/routers/discord.py index 9be1dbf..040c4c8 100644 --- a/drb-c2-core/app/routers/discord.py +++ b/drb-c2-core/app/routers/discord.py @@ -27,6 +27,7 @@ from pydantic import BaseModel, Field from app.config import settings from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal import discord_streams as streams from app.internal import firestore as fstore from app.internal.audit import write_audit @@ -244,7 +245,7 @@ async def remove_guild(guild_id: str, leave_radio_bots: bool = False, return _error(e) left: list = [] if leave_radio_bots: - left = await leases.leave_guild(guild_id, out["previous_org_id"] or acting_org) + left = await discord_rest.leave_guild(guild_id, out["previous_org_id"] or acting_org) await write_audit(uid, email, "discord.guild.remove", details={ "guild_id": guild_id, "previous_org_id": out["previous_org_id"], "acting_org_id": acting_org, "streams_removed": out["streams_removed"], "leave_radio_bots": leave_radio_bots, diff --git a/drb-c2-core/app/routers/tokens.py b/drb-c2-core/app/routers/tokens.py index 729bd07..d3310a1 100644 --- a/drb-c2-core/app/routers/tokens.py +++ b/drb-c2-core/app/routers/tokens.py @@ -4,6 +4,7 @@ from pydantic import BaseModel from typing import Optional from app.internal import firestore as fstore from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal.auth import require_admin_token from app.internal.logger import logger @@ -47,8 +48,8 @@ async def add_token(body: TokenCreate, _: dict = Depends(require_admin_token)): (server-26#204 item 5): a string Discord answers 401 for is refused; if Discord is merely unreachable the token is stored with membership unknown.""" try: - info = await leases.inspect_token(body.token.strip()) - except leases.TokenRejected as e: + info = await discord_rest.inspect_token(body.token.strip()) + except discord_rest.TokenRejected as e: raise HTTPException(400, str(e)) token_id = str(uuid.uuid4()) doc = { @@ -73,7 +74,7 @@ async def refresh_token(token_id: str, _: dict = Depends(require_admin_token)): """Re-check a bot with Discord (valid? which servers is it in?).""" if not await fstore.doc_get("bot_tokens", token_id): raise HTTPException(404, "Token not found.") - await leases.refresh_membership([token_id]) + await discord_rest.refresh_membership([token_id]) return _masked(await fstore.doc_get("bot_tokens", token_id) or {}) diff --git a/drb-c2-core/tests/test_discord_leases.py b/drb-c2-core/tests/test_discord_leases.py index 9ae9ba9..c12cf20 100644 --- a/drb-c2-core/tests/test_discord_leases.py +++ b/drb-c2-core/tests/test_discord_leases.py @@ -13,6 +13,7 @@ from datetime import datetime, timezone, timedelta import pytest from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal import discord_streams as streams from tests.fake_firestore import FakeFirestore @@ -38,7 +39,7 @@ def fs(monkeypatch): async def _no_discord(token_ids): return 0 - monkeypatch.setattr(leases, "refresh_membership", _no_discord) + monkeypatch.setattr(discord_rest, "refresh_membership", _no_discord) return fake diff --git a/drb-c2-core/tests/test_discord_multi_stream.py b/drb-c2-core/tests/test_discord_multi_stream.py index ad5bb17..cb5cfa3 100644 --- a/drb-c2-core/tests/test_discord_multi_stream.py +++ b/drb-c2-core/tests/test_discord_multi_stream.py @@ -17,6 +17,7 @@ from fastapi.testclient import TestClient from app.config import settings from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal import discord_streams as streams from app.internal.auth import require_admin_token, require_service_key_or_admin from app.main import app @@ -42,7 +43,7 @@ def fs(monkeypatch): async def _no_discord(token_ids): return 0 - monkeypatch.setattr(leases, "refresh_membership", _no_discord) + monkeypatch.setattr(discord_rest, "refresh_membership", _no_discord) monkeypatch.setattr(settings, "discord_unlinked_guild_org_id", "") monkeypatch.setattr(settings, "discord_guild_max_streams", 4) monkeypatch.setattr(settings, "discord_node_max_streams", 4) diff --git a/drb-c2-core/tests/test_discord_remove_and_logging.py b/drb-c2-core/tests/test_discord_remove_and_logging.py index 7b2b910..2858425 100644 --- a/drb-c2-core/tests/test_discord_remove_and_logging.py +++ b/drb-c2-core/tests/test_discord_remove_and_logging.py @@ -15,6 +15,7 @@ from fastapi.testclient import TestClient from app.config import settings from app.internal import discord_guild_link as glink from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal import discord_streams as streams from app.internal.auth import require_admin_token, require_firebase_token from app.main import app @@ -57,7 +58,7 @@ def fs(monkeypatch): async def _no_discord(token_ids): return 0 - monkeypatch.setattr(leases, "refresh_membership", _no_discord) + monkeypatch.setattr(discord_rest, "refresh_membership", _no_discord) monkeypatch.setattr(settings, "discord_unlinked_guild_org_id", "") for nid, org in (("n1", ORG), ("n2", ORG), ("nx", OTHER)): @@ -340,32 +341,21 @@ def test_remove_guild_can_make_pool_and_own_radio_bots_leave(fs, monkeypatch): "org_id": None, "guild_ids": ["300"], "valid": True}) calls = [] - class _Resp: - def __init__(self, code): - self.status_code = code + def _handler(request): + auth = request.headers["Authorization"] + calls.append((auth, request.method, request.url.path)) + return httpx.Response(500 if auth == "Bot secret-t2" else 204) - class _Client: - def __init__(self, *, headers, **kw): - self.auth = headers["Authorization"] - - async def __aenter__(self): - return self - - async def __aexit__(self, *exc): - return False - - async def delete(self, url): - calls.append((self.auth, url)) - return _Resp(500 if self.auth == "Bot secret-t2" else 204) - - monkeypatch.setattr(httpx, "AsyncClient", _Client) + discord_rest.reset() + monkeypatch.setattr(discord_rest, "_transport", httpx.MockTransport(_handler)) + monkeypatch.setattr(discord_rest, "_client_obj", None) app.dependency_overrides[require_admin_token] = lambda: ADMIN r = client.delete("/discord/guilds/100?leave_radio_bots=true") assert r.status_code == 200 - tried = sorted(a for a, _ in calls) + tried = sorted(a for a, _, _ in calls) # pool bots in the server + this org's BYO; never another org's bot, never a bot not in it assert tried == ["Bot secret-mine", "Bot secret-t1", "Bot secret-t2"] - assert all(u == "/users/@me/guilds/100" for _, u in calls) + assert all(m == "DELETE" and u == "/api/v10/users/@me/guilds/100" for _, m, u in calls) rows = {x["name"]: x for x in r.json()["radio_bots"]} assert rows["Radio t1"]["left"] is True and rows["Radio t2"]["left"] is False assert "secret" not in r.text diff --git a/drb-c2-core/tests/test_discord_rest.py b/drb-c2-core/tests/test_discord_rest.py new file mode 100644 index 0000000..066c71d --- /dev/null +++ b/drb-c2-core/tests/test_discord_rest.py @@ -0,0 +1,442 @@ +""" +Discord REST client, cache and settings (server-26#225 slice S2). + +Every request goes through httpx.MockTransport: nothing here talks to Discord. +""" +import asyncio +import logging +from types import SimpleNamespace + +import httpx +import pytest + +from app.config import settings +from app.internal import discord_rest as rest +from tests.fake_firestore import FakeFirestore + +TOKEN = "Nzk5.SECRET-TOKEN-VALUE" +G = "100000000000000001" +BOT = "900000000000000009" + + +class Clock: + def __init__(self): + self.t = 1000.0 + + def __call__(self): + return self.t + + +class Discord: + """A scripted Discord: ``routes[(method, path)]`` is a Response or a + callable(request) -> Response. Records every request.""" + + def __init__(self): + self.routes: dict = {} + self.calls: list[httpx.Request] = [] + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.calls.append(request) + path = request.url.path.removeprefix("/api/v10") + handler = self.routes.get((request.method, path)) + if handler is None: + return httpx.Response(404, json={"code": 0, "message": "no route in test"}) + return handler(request) if callable(handler) else handler + + def paths(self): + return [r.url.path.removeprefix("/api/v10") for r in self.calls] + + +@pytest.fixture +def discord(monkeypatch): + d = Discord() + rest.reset() + monkeypatch.setattr(rest, "_transport", httpx.MockTransport(d)) + monkeypatch.setattr(rest, "_client_obj", None) + clock = Clock() + monkeypatch.setattr(rest, "time", SimpleNamespace(monotonic=clock, time=clock)) + slept: list = [] + + async def _sleep(s): + slept.append(s) + clock.t += s + monkeypatch.setattr(rest, "_sleep", _sleep) + d.clock, d.slept = clock, slept + yield d + rest.reset() + + +@pytest.fixture +def fs(monkeypatch): + return FakeFirestore.install(monkeypatch) + + +def _token_doc(fs, token_id="t1", *, guilds=None, valid=True, bot_user_id=BOT): + doc = {"token_id": token_id, "name": "Radio 1", "token": TOKEN, "org_id": None, + "guild_ids": [G] if guilds is None else guilds, "valid": valid, "bot_user_id": bot_user_id} + fs.put("bot_tokens", token_id, doc) + return dict(doc) + + +def _guild_routes(d, guild=G): + d.routes[("GET", f"/guilds/{guild}")] = httpx.Response(200, json={"id": guild, "owner_id": "1", "roles": []}) + d.routes[("GET", f"/guilds/{guild}/members/{BOT}")] = httpx.Response(200, json={"roles": []}) + d.routes[("GET", f"/guilds/{guild}/channels")] = httpx.Response(200, json=[{"id": "5", "type": 2}]) + + +# --- client hygiene ---------------------------------------------------------- + +async def test_requests_carry_user_agent_and_bot_auth(discord): + discord.routes[("GET", "/users/@me")] = httpx.Response(200, json={"id": BOT}) + r = await rest.request(TOKEN, "GET", "/users/@me") + assert r.ok and r.data == {"id": BOT} + (req,) = discord.calls + assert req.headers["User-Agent"] == "DiscordBot (https://drb.cusano.net, 1)" + assert req.headers["Authorization"] == f"Bot {TOKEN}" + + +async def test_token_never_logged_or_in_errors(discord, fs, caplog): + caplog.set_level(logging.DEBUG) + discord.routes[("GET", "/users/@me")] = httpx.Response(500, json={"message": "boom"}) + info = await rest.inspect_token(TOKEN) + assert info["valid"] is None and info["check_error"] + + def _boom(request): + raise httpx.ConnectError("connection refused", request=request) + discord.routes[("GET", f"/guilds/{G}")] = _boom + snap = await rest.guild_snapshot(_token_doc(fs), G) + assert snap.status == rest.UNKNOWN + assert TOKEN not in (info["check_error"] or "") and TOKEN not in (snap.error or "") + assert TOKEN not in caplog.text and "SECRET" not in caplog.text + + +# --- /users/@me/guilds pagination -------------------------------------------- + +def _guild_pages(d, total): + ids = [str(10**17 + i) for i in range(total)] + + def _page(request): + after = request.url.params.get("after") + limit = int(request.url.params["limit"]) + rest_ids = [i for i in ids if after is None or int(i) > int(after)] + return httpx.Response(200, json=[{"id": i} for i in rest_ids[:limit]]) + d.routes[("GET", "/users/@me/guilds")] = _page + return ids + + +async def test_guild_list_follows_after_pages(discord): + ids = _guild_pages(discord, 450) + r = await rest.list_guild_ids(TOKEN) + assert r.ok and r.data == ids + pages = [c for c in discord.calls if c.url.path.endswith("/users/@me/guilds")] + assert len(pages) == 3 + assert pages[0].url.params.get("after") is None + assert pages[1].url.params["after"] == ids[199] and pages[2].url.params["after"] == ids[399] + + +async def test_guild_list_exact_multiple_of_page_size_stops_on_empty_page(discord): + ids = _guild_pages(discord, 400) + r = await rest.list_guild_ids(TOKEN) + assert r.data == ids and len(discord.calls) == 3 + + +async def test_failed_page_is_unknown_not_a_truncated_list(discord): + ids = [str(10**17 + i) for i in range(200)] + + def _page(request): + if request.url.params.get("after"): + return httpx.Response(502) + return httpx.Response(200, json=[{"id": i} for i in ids]) + discord.routes[("GET", "/users/@me/guilds")] = _page + r = await rest.list_guild_ids(TOKEN) + assert not r.ok and r.unknown and r.data is None + + +async def test_inspect_token_reports_all_servers(discord): + ids = _guild_pages(discord, 250) + discord.routes[("GET", "/users/@me")] = httpx.Response(200, json={"id": BOT, "username": "DRB Radio 1"}) + discord.routes[("GET", "/oauth2/applications/@me")] = httpx.Response(200, json={"id": "777"}) + info = await rest.inspect_token(TOKEN) + assert info["valid"] is True and info["guild_ids"] == ids + assert info["bot_user_id"] == BOT and info["bot_name"] == "DRB Radio 1" and info["application_id"] == "777" + + +async def test_inspect_token_401_raises_and_5xx_is_unknown(discord): + discord.routes[("GET", "/users/@me")] = httpx.Response(401, json={"code": 0}) + with pytest.raises(rest.TokenRejected): + await rest.inspect_token(TOKEN) + discord.routes[("GET", "/users/@me")] = httpx.Response(503) + info = await rest.inspect_token(TOKEN) + assert info["valid"] is None and info["guild_ids"] is None and "503" in info["check_error"] + + +async def test_refresh_membership_writes_full_list_and_drops_cache(discord, fs): + doc = _token_doc(fs, guilds=["1"]) + ids = _guild_pages(discord, 201) + discord.routes[("GET", "/users/@me")] = httpx.Response(200, json={"id": BOT, "username": "r"}) + rest.access_cache._entries[("t1", G, "snapshot")] = (10**9, "cached") + assert await rest.refresh_membership(["t1"]) == 1 + assert fs.get("bot_tokens", "t1")["guild_ids"] == ids + assert rest.access_cache.peek(("t1", G, "snapshot")) is None + assert doc["token"] == fs.get("bot_tokens", "t1")["token"] + + +# --- rate limits --------------------------------------------------------------- + +async def test_short_429_is_retried_once(discord): + seq = [httpx.Response(429, json={"retry_after": 1.5, "global": False}), + httpx.Response(200, json={"id": BOT})] + discord.routes[("GET", "/users/@me")] = lambda req: seq.pop(0) + r = await rest.request(TOKEN, "GET", "/users/@me") + assert r.ok and discord.slept == [1.5] and len(discord.calls) == 2 + + +async def test_second_429_or_long_retry_after_is_unknown(discord): + discord.routes[("GET", "/users/@me")] = httpx.Response(429, json={"retry_after": 1.0}) + r = await rest.request(TOKEN, "GET", "/users/@me") + assert r.status == 429 and r.unknown and len(discord.calls) == 2 + + discord.calls.clear() + discord.slept.clear() + discord.routes[("GET", "/users/@me")] = httpx.Response(429, json={"retry_after": 30}) + r = await rest.request(TOKEN, "GET", "/users/@me") + assert r.unknown and len(discord.calls) == 1 and discord.slept == [] + + +async def test_global_429_pauses_only_that_token(discord): + discord.routes[("GET", "/users/@me")] = httpx.Response( + 429, json={"retry_after": 60, "global": True}, headers={"X-RateLimit-Global": "true"}) + assert (await rest.request(TOKEN, "GET", "/users/@me")).unknown + discord.routes[("GET", "/users/@me")] = httpx.Response(200, json={"id": BOT}) + r = await rest.request(TOKEN, "GET", "/users/@me") + assert r.status is None and "rate limited" in r.error and len(discord.calls) == 1 + assert (await rest.request("other-token", "GET", "/users/@me")).ok + discord.clock.t += 61 + assert (await rest.request(TOKEN, "GET", "/users/@me")).ok + + +async def test_exhausted_bucket_is_per_token_and_guild(discord): + other = "200000000000000002" + hdrs = {"X-RateLimit-Bucket": "abc", "X-RateLimit-Remaining": "0", "X-RateLimit-Reset-After": "60"} + discord.routes[("GET", f"/guilds/{G}/channels")] = httpx.Response(200, json=[], headers=hdrs) + discord.routes[("GET", f"/guilds/{other}/channels")] = httpx.Response(200, json=[], headers=hdrs) + assert (await rest.request(TOKEN, "GET", f"/guilds/{G}/channels")).ok + blocked = await rest.request(TOKEN, "GET", f"/guilds/{G}/channels") + assert blocked.status is None and len(discord.calls) == 1 # not sent + assert (await rest.request(TOKEN, "GET", f"/guilds/{other}/channels")).ok # other guild + assert (await rest.request("tok-2", "GET", f"/guilds/{G}/channels")).ok # other token + discord.clock.t += 61 + assert (await rest.request(TOKEN, "GET", f"/guilds/{G}/channels")).ok + + +async def test_short_bucket_wait_is_slept_not_refused(discord): + hdrs = {"X-RateLimit-Bucket": "abc", "X-RateLimit-Remaining": "0", "X-RateLimit-Reset-After": "2"} + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(200, json={}, headers=hdrs) + await rest.request(TOKEN, "GET", f"/guilds/{G}") + assert (await rest.request(TOKEN, "GET", f"/guilds/{G}")).ok and discord.slept == [2.0] + + +# --- invalid-request breaker ----------------------------------------------------- + +async def test_breaker_stops_all_requests_and_drains(discord, monkeypatch): + monkeypatch.setattr(settings, "discord_rest_invalid_budget", 3) + discord.routes[("GET", "/users/@me")] = httpx.Response(401) + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response( + 429, json={"retry_after": 99}, headers={"X-RateLimit-Scope": "shared"}) + for _ in range(4): # shared 429s never count + await rest.request("t-x", "GET", f"/guilds/{G}") + for _ in range(3): + assert (await rest.request(TOKEN, "GET", "/users/@me")).status == 401 + sent = len(discord.calls) + r = await rest.request("a-fresh-token", "GET", "/users/@me") + assert r.status is None and "paused" in r.error and len(discord.calls) == sent + discord.clock.t += settings.discord_rest_invalid_window_seconds + 1 + assert (await rest.request("a-fresh-token", "GET", "/users/@me")).status == 401 + + +async def test_403_and_unshared_429_count_toward_breaker(discord, monkeypatch): + monkeypatch.setattr(settings, "discord_rest_invalid_budget", 2) + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(403, json={"code": 50013}) + discord.routes[("GET", "/users/@me")] = httpx.Response(429, json={"retry_after": 99}) + await rest.request(TOKEN, "GET", f"/guilds/{G}") + await rest.request("t2", "GET", "/users/@me") + assert (await rest.request("t3", "GET", "/users/@me")).error.startswith("Discord REST paused") + + +# --- per-(token, guild) snapshot + cache --------------------------------------------- + +async def test_snapshot_reads_three_endpoints_and_caches(discord, fs, monkeypatch): + _guild_routes(discord) + tok = _token_doc(fs) + snap = await rest.guild_snapshot(tok, G) + assert snap.status == rest.OK and snap.channels == [{"id": "5", "type": 2}] and snap.member == {"roles": []} + assert sorted(discord.paths()) == sorted([f"/guilds/{G}", f"/guilds/{G}/members/{BOT}", f"/guilds/{G}/channels"]) + assert (await rest.guild_snapshot(tok, G)) is snap and len(discord.calls) == 3 + discord.clock.t += settings.discord_access_ttl_seconds + 1 + await rest.guild_snapshot(tok, G) + assert len(discord.calls) == 6 + assert rest.invalidate(guild_id=G) == 1 + await rest.guild_snapshot(tok, G) + assert len(discord.calls) == 9 + await rest.guild_snapshot(tok, G, refresh=True) + assert len(discord.calls) == 12 + + +async def test_snapshot_learns_and_stores_bot_user_id(discord, fs): + _guild_routes(discord) + discord.routes[("GET", "/users/@me")] = httpx.Response(200, json={"id": BOT}) + snap = await rest.guild_snapshot(_token_doc(fs, bot_user_id=None), G) + assert snap.status == rest.OK and snap.bot_user_id == BOT + assert fs.get("bot_tokens", "t1")["bot_user_id"] == BOT + + +async def test_snapshot_401_marks_token_invalid_and_is_not_retried(discord, fs): + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(401, json={"code": 0}) + tok = _token_doc(fs) + snap = await rest.guild_snapshot(tok, G) + assert snap.status == rest.INVALID and fs.get("bot_tokens", "t1")["valid"] is False + assert len(discord.calls) == 1 + discord.clock.t += settings.discord_access_negative_ttl_seconds - 1 + assert (await rest.guild_snapshot(tok, G)).status == rest.INVALID and len(discord.calls) == 1 + # an invalid token doc is never sent to Discord at all + rest.invalidate() + assert (await rest.guild_snapshot({**tok, "valid": False}, G)).status == rest.INVALID + assert len(discord.calls) == 1 + + +@pytest.mark.parametrize("code,path", [(10004, f"/guilds/{G}"), (10007, f"/guilds/{G}/members/{BOT}")]) +async def test_unknown_guild_or_member_drops_membership(discord, fs, code, path): + _guild_routes(discord) + discord.routes[("GET", path)] = httpx.Response(404, json={"code": code}) + snap = await rest.guild_snapshot(_token_doc(fs, guilds=[G, "7"]), G) + assert snap.status == rest.NOT_MEMBER + assert fs.get("bot_tokens", "t1")["guild_ids"] == ["7"] + + +async def test_missing_access_confirms_with_one_membership_read(discord, fs): + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(403, json={"code": 50001}) + discord.routes[("GET", "/users/@me/guilds")] = httpx.Response(200, json=[{"id": "7"}]) + snap = await rest.guild_snapshot(_token_doc(fs, guilds=[G, "7"]), G) + assert snap.status == rest.NOT_MEMBER + assert fs.get("bot_tokens", "t1")["guild_ids"] == ["7"] + assert discord.paths().count("/users/@me/guilds") == 1 + + +async def test_missing_access_alone_never_changes_membership(discord, fs): + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(403, json={"code": 50001}) + discord.routes[("GET", "/users/@me/guilds")] = httpx.Response(200, json=[{"id": G}]) + snap = await rest.guild_snapshot(_token_doc(fs, guilds=[G]), G) + assert snap.status == rest.UNKNOWN and fs.get("bot_tokens", "t1")["guild_ids"] == [G] + + rest.invalidate() + discord.routes[("GET", "/users/@me/guilds")] = httpx.Response(503) # confirmation unknown too + snap = await rest.guild_snapshot(_token_doc(fs, guilds=[G]), G) + assert snap.status == rest.UNKNOWN and fs.get("bot_tokens", "t1")["guild_ids"] == [G] + + +async def test_non_discord_failure_is_unknown_not_raised(discord, fs, monkeypatch): + # A Firestore error while classifying (here: dropping membership) must + # fail open — S3's start_stream relies on the snapshot never raising. + _guild_routes(discord) + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(404, json={"code": 10004}) + tok = _token_doc(fs) + + async def _boom(*a, **kw): + raise RuntimeError("firestore down") + monkeypatch.setattr(rest, "_drop_membership", _boom) + snap = await rest.guild_snapshot(tok, G) + assert snap.status == rest.UNKNOWN and "RuntimeError" in snap.error + assert TOKEN not in snap.error + + +async def test_unknown_result_is_cached_briefly(discord, fs): + discord.routes[("GET", f"/guilds/{G}")] = httpx.Response(503) + tok = _token_doc(fs) + assert (await rest.guild_snapshot(tok, G)).status == rest.UNKNOWN + await rest.guild_snapshot(tok, G) + assert len(discord.calls) == 1 + discord.clock.t += settings.discord_access_unknown_ttl_seconds + 1 + await rest.guild_snapshot(tok, G) + assert len(discord.calls) == 2 + + +async def test_concurrent_loads_share_one_flight(discord, fs, monkeypatch): + gate = asyncio.Event() + _guild_routes(discord) + tok = _token_doc(fs) + real = rest._load_snapshot + loads = [] + + async def _slow(token_doc, guild_id): + loads.append(guild_id) + await gate.wait() + return await real(token_doc, guild_id) + monkeypatch.setattr(rest, "_load_snapshot", _slow) + tasks = [asyncio.create_task(rest.guild_snapshot(tok, G)) for _ in range(5)] + await asyncio.sleep(0) + gate.set() + snaps = await asyncio.gather(*tasks) + assert loads == [G] and len(discord.calls) == 3 + assert all(s is snaps[0] for s in snaps) + + +async def test_invalidate_during_a_load_is_not_undone(discord, fs, monkeypatch): + gate = asyncio.Event() + _guild_routes(discord) + tok = _token_doc(fs) + real = rest._load_snapshot + + async def _slow(token_doc, guild_id): + await gate.wait() + return await real(token_doc, guild_id) + monkeypatch.setattr(rest, "_load_snapshot", _slow) + task = asyncio.create_task(rest.guild_snapshot(tok, G)) + await asyncio.sleep(0) + rest.invalidate(token_id="t1") + gate.set() + assert (await task).status == rest.OK + assert rest.access_cache.peek(("t1", G, "snapshot")) is None + + +def test_invalidate_scopes(): + c = rest.AccessCache() + for k in [("t1", "g1", "s"), ("t1", "g2", "s"), ("t2", "g1", "s")]: + c._entries[k] = (10**12, k) + assert c.invalidate(token_id="t1", guild_id="g1") == 1 + assert c.invalidate(guild_id="g1") == 1 + assert c.invalidate(token_id="t1") == 1 + assert c.invalidate() == 0 + + +def test_ttl_for_outcomes(): + assert rest.ttl_for(rest.OK) == 120 + assert rest.ttl_for(rest.DENIED) == rest.ttl_for(rest.NOT_MEMBER) == rest.ttl_for(rest.INVALID) == 600 + assert rest.ttl_for(rest.UNKNOWN) == 30 + + +def test_225_settings_defaults(): + assert settings.discord_perm_check_enabled is True + assert settings.discord_access_ttl_seconds == 120 + assert settings.discord_rest_invalid_budget == 4000 + assert settings.discord_oauth_redirect_uri == "" + assert settings.discord_pool_bot_guild_cap == 90 + + +# --- leave_guild on the shared client --------------------------------------------- + +async def test_leave_guild_drops_cache_and_reports_unknown(discord, fs): + _token_doc(fs, "t1", guilds=[G]) + fs.put("bot_tokens", "t2", {"token_id": "t2", "name": "R2", "token": "tok-two", "org_id": None, + "guild_ids": None, "valid": True}) + rest.access_cache._entries[("t1", G, "snapshot")] = (10**12, "x") + + def _delete(request): + if request.headers["Authorization"] == "Bot tok-two": + raise httpx.ReadTimeout("timed out", request=request) + return httpx.Response(204) + discord.routes[("DELETE", f"/users/@me/guilds/{G}")] = _delete + rows = {r["token_id"]: r for r in await rest.leave_guild(G, None)} + assert rows["t1"]["left"] is True and rows["t2"]["left"] is False + assert "ReadTimeout" in rows["t2"]["error"] and "tok-two" not in rows["t2"]["error"] + assert fs.get("bot_tokens", "t1")["guild_ids"] == [] + assert rest.access_cache.peek(("t1", G, "snapshot")) is None diff --git a/drb-c2-core/tests/test_discord_web_picker.py b/drb-c2-core/tests/test_discord_web_picker.py index 24eb44c..708c0a9 100644 --- a/drb-c2-core/tests/test_discord_web_picker.py +++ b/drb-c2-core/tests/test_discord_web_picker.py @@ -14,6 +14,7 @@ from fastapi import HTTPException from app.config import settings from app.internal import discord_leases as leases +from app.internal import discord_rest from app.internal import discord_streams as streams from app.models import CommandPayload from app.routers import discord_web as web @@ -43,7 +44,7 @@ def fs(monkeypatch): async def _no_discord(token_ids): return 0 - monkeypatch.setattr(leases, "refresh_membership", _no_discord) + monkeypatch.setattr(discord_rest, "refresh_membership", _no_discord) monkeypatch.setattr(settings, "discord_unlinked_guild_org_id", "") fake.published = published