Build JWT auth with refresh-token rotation and Redis sessions

Loading diagram...

• Lines 1-5: Defines a Mermaid sequence diagram with four participants: Client, FastAPI (labeled as Auth Service), Redis (for token storage), and LLM (for inference).
• Lines 7-10: Depicts the login flow — the client sends credentials via POST /auth/login, FastAPI validates them and generates an RS256-signed JWT, stores the refresh token in Redis with a 7-day TTL, and returns both access and refresh tokens to the client.
• Lines 12-14: Shows an authenticated API call — the client sends a GET /conversations request with a Bearer access token, FastAPI decodes and verifies the JWT using the RS256 public key, and returns the conversations list with a 200 status.
• Lines 24-27: Shows an authenticated LLM inference request — the client sends a message via POST /conversations/{id}/message with the new access token, FastAPI validates the new JWT, forwards the prompt to the LLM inference service, and the LLM streams the response back to the client via Server-Sent Events (SSE).

1import jwt
2import uuid
3from datetime import datetime, timedelta, timezone
4from dataclasses import dataclass
5
6@dataclass
7class TokenPair:
8    access_token: str
9    refresh_token: str
10    token_family: str
11
12class JWTTokenService:
13    def __init__(self, private_key: str, public_key: str):
14        self._private_key = private_key
15        self._public_key = public_key
16        self.access_ttl = timedelta(minutes=15)
17        self.refresh_ttl = timedelta(days=7)
18
19    def create_token_pair(self, user_id: str, roles: list[str]) -> TokenPair:
20        family_id = str(uuid.uuid4())
21        now = datetime.now(timezone.utc)
22
23        access_payload = {
24            "sub": user_id,
25            "roles": roles,
26            "type": "access",
27            "iat": now,
28            "exp": now + self.access_ttl,
29            "jti": str(uuid.uuid4()),
30        }
31        refresh_payload = {
32            "sub": user_id,
33            "type": "refresh",
34            "family": family_id,
35            "iat": now,
36            "exp": now + self.refresh_ttl,
37            "jti": str(uuid.uuid4()),
38        }
39
40        access_token = jwt.encode(access_payload, self._private_key, algorithm="RS256")
41        refresh_token = jwt.encode(refresh_payload, self._private_key, algorithm="RS256")
42
43        return TokenPair(
44            access_token=access_token,
45            refresh_token=refresh_token,
46            token_family=family_id,
47        )
48
49    def verify_token(self, token: str, expected_type: str = "access") -> dict:
50        try:
51            payload = jwt.decode(token, self._public_key, algorithms=["RS256"])
52        except jwt.ExpiredSignatureError:
53            raise AuthError("Token has expired", code=401)
54        except jwt.InvalidTokenError as exc:
55            raise AuthError(f"Invalid token: {exc}", code=401)
56
57        if payload.get("type") != expected_type:
58            raise AuthError(f"Expected {expected_type} token", code=401)
59
60        return payload

• Lines 1–4: Import PyJWT as jwt, the uuid module for generating unique token identifiers, datetime utilities for expiration calculations, and dataclass for the TokenPair return type.
• Lines 7–10: The TokenPair dataclass bundles both tokens with their token_family identifier, which links refresh tokens together for rotation tracking. A token family represents all refresh tokens descended from a single login event.
• Lines 13–17: The constructor accepts PEM-formatted RSA keys as strings. Access tokens live for 15 minutes—long enough for a multi-turn conversation but short enough to limit damage from token theft. Refresh tokens live for 7 days, matching typical "remember me" session expectations.
• Lines 53–56: Catching jwt.ExpiredSignatureError separately from jwt.InvalidTokenError allows the client to distinguish "you need to refresh" from "this token is fundamentally invalid." The type check on line 56 prevents a refresh token from being used as an access token at regular endpoints.

