Coverage for src/ecnl/security/jwks.py: 94%

1"""JWKS client for fetching auth-server's signing keys. 

2 

3Per identity-platform-go ADR-0008, the auth-server publishes its RSA 

4public keys at ``/.well-known/jwks.json``. This module wraps that 

5endpoint with a small caching layer keyed on ``kid`` (RFC 7517 §4.5) 

6so a token whose JOSE header references a known key validates without 

7a network round-trip on every request. 

8 

9The cache TTL is intentionally long (default 1 hour). A key rotation 

10takes effect on the next refresh — auth-server's rotation cadence 

11(ADR-0008) is operator-controlled at a scale of days, so a one-hour 

12staleness window is harmless. A token signed with a not-yet-cached 

13key triggers an immediate refetch so newly-promoted keys validate 

14without waiting for the TTL. 

15""" 

16 

17from __future__ import annotations 

18 

19import asyncio 

20import logging 

21import time 

22from dataclasses import dataclass, field 

23 

24import httpx 

25import jwt 

26from jwt import PyJWK 

27 

28logger = logging.getLogger(__name__) 

29 

30 

31@dataclass 

32class JWKSCache: 

33 """In-memory cache of JWKS entries keyed by ``kid``. 

34 

35 The cache is intentionally per-process — every MCP server replica 

36 fetches its own copy. Sharing across replicas via Redis would add 

37 operational complexity for negligible gain at the volumes the MCP 

38 deployment sees today. 

39 """ 

40 

41 issuer_url: str 

42 ttl_seconds: int = 3600 

43 http_client: httpx.AsyncClient | None = None 

44 

45 _keys: dict[str, PyJWK] = field(default_factory=dict, init=False) 

46 _expires_at: float = field(default=0.0, init=False) 

47 _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False) 

48 

49 async def get(self, kid: str) -> PyJWK: 

50 """Return the JWK for ``kid``. 

51 

52 Triggers a refresh when the cache is stale or the requested 

53 ``kid`` is unknown. Raises :class:`KeyError` when the kid is 

54 still missing after a fresh fetch — auth-server has rotated 

55 beyond this key or the caller is using a forged header. 

56 """ 

57 async with self._lock: 

58 if self._stale() or kid not in self._keys: 

59 await self._refresh() 

60 try: 

61 return self._keys[kid] 

62 except KeyError as exc: 

63 raise KeyError(f"unknown JWKS kid: {kid}") from exc 

64 

65 def _stale(self) -> bool: 

66 return time.time() >= self._expires_at 

67 

68 async def _refresh(self) -> None: 

69 client = self.http_client or httpx.AsyncClient(timeout=5.0) 

70 owns_client = self.http_client is None 

71 url = self.issuer_url.rstrip("/") + "/.well-known/jwks.json" 

72 try: 

73 resp = await client.get(url) 

74 resp.raise_for_status() 

75 payload = resp.json() 

76 finally: 

77 if owns_client: 

78 await client.aclose() 

79 

80 keys: dict[str, PyJWK] = {} 

81 for raw in payload.get("keys", []): 

82 try: 

83 kid = raw["kid"] 

84 except KeyError: 

85 logger.warning("JWKS entry missing kid; skipping: %s", raw) 

86 continue 

87 keys[kid] = PyJWK(raw) 

88 self._keys = keys 

89 self._expires_at = time.time() + self.ttl_seconds 

90 logger.debug("JWKS cache refreshed: %d keys, ttl=%ds", len(keys), self.ttl_seconds) 

91 

92 

93def jwks_url(issuer_url: str) -> str: 

94 """Compose the well-known JWKS URL from an issuer base. 

95 

96 Public so server-side wiring code can log the resolved URL at 

97 startup without re-deriving it; the join logic is trivial but 

98 centralizing it keeps the contract with auth-server in one spot. 

99 """ 

100 return issuer_url.rstrip("/") + "/.well-known/jwks.json" 

101 

102 

103# Tiny PyJWT alias re-exported so test fixtures can construct fake 

104# tokens without importing PyJWT directly — keeps the test surface 

105# coupled to this module rather than the underlying library. 

106encode = jwt.encode