Coverage for src/lumen/auth/jwks.py: 98%

1""" 

2This file manages caching of the JSON Web Key Sets (JWKS) from IdP endpoints. 

3It fetches and updates public keys with lock protection to prevent concurrent re-fetches. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import json 

10import time 

11from typing import Any 

12 

13import httpx 

14from jwt.algorithms import RSAAlgorithm 

15 

16 

17class JwksCache: 

18 def __init__(self, ttl: int = 300) -> None: 

19 self.ttl = ttl 

20 # jwks_uri -> (keys, expires_at) 

21 self._cache: dict[str, tuple[list[dict[str, Any]], float]] = {} 

22 # One lock per JWKS URI to avoid concurrent re-fetches. 

23 self._locks: dict[str, asyncio.Lock] = {} 

24 

25 async def _get_lock(self, key: str) -> asyncio.Lock: 

26 if key not in self._locks: 

27 self._locks[key] = asyncio.Lock() 

28 return self._locks[key] 

29 

30 async def _fetch_jwks(self, jwks_uri: str) -> list[dict[str, Any]]: 

31 async with httpx.AsyncClient(timeout=10) as client: 

32 resp = await client.get(jwks_uri) 

33 resp.raise_for_status() 

34 return resp.json()["keys"] # type: ignore[no-any-return] 

35 

36 async def get_keys(self, jwks_uri: str) -> list[dict[str, Any]]: 

37 now = time.monotonic() 

38 cached = self._cache.get(jwks_uri) 

39 if cached and now < cached[1]: 

40 return cached[0] 

41 

42 lock = await self._get_lock(jwks_uri) 

43 async with lock: 

44 # Re-check after acquiring the lock (another coroutine may have fetched). 

45 cached = self._cache.get(jwks_uri) 

46 if cached and now < cached[1]: 

47 return cached[0] 

48 

49 keys = await self._fetch_jwks(jwks_uri) 

50 self._cache[jwks_uri] = (keys, now + self.ttl) 

51 return keys 

52 

53 async def get_public_key(self, jwks_uri: str, kid: str) -> Any: 

54 keys = await self.get_keys(jwks_uri) 

55 for key in keys: 

56 if key.get("kid") == kid: 

57 return RSAAlgorithm.from_jwk(json.dumps(key)) 

58 raise KeyError(f"kid '{kid}' not found in JWKS at {jwks_uri}") 

59 

60 def invalidate(self, jwks_uri: str) -> None: 

61 self._cache.pop(jwks_uri, None) 

62 

63 

64async def discover_oidc_config(issuer: str, *, fetch_base: str | None = None) -> dict[str, Any]: 

65 """Fetch and return the full OIDC discovery document for an issuer. 

66 

67 `fetch_base` overrides the base URL used for the HTTP request (useful when 

68 Lumen runs in Docker and the external issuer URL is not reachable internally). 

69 The returned document is served as-is from the IdP — callers should use 

70 `issuer` for tenant/iss matching and `fetch_base` only for the network call. 

71 """ 

72 base = (fetch_base or issuer).rstrip("/") 

73 discovery_url = f"{base}/.well-known/openid-configuration" 

74 async with httpx.AsyncClient(timeout=10) as client: 

75 resp = await client.get(discovery_url) 

76 resp.raise_for_status() 

77 return resp.json() # type: ignore[no-any-return] 

78 

79 

80async def discover_jwks_uri(issuer: str, *, fetch_base: str | None = None) -> str: 

81 """Resolve the JWKS URI from an OIDC discovery document.""" 

82 config = await discover_oidc_config(issuer, fetch_base=fetch_base) 

83 return config["jwks_uri"] # type: ignore[no-any-return] 

84 

85 

86# Module-level singleton — one cache per process. 

87_cache: JwksCache | None = None 

88 

89 

90def get_jwks_cache(ttl: int = 300) -> JwksCache: 

91 global _cache 

92 if _cache is None: 

93 _cache = JwksCache(ttl=ttl) 

94 return _cache 

95 

96 

97def reset_jwks_cache() -> None: 

98 """Replace the singleton with a fresh cache. Intended for tests only.""" 

99 global _cache 

100 _cache = None