Skip to content

Commit a1e6fc2

Browse files
larrysingleton007larrysingleton007
authored
fix: Reuse the OIDC JWKS client across requests (#6683)
* fix: Reuse the OIDC JWKS client across requests _decode_token constructed a fresh PyJWKClient on every call. The JWK-set cache lives on the instance, so every authenticated request started cold and paid a full HTTPS fetch of the JWKS document (TLS handshake included) before signature verification. That blocking I/O also runs on the event loop, so concurrent requests serialized behind it. Build the client lazily once per parser instead (the parser is constructed once per process by init_auth_manager), letting PyJWT's built-in JWK-set cache (default lifespan 300s) do its job. IdP key rotation stays safe: PyJWKClient.get_signing_key refreshes the set and retries once whenever it sees an unknown kid. Local benchmark, 25 authenticated decodes against a mock JWKS endpoint: 25 JWKS fetches before, 1 after. Signed-off-by: Larry Singleton <[email protected]> * fix: Harden JWKS client reuse from review findings Floor pyjwt at 2.13.0: the reuse makes PyJWT's JWK-set cache semantics load-bearing, and only >=2.13 keeps the cache through transient fetch failures (2.5-2.12 wipe it; pre-2.5 has no cache, silently reverting the fix). Pass lifespan and timeout explicitly so upgrades cannot change the documented staleness window and a hung IdP bounds how long a fetch can block the serving path. Scope the rotation docstring honestly: new-kid rotations recover immediately, but a removed key keeps validating and a reused-kid rotation keeps failing for up to the cache lifespan. Test hardening from mutation results: pin the SSL context per verify_ssl config (both dropping ssl_context and inverting the guard previously left all tests green), pin laziness (no construction before the first request), pin per-parser scoping (second parser builds its own client), and pin construction-failure recovery (a failed first build must not wedge the process-singleton parser). Signed-off-by: Larry Singleton <[email protected]> * feat: Make the JWKS cache lifespan and fetch timeout configurable Review feedback on #6683: expose the two PyJWKClient tunables as OidcAuthConfig options rather than hardcoding them. jwks_cache_lifespan_seconds (default 300) also bounds how long a key the IdP has revoked keeps validating tokens, so operators whose provider rotates or revokes aggressively can tighten it at the cost of more JWKS fetches. jwks_request_timeout_seconds (default 10) bounds how long an unresponsive IdP blocks the serving path. Both reject non-positive values at config load: a zero or negative lifespan would expire the cache immediately and silently restore a JWKS fetch per request, undoing this PR. Signed-off-by: Larry Singleton <[email protected]> * chore: Update pixi.lock for the pyjwt floor CI runs pixi install --locked, which fails at setup when pyproject.toml and pixi.lock disagree, so adding the pyjwt>=2.13.0 bound took out all five pixi-based integration jobs before any test ran. Regenerate the lock; the only change is the corresponding requires_dist entry. Signed-off-by: Larry Singleton <[email protected]> --------- Signed-off-by: Larry Singleton <[email protected]>
1 parent 3a6a103 commit a1e6fc2

6 files changed

Lines changed: 236 additions & 18 deletions

File tree

‎docs/getting-started/components/authz_manager.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,19 @@ A token whose `aud` (or `iss`) claim does not match is rejected at authenticatio
122122
Set these to the values your IdP puts **in the token itself**, which are not always the ones in the discovery document. For example, Microsoft Entra ID commonly issues v1.0 tokens (`iss: https://sts.windows.net/<tenant-id>/`, `aud: api://<app-id-uri>`) even when `auth_discovery_url` points at the v2.0 endpoint. That setup keeps working with these options unset, or set to the v1.0 values — but copying the v2.0 issuer from the discovery document would reject every v1.0 token.
123123
{% endhint %}
124124

125+
To validate token signatures the server fetches the provider's JWKS document and caches it, refetching when the cache expires or when a token presents an unknown key id. Two options tune that behavior:
126+
127+
```yaml
128+
auth:
129+
type: oidc
130+
client_id: _CLIENT_ID_
131+
auth_discovery_url: https://login.example.com/.well-known/openid-configuration
132+
jwks_cache_lifespan_seconds: 300 # default; how long the fetched key set is reused
133+
jwks_request_timeout_seconds: 10 # default; network timeout for the JWKS fetch
134+
```
135+
136+
`jwks_cache_lifespan_seconds` also bounds how long a key the provider has **revoked** continues to validate tokens, so lower it if your provider rotates or revokes aggressively; each reduction costs proportionally more JWKS fetches. Key rotations that introduce a new key id are picked up immediately regardless of this setting, because an unknown key id triggers a refetch. `jwks_request_timeout_seconds` bounds how long an unresponsive provider can block request serving. Both must be greater than zero.
137+
125138
#### Client-Side Configuration
126139

127140
The client supports multiple token source modes. The SDK resolves tokens in the following priority order:

‎pixi.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,9 @@ dependencies = [
4343
"prometheus_client>=0.20.0,<0.25.0",
4444
"psutil",
4545
"bigtree>=0.19.2",
46-
"pyjwt",
46+
# >=2.13 for PyJWKClient's JWK-set cache surviving transient fetch
47+
# failures; the OIDC token parser reuses one client and relies on it.
48+
"pyjwt>=2.13.0",
4749
]
4850

4951
[project.optional-dependencies]

‎sdk/python/feast/permissions/auth/oidc_token_parser.py‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,41 @@ def __init__(self, auth_config: OidcAuthConfig):
4141
ca_cert_path=self._auth_config.ca_cert_path,
4242
)
4343
self._k8s_auth_api = None
44+
self._jwks_client: Optional[PyJWKClient] = None # Initialize it lazily.
45+
46+
def _get_jwks_client(self) -> PyJWKClient:
47+
"""Lazily build and cache a parser-lifetime ``PyJWKClient``.
48+
49+
A per-request client starts with a cold JWK-set cache, forcing a
50+
full HTTPS fetch of the JWKS document on every authenticated
51+
request. Reusing one client lets PyJWT cache the JWK set for
52+
``jwks_cache_lifespan_seconds``, which also bounds two staleness
53+
windows: a key the IdP has removed keeps validating, and a
54+
rotation that reuses an existing ``kid`` keeps failing, for at
55+
most that long. Rotations that introduce a new ``kid`` recover
56+
immediately (``PyJWKClient.get_signing_key`` refreshes and
57+
retries once on a cache miss).
58+
"""
59+
if self._jwks_client is None:
60+
ssl_ctx = ssl.create_default_context()
61+
if not self._auth_config.verify_ssl:
62+
ssl_ctx.check_hostname = False
63+
ssl_ctx.verify_mode = ssl.CERT_NONE
64+
elif self._auth_config.ca_cert_path and os.path.exists(
65+
self._auth_config.ca_cert_path
66+
):
67+
ssl_ctx.load_verify_locations(self._auth_config.ca_cert_path)
68+
self._jwks_client = PyJWKClient(
69+
self.oidc_discovery_service.get_jwks_url(),
70+
headers={"User-agent": "custom-user-agent"},
71+
ssl_context=ssl_ctx,
72+
# Explicit so upgrades cannot silently change the staleness
73+
# window documented above, and so a hung IdP bounds how long
74+
# a fetch can block the serving path.
75+
lifespan=self._auth_config.jwks_cache_lifespan_seconds,
76+
timeout=self._auth_config.jwks_request_timeout_seconds,
77+
)
78+
return self._jwks_client
4479

