refactor: move pop replay cache to server

src/authsome/identity/__init__.py

@@ -20,7 +20,6 @@ from authsome.identity.proof import (
     POP_AUTH_SCHEME,
     ProofClaims,
     ProofValidationError,
-    ReplayCache,
     create_proof_jwt,
     validate_proof_jwt,
 )
@@ -36,7 +35,6 @@ __all__ = [
     "POP_AUTH_SCHEME",
     "ProofClaims",
     "ProofValidationError",
-    "ReplayCache",
     "create_identity_material",
     "create_proof_jwt",
     "generate_handle",

src/authsome/identity/proof.py

@@ -28,20 +28,6 @@ class ProofClaims:
     jwt_id: str
 
 
-class ReplayCache:
-    """Small in-memory jti replay cache."""
-
-    def __init__(self) -> None:
-        self._seen: dict[str, int] = {}
-
-    def check_and_store(self, jti: str, exp: int) -> None:
-        now = int(time.time())
-        self._seen = {key: value for key, value in self._seen.items() if value > now}
-        if jti in self._seen:
-            raise ProofValidationError("Proof JWT was already used")
-        self._seen[jti] = exp
-
-
 def body_sha256(body: bytes) -> str:
     return hashlib.sha256(body).hexdigest()
 
@@ -78,7 +64,6 @@ def validate_proof_jwt(  # noqa: PLR0913
     method: str,
     path_query: str,
     body: bytes,
-    replay_cache: ReplayCache | None = None,
     audience: str = DEFAULT_AUDIENCE,
 ) -> ProofClaims:
     unverified = _unverified_claims(token)
@@ -101,8 +86,6 @@ def validate_proof_jwt(  # noqa: PLR0913
     exp = claims.get("exp")
     if not isinstance(exp, int):
         raise ProofValidationError("Proof JWT exp must be an integer")
-    if replay_cache is not None:
-        replay_cache.check_and_store(jwt_id, exp)
     return ProofClaims(issuer=issuer, subject=subject, expires_at=exp, jwt_id=jwt_id)
 
 

src/authsome/server/app.py

@@ -9,7 +9,6 @@ from fastapi.staticfiles import StaticFiles
 
 from authsome.auth.sessions import AuthSessionStore
 from authsome.errors import AuthsomeError
-from authsome.identity.proof import ReplayCache
 from authsome.server.analytics import init_posthog, shutdown_posthog
 from authsome.server.dependencies import (
     create_account_auth_service,
@@ -21,6 +20,7 @@ from authsome.server.dependencies import (
     load_server_config,
 )
 from authsome.server.provider_repository import ProviderRepository
+from authsome.server.replay_cache import MemoryReplayCache
 from authsome.server.routes.audit import router as audit_router
 from authsome.server.routes.auth import browser_router as auth_browser_router
 from authsome.server.routes.auth import router as auth_router
@@ -46,7 +46,7 @@ async def lifespan(app: FastAPI):
     app.state.vault = await create_vault(app.state.store.home)
     app.state.auth_sessions = AuthSessionStore()
     app.state.ui_sessions = UiSessionStore(load_ui_session_signing_secret(app.state.store.home))
-    app.state.proof_replay_cache = ReplayCache()
+    app.state.proof_replay_cache = MemoryReplayCache()
     app.state.provider_repository = ProviderRepository(app.state.store.provider_definitions)
     app.state.account_auth_service = create_account_auth_service(app.state.store, app.state.ui_sessions)
     app.state.server_base_url = get_server_base_url()

src/authsome/server/replay_cache.py

@@ -0,0 +1,41 @@
+"""Server-owned PoP replay caches."""
+
+import time
+from typing import Protocol
+
+from authsome.identity.proof import ProofValidationError
+
+
+class ReplayCache(Protocol):
+    async def check_and_store(self, jti: str, exp: int) -> None:
+        """Store a JTI until expiry or raise when it has already been used."""
+
+
+class MemoryReplayCache:
+    """Process-local replay cache for local development and tests."""
+
+    def __init__(self) -> None:
+        self._seen: dict[str, int] = {}
+
+    async def check_and_store(self, jti: str, exp: int) -> None:
+        now = int(time.time())
+        self._seen = {key: value for key, value in self._seen.items() if value > now}
+        if jti in self._seen:
+            raise ProofValidationError("Proof JWT was already used")
+        if exp > now:
+            self._seen[jti] = exp
+
+
+class RedisReplayCache:
+    """Redis-backed replay cache shared across server replicas."""
+
+    def __init__(self, client, *, key_prefix: str = "authsome:pop:jti") -> None:
+        self._client = client
+        self._key_prefix = key_prefix.rstrip(":")
+
+    async def check_and_store(self, jti: str, exp: int) -> None:
+        ttl = max(exp - int(time.time()), 1)
+        key = f"{self._key_prefix}:{jti}"
+        stored = await self._client.set(key, "1", ex=ttl, nx=True)
+        if not stored:
+            raise ProofValidationError("Proof JWT was already used")

src/authsome/server/routes/_deps.py

@@ -143,11 +143,12 @@ async def verify_pop_caller(request: Request) -> ResolvedOwnership:
             method=request.method,
             path_query=path_query,
             body=body,
-            replay_cache=request.app.state.proof_replay_cache,
         )
     except (ProofValidationError, ValueError) as exc:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=str(exc)) from exc
 
+    await request.app.state.proof_replay_cache.check_and_store(claims.jwt_id, claims.expires_at)
+
     registration = await request.app.state.store.identity_registry.resolve(claims.subject)
     if registration is None:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Unknown identity handle")

tests/identity/test_proof.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import pytest
 
 from authsome.cli.identity import RuntimeIdentity
-from authsome.identity.proof import ReplayCache, create_proof_jwt, validate_proof_jwt
+from authsome.identity.proof import create_proof_jwt, validate_proof_jwt
 
 
 def _token(tmp_path: Path, *, method: str = "POST", path: str = "/connections", body: bytes = b"{}") -> str:
@@ -42,10 +42,10 @@ def test_validate_proof_jwt_rejects_wrong_body(tmp_path: Path) -> None:
         validate_proof_jwt(token=token, method="POST", path_query="/connections", body=b'{"x":1}')
 
 
-def test_validate_proof_jwt_rejects_replay(tmp_path: Path) -> None:
+def test_validate_proof_jwt_returns_jti_for_server_replay_check(tmp_path: Path) -> None:
     token = _token(tmp_path)
-    cache = ReplayCache()
 
-    validate_proof_jwt(token=token, method="POST", path_query="/connections", body=b"{}", replay_cache=cache)
-    with pytest.raises(ValueError, match="already used"):
-        validate_proof_jwt(token=token, method="POST", path_query="/connections", body=b"{}", replay_cache=cache)
+    claims = validate_proof_jwt(token=token, method="POST", path_query="/connections", body=b"{}")
+
+    assert claims.jwt_id
+    assert claims.expires_at > 0

tests/server/test_replay_cache.py

@@ -0,0 +1,25 @@
+import time
+
+import pytest
+
+from authsome.identity.proof import ProofValidationError
+from authsome.server.replay_cache import MemoryReplayCache
+
+
+@pytest.mark.asyncio
+async def test_memory_replay_cache_rejects_duplicate_jti() -> None:
+    cache = MemoryReplayCache()
+    exp = int(time.time()) + 60
+
+    await cache.check_and_store("jti-1", exp)
+
+    with pytest.raises(ProofValidationError, match="already used"):
+        await cache.check_and_store("jti-1", exp)
+
+
+@pytest.mark.asyncio
+async def test_memory_replay_cache_drops_expired_entries() -> None:
+    cache = MemoryReplayCache()
+
+    await cache.check_and_store("jti-1", int(time.time()) - 1)
+    await cache.check_and_store("jti-1", int(time.time()) + 60)