1import redis.asyncio as aioredis
2import json
3from datetime import timedelta
4from starlette.middleware.base import BaseHTTPMiddleware
5from starlette.requests import Request
6from starlette.responses import JSONResponse
7
8class RedisSessionManager:
9    def __init__(self, redis_url: str = "redis://localhost:6379/0"):
10        self._redis = aioredis.from_url(redis_url, decode_responses=True)
11        self._refresh_prefix = "refresh:"
12        self._session_prefix = "session:"
13        self._family_prefix = "family:"
14
15    async def store_refresh_token(
16        self, jti: str, family: str, user_id: str, ttl: timedelta
17    ) -> None:
18        pipe = self._redis.pipeline()
19        pipe.setex(f"{self._refresh_prefix}{jti}", ttl, json.dumps({
20            "family": family, "user_id": user_id
21        }))
22        pipe.setex(f"{self._family_prefix}{family}", ttl, "active")
23        await pipe.execute()
24
25    async def rotate_refresh_token(
26        self, old_jti: str, old_family: str
27    ) -> bool:
28        stored = await self._redis.get(f"{self._refresh_prefix}{old_jti}")
29        if stored is None:
30            family_status = await self._redis.get(
31                f"{self._family_prefix}{old_family}"
32            )
33            if family_status == "revoked":
34                return False  # Replay attack detected
35            await self._redis.setex(
36                f"{self._family_prefix}{old_family}",
37                timedelta(days=7),
38                "revoked",
39            )
40            return False  # Token already used — revoke entire family
41
42        await self._redis.delete(f"{self._refresh_prefix}{old_jti}")
43        return True
44
45    async def create_session(
46        self, user_id: str, metadata: dict, ttl: timedelta = timedelta(hours=8)
47    ) -> str:
48        session_key = f"{self._session_prefix}{user_id}"
49        session_data = {
50            "user_id": user_id,
51            "active_conversations": [],
52            "request_count": 0,
53            **metadata,
54        }
55        await self._redis.setex(session_key, ttl, json.dumps(session_data))
56        return session_key
57
58    async def get_session(self, user_id: str) -> dict | None:
59        data = await self._redis.get(f"{self._session_prefix}{user_id}")
60        if data is None:
61            return None
62        return json.loads(data)
63
64    async def increment_request_count(self, user_id: str) -> int:
65        session_key = f"{self._session_prefix}{user_id}"
66        raw = await self._redis.get(session_key)
67        if raw is None:
68            raise SessionExpiredError(f"No session for {user_id}")
69        session = json.loads(raw)
70        session["request_count"] += 1
71        ttl = await self._redis.ttl(session_key)
72        if ttl > 0:
73            await self._redis.setex(session_key, ttl, json.dumps(session))
74        return session["request_count"]
75
76PUBLIC_PATHS = {"/auth/login", "/auth/register", "/health", "/docs", "/openapi.json"}
77
78class JWTAuthMiddleware(BaseHTTPMiddleware):
79    def __init__(self, app, token_service: JWTTokenService,
80                 session_manager: RedisSessionManager):
81        super().__init__(app)
82        self.token_service = token_service
83        self.session_manager = session_manager
84
85    async def dispatch(self, request: Request, call_next):
86        if request.url.path in PUBLIC_PATHS:
87            return await call_next(request)
88
89        auth_header = request.headers.get("Authorization")
90        if auth_header is None or not auth_header.startswith("Bearer "):
91            return JSONResponse(
92                {"detail": "Missing or malformed Authorization header"},
93                status_code=401,
94            )
95
96        token = auth_header[7:]  # Strip "Bearer " prefix
97        try:
98            payload = self.token_service.verify_token(token, expected_type="access")
99        except AuthError as exc:
100            return JSONResponse({"detail": exc.message}, status_code=exc.code)
101
102        user_id = payload["sub"]
103        session = await self.session_manager.get_session(user_id)
104        if session is None:
105            session = await self.session_manager.create_session(
106                user_id, metadata={"roles": payload.get("roles", [])}
107            )
108
109        request.state.user_id = user_id
110        request.state.user_roles = payload.get("roles", [])
111        request.state.token_jti = payload.get("jti")
112
113        await self.session_manager.increment_request_count(user_id)
114        return await call_next(request)

• Lines 1–3: The middleware inherits from Starlette's BaseHTTPMiddleware, which provides the dispatch() hook for intercepting every request/response cycle. This works identically in FastAPI since FastAPI is built on Starlette.
• Line 6: The PUBLIC_PATHS set defines endpoints that do not require authentication. Checking membership in a set is O(1), which matters when this check runs on every single request including health checks from load balancers.
• Lines 9–14: The constructor accepts pre-configured JWTTokenService and RedisSessionManager instances, following dependency injection. This makes testing straightforward—pass mock implementations during unit tests.
• Lines 45–46: Incrementing the request count on every authenticated request enables per-user rate limiting at the session level. For GenAI applications, this is essential—a single LLM inference call can cost $0.01–$0.10, so unthrottled access represents a direct financial risk.