diff --git a/openrag/components/auth/__init__.py b/openrag/components/auth/__init__.py index a1eb42365..df560c692 100644 --- a/openrag/components/auth/__init__.py +++ b/openrag/components/auth/__init__.py @@ -1,7 +1,8 @@ -from components.auth.deps import get_oidc_client, reset_oidc_client -from components.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle -from components.auth.session_tokens import decrypt_token, encrypt_token, hash_session_token, issue_session_token -from components.auth.state_cookie import StateCookiePayload, StateCookieSerializer +# Adapter shim -- canonical exports moved to services.auth (Phase 6F). +from services.auth.deps import get_oidc_client, reset_oidc_client +from services.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle +from services.auth.session_tokens import decrypt_token, encrypt_token, hash_session_token, issue_session_token +from services.auth.state_cookie import StateCookiePayload, StateCookieSerializer __all__ = [ "OIDCClient", diff --git a/openrag/components/auth/deps.py b/openrag/components/auth/deps.py index c68d50e21..7905e92b0 100644 --- a/openrag/components/auth/deps.py +++ b/openrag/components/auth/deps.py @@ -1,83 +1,4 @@ -"""Lazy, process-local singleton for the OIDCClient. +# Adapter shim -- canonical code moved to services.auth.deps (Phase 6F). +from services.auth.deps import get_oidc_client, reset_oidc_client -Kept in a dedicated module to avoid circular imports between the router -(``openrag/routers/auth.py``) and the application entry point (``openrag/api.py``). - -The OIDC config env vars are resolved here via ``os.getenv`` — the same values -that ``openrag/api.py`` validates at startup. In ``AUTH_MODE=oidc`` mode, these -are guaranteed to be non-empty (api.py refuses to start otherwise), so this -module simply trusts them. -""" - -from __future__ import annotations - -import os -from threading import Lock - -from components.auth.oidc_client import OIDCClient - -_client: OIDCClient | None = None -_lock = Lock() - - -def get_oidc_client() -> OIDCClient: - """Return the shared OIDCClient instance, creating it on first call. - - The instance caches the discovery doc and JWKS, so a single shared client - per worker process is both correct and more efficient than one-per-request. - - Env vars read (all required in AUTH_MODE=oidc): - - OIDC_ENDPOINT - - OIDC_CLIENT_ID - - OIDC_CLIENT_SECRET - - OIDC_REDIRECT_URI - - OIDC_SCOPES (default ``openid email profile offline_access``) - """ - global _client - if _client is not None: - return _client - with _lock: - if _client is not None: - return _client - issuer = os.environ["OIDC_ENDPOINT"] - client_id = os.environ["OIDC_CLIENT_ID"] - client_secret = os.environ["OIDC_CLIENT_SECRET"] - redirect_uri = os.environ["OIDC_REDIRECT_URI"] - scopes = os.getenv("OIDC_SCOPES", "openid email profile offline_access") - _client = OIDCClient( - issuer=issuer, - client_id=client_id, - client_secret=client_secret, - redirect_uri=redirect_uri, - scopes=scopes, - ) - return _client - - -def reset_oidc_client() -> None: - """Test hook — drops the cached client so the next call rebuilds from env. - - Best-effort closes the underlying httpx.AsyncClient to avoid "Unclosed - client session" warnings and leaking connections when tests repeatedly - reset the singleton. If no event loop is running we skip the close call - — the GC will eventually reclaim the socket. - """ - global _client - with _lock: - old = _client - _client = None - if old is None: - return - try: - import asyncio - - loop = asyncio.get_event_loop_policy().get_event_loop() - if loop.is_running(): - # Schedule close on the running loop without awaiting — caller - # doesn't need to be async. - loop.create_task(old.aclose()) - else: - loop.run_until_complete(old.aclose()) - except Exception: - # Closing is best-effort; never let a reset blow up the caller. - pass +__all__ = ["get_oidc_client", "reset_oidc_client"] diff --git a/openrag/components/auth/oidc_client.py b/openrag/components/auth/oidc_client.py index 50b449b9e..622228223 100644 --- a/openrag/components/auth/oidc_client.py +++ b/openrag/components/auth/oidc_client.py @@ -1,376 +1,4 @@ -"""Lightweight OIDC Relying Party client for OpenRAG. +# Adapter shim -- canonical code moved to services.auth.oidc_client (Phase 6F). +from services.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle -Wraps Authlib's JWT/JWK primitives with: -- Discovery endpoint caching (1 h TTL) -- JWKS caching with automatic refresh on kid-miss -- PKCE pair generation (S256) -- Authorization URL builder -- Code exchange with ID token verification -- Token refresh (lazy, called by middleware when access_token near expiry) -- Userinfo fetch -- Back-channel logout token verification - -One instance per (issuer, client_id, client_secret) tuple. -The instance is not thread-safe for writes but safe for concurrent reads once -the metadata and JWKS caches are populated. -""" - -import base64 -import hashlib -import secrets -import time -from dataclasses import dataclass -from typing import Any -from urllib.parse import urlencode - -import httpx -from authlib.jose import JsonWebKey, JsonWebToken -from authlib.jose.errors import JoseError - - -@dataclass -class TokenBundle: - """Holds the token set returned by the IdP together with verified ID token claims.""" - - id_token: str - access_token: str - refresh_token: str | None - expires_in: int # seconds - token_type: str # usually "Bearer" - claims: dict[str, Any] # verified claims from id_token - - -@dataclass -class LogoutTokenClaims: - """Verified claims from a back-channel logout token.""" - - iss: str - aud: str | list[str] - sub: str | None - sid: str | None - iat: int - jti: str | None - - -class OIDCClient: - """Lightweight OIDC Relying Party client. - - One instance per (issuer, client_id, client_secret) tuple. - """ - - _DISCOVERY_TTL = 3600 # 1 hour - _JWKS_TTL = 3600 # 1 hour - - def __init__( - self, - *, - issuer: str, - client_id: str, - client_secret: str, - redirect_uri: str, - scopes: str, - http_client: httpx.AsyncClient | None = None, - ): - # Keep the issuer string verbatim (including any trailing "/") — the OIDC - # spec mandates strict byte-for-byte equality between ``self.issuer``, the - # issuer advertised by the discovery document, and the ``iss`` claim in - # tokens. Operators must configure ``OIDC_ENDPOINT`` to match EXACTLY - # what the IdP returns. - self.issuer = issuer - self.client_id = client_id - self.client_secret = client_secret - self.redirect_uri = redirect_uri - self.scopes = scopes - self._http = http_client or httpx.AsyncClient(timeout=10.0) - self._metadata: dict | None = None - self._metadata_fetched_at: float = 0.0 - self._jwks: JsonWebKey | None = None - self._jwks_fetched_at: float = 0.0 - - # ------------------------------------------------------------------ - # Discovery - # ------------------------------------------------------------------ - - async def discover(self) -> dict: - """Fetch and cache the OIDC discovery document. - - Returns the cached document if it is less than _DISCOVERY_TTL seconds old. - Raises ValueError if the returned issuer does not match the configured one. - """ - if self._metadata and (time.time() - self._metadata_fetched_at) < self._DISCOVERY_TTL: - return self._metadata - url = f"{self.issuer.rstrip('/')}/.well-known/openid-configuration" - resp = await self._http.get(url) - resp.raise_for_status() - self._metadata = resp.json() - self._metadata_fetched_at = time.time() - if self._metadata.get("issuer") != self.issuer: - raise ValueError(f"Issuer mismatch: configured {self.issuer!r}, got {self._metadata.get('issuer')!r}") - return self._metadata - - # ------------------------------------------------------------------ - # JWKS - # ------------------------------------------------------------------ - - async def _load_jwks(self, force: bool = False) -> JsonWebKey: - meta = await self.discover() - if not force and self._jwks and (time.time() - self._jwks_fetched_at) < self._JWKS_TTL: - return self._jwks - resp = await self._http.get(meta["jwks_uri"]) - resp.raise_for_status() - self._jwks = JsonWebKey.import_key_set(resp.json()) - self._jwks_fetched_at = time.time() - return self._jwks - - # ------------------------------------------------------------------ - # PKCE helpers - # ------------------------------------------------------------------ - - @staticmethod - def generate_pkce_pair() -> tuple[str, str]: - """Generate a PKCE (code_verifier, code_challenge) pair using S256. - - Returns: - (verifier, challenge) — verifier is 128 url-safe chars, - challenge is the base64url-encoded SHA-256 of the verifier. - """ - verifier = secrets.token_urlsafe(96)[:128] - digest = hashlib.sha256(verifier.encode()).digest() - challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() - return verifier, challenge - - @staticmethod - def generate_state_and_nonce() -> tuple[str, str]: - """Generate cryptographically random state and nonce values.""" - return secrets.token_urlsafe(32), secrets.token_urlsafe(32) - - # ------------------------------------------------------------------ - # Authorization URL - # ------------------------------------------------------------------ - - async def build_authorization_url(self, *, state: str, nonce: str, code_challenge: str) -> str: - """Build the full authorization URL to redirect the browser to.""" - meta = await self.discover() - params = { - "response_type": "code", - "client_id": self.client_id, - "redirect_uri": self.redirect_uri, - "scope": self.scopes, - "state": state, - "nonce": nonce, - "code_challenge": code_challenge, - "code_challenge_method": "S256", - } - return f"{meta['authorization_endpoint']}?{urlencode(params)}" - - # ------------------------------------------------------------------ - # Code exchange - # ------------------------------------------------------------------ - - async def exchange_code(self, *, code: str, code_verifier: str, expected_nonce: str) -> TokenBundle: - """Exchange an authorization code for tokens. - - Verifies the returned id_token (signature, iss, aud, exp, nonce). - - Args: - code: The authorization code from the IdP callback. - code_verifier: The PKCE verifier corresponding to the challenge sent earlier. - expected_nonce: The nonce value that was sent in the authorization request. - - Returns: - A TokenBundle with verified claims. - """ - meta = await self.discover() - data = { - "grant_type": "authorization_code", - "code": code, - "redirect_uri": self.redirect_uri, - "client_id": self.client_id, - "client_secret": self.client_secret, - "code_verifier": code_verifier, - } - resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) - resp.raise_for_status() - payload = resp.json() - id_token = payload["id_token"] - claims = await self._verify_id_token(id_token, expected_nonce=expected_nonce) - return TokenBundle( - id_token=id_token, - access_token=payload["access_token"], - refresh_token=payload.get("refresh_token"), - expires_in=int(payload.get("expires_in", 0)), - token_type=payload.get("token_type", "Bearer"), - claims=claims, - ) - - # ------------------------------------------------------------------ - # Token refresh - # ------------------------------------------------------------------ - - async def refresh_access_token(self, refresh_token: str) -> TokenBundle: - """Use the refresh_token to obtain a new access_token. - - If the IdP returns a new id_token, it is re-verified (nonce check skipped - per RFC 8252 §8.2 — nonce is only required during the initial code exchange). - If the IdP omits the refresh_token in the response, the caller's existing - refresh_token is preserved. - - Returns: - A new TokenBundle. - """ - meta = await self.discover() - data = { - "grant_type": "refresh_token", - "refresh_token": refresh_token, - "client_id": self.client_id, - "client_secret": self.client_secret, - } - resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) - resp.raise_for_status() - payload = resp.json() - new_id_token = payload.get("id_token") - claims: dict[str, Any] = {} - if new_id_token: - claims = await self._verify_id_token(new_id_token, expected_nonce=None) - return TokenBundle( - id_token=new_id_token or "", - access_token=payload["access_token"], - # Some IdPs omit the refresh_token on rotation — keep the old one. - refresh_token=payload.get("refresh_token", refresh_token), - expires_in=int(payload.get("expires_in", 0)), - token_type=payload.get("token_type", "Bearer"), - claims=claims, - ) - - # ------------------------------------------------------------------ - # Userinfo - # ------------------------------------------------------------------ - - async def fetch_userinfo(self, access_token: str) -> dict: - """Fetch the userinfo endpoint with the given access token.""" - meta = await self.discover() - resp = await self._http.get( - meta["userinfo_endpoint"], - headers={"Authorization": f"Bearer {access_token}"}, - ) - resp.raise_for_status() - return resp.json() - - # ------------------------------------------------------------------ - # ID token verification - # ------------------------------------------------------------------ - - async def _verify_id_token(self, token: str, *, expected_nonce: str | None) -> dict[str, Any]: - """Verify an ID token's signature and standard claims. - - Retries with a fresh JWKS fetch on kid-miss (covers IdP key rotation). - Raises JoseError / ValueError on any validation failure. - """ - jwks = await self._load_jwks() - jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) - try: - claims = jwt.decode(token, jwks) - except JoseError: - # Force JWKS refresh in case of kid rotation; retry once. - jwks = await self._load_jwks(force=True) - claims = jwt.decode(token, jwks) - - # Manual validation — avoids authlib version differences around claims.params - decoded: dict[str, Any] = dict(claims) - now = int(time.time()) - - if decoded.get("iss") != self.issuer: - raise ValueError(f"ID token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") - - aud = decoded.get("aud") - if isinstance(aud, list): - if self.client_id not in aud: - raise ValueError(f"ID token aud {aud!r} does not contain client_id {self.client_id!r}") - elif aud != self.client_id: - raise ValueError(f"ID token aud {aud!r} != client_id {self.client_id!r}") - - if "exp" not in decoded: - raise ValueError("ID token missing exp claim") - if int(decoded["exp"]) < now: - raise ValueError("ID token has expired") - - if "iat" not in decoded: - raise ValueError("ID token missing iat claim") - - if expected_nonce is not None: - if decoded.get("nonce") != expected_nonce: - raise ValueError("OIDC nonce mismatch") - - return decoded - - # ------------------------------------------------------------------ - # Back-channel logout token verification - # ------------------------------------------------------------------ - - async def verify_logout_token(self, token: str) -> LogoutTokenClaims: - """Verify an OIDC back-channel logout token. - - Validates: - - Signature (with JWKS kid-miss retry) - - Standard claims (iss, aud, iat) - - events claim contains the back-channel-logout URI key - - nonce must NOT be present (spec requirement) - - At least one of sub or sid must be present - - Returns: - LogoutTokenClaims with the verified values. - Raises: - ValueError: on any spec violation. - """ - jwks = await self._load_jwks() - jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) - try: - claims = jwt.decode(token, jwks) - except JoseError: - jwks = await self._load_jwks(force=True) - claims = jwt.decode(token, jwks) - - decoded: dict[str, Any] = dict(claims) - now = int(time.time()) - - if decoded.get("iss") != self.issuer: - raise ValueError(f"logout_token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") - - aud = decoded.get("aud") - if isinstance(aud, list): - if self.client_id not in aud: - raise ValueError(f"logout_token aud {aud!r} does not contain client_id {self.client_id!r}") - elif aud != self.client_id: - raise ValueError(f"logout_token aud {aud!r} != client_id {self.client_id!r}") - - if "iat" not in decoded: - raise ValueError("logout_token missing iat claim") - if int(decoded.get("exp", now + 1)) < now: - raise ValueError("logout_token has expired") - - events = decoded.get("events") or {} - if "http://schemas.openid.net/event/backchannel-logout" not in events: - raise ValueError("logout_token missing required back-channel-logout event claim") - - if decoded.get("nonce"): - raise ValueError("logout_token must not contain nonce") - - if not decoded.get("sub") and not decoded.get("sid"): - raise ValueError("logout_token must contain sub or sid") - - return LogoutTokenClaims( - iss=decoded["iss"], - aud=decoded["aud"], - sub=decoded.get("sub"), - sid=decoded.get("sid"), - iat=int(decoded["iat"]), - jti=decoded.get("jti"), - ) - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ - - async def aclose(self) -> None: - """Close the underlying HTTP client.""" - await self._http.aclose() +__all__ = ["OIDCClient", "TokenBundle", "LogoutTokenClaims"] diff --git a/openrag/components/auth/refresh.py b/openrag/components/auth/refresh.py index cd02e80f0..cc52a6845 100644 --- a/openrag/components/auth/refresh.py +++ b/openrag/components/auth/refresh.py @@ -1,176 +1,4 @@ -"""Lazy refresh helper for OIDC access tokens. +# Adapter shim -- canonical code moved to services.auth.refresh (Phase 6F). +from services.auth.refresh import refresh_session_if_needed -Extracted from ``AuthMiddleware`` (Phase 5) to keep ``api.py`` small and -independently testable. Called per-request when a valid cookie session is -found; a no-op when the access token is still fresh. - -Timezone policy ---------------- -Phase 2 stores all OIDC session timestamps as **naive local time** via -``datetime.now()`` (see ``test_oidc_sessions.py`` and -``PartitionFileManager.get_oidc_session_by_token``). We match that style -everywhere in this module to avoid tz-mismatch bugs when comparing -``access_token_expires_at`` against "now". - -Refresh-token stampede guard (M1) ---------------------------------- -IdPs with refresh_token rotation enabled invalidate the old refresh_token the -first time it is redeemed. Under concurrency, multiple requests can each notice -"my access_token is about to expire" at the same time and race each other to -the token endpoint. The second attempt fails with ``invalid_grant`` and -(without a guard) its session would be revoked mid-flight. - -We mitigate that with two cooperating mechanisms: - -1. A **short-circuit** here: if ``last_refresh_at`` was bumped less than 5 - seconds ago, we assume a sibling request already rotated the tokens, - re-read the row, and reuse those freshly rotated tokens instead of calling - the IdP. -2. A **row-level write lock** in :meth:`PartitionFileManager.update_oidc_session_tokens` - (``SELECT ... FOR UPDATE``) so that only one writer commits at a time on - Postgres. -3. An **error-recovery branch** here: if the IdP does reject our refresh_token - (typically because a sibling raced us and won), we re-read the row once - more and, if the tokens were advanced meanwhile, return the fresh session - rather than giving up. -""" - -from __future__ import annotations - -from datetime import datetime, timedelta -from typing import Any - -from components.auth.deps import get_oidc_client -from components.auth.session_tokens import decrypt_token, encrypt_token -from utils.logger import get_logger - -_REFRESH_BUFFER = timedelta(seconds=60) -_STAMPEDE_WINDOW = timedelta(seconds=5) - -logger = get_logger() - - -def _to_dt(val: Any) -> datetime: - """Coerce a datetime-or-ISO-string into a ``datetime``. - - Ray occasionally ships values across actors in serialised form; accept - either shape so callers never have to care about the transport. - """ - if isinstance(val, datetime): - return val - if isinstance(val, str): - return datetime.fromisoformat(val) - raise TypeError(f"Expected datetime or ISO string, got {type(val).__name__}") - - -async def refresh_session_if_needed( - *, - session: dict[str, Any], - enc_key: str, - vectordb: Any, -) -> dict[str, Any] | None: - """Refresh the IdP access_token if it is within ``_REFRESH_BUFFER`` of expiry. - - Behaviour: - - If the access_token is still valid with the 60s buffer → return ``session`` unchanged. - - Stampede guard: if another request has just refreshed this session - (``last_refresh_at`` within 5s), re-read the row and reuse the fresh - tokens without calling the IdP. - - If near/past expiry AND a ``refresh_token_encrypted`` blob is stored → - call the IdP, persist rotated tokens, return an updated session dict. - - If near/past expiry AND no refresh_token is stored → return ``session`` as-is - when still formally valid, or ``None`` when already expired (caller should - treat as a revoked session). - - If the refresh call raises (typically because a sibling already rotated - the tokens and the IdP now rejects ours) → re-read the row; if a sibling - succeeded, return their fresh session; otherwise ``None``. - - The session dict returned mirrors the DB row shape produced by - ``PartitionFileManager._oidc_session_to_dict``. - """ - now = datetime.now() - access_exp = _to_dt(session["access_token_expires_at"]) - - if access_exp > now + _REFRESH_BUFFER: - return session - - # --- Stampede short-circuit ------------------------------------------- - # If a sibling request just refreshed this same session, re-read the row - # and reuse the freshly rotated tokens. This avoids racing the IdP with a - # refresh_token that the sibling's success has already invalidated. - last_refresh_at = session.get("last_refresh_at") - if last_refresh_at is not None: - try: - last_refresh_at_dt = _to_dt(last_refresh_at) - except TypeError: - last_refresh_at_dt = None - if last_refresh_at_dt is not None and (now - last_refresh_at_dt) < _STAMPEDE_WINDOW: - try: - fresh = await vectordb.get_oidc_session_by_id.remote(session["id"]) - except Exception as e: - logger.bind(session_id=session.get("id"), error=str(e)).warning( - "Stampede-guard re-read failed; falling through to refresh" - ) - fresh = None - if fresh is not None: - fresh_exp = _to_dt(fresh["access_token_expires_at"]) - if fresh_exp > now + _REFRESH_BUFFER: - return fresh - - refresh_enc = session.get("refresh_token_encrypted") - if not refresh_enc: - # No refresh_token available. - # - If still formally valid (within the 60s buffer window but not yet past exp), - # keep using it. - # - If already expired, caller should treat the session as dead. - return session if access_exp > now else None - - try: - refresh_token = decrypt_token(refresh_enc, enc_key) - client = get_oidc_client() - bundle = await client.refresh_access_token(refresh_token) - except Exception as e: - # Maybe a sibling refreshed between our staleness check and the IdP call - # and the IdP has already invalidated our refresh_token. Re-read the - # row once before giving up: if the tokens were rotated meanwhile, - # treat this as a successful refresh (the sibling's). - logger.bind(session_id=session.get("id"), error=str(e)).warning( - "OIDC refresh_token exchange failed — re-reading session for stampede recovery" - ) - try: - fresh = await vectordb.get_oidc_session_by_id.remote(session["id"]) - except Exception as re: - logger.bind(session_id=session.get("id"), error=str(re)).error( - "Post-failure re-read of OIDC session failed — invalidating" - ) - return None - if fresh is not None: - fresh_exp = _to_dt(fresh["access_token_expires_at"]) - if fresh_exp > now + _REFRESH_BUFFER: - return fresh - return None - - new_access_exp = now + timedelta(seconds=max(int(bundle.expires_in or 0), 60)) - new_access_enc = encrypt_token(bundle.access_token, enc_key) - new_refresh_enc = encrypt_token(bundle.refresh_token, enc_key) if bundle.refresh_token else refresh_enc - - try: - await vectordb.update_oidc_session_tokens.remote( - session_id=session["id"], - access_token_encrypted=new_access_enc, - refresh_token_encrypted=new_refresh_enc, - access_token_expires_at=new_access_exp, - ) - except Exception as e: - logger.bind(session_id=session.get("id"), error=str(e)).error( - "Failed to persist refreshed OIDC tokens — invalidating session" - ) - return None - - return { - **session, - "access_token_encrypted": new_access_enc, - "access_token_expires_at": new_access_exp, - "refresh_token_encrypted": new_refresh_enc, - "last_refresh_at": now, - } +__all__ = ["refresh_session_if_needed"] diff --git a/openrag/components/auth/session_tokens.py b/openrag/components/auth/session_tokens.py index 79cc085e1..1d611e2b2 100644 --- a/openrag/components/auth/session_tokens.py +++ b/openrag/components/auth/session_tokens.py @@ -1,64 +1,9 @@ -"""Session token utilities for OpenRAG OIDC sessions. - -Opaque session tokens are issued at callback and stored hashed (SHA-256) in the DB. -IdP tokens (access_token, refresh_token) are encrypted with Fernet before storage. - -The Fernet key is provided via the OIDC_TOKEN_ENCRYPTION_KEY environment variable. -Generate one with: - python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())' -""" - -import hashlib -import secrets - -from cryptography.fernet import Fernet, InvalidToken - - -def issue_session_token() -> tuple[str, str]: - """Generate a new session token. - - Returns: - (plaintext, sha256_hex) — the plaintext is set in the cookie, - the hash is stored in the database. - """ - plain = secrets.token_urlsafe(32) # 43 chars, >= 256 bits entropy - return plain, hash_session_token(plain) - - -def hash_session_token(token: str) -> str: - """Return the SHA-256 hex digest of the session token.""" - return hashlib.sha256(token.encode("utf-8")).hexdigest() - - -def _fernet(key: str | bytes) -> Fernet: - try: - return Fernet(key.encode("utf-8") if isinstance(key, str) else key) - except Exception as e: - raise ValueError( - "OIDC_TOKEN_ENCRYPTION_KEY is not a valid Fernet key. " - "Generate one with: python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())'" - ) from e - - -def encrypt_token(plaintext: str | None, key: str) -> bytes | None: - """Encrypt a plaintext token string. - - Returns None if plaintext is None (refresh_token may be absent). - """ - if plaintext is None: - return None - return _fernet(key).encrypt(plaintext.encode("utf-8")) - - -def decrypt_token(ciphertext: bytes | None, key: str) -> str | None: - """Decrypt a Fernet-encrypted token. - - Returns None if ciphertext is None. - Raises ValueError on key mismatch or data corruption. - """ - if ciphertext is None: - return None - try: - return _fernet(key).decrypt(ciphertext).decode("utf-8") - except InvalidToken as e: - raise ValueError("Failed to decrypt stored OIDC token — key mismatch or corruption") from e +# Adapter shim -- canonical code moved to services.auth.session_tokens (Phase 6F). +from services.auth.session_tokens import ( + decrypt_token, + encrypt_token, + hash_session_token, + issue_session_token, +) + +__all__ = ["issue_session_token", "hash_session_token", "encrypt_token", "decrypt_token"] diff --git a/openrag/components/auth/state_cookie.py b/openrag/components/auth/state_cookie.py index c6e5c566f..6a462a97a 100644 --- a/openrag/components/auth/state_cookie.py +++ b/openrag/components/auth/state_cookie.py @@ -1,55 +1,4 @@ -"""Signed state cookie for OIDC Authorization Code + PKCE flow. +# Adapter shim -- canonical code moved to services.auth.state_cookie (Phase 6F). +from services.auth.state_cookie import StateCookiePayload, StateCookieSerializer -The cookie transports state/nonce/code_verifier between /auth/login and /auth/callback. -It is signed (not encrypted) using itsdangerous.URLSafeTimedSerializer with HMAC-SHA1. - -The signing key is the OIDC_TOKEN_ENCRYPTION_KEY (a Fernet base64url key, which is -valid arbitrary bytes for HMAC). The consuming code (phase 4 router) will pass the -key to StateCookieSerializer(key). Using the same key for both Fernet encryption and -HMAC signing is safe since itsdangerous derives separate subkeys via HMAC. - -TTL defaults to 600 s (10 minutes) — long enough for a slow user at the IdP login page. -""" - -from dataclasses import asdict, dataclass - -from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer - - -@dataclass -class StateCookiePayload: - state: str - nonce: str - code_verifier: str - next_url: str = "/" - - -class StateCookieSerializer: - """Signs/verifies the short-lived cookie holding OIDC state/nonce/code_verifier. - - TTL defaults to 600 s (10 minutes) — long enough for a slow user at the IdP. - """ - - COOKIE_NAME = "openrag_oidc_state" - DEFAULT_TTL_SECONDS = 600 - - def __init__(self, secret_key: str, salt: str = "openrag-oidc-state-v1"): - self._serializer = URLSafeTimedSerializer(secret_key, salt=salt) - - def dumps(self, payload: StateCookiePayload) -> str: - """Serialize and sign the payload, returning an opaque cookie value.""" - return self._serializer.dumps(asdict(payload)) - - def loads(self, token: str, max_age: int = DEFAULT_TTL_SECONDS) -> StateCookiePayload: - """Verify and deserialize the cookie value. - - Raises: - ValueError: if the cookie is expired or the signature is invalid. - """ - try: - data = self._serializer.loads(token, max_age=max_age) - except SignatureExpired as e: - raise ValueError("OIDC state cookie expired") from e - except BadSignature as e: - raise ValueError("OIDC state cookie signature invalid") from e - return StateCookiePayload(**data) +__all__ = ["StateCookiePayload", "StateCookieSerializer"] diff --git a/openrag/components/auth/test_middleware.py b/openrag/components/auth/test_middleware.py index 67092d073..2896cb7c3 100644 --- a/openrag/components/auth/test_middleware.py +++ b/openrag/components/auth/test_middleware.py @@ -390,8 +390,8 @@ async def test_expired_no_refresh_token_returns_none(self): async def test_refresh_short_circuit_when_last_refresh_recent(self): """If another request refreshed <5s ago, reuse the fresh row; do NOT hit the IdP again with a refresh_token that has already been rotated.""" - from components.auth import refresh as refresh_mod - from components.auth.refresh import refresh_session_if_needed + from services.auth import refresh as refresh_mod + from services.auth.refresh import refresh_session_if_needed now = datetime.now() fresh_exp = now + timedelta(minutes=30) @@ -432,8 +432,8 @@ async def test_refresh_short_circuit_when_last_refresh_recent(self): async def test_refresh_recovers_when_idp_rejects_stale_refresh_token(self): """IdP rejects our refresh_token (sibling already rotated it); the helper re-reads the session and returns the sibling's fresh tokens.""" - from components.auth import refresh as refresh_mod - from components.auth.refresh import refresh_session_if_needed + from services.auth import refresh as refresh_mod + from services.auth.refresh import refresh_session_if_needed now = datetime.now() stale_session = { @@ -470,8 +470,8 @@ async def test_refresh_recovers_when_idp_rejects_stale_refresh_token(self): @pytest.mark.asyncio async def test_refresh_returns_none_when_idp_rejects_and_no_concurrent_refresh(self): """IdP rejects us and no sibling rotated the tokens → invalidate session.""" - from components.auth import refresh as refresh_mod - from components.auth.refresh import refresh_session_if_needed + from services.auth import refresh as refresh_mod + from services.auth.refresh import refresh_session_if_needed now = datetime.now() stale_session = { diff --git a/openrag/components/auth/test_oidc_client.py b/openrag/components/auth/test_oidc_client.py index 8a2e79198..531d109a1 100644 --- a/openrag/components/auth/test_oidc_client.py +++ b/openrag/components/auth/test_oidc_client.py @@ -7,7 +7,7 @@ import pytest_asyncio import respx from authlib.jose import JsonWebKey -from components.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle +from services.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle # --------------------------------------------------------------------------- # Helpers — RSA test key + JWT factory diff --git a/openrag/components/auth/test_session_tokens.py b/openrag/components/auth/test_session_tokens.py index 191688962..59335e487 100644 --- a/openrag/components/auth/test_session_tokens.py +++ b/openrag/components/auth/test_session_tokens.py @@ -1,7 +1,7 @@ """Unit tests for session_tokens.py.""" import pytest -from components.auth.session_tokens import ( +from services.auth.session_tokens import ( decrypt_token, encrypt_token, hash_session_token, diff --git a/openrag/components/auth/test_state_cookie.py b/openrag/components/auth/test_state_cookie.py index 4a5b8d2a2..7faa91241 100644 --- a/openrag/components/auth/test_state_cookie.py +++ b/openrag/components/auth/test_state_cookie.py @@ -3,7 +3,7 @@ import time import pytest -from components.auth.state_cookie import StateCookiePayload, StateCookieSerializer +from services.auth.state_cookie import StateCookiePayload, StateCookieSerializer SECRET = "test-secret-key-for-state-cookie" diff --git a/openrag/services/auth/__init__.py b/openrag/services/auth/__init__.py index e69de29bb..43377f9fd 100644 --- a/openrag/services/auth/__init__.py +++ b/openrag/services/auth/__init__.py @@ -0,0 +1,18 @@ +from services.auth.deps import get_oidc_client, reset_oidc_client +from services.auth.oidc_client import LogoutTokenClaims, OIDCClient, TokenBundle +from services.auth.session_tokens import decrypt_token, encrypt_token, hash_session_token, issue_session_token +from services.auth.state_cookie import StateCookiePayload, StateCookieSerializer + +__all__ = [ + "OIDCClient", + "TokenBundle", + "LogoutTokenClaims", + "issue_session_token", + "encrypt_token", + "decrypt_token", + "hash_session_token", + "StateCookieSerializer", + "StateCookiePayload", + "get_oidc_client", + "reset_oidc_client", +] diff --git a/openrag/services/auth/deps.py b/openrag/services/auth/deps.py new file mode 100644 index 000000000..38927ff60 --- /dev/null +++ b/openrag/services/auth/deps.py @@ -0,0 +1,80 @@ +"""Lazy, process-local singleton for the OIDCClient. + +Kept in a dedicated module to avoid circular imports between the router +(``openrag/routers/auth.py``) and the application entry point (``openrag/api.py``). + +The OIDC config env vars are resolved here via ``os.getenv`` -- the same values +that ``openrag/api.py`` validates at startup. In ``AUTH_MODE=oidc`` mode, these +are guaranteed to be non-empty (api.py refuses to start otherwise), so this +module simply trusts them. +""" + +from __future__ import annotations + +import os +from threading import Lock + +from services.auth.oidc_client import OIDCClient + +_client: OIDCClient | None = None +_lock = Lock() + + +def get_oidc_client() -> OIDCClient: + """Return the shared OIDCClient instance, creating it on first call. + + The instance caches the discovery doc and JWKS, so a single shared client + per worker process is both correct and more efficient than one-per-request. + + Env vars read (all required in AUTH_MODE=oidc): + - OIDC_ENDPOINT + - OIDC_CLIENT_ID + - OIDC_CLIENT_SECRET + - OIDC_REDIRECT_URI + - OIDC_SCOPES (default ``openid email profile offline_access``) + """ + global _client + if _client is not None: + return _client + with _lock: + if _client is not None: + return _client + issuer = os.environ["OIDC_ENDPOINT"] + client_id = os.environ["OIDC_CLIENT_ID"] + client_secret = os.environ["OIDC_CLIENT_SECRET"] + redirect_uri = os.environ["OIDC_REDIRECT_URI"] + scopes = os.getenv("OIDC_SCOPES", "openid email profile offline_access") + _client = OIDCClient( + issuer=issuer, + client_id=client_id, + client_secret=client_secret, + redirect_uri=redirect_uri, + scopes=scopes, + ) + return _client + + +def reset_oidc_client() -> None: + """Test hook -- drops the cached client so the next call rebuilds from env. + + Best-effort closes the underlying httpx.AsyncClient to avoid "Unclosed + client session" warnings and leaking connections when tests repeatedly + reset the singleton. If no event loop is running we skip the close call + -- the GC will eventually reclaim the socket. + """ + global _client + with _lock: + old = _client + _client = None + if old is None: + return + try: + import asyncio + + loop = asyncio.get_event_loop_policy().get_event_loop() + if loop.is_running(): + loop.create_task(old.aclose()) + else: + loop.run_until_complete(old.aclose()) + except Exception: + pass diff --git a/openrag/services/auth/oidc_client.py b/openrag/services/auth/oidc_client.py new file mode 100644 index 000000000..1d21a59e0 --- /dev/null +++ b/openrag/services/auth/oidc_client.py @@ -0,0 +1,376 @@ +"""Lightweight OIDC Relying Party client for OpenRAG. + +Wraps Authlib's JWT/JWK primitives with: +- Discovery endpoint caching (1 h TTL) +- JWKS caching with automatic refresh on kid-miss +- PKCE pair generation (S256) +- Authorization URL builder +- Code exchange with ID token verification +- Token refresh (lazy, called by middleware when access_token near expiry) +- Userinfo fetch +- Back-channel logout token verification + +One instance per (issuer, client_id, client_secret) tuple. +The instance is not thread-safe for writes but safe for concurrent reads once +the metadata and JWKS caches are populated. +""" + +import base64 +import hashlib +import secrets +import time +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlencode + +import httpx +from authlib.jose import JsonWebKey, JsonWebToken +from authlib.jose.errors import JoseError + + +@dataclass +class TokenBundle: + """Holds the token set returned by the IdP together with verified ID token claims.""" + + id_token: str + access_token: str + refresh_token: str | None + expires_in: int # seconds + token_type: str # usually "Bearer" + claims: dict[str, Any] # verified claims from id_token + + +@dataclass +class LogoutTokenClaims: + """Verified claims from a back-channel logout token.""" + + iss: str + aud: str | list[str] + sub: str | None + sid: str | None + iat: int + jti: str | None + + +class OIDCClient: + """Lightweight OIDC Relying Party client. + + One instance per (issuer, client_id, client_secret) tuple. + """ + + _DISCOVERY_TTL = 3600 # 1 hour + _JWKS_TTL = 3600 # 1 hour + + def __init__( + self, + *, + issuer: str, + client_id: str, + client_secret: str, + redirect_uri: str, + scopes: str, + http_client: httpx.AsyncClient | None = None, + ): + # Keep the issuer string verbatim (including any trailing "/") — the OIDC + # spec mandates strict byte-for-byte equality between ``self.issuer``, the + # issuer advertised by the discovery document, and the ``iss`` claim in + # tokens. Operators must configure ``OIDC_ENDPOINT`` to match EXACTLY + # what the IdP returns. + self.issuer = issuer + self.client_id = client_id + self.client_secret = client_secret + self.redirect_uri = redirect_uri + self.scopes = scopes + self._http = http_client or httpx.AsyncClient(timeout=10.0) + self._metadata: dict | None = None + self._metadata_fetched_at: float = 0.0 + self._jwks: JsonWebKey | None = None + self._jwks_fetched_at: float = 0.0 + + # ------------------------------------------------------------------ + # Discovery + # ------------------------------------------------------------------ + + async def discover(self) -> dict: + """Fetch and cache the OIDC discovery document. + + Returns the cached document if it is less than _DISCOVERY_TTL seconds old. + Raises ValueError if the returned issuer does not match the configured one. + """ + if self._metadata and (time.time() - self._metadata_fetched_at) < self._DISCOVERY_TTL: + return self._metadata + url = f"{self.issuer.rstrip('/')}/.well-known/openid-configuration" + resp = await self._http.get(url) + resp.raise_for_status() + self._metadata = resp.json() + self._metadata_fetched_at = time.time() + if self._metadata.get("issuer") != self.issuer: + raise ValueError(f"Issuer mismatch: configured {self.issuer!r}, got {self._metadata.get('issuer')!r}") + return self._metadata + + # ------------------------------------------------------------------ + # JWKS + # ------------------------------------------------------------------ + + async def _load_jwks(self, force: bool = False) -> JsonWebKey: + meta = await self.discover() + if not force and self._jwks and (time.time() - self._jwks_fetched_at) < self._JWKS_TTL: + return self._jwks + resp = await self._http.get(meta["jwks_uri"]) + resp.raise_for_status() + self._jwks = JsonWebKey.import_key_set(resp.json()) + self._jwks_fetched_at = time.time() + return self._jwks + + # ------------------------------------------------------------------ + # PKCE helpers + # ------------------------------------------------------------------ + + @staticmethod + def generate_pkce_pair() -> tuple[str, str]: + """Generate a PKCE (code_verifier, code_challenge) pair using S256. + + Returns: + (verifier, challenge) — verifier is 128 url-safe chars, + challenge is the base64url-encoded SHA-256 of the verifier. + """ + verifier = secrets.token_urlsafe(96)[:128] + digest = hashlib.sha256(verifier.encode()).digest() + challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() + return verifier, challenge + + @staticmethod + def generate_state_and_nonce() -> tuple[str, str]: + """Generate cryptographically random state and nonce values.""" + return secrets.token_urlsafe(32), secrets.token_urlsafe(32) + + # ------------------------------------------------------------------ + # Authorization URL + # ------------------------------------------------------------------ + + async def build_authorization_url(self, *, state: str, nonce: str, code_challenge: str) -> str: + """Build the full authorization URL to redirect the browser to.""" + meta = await self.discover() + params = { + "response_type": "code", + "client_id": self.client_id, + "redirect_uri": self.redirect_uri, + "scope": self.scopes, + "state": state, + "nonce": nonce, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + return f"{meta['authorization_endpoint']}?{urlencode(params)}" + + # ------------------------------------------------------------------ + # Code exchange + # ------------------------------------------------------------------ + + async def exchange_code(self, *, code: str, code_verifier: str, expected_nonce: str) -> TokenBundle: + """Exchange an authorization code for tokens. + + Verifies the returned id_token (signature, iss, aud, exp, nonce). + + Args: + code: The authorization code from the IdP callback. + code_verifier: The PKCE verifier corresponding to the challenge sent earlier. + expected_nonce: The nonce value that was sent in the authorization request. + + Returns: + A TokenBundle with verified claims. + """ + meta = await self.discover() + data = { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": self.redirect_uri, + "client_id": self.client_id, + "client_secret": self.client_secret, + "code_verifier": code_verifier, + } + resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) + resp.raise_for_status() + payload = resp.json() + id_token = payload["id_token"] + claims = await self._verify_id_token(id_token, expected_nonce=expected_nonce) + return TokenBundle( + id_token=id_token, + access_token=payload["access_token"], + refresh_token=payload.get("refresh_token"), + expires_in=int(payload.get("expires_in", 0)), + token_type=payload.get("token_type", "Bearer"), + claims=claims, + ) + + # ------------------------------------------------------------------ + # Token refresh + # ------------------------------------------------------------------ + + async def refresh_access_token(self, refresh_token: str) -> TokenBundle: + """Use the refresh_token to obtain a new access_token. + + If the IdP returns a new id_token, it is re-verified (nonce check skipped + per RFC 8252 S8.2 -- nonce is only required during the initial code exchange). + If the IdP omits the refresh_token in the response, the caller's existing + refresh_token is preserved. + + Returns: + A new TokenBundle. + """ + meta = await self.discover() + data = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": self.client_id, + "client_secret": self.client_secret, + } + resp = await self._http.post(meta["token_endpoint"], data=data, headers={"Accept": "application/json"}) + resp.raise_for_status() + payload = resp.json() + new_id_token = payload.get("id_token") + claims: dict[str, Any] = {} + if new_id_token: + claims = await self._verify_id_token(new_id_token, expected_nonce=None) + return TokenBundle( + id_token=new_id_token or "", + access_token=payload["access_token"], + # Some IdPs omit the refresh_token on rotation -- keep the old one. + refresh_token=payload.get("refresh_token", refresh_token), + expires_in=int(payload.get("expires_in", 0)), + token_type=payload.get("token_type", "Bearer"), + claims=claims, + ) + + # ------------------------------------------------------------------ + # Userinfo + # ------------------------------------------------------------------ + + async def fetch_userinfo(self, access_token: str) -> dict: + """Fetch the userinfo endpoint with the given access token.""" + meta = await self.discover() + resp = await self._http.get( + meta["userinfo_endpoint"], + headers={"Authorization": f"Bearer {access_token}"}, + ) + resp.raise_for_status() + return resp.json() + + # ------------------------------------------------------------------ + # ID token verification + # ------------------------------------------------------------------ + + async def _verify_id_token(self, token: str, *, expected_nonce: str | None) -> dict[str, Any]: + """Verify an ID token's signature and standard claims. + + Retries with a fresh JWKS fetch on kid-miss (covers IdP key rotation). + Raises JoseError / ValueError on any validation failure. + """ + jwks = await self._load_jwks() + jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) + try: + claims = jwt.decode(token, jwks) + except JoseError: + # Force JWKS refresh in case of kid rotation; retry once. + jwks = await self._load_jwks(force=True) + claims = jwt.decode(token, jwks) + + # Manual validation -- avoids authlib version differences around claims.params + decoded: dict[str, Any] = dict(claims) + now = int(time.time()) + + if decoded.get("iss") != self.issuer: + raise ValueError(f"ID token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") + + aud = decoded.get("aud") + if isinstance(aud, list): + if self.client_id not in aud: + raise ValueError(f"ID token aud {aud!r} does not contain client_id {self.client_id!r}") + elif aud != self.client_id: + raise ValueError(f"ID token aud {aud!r} != client_id {self.client_id!r}") + + if "exp" not in decoded: + raise ValueError("ID token missing exp claim") + if int(decoded["exp"]) < now: + raise ValueError("ID token has expired") + + if "iat" not in decoded: + raise ValueError("ID token missing iat claim") + + if expected_nonce is not None: + if decoded.get("nonce") != expected_nonce: + raise ValueError("OIDC nonce mismatch") + + return decoded + + # ------------------------------------------------------------------ + # Back-channel logout token verification + # ------------------------------------------------------------------ + + async def verify_logout_token(self, token: str) -> LogoutTokenClaims: + """Verify an OIDC back-channel logout token. + + Validates: + - Signature (with JWKS kid-miss retry) + - Standard claims (iss, aud, iat) + - events claim contains the back-channel-logout URI key + - nonce must NOT be present (spec requirement) + - At least one of sub or sid must be present + + Returns: + LogoutTokenClaims with the verified values. + Raises: + ValueError: on any spec violation. + """ + jwks = await self._load_jwks() + jwt = JsonWebToken(["RS256", "ES256", "EdDSA", "RS384", "RS512"]) + try: + claims = jwt.decode(token, jwks) + except JoseError: + jwks = await self._load_jwks(force=True) + claims = jwt.decode(token, jwks) + + decoded: dict[str, Any] = dict(claims) + now = int(time.time()) + + if decoded.get("iss") != self.issuer: + raise ValueError(f"logout_token iss mismatch: expected {self.issuer!r}, got {decoded.get('iss')!r}") + + aud = decoded.get("aud") + if isinstance(aud, list): + if self.client_id not in aud: + raise ValueError(f"logout_token aud {aud!r} does not contain client_id {self.client_id!r}") + elif aud != self.client_id: + raise ValueError(f"logout_token aud {aud!r} != client_id {self.client_id!r}") + + if "iat" not in decoded: + raise ValueError("logout_token missing iat claim") + if int(decoded.get("exp", now + 1)) < now: + raise ValueError("logout_token has expired") + + events = decoded.get("events") or {} + if "http://schemas.openid.net/event/backchannel-logout" not in events: + raise ValueError("logout_token missing required back-channel-logout event claim") + + if decoded.get("nonce"): + raise ValueError("logout_token must not contain nonce") + + if not decoded.get("sub") and not decoded.get("sid"): + raise ValueError("logout_token must contain sub or sid") + + return LogoutTokenClaims( + iss=decoded["iss"], + aud=decoded["aud"], + sub=decoded.get("sub"), + sid=decoded.get("sid"), + iat=int(decoded["iat"]), + jti=decoded.get("jti"), + ) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def aclose(self) -> None: + """Close the underlying HTTP client.""" + await self._http.aclose() diff --git a/openrag/services/auth/refresh.py b/openrag/services/auth/refresh.py new file mode 100644 index 000000000..a3715d589 --- /dev/null +++ b/openrag/services/auth/refresh.py @@ -0,0 +1,176 @@ +"""Lazy refresh helper for OIDC access tokens. + +Extracted from ``AuthMiddleware`` (Phase 5) to keep ``api.py`` small and +independently testable. Called per-request when a valid cookie session is +found; a no-op when the access token is still fresh. + +Timezone policy +--------------- +Phase 2 stores all OIDC session timestamps as **naive local time** via +``datetime.now()`` (see ``test_oidc_sessions.py`` and +``PartitionFileManager.get_oidc_session_by_token``). We match that style +everywhere in this module to avoid tz-mismatch bugs when comparing +``access_token_expires_at`` against "now". + +Refresh-token stampede guard (M1) +--------------------------------- +IdPs with refresh_token rotation enabled invalidate the old refresh_token the +first time it is redeemed. Under concurrency, multiple requests can each notice +"my access_token is about to expire" at the same time and race each other to +the token endpoint. The second attempt fails with ``invalid_grant`` and +(without a guard) its session would be revoked mid-flight. + +We mitigate that with two cooperating mechanisms: + +1. A **short-circuit** here: if ``last_refresh_at`` was bumped less than 5 + seconds ago, we assume a sibling request already rotated the tokens, + re-read the row, and reuse those freshly rotated tokens instead of calling + the IdP. +2. A **row-level write lock** in :meth:`PartitionFileManager.update_oidc_session_tokens` + (``SELECT ... FOR UPDATE``) so that only one writer commits at a time on + Postgres. +3. An **error-recovery branch** here: if the IdP does reject our refresh_token + (typically because a sibling raced us and won), we re-read the row once + more and, if the tokens were advanced meanwhile, return the fresh session + rather than giving up. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta +from typing import Any + +from services.auth.deps import get_oidc_client +from services.auth.session_tokens import decrypt_token, encrypt_token +from utils.logger import get_logger + +_REFRESH_BUFFER = timedelta(seconds=60) +_STAMPEDE_WINDOW = timedelta(seconds=5) + +logger = get_logger() + + +def _to_dt(val: Any) -> datetime: + """Coerce a datetime-or-ISO-string into a ``datetime``. + + Ray occasionally ships values across actors in serialised form; accept + either shape so callers never have to care about the transport. + """ + if isinstance(val, datetime): + return val + if isinstance(val, str): + return datetime.fromisoformat(val) + raise TypeError(f"Expected datetime or ISO string, got {type(val).__name__}") + + +async def refresh_session_if_needed( + *, + session: dict[str, Any], + enc_key: str, + vectordb: Any, +) -> dict[str, Any] | None: + """Refresh the IdP access_token if it is within ``_REFRESH_BUFFER`` of expiry. + + Behaviour: + - If the access_token is still valid with the 60s buffer -> return ``session`` unchanged. + - Stampede guard: if another request has just refreshed this session + (``last_refresh_at`` within 5s), re-read the row and reuse the fresh + tokens without calling the IdP. + - If near/past expiry AND a ``refresh_token_encrypted`` blob is stored -> + call the IdP, persist rotated tokens, return an updated session dict. + - If near/past expiry AND no refresh_token is stored -> return ``session`` as-is + when still formally valid, or ``None`` when already expired (caller should + treat as a revoked session). + - If the refresh call raises (typically because a sibling already rotated + the tokens and the IdP now rejects ours) -> re-read the row; if a sibling + succeeded, return their fresh session; otherwise ``None``. + + The session dict returned mirrors the DB row shape produced by + ``PartitionFileManager._oidc_session_to_dict``. + """ + now = datetime.now() + access_exp = _to_dt(session["access_token_expires_at"]) + + if access_exp > now + _REFRESH_BUFFER: + return session + + # --- Stampede short-circuit ------------------------------------------- + # If a sibling request just refreshed this same session, re-read the row + # and reuse the freshly rotated tokens. This avoids racing the IdP with a + # refresh_token that the sibling's success has already invalidated. + last_refresh_at = session.get("last_refresh_at") + if last_refresh_at is not None: + try: + last_refresh_at_dt = _to_dt(last_refresh_at) + except TypeError: + last_refresh_at_dt = None + if last_refresh_at_dt is not None and (now - last_refresh_at_dt) < _STAMPEDE_WINDOW: + try: + fresh = await vectordb.get_oidc_session_by_id.remote(session["id"]) + except Exception as e: + logger.bind(session_id=session.get("id"), error=str(e)).warning( + "Stampede-guard re-read failed; falling through to refresh" + ) + fresh = None + if fresh is not None: + fresh_exp = _to_dt(fresh["access_token_expires_at"]) + if fresh_exp > now + _REFRESH_BUFFER: + return fresh + + refresh_enc = session.get("refresh_token_encrypted") + if not refresh_enc: + # No refresh_token available. + # - If still formally valid (within the 60s buffer window but not yet past exp), + # keep using it. + # - If already expired, caller should treat the session as dead. + return session if access_exp > now else None + + try: + refresh_token = decrypt_token(refresh_enc, enc_key) + client = get_oidc_client() + bundle = await client.refresh_access_token(refresh_token) + except Exception as e: + # Maybe a sibling refreshed between our staleness check and the IdP call + # and the IdP has already invalidated our refresh_token. Re-read the + # row once before giving up: if the tokens were rotated meanwhile, + # treat this as a successful refresh (the sibling's). + logger.bind(session_id=session.get("id"), error=str(e)).warning( + "OIDC refresh_token exchange failed -- re-reading session for stampede recovery" + ) + try: + fresh = await vectordb.get_oidc_session_by_id.remote(session["id"]) + except Exception as re: + logger.bind(session_id=session.get("id"), error=str(re)).error( + "Post-failure re-read of OIDC session failed -- invalidating" + ) + return None + if fresh is not None: + fresh_exp = _to_dt(fresh["access_token_expires_at"]) + if fresh_exp > now + _REFRESH_BUFFER: + return fresh + return None + + new_access_exp = now + timedelta(seconds=max(int(bundle.expires_in or 0), 60)) + new_access_enc = encrypt_token(bundle.access_token, enc_key) + new_refresh_enc = encrypt_token(bundle.refresh_token, enc_key) if bundle.refresh_token else refresh_enc + + try: + await vectordb.update_oidc_session_tokens.remote( + session_id=session["id"], + access_token_encrypted=new_access_enc, + refresh_token_encrypted=new_refresh_enc, + access_token_expires_at=new_access_exp, + ) + except Exception as e: + logger.bind(session_id=session.get("id"), error=str(e)).error( + "Failed to persist refreshed OIDC tokens -- invalidating session" + ) + return None + + return { + **session, + "access_token_encrypted": new_access_enc, + "access_token_expires_at": new_access_exp, + "refresh_token_encrypted": new_refresh_enc, + "last_refresh_at": now, + } diff --git a/openrag/services/auth/session_tokens.py b/openrag/services/auth/session_tokens.py new file mode 100644 index 000000000..68b5c18d7 --- /dev/null +++ b/openrag/services/auth/session_tokens.py @@ -0,0 +1,64 @@ +"""Session token utilities for OpenRAG OIDC sessions. + +Opaque session tokens are issued at callback and stored hashed (SHA-256) in the DB. +IdP tokens (access_token, refresh_token) are encrypted with Fernet before storage. + +The Fernet key is provided via the OIDC_TOKEN_ENCRYPTION_KEY environment variable. +Generate one with: + python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())' +""" + +import hashlib +import secrets + +from cryptography.fernet import Fernet, InvalidToken + + +def issue_session_token() -> tuple[str, str]: + """Generate a new session token. + + Returns: + (plaintext, sha256_hex) -- the plaintext is set in the cookie, + the hash is stored in the database. + """ + plain = secrets.token_urlsafe(32) # 43 chars, >= 256 bits entropy + return plain, hash_session_token(plain) + + +def hash_session_token(token: str) -> str: + """Return the SHA-256 hex digest of the session token.""" + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def _fernet(key: str | bytes) -> Fernet: + try: + return Fernet(key.encode("utf-8") if isinstance(key, str) else key) + except Exception as e: + raise ValueError( + "OIDC_TOKEN_ENCRYPTION_KEY is not a valid Fernet key. " + "Generate one with: python -c 'from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())'" + ) from e + + +def encrypt_token(plaintext: str | None, key: str) -> bytes | None: + """Encrypt a plaintext token string. + + Returns None if plaintext is None (refresh_token may be absent). + """ + if plaintext is None: + return None + return _fernet(key).encrypt(plaintext.encode("utf-8")) + + +def decrypt_token(ciphertext: bytes | None, key: str) -> str | None: + """Decrypt a Fernet-encrypted token. + + Returns None if ciphertext is None. + Raises ValueError on key mismatch or data corruption. + """ + if ciphertext is None: + return None + try: + return _fernet(key).decrypt(ciphertext).decode("utf-8") + except InvalidToken as e: + raise ValueError("Failed to decrypt stored OIDC token -- key mismatch or corruption") from e diff --git a/openrag/services/auth/state_cookie.py b/openrag/services/auth/state_cookie.py new file mode 100644 index 000000000..d22db7006 --- /dev/null +++ b/openrag/services/auth/state_cookie.py @@ -0,0 +1,55 @@ +"""Signed state cookie for OIDC Authorization Code + PKCE flow. + +The cookie transports state/nonce/code_verifier between /auth/login and /auth/callback. +It is signed (not encrypted) using itsdangerous.URLSafeTimedSerializer with HMAC-SHA1. + +The signing key is the OIDC_TOKEN_ENCRYPTION_KEY (a Fernet base64url key, which is +valid arbitrary bytes for HMAC). The consuming code (phase 4 router) will pass the +key to StateCookieSerializer(key). Using the same key for both Fernet encryption and +HMAC signing is safe since itsdangerous derives separate subkeys via HMAC. + +TTL defaults to 600 s (10 minutes) -- long enough for a slow user at the IdP login page. +""" + +from dataclasses import asdict, dataclass + +from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer + + +@dataclass +class StateCookiePayload: + state: str + nonce: str + code_verifier: str + next_url: str = "/" + + +class StateCookieSerializer: + """Signs/verifies the short-lived cookie holding OIDC state/nonce/code_verifier. + + TTL defaults to 600 s (10 minutes) -- long enough for a slow user at the IdP. + """ + + COOKIE_NAME = "openrag_oidc_state" + DEFAULT_TTL_SECONDS = 600 + + def __init__(self, secret_key: str, salt: str = "openrag-oidc-state-v1"): + self._serializer = URLSafeTimedSerializer(secret_key, salt=salt) + + def dumps(self, payload: StateCookiePayload) -> str: + """Serialize and sign the payload, returning an opaque cookie value.""" + return self._serializer.dumps(asdict(payload)) + + def loads(self, token: str, max_age: int = DEFAULT_TTL_SECONDS) -> StateCookiePayload: + """Verify and deserialize the cookie value. + + Raises: + ValueError: if the cookie is expired or the signature is invalid. + """ + try: + data = self._serializer.loads(token, max_age=max_age) + except SignatureExpired as e: + raise ValueError("OIDC state cookie expired") from e + except BadSignature as e: + raise ValueError("OIDC state cookie signature invalid") from e + return StateCookiePayload(**data)