Skip to content

Commit 4b349fc

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

2 files changed

Lines changed: 61 additions & 15 deletions

File tree

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

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,34 @@ 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, 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).
56+
"""
57+
if self._jwks_client is None:
58+
ssl_ctx = ssl.create_default_context()
59+
if not self._auth_config.verify_ssl:
60+
ssl_ctx.check_hostname = False
61+
ssl_ctx.verify_mode = ssl.CERT_NONE
62+
elif self._auth_config.ca_cert_path and os.path.exists(
63+
self._auth_config.ca_cert_path
64+
):
65+
ssl_ctx.load_verify_locations(self._auth_config.ca_cert_path)
66+
self._jwks_client = PyJWKClient(
67+
self.oidc_discovery_service.get_jwks_url(),
68+
headers={"User-agent": "custom-user-agent"},
69+
ssl_context=ssl_ctx,
70+
)
71+
return self._jwks_client
4472

4573
async def _validate_token(self, access_token: str):
4674
"""
@@ -125,21 +153,7 @@ def _decode_token(self, access_token: str) -> dict:
125153
metadata (e.g. Entra ID v1.0 tokens validated against a v2.0
126154
discovery document).
127155
"""
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)
156+
signing_key = self._get_jwks_client().get_signing_key_from_jwt(access_token)
143157
expected_audience = self._auth_config.audience
144158
expected_issuer = self._auth_config.issuer
145159
return jwt.decode(

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -671,6 +671,38 @@ def test_oidc_default_supports_v1_tokens_against_v2_discovery(
671671
assertpy.assert_that(user.roles).is_equal_to(["reader"])
672672

673673

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+
674706
# TODO RBAC: Move role bindings to a reusable fixture
675707
@patch("feast.permissions.auth.kubernetes_token_parser.config.load_incluster_config")
676708
@patch("feast.permissions.auth.kubernetes_token_parser.jwt.decode")

0 commit comments

Comments
 (0)