4580
async def _validate_token(self, access_token: str):
4681
"""
@@ -125,21 +160,7 @@ def _decode_token(self, access_token: str) -> dict:
125160
metadata (e.g. Entra ID v1.0 tokens validated against a v2.0
126161
discovery document).
127162
"""
128-
optional_custom_headers = {"User-agent": "custom-user-agent"}
129-
ssl_ctx = ssl.create_default_context()
130-
if not self._auth_config.verify_ssl:
131-
ssl_ctx.check_hostname = False
132-
ssl_ctx.verify_mode = ssl.CERT_NONE
133-
elif self._auth_config.ca_cert_path and os.path.exists(
134-
self._auth_config.ca_cert_path
135-
):
136-
ssl_ctx.load_verify_locations(self._auth_config.ca_cert_path)
137-
jwks_client = PyJWKClient(
138-
self.oidc_discovery_service.get_jwks_url(),
139-
headers=optional_custom_headers,
140-
ssl_context=ssl_ctx,
141-
)
142-
signing_key = jwks_client.get_signing_key_from_jwt(access_token)
163+
signing_key = self._get_jwks_client().get_signing_key_from_jwt(access_token)
143164
expected_audience = self._auth_config.audience
144165
expected_issuer = self._auth_config.issuer
145166
return jwt.decode(

‎sdk/python/feast/permissions/auth_model.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from typing import Literal, Optional, Tuple
44

5-
from pydantic import ConfigDict, model_validator
5+
from pydantic import ConfigDict, Field, model_validator
66

77
from feast.repo_config import FeastConfigBaseModel
88

@@ -47,6 +47,15 @@ class OidcAuthConfig(AuthConfig):
4747
# against a v2.0 discovery URL).
4848
audience: Optional[str] = None
4949
issuer: Optional[str] = None
50+
# How long the fetched JWK set is reused before the server refetches it.
51+
# This also bounds how long a key the IdP has revoked keeps validating
52+
# tokens, so lower it if your provider rotates or revokes aggressively;
53+
# every reduction costs a corresponding increase in JWKS fetches.
54+
jwks_cache_lifespan_seconds: int = Field(default=300, gt=0)
55+
# Network timeout for the JWKS fetch. This fetch happens inline on the
56+
# request path, so an unresponsive IdP blocks serving for at most this
57+
# long.
58+
jwks_request_timeout_seconds: float = Field(default=10, gt=0)
5059

5160

5261
class OidcClientAuthConfig(OidcAuthConfig):

‎sdk/python/tests/unit/permissions/auth/test_token_parser.py‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
import asyncio
22
import os
3+
import ssl
34
import time
45
from unittest import mock
56
from unittest.mock import MagicMock, patch
67

78
import assertpy
89
import jwt
910
import pytest
11+
from pydantic import ValidationError
1012
from starlette.authentication import (
1113
AuthenticationError,
1214
)
@@ -472,6 +474,177 @@ async def mock_oath2(self, request):
472474
assertpy.assert_that(user.has_matching_role(["updater"])).is_false()
473475

474476

477+
# ---------------------------------------------------------------------------
478+
# JWKS client lifecycle (one lazy client per parser)
479+
# ---------------------------------------------------------------------------
480+
481+
482+
@patch(
483+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
484+
)
485+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
486+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
487+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
488+
def test_oidc_jwks_client_is_lazy_and_reused_per_parser(
489+
mock_discovery_data,
490+
mock_jwks_client_cls,
491+
mock_jwt,
492+
mock_oauth2,
493+
oidc_config,
494+
discovery_data,
495+
):
496+
"""One JWKS client per parser: built on the first request (not at
497+
construction, which would move a blocking discovery fetch into server
498+
startup), reused across requests, and scoped to the parser instance."""
499+
mock_discovery_data.return_value = discovery_data
500+
mock_jwt.return_value = {"preferred_username": "my-name"}
501+
502+
token_parser = OidcTokenParser(auth_config=oidc_config)
503+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(0)
504+
505+
for _ in range(3):
506+
asyncio.run(
507+
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
508+
)
509+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(1)
510+
511+
# A second parser must not see the first parser's client: each parser
512+
# verifies against the JWKS of its own configured provider.
513+
other_parser = OidcTokenParser(auth_config=oidc_config)
514+
asyncio.run(other_parser.user_details_from_access_token(access_token="aaa-bbb-ccc"))
515+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(2)
516+
517+
518+
@pytest.mark.parametrize("verify_ssl", [True, False])
519+
@patch(
520+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
521+
)
522+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
523+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
524+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
525+
def test_oidc_jwks_client_ssl_context_follows_config(
526+
mock_discovery_data,
527+
mock_jwks_client_cls,
528+
mock_jwt,
529+
mock_oauth2,
530+
verify_ssl,
531+
discovery_data,
532+
):
533+
"""The client is built with the discovery JWKS URL and an SSL context
534+
matching verify_ssl: default configs must keep certificate verification
535+
on, and verify_ssl=False must be the only way to turn it off."""
536+
mock_discovery_data.return_value = discovery_data
537+
mock_jwt.return_value = {"preferred_username": "my-name"}
538+
539+
token_parser = OidcTokenParser(auth_config=_oidc_config_with(verify_ssl=verify_ssl))
540+
asyncio.run(token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc"))
541+
542+
call = mock_jwks_client_cls.call_args
543+
assertpy.assert_that(call.args[0]).is_equal_to(discovery_data["jwks_uri"])
544+
ssl_ctx = call.kwargs["ssl_context"]
545+
if verify_ssl:
546+
assertpy.assert_that(ssl_ctx.verify_mode).is_equal_to(ssl.CERT_REQUIRED)
547+
assertpy.assert_that(ssl_ctx.check_hostname).is_true()
548+
else:
549+
assertpy.assert_that(ssl_ctx.verify_mode).is_equal_to(ssl.CERT_NONE)
550+
assertpy.assert_that(ssl_ctx.check_hostname).is_false()
551+
552+
553+
@pytest.mark.parametrize(
554+
"overrides,expected_lifespan,expected_timeout",
555+
[
556+
({}, 300, 10),
557+
(
558+
{
559+
"jwks_cache_lifespan_seconds": 60,
560+
"jwks_request_timeout_seconds": 2.5,
561+
},
562+
60,
563+
2.5,
564+
),
565+
],
566+
)
567+
@patch(
568+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
569+
)
570+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
571+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
572+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
573+
def test_oidc_jwks_client_cache_and_timeout_follow_config(
574+
mock_discovery_data,
575+
mock_jwks_client_cls,
576+
mock_jwt,
577+
mock_oauth2,
578+
overrides,
579+
expected_lifespan,
580+
expected_timeout,
581+
discovery_data,
582+
):
583+
"""The JWK-set cache lifespan and the fetch timeout are operator-tunable:
584+
the lifespan bounds how long a revoked key keeps validating, and the
585+
timeout bounds how long an unresponsive IdP blocks the serving path."""
586+
mock_discovery_data.return_value = discovery_data
587+
mock_jwt.return_value = {"preferred_username": "my-name"}
588+
589+
token_parser = OidcTokenParser(auth_config=_oidc_config_with(**overrides))
590+
asyncio.run(token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc"))
591+
592+
kwargs = mock_jwks_client_cls.call_args.kwargs
593+
assertpy.assert_that(kwargs["lifespan"]).is_equal_to(expected_lifespan)
594+
assertpy.assert_that(kwargs["timeout"]).is_equal_to(expected_timeout)
595+
596+
597+
@pytest.mark.parametrize(
598+
"overrides",
599+
[
600+
{"jwks_cache_lifespan_seconds": 0},
601+
{"jwks_cache_lifespan_seconds": -1},
602+
{"jwks_request_timeout_seconds": 0},
603+
{"jwks_request_timeout_seconds": -1},
604+
],
605+
)
606+
def test_oidc_jwks_tunables_reject_non_positive_values(overrides):
607+
"""A non-positive lifespan would expire the cache immediately, silently
608+
restoring a JWKS fetch per request; a non-positive timeout is equally
609+
meaningless. Reject both at config load rather than at serving time."""
610+
with pytest.raises(ValidationError):
611+
_oidc_config_with(**overrides)
612+
613+
614+
@patch(
615+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
616+
)
617+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
618+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
619+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
620+
def test_oidc_jwks_client_construction_failure_is_retried(
621+
mock_discovery_data,
622+
mock_jwks_client_cls,
623+
mock_jwt,
624+
mock_oauth2,
625+
oidc_config,
626+
discovery_data,
627+
):
628+
"""A failed first construction must leave the parser able to retry on
629+
the next request: the parser is a process singleton, so caching a failed
630+
or half-built client would wedge authentication until restart."""
631+
mock_discovery_data.return_value = discovery_data
632+
mock_jwt.return_value = {"preferred_username": "my-name"}
633+
mock_jwks_client_cls.side_effect = [RuntimeError("IdP unreachable"), MagicMock()]
634+
635+
token_parser = OidcTokenParser(auth_config=oidc_config)
636+
with pytest.raises(RuntimeError):
637+
asyncio.run(
638+
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
639+
)
640+
641+
user = asyncio.run(
642+
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
643+
)
644+
assertpy.assert_that(user.username).is_equal_to("my-name")
645+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(2)
646+
647+
475648
# ---------------------------------------------------------------------------
476649
# Optional audience / issuer verification (opt-in via OidcAuthConfig)
477650
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)