Skip to content

Commit 602feb0

Browse files
larrysingleton007larrysingleton007
committed
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 <166439969+larrysingleton007@users.noreply.github.com>
1 parent 4b349fc commit 602feb0

3 files changed

Lines changed: 128 additions & 40 deletions

File tree

‎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: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -46,13 +46,15 @@ def __init__(self, auth_config: OidcAuthConfig):
4646
def _get_jwks_client(self) -> PyJWKClient:
4747
"""Lazily build and cache a parser-lifetime ``PyJWKClient``.
4848
49-
A per-request client starts with a cold JWK-set cache, so every
50-
authenticated request pays a full HTTPS fetch of the JWKS document
51-
(TLS handshake included) before signature verification, and
52-
concurrent requests serialize behind that blocking I/O. Reusing one
53-
client lets PyJWT's built-in JWK-set cache do its job; on IdP key
54-
rotation the client refetches automatically when it sees an unknown
55-
``kid`` (``PyJWKClient.get_signing_key`` refreshes and retries once).
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+
``lifespan`` seconds, which also bounds two staleness windows: a
53+
key the IdP has removed keeps validating, and a rotation that
54+
reuses an existing ``kid`` keeps failing, for at most that long.
55+
Rotations that introduce a new ``kid`` recover immediately
56+
(``PyJWKClient.get_signing_key`` refreshes and retries once on a
57+
cache miss).
5658
"""
5759
if self._jwks_client is None:
5860
ssl_ctx = ssl.create_default_context()
@@ -67,6 +69,11 @@ def _get_jwks_client(self) -> PyJWKClient:
6769
self.oidc_discovery_service.get_jwks_url(),
6870
headers={"User-agent": "custom-user-agent"},
6971
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=300,
76+
timeout=10,
7077
)
7178
return self._jwks_client
7279

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

Lines changed: 111 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import asyncio
22
import os
3+
import ssl
34
import time
45
from unittest import mock
56
from unittest.mock import MagicMock, patch
@@ -472,6 +473,116 @@ async def mock_oath2(self, request):
472473
assertpy.assert_that(user.has_matching_role(["updater"])).is_false()
473474

474475

