From 56045a96e6f35a21d06d1bf7f32339dee97cea91 Mon Sep 17 00:00:00 2001 From: tri-pti Date: Thu, 23 Jul 2026 11:27:59 +0000 Subject: [PATCH] Add Redis-based token denylist implementation --- .env.example | 8 ++ backend/app/token_denylist.py | 117 ++++++++++++++++--- backend/requirements.txt | 1 + backend/tests/test_token_denylist.py | 162 +++++++++++++++++++++++++++ deploy/k8s/deployment.example.yaml | 10 ++ docker-compose.yml | 10 ++ 6 files changed, 291 insertions(+), 17 deletions(-) create mode 100644 backend/tests/test_token_denylist.py diff --git a/.env.example b/.env.example index adef860b..5ecca913 100644 --- a/.env.example +++ b/.env.example @@ -37,6 +37,14 @@ ACCESS_TOKEN_MINUTES=720 # ── Database ─────────────────────────────────────────────────── DATABASE_URL=sqlite:///./assistant.db +# ── Token Denylist ──────────────────────────────────────────── +# Required when running more than one backend replica/worker (e.g. the +# example Kubernetes manifest's replicas: 2, or uvicorn --workers N): without +# it, a token revoked via /logout on one replica is still accepted by the +# others until it naturally expires. Leave blank for single-process local +# dev — the denylist then falls back to an in-memory store. +# REDIS_URL=redis://localhost:6379/0 + # ── Observability (Prometheus metrics + health probes) ───────── # Set METRICS_ENABLED=false to disable /metrics and skip the metrics # middleware entirely (zero overhead). diff --git a/backend/app/token_denylist.py b/backend/app/token_denylist.py index f1f72826..daa74278 100644 --- a/backend/app/token_denylist.py +++ b/backend/app/token_denylist.py @@ -1,4 +1,4 @@ -"""In-memory revocation store for JWT identifiers (``jti``). +"""Revocation store for JWT identifiers (``jti``). Access tokens are revoked by recording their ``jti`` here. ``get_current_user`` consults this store on every authenticated request and rejects any token whose @@ -6,27 +6,46 @@ — for example, a token a user explicitly invalidated by logging out can no longer be reused even though its signature and ``exp`` are still valid. -Design notes: - -* Each entry carries the revoked token's own expiry, so a ``jti`` only needs to - be remembered until the moment the token would have expired anyway. After that - the standard signature/``exp`` check rejects it for free, so the entry is - purged to keep memory bounded. -* The store is process-local and guarded by a lock so it is safe to call from - the threadpool FastAPI uses for synchronous route handlers. -* It is intentionally a thin seam: the same ``revoke`` / ``is_revoked`` API can - later be backed by Redis (the project already exposes ``settings.redis_url``) - without touching any call sites. +Two implementations share the same ``revoke`` / ``is_revoked`` / ``clear`` +interface: + +* ``InMemoryTokenDenylist`` keeps revocations in a process-local dict guarded + by a lock. It's the only option in single-process/single-replica setups + (e.g. local dev via ``docker-compose.yml``), but a revocation made on one + worker/replica is invisible to the others. +* ``RedisTokenDenylist`` stores revocations in Redis (``SETEX``/``EXISTS``), + so a revocation is immediately visible to every replica behind the load + balancer — required once you run more than one backend process, as the + example Kubernetes manifest does (``deploy/k8s/deployment.example.yaml``). + +The module picks between them at import time based on ``settings.redis_url``, +falling back to the in-memory store (with a warning) if Redis is configured +but unreachable, so a Redis outage degrades the deployment rather than +breaking every login. """ from __future__ import annotations +import logging import threading import time +from typing import Protocol + +from .config import settings + +logger = logging.getLogger(__name__) + + +class SupportsTokenDenylist(Protocol): + def revoke(self, jti: str, expires_at: float) -> None: ... + + def is_revoked(self, jti: str) -> bool: ... + def clear(self) -> None: ... -class TokenDenylist: - """A TTL-bounded set of revoked JWT ``jti`` values.""" + +class InMemoryTokenDenylist: + """A TTL-bounded set of revoked JWT ``jti`` values, local to this process.""" def __init__(self) -> None: # Maps jti -> epoch-seconds expiry of the revoked token. @@ -75,6 +94,70 @@ def clear(self) -> None: self._revoked.clear() -# Process-wide singleton. Swap for a Redis-backed implementation behind the same -# interface to share revocations across workers. -token_denylist = TokenDenylist() +class RedisTokenDenylist: + """Redis-backed revoked-``jti`` store, shared across all replicas/workers. + + A revoked ``jti`` is stored as a key with a TTL matching the remaining + lifetime of the token it belongs to, so Redis expires the bookkeeping + entry for free once the token would have expired anyway — mirroring the + self-pruning behaviour of ``InMemoryTokenDenylist``. + """ + + _KEY_PREFIX = "token_denylist:" + # Without an explicit timeout, a host that silently drops packets (a + # common failure mode for k8s NetworkPolicies/security groups, as opposed + # to one that actively refuses the connection) falls back to the OS's + # default TCP connect timeout - tens of seconds to minutes - which would + # block application startup and defeat the "fail fast" behaviour below. + _CONNECT_TIMEOUT_SECONDS = 5 + + def __init__(self, redis_url: str) -> None: + import redis + + self._client = redis.Redis.from_url( + redis_url, + decode_responses=True, + socket_connect_timeout=self._CONNECT_TIMEOUT_SECONDS, + socket_timeout=self._CONNECT_TIMEOUT_SECONDS, + ) + # Fail fast at construction time so callers can fall back to the + # in-memory store instead of discovering the outage on first request. + self._client.ping() + + def revoke(self, jti: str, expires_at: float) -> None: + if not jti: + return + ttl_seconds = max(1, int(expires_at - time.time())) + self._client.set(self._KEY_PREFIX + jti, "1", ex=ttl_seconds) + + def is_revoked(self, jti: str) -> bool: + if not jti: + return False + return bool(self._client.exists(self._KEY_PREFIX + jti)) + + def clear(self) -> None: + """Forget all revoked tokens. Primarily a test helper.""" + keys = list(self._client.scan_iter(f"{self._KEY_PREFIX}*")) + if keys: + self._client.delete(*keys) + + +def _build_token_denylist() -> SupportsTokenDenylist: + if not settings.redis_url: + return InMemoryTokenDenylist() + try: + return RedisTokenDenylist(settings.redis_url) + except Exception: + logger.warning( + "REDIS_URL is set but Redis is unreachable; falling back to an " + "in-memory token denylist. Revocations will NOT be shared across " + "replicas until Redis is reachable.", + exc_info=True, + ) + return InMemoryTokenDenylist() + + +# Process-wide singleton, backed by Redis when settings.redis_url is set so +# revocations are visible to every replica/worker, otherwise an in-memory +# fallback for single-process setups. +token_denylist: SupportsTokenDenylist = _build_token_denylist() diff --git a/backend/requirements.txt b/backend/requirements.txt index aea3337d..225294dd 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -10,6 +10,7 @@ python-multipart>=0.0.9 pytest>=8.0.0 pytest-asyncio>=0.23.0 apscheduler>=3.10.0 +redis>=5.0.0 python-magic-bin>=0.4.14; platform_system == "Windows" python-magic>=0.4.27; platform_system != "Windows" prometheus-client>=0.20.0 diff --git a/backend/tests/test_token_denylist.py b/backend/tests/test_token_denylist.py new file mode 100644 index 00000000..88f1279a --- /dev/null +++ b/backend/tests/test_token_denylist.py @@ -0,0 +1,162 @@ +import time + +import pytest +import redis as redis_lib +from app.config import settings +from app.token_denylist import ( + InMemoryTokenDenylist, + RedisTokenDenylist, + _build_token_denylist, +) + +# A separate DB index from whatever local dev might use, so these tests never +# collide with data a developer is inspecting by hand. +_TEST_REDIS_URL = "redis://localhost:6379/15" + + +def _redis_available() -> bool: + try: + redis_lib.Redis.from_url(_TEST_REDIS_URL, socket_connect_timeout=1).ping() + return True + except Exception: + return False + + +requires_redis = pytest.mark.skipif( + not _redis_available(), reason="no local Redis reachable on db 15" +) + + +def test_in_memory_revoke_then_is_revoked(): + denylist = InMemoryTokenDenylist() + + denylist.revoke("jti-1", expires_at=time.time() + 60) + + assert denylist.is_revoked("jti-1") is True + + +def test_in_memory_unrevoked_jti_is_not_revoked(): + denylist = InMemoryTokenDenylist() + + assert denylist.is_revoked("never-seen") is False + + +def test_in_memory_expired_entry_is_treated_as_not_revoked(): + denylist = InMemoryTokenDenylist() + denylist.revoke("jti-expired", expires_at=time.time() - 1) + + assert denylist.is_revoked("jti-expired") is False + # The stale bookkeeping entry should also have been purged. + assert "jti-expired" not in denylist._revoked + + +def test_in_memory_falsy_jti_is_ignored(): + denylist = InMemoryTokenDenylist() + + denylist.revoke("", expires_at=time.time() + 60) + + assert denylist.is_revoked("") is False + assert denylist._revoked == {} + + +def test_in_memory_clear_forgets_all_revocations(): + denylist = InMemoryTokenDenylist() + denylist.revoke("jti-1", expires_at=time.time() + 60) + + denylist.clear() + + assert denylist.is_revoked("jti-1") is False + + +@requires_redis +def test_redis_revoke_then_is_revoked(): + denylist = RedisTokenDenylist(_TEST_REDIS_URL) + denylist.clear() + + denylist.revoke("jti-1", expires_at=time.time() + 60) + + assert denylist.is_revoked("jti-1") is True + denylist.clear() + + +@requires_redis +def test_redis_unrevoked_jti_is_not_revoked(): + denylist = RedisTokenDenylist(_TEST_REDIS_URL) + denylist.clear() + + assert denylist.is_revoked("never-seen") is False + + +@requires_redis +def test_redis_key_ttl_matches_remaining_token_lifetime(): + denylist = RedisTokenDenylist(_TEST_REDIS_URL) + denylist.clear() + + denylist.revoke("jti-1", expires_at=time.time() + 30) + + ttl = denylist._client.ttl(denylist._KEY_PREFIX + "jti-1") + assert 0 < ttl <= 30 + denylist.clear() + + +@requires_redis +def test_redis_falsy_jti_is_ignored(): + denylist = RedisTokenDenylist(_TEST_REDIS_URL) + denylist.clear() + + denylist.revoke("", expires_at=time.time() + 60) + + assert denylist.is_revoked("") is False + + +@requires_redis +def test_redis_clear_forgets_all_revocations(): + denylist = RedisTokenDenylist(_TEST_REDIS_URL) + denylist.revoke("jti-1", expires_at=time.time() + 60) + + denylist.clear() + + assert denylist.is_revoked("jti-1") is False + + +@requires_redis +def test_redis_denylist_is_shared_across_instances(): + """Two separate clients pointed at the same Redis see each other's revocations, + which is the multi-replica scenario this store exists to fix.""" + pod_a = RedisTokenDenylist(_TEST_REDIS_URL) + pod_a.clear() + pod_b = RedisTokenDenylist(_TEST_REDIS_URL) + + pod_a.revoke("jti-1", expires_at=time.time() + 60) + + assert pod_b.is_revoked("jti-1") is True + pod_a.clear() + + +def test_build_token_denylist_uses_in_memory_when_redis_url_unset(monkeypatch): + monkeypatch.setattr(settings, "redis_url", None) + + denylist = _build_token_denylist() + + assert isinstance(denylist, InMemoryTokenDenylist) + + +@requires_redis +def test_build_token_denylist_uses_redis_when_reachable(monkeypatch): + monkeypatch.setattr(settings, "redis_url", _TEST_REDIS_URL) + + denylist = _build_token_denylist() + + assert isinstance(denylist, RedisTokenDenylist) + + +def test_build_token_denylist_falls_back_when_redis_unreachable(monkeypatch, caplog): + # Port 1 on localhost is not listening, so the connection is refused + # immediately rather than hanging for the configured connect timeout. + monkeypatch.setattr(settings, "redis_url", "redis://localhost:1/0") + + with caplog.at_level("WARNING"): + denylist = _build_token_denylist() + + assert isinstance(denylist, InMemoryTokenDenylist) + assert "falling back to an in-memory token denylist" in caplog.text diff --git a/deploy/k8s/deployment.example.yaml b/deploy/k8s/deployment.example.yaml index 252f4701..dcb69629 100644 --- a/deploy/k8s/deployment.example.yaml +++ b/deploy/k8s/deployment.example.yaml @@ -38,6 +38,16 @@ spec: value: "true" - name: RATE_LIMIT_PER_MINUTE value: "120" + # Required with replicas > 1 (as below): without it, each pod + # revokes tokens into its own process-local denylist, so a token + # revoked via /logout on one pod is still accepted by the others + # until it naturally expires. Point this at a Redis reachable + # from every pod (e.g. a managed Redis or an in-cluster Service). + - name: REDIS_URL + valueFrom: + secretKeyRef: + name: qyverixai-redis + key: url # Liveness — restarts the container when the process can no longer # respond. Must NOT depend on external services. livenessProbe: diff --git a/docker-compose.yml b/docker-compose.yml index d6a62f56..ead6d964 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -43,5 +43,15 @@ services: volumes: - postgres_data:/var/lib/postgresql/data + # Optional: only needed to exercise the Redis-backed token denylist locally + # (set REDIS_URL=redis://redis:6379/0 in .env to opt in). The backend falls + # back to an in-memory denylist when REDIS_URL is unset, which is fine for + # this single-replica compose setup, so nothing depends_on this service. + redis: + image: redis:7-alpine + container_name: ai-dev-redis + ports: + - "6379:6379" + volumes: postgres_data: