logan/server-26
1
0
Fork 0
Code Issues 84 Pull Requests Actions Packages Projects Releases Wiki Activity

Discord REST client, cache and #225 settings (server-26#225 S2)

Browse Source 
New app/internal/discord_rest.py: one shared httpx client with the
DiscordBot User-Agent, per-token rate-limit buckets (X-RateLimit-Bucket +
guild), one retry on a 429 of <=3s, global-429 pause per token, and a
process-wide breaker at 4000 401/403/unshared-429 responses per 10 min so
C2 can never get the shared egress IP banned. Tokens are never logged.

inspect_token / refresh_membership / leave_guild / TokenRejected move off
discord_leases onto it. /users/@me/guilds now follows after= pages (it read
one page of 200 and silently dropped the rest); a failed page is unknown,
never a truncated list.

guild_snapshot(token_doc, guild_id): cached, single-flight read of
/guilds/{g}, the bot's own member and /channels for S3's permission check.
401 marks the token invalid; 10004/10007 drop the server from guild_ids;
50001 is confirmed with one membership read and never changes guild_ids
on its own. TTLs: ok 120s, not-member/invalid 600s, unknown 30s.

config.py gets every setting the #225 slices use (perm-check kill switch,
TTLs, breaker budget, OAuth redirect URI, discord_pool_bot_guild_cap=90).
No behaviour change beyond pagination, User-Agent and the breaker.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Logan Cusano
2026-10-05 01:57:27 -04:00
co-authored by Claude Opus 5.5
parent c4e1f9aa61
commit 2b3b875cdc
11 changed files with 1139 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
drb-c2-core/app/config.py
+38
View File
@@ -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"
drb-c2-core/app/internal/discord_leases.py
+3 -103
View File
@@ -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
drb-c2-core/app/internal/discord_rest.py
+645
View File
@@ -0,0 +1,645 @@
"""
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)
key = (str(token_doc.get("token_id")), guild_id, "snapshot")
return await access_cache.get_or_load(
key, lambda: _load_snapshot(token_doc, guild_id),
lambda s: ttl_for(s.status), refresh=refresh)
drb-c2-core/app/internal/discord_streams.py
+3 -2
View File
@@ -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 [])}
drb-c2-core/app/routers/discord.py
+2 -1
View File
@@ -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,
drb-c2-core/app/routers/tokens.py
+4 -3
View File
@@ -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 {})
drb-c2-core/tests/test_discord_leases.py
+2 -1
View File
@@ -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
drb-c2-core/tests/test_discord_multi_stream.py
+2 -1
View File
@@ -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)
drb-c2-core/tests/test_discord_remove_and_logging.py
+11 -21
View File
@@ -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
drb-c2-core/tests/test_discord_rest.py
+427
View File
@@ -0,0 +1,427 @@
"""
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_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
drb-c2-core/tests/test_discord_web_picker.py
+2 -1
View File
@@ -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