476+
# ---------------------------------------------------------------------------
477+
# JWKS client lifecycle (one lazy client per parser)
478+
# ---------------------------------------------------------------------------
479+
480+
481+
@patch(
482+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
483+
)
484+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
485+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
486+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
487+
def test_oidc_jwks_client_is_lazy_and_reused_per_parser(
488+
mock_discovery_data,
489+
mock_jwks_client_cls,
490+
mock_jwt,
491+
mock_oauth2,
492+
oidc_config,
493+
discovery_data,
494+
):
495+
"""One JWKS client per parser: built on the first request (not at
496+
construction, which would move a blocking discovery fetch into server
497+
startup), reused across requests, and scoped to the parser instance."""
498+
mock_discovery_data.return_value = discovery_data
499+
mock_jwt.return_value = {"preferred_username": "my-name"}
500+
501+
token_parser = OidcTokenParser(auth_config=oidc_config)
502+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(0)
503+
504+
for _ in range(3):
505+
asyncio.run(
506+
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
507+
)
508+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(1)
509+
510+
# A second parser must not see the first parser's client: each parser
511+
# verifies against the JWKS of its own configured provider.
512+
other_parser = OidcTokenParser(auth_config=oidc_config)
513+
asyncio.run(other_parser.user_details_from_access_token(access_token="aaa-bbb-ccc"))
514+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(2)
515+
516+
517+
@pytest.mark.parametrize("verify_ssl", [True, False])
518+
@patch(
519+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
520+
)
521+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
522+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
523+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
524+
def test_oidc_jwks_client_ssl_context_follows_config(
525+
mock_discovery_data,
526+
mock_jwks_client_cls,
527+
mock_jwt,
528+
mock_oauth2,
529+
verify_ssl,
530+
discovery_data,
531+
):
532+
"""The client is built with the discovery JWKS URL and an SSL context
533+
matching verify_ssl: default configs must keep certificate verification
534+
on, and verify_ssl=False must be the only way to turn it off."""
535+
mock_discovery_data.return_value = discovery_data
536+
mock_jwt.return_value = {"preferred_username": "my-name"}
537+
538+
token_parser = OidcTokenParser(auth_config=_oidc_config_with(verify_ssl=verify_ssl))
539+
asyncio.run(token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc"))
540+
541+
call = mock_jwks_client_cls.call_args
542+
assertpy.assert_that(call.args[0]).is_equal_to(discovery_data["jwks_uri"])
543+
ssl_ctx = call.kwargs["ssl_context"]
544+
if verify_ssl:
545+
assertpy.assert_that(ssl_ctx.verify_mode).is_equal_to(ssl.CERT_REQUIRED)
546+
assertpy.assert_that(ssl_ctx.check_hostname).is_true()
547+
else:
548+
assertpy.assert_that(ssl_ctx.verify_mode).is_equal_to(ssl.CERT_NONE)
549+
assertpy.assert_that(ssl_ctx.check_hostname).is_false()
550+
551+
552+
@patch(
553+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
554+
)
555+
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
556+
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
557+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
558+
def test_oidc_jwks_client_construction_failure_is_retried(
559+
mock_discovery_data,
560+
mock_jwks_client_cls,
561+
mock_jwt,
562+
mock_oauth2,
563+
oidc_config,
564+
discovery_data,
565+
):
566+
"""A failed first construction must leave the parser able to retry on
567+
the next request: the parser is a process singleton, so caching a failed
568+
or half-built client would wedge authentication until restart."""
569+
mock_discovery_data.return_value = discovery_data
570+
mock_jwt.return_value = {"preferred_username": "my-name"}
571+
mock_jwks_client_cls.side_effect = [RuntimeError("IdP unreachable"), MagicMock()]
572+
573+
token_parser = OidcTokenParser(auth_config=oidc_config)
574+
with pytest.raises(RuntimeError):
575+
asyncio.run(
576+
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
577+
)
578+
579+
user = asyncio.run(
580+
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
581+
)
582+
assertpy.assert_that(user.username).is_equal_to("my-name")
583+
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(2)
584+
585+
475586
# ---------------------------------------------------------------------------
476587
# Optional audience / issuer verification (opt-in via OidcAuthConfig)
477588
# ---------------------------------------------------------------------------
@@ -671,38 +782,6 @@ def test_oidc_default_supports_v1_tokens_against_v2_discovery(
671782
assertpy.assert_that(user.roles).is_equal_to(["reader"])
672783

673784

674-
@patch(
675-
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__"
676-
)
677-
@patch("feast.permissions.auth.oidc_token_parser.jwt.decode")
678-
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
679-
@patch("feast.permissions.auth.oidc_token_parser.PyJWKClient")
680-
def test_oidc_jwks_client_is_reused_across_requests(
681-
mock_jwks_client_cls,
682-
mock_discovery_data,
683-
mock_jwt,
684-
mock_oauth2,
685-
oidc_config,
686-
discovery_data,
687-
):
688-
"""The JWKS client must be built once per parser, not once per request:
689-
a fresh client starts with a cold JWK-set cache, which forces a full
690-
HTTPS fetch of the JWKS document on every authenticated call."""
691-
mock_discovery_data.return_value = discovery_data
692-
mock_jwt.return_value = {"preferred_username": "my-name"}
693-
694-
token_parser = OidcTokenParser(auth_config=oidc_config)
695-
for _ in range(3):
696-
asyncio.run(
697-
token_parser.user_details_from_access_token(access_token="aaa-bbb-ccc")
698-
)
699-
700-
assertpy.assert_that(mock_jwks_client_cls.call_count).is_equal_to(1)
701-
assertpy.assert_that(
702-
mock_jwks_client_cls.return_value.get_signing_key_from_jwt.call_count
703-
).is_equal_to(3)
704-
705-
706785
# TODO RBAC: Move role bindings to a reusable fixture
707786
@patch("feast.permissions.auth.kubernetes_token_parser.config.load_incluster_config")
708787
@patch("feast.permissions.auth.kubernetes_token_parser.jwt.decode")

0 commit comments

Comments
 (0)