Skip to content

Commit 602d752

Browse files
larrysingleton007ntkathole
larrysingleton007
authored and
ntkathole
committed
fix: Reuse IdP-issued client tokens until near expiry
On the client_secret branch every outbound RPC built a fresh auth-token factory, manager and OIDCDiscoveryService, paying a discovery GET plus a token POST per call and discarding the token. All three auth interceptors invoke it per RPC, so a batch materialization loop multiplied IdP load by request count and could trip IdP rate limits. Measured before: 5 calls = 5 discovery GETs + 5 token POSTs. After: 1 and 1. Caches IdP tokens in a module-level dict keyed by the token-request identity, since the interceptors build a fresh manager per call and instance state would not survive. Expiry comes from the token's own exp claim, falling back to the token endpoint's expires_in; a token whose expiry is unknowable is not cached, preserving per-call behaviour for opaque tokens. The cache stores true expiry and applies the caller's token_refresh_margin_seconds on read, rather than storing a deadline. The margin is not part of the key, so baking it in let a config with a wider margin reuse a token past its own safety window when another config sharing the same credentials had written the entry. token_refresh_margin_seconds is configurable on OidcClientAuthConfig (default 30, gt=0) rather than hardcoded. Inserting on a miss prunes entries whose stored expiry has passed. Keys are credential identities so the cache is bounded by distinct configs, but a long-lived process rotating credentials would otherwise retain every retired identity. Reuse means a token the IdP revokes mid-life keeps being presented until its own expiry, where fetching per call self-corrected. Adds OidcAuthClientManager.invalidate_token and a transport-agnostic invalidate_auth_token(auth_config), wired into the gRPC interceptor: an UNAUTHENTICATED response drops the cached token so the next call refetches, bounding staleness to the rejected request. The call is not retried, because all four interceptor methods share that path and a stream's request_iterator may already be consumed. 329 permissions tests pass. Signed-off-by: Larry Singleton <166439969+larrysingleton007@users.noreply.github.com>
1 parent 5fd7af7 commit 602d752

6 files changed

Lines changed: 414 additions & 9 deletions

File tree

‎.secrets.baseline‎

Lines changed: 3 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,11 @@ class OidcClientAuthConfig(OidcAuthConfig):
6767
client_secret: Optional[str] = None
6868
token: Optional[str] = None
6969
token_env_var: Optional[str] = None
70+
# Stop reusing an IdP-issued token this many seconds before it expires,
71+
# so a reused token still has life left when the server validates it.
72+
# Raise it if clients see sporadic 401s from clock skew or slow calls;
73+
# lower it to squeeze more reuse out of short-lived tokens.
74+
token_refresh_margin_seconds: float = Field(default=30, gt=0)
7075

7176
@model_validator(mode="after")
7277
def _validate_credentials(self):

‎sdk/python/feast/permissions/client/client_auth_token.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,16 @@ def get_auth_token(auth_config: AuthConfig) -> str:
1212
.get_auth_client_manager()
1313
.get_token()
1414
)
15+
16+
17+
def invalidate_auth_token(auth_config: AuthConfig) -> bool:
18+
"""Drop any cached token for *auth_config*, returning whether one was held.
19+
20+
Only the OIDC client manager caches, so this is a no-op for the other auth
21+
types. Callers that can observe an authentication failure should use it: a
22+
token the IdP revokes mid-life still looks valid to the client until its
23+
own expiry, and dropping it bounds that to a single rejected request.
24+
"""
25+
manager = AuthenticationClientManagerFactory(auth_config).get_auth_client_manager()
26+
invalidate = getattr(manager, "invalidate_token", None)
27+
return bool(invalidate()) if callable(invalidate) else False

‎sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@
55
from feast.errors import FeastError
66
from feast.permissions.auth.auth_type import AuthType
77
from feast.permissions.auth_model import AuthConfig
8-
from feast.permissions.client.client_auth_token import get_auth_token
8+
from feast.permissions.client.client_auth_token import (
9+
get_auth_token,
10+
invalidate_auth_token,
11+
)
912

1013
logger = logging.getLogger(__name__)
1114

@@ -44,11 +47,37 @@ def _handle_call(self, continuation, client_call_details, request_iterator):
4447
client_call_details = self._append_auth_header_metadata(client_call_details)
4548
result = continuation(client_call_details, request_iterator)
4649
if result.exception() is not None:
50+
self._invalidate_token_if_rejected(result)
4751
mapped_error = FeastError.from_error_detail(result.exception().details())
4852
if mapped_error is not None:
4953
raise mapped_error
5054
return result
5155

56+
def _invalidate_token_if_rejected(self, result) -> None:
57+
"""Drop the cached token when the server rejects it as unauthenticated.
58+
59+
Tokens are reused until near expiry, so one the IdP revoked mid-life
60+
would otherwise keep being presented for the rest of its lifetime.
61+
Dropping it here bounds that to the single request that was rejected;
62+
the next call fetches a fresh token.
63+
64+
The call is deliberately not retried. All four interceptor methods
65+
share this path, and a stream's ``request_iterator`` may already be
66+
consumed, so retrying here could replay a partially-sent stream.
67+
"""
68+
if self._auth_config.type == AuthType.NONE.value:
69+
return
70+
try:
71+
if result.code() != grpc.StatusCode.UNAUTHENTICATED:
72+
return
73+
except Exception: # pragma: no cover - result without a status code
74+
return
75+
if invalidate_auth_token(self._auth_config):
76+
logger.debug(
77+
"Server rejected the cached auth token; dropped it so the next "
78+
"call fetches a fresh one."
79+
)
80+
5281
def _append_auth_header_metadata(self, client_call_details):
5382
logger.debug(
5483
"Intercepted the grpc api method call to inject Authorization header "

‎sdk/python/feast/permissions/client/oidc_authentication_client_manager.py‎

Lines changed: 110 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import logging
22
import os
3-
from typing import Optional
3+
import threading
4+
import time
5+
from typing import Dict, Optional, Tuple
46

57
import jwt
68
import requests
@@ -13,6 +15,17 @@
1315

1416
SA_TOKEN_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/token"
1517

18+
# IdP-issued tokens keyed by the token-request identity, stored with their
19+
# true expiry. The auth interceptors build a fresh manager for every outbound
20+
# RPC, so instance state would not survive between calls; without this
21+
# cache every RPC pays a discovery GET plus a token POST against the IdP.
22+
# The refresh margin is applied on read, not baked into the stored value, so
23+
# configs sharing IdP credentials but setting different margins can share
24+
# tokens while each still honors its own margin.
25+
# Concurrent misses may fetch in parallel (benign: last write wins).
26+
_token_cache: Dict[Tuple, Tuple[str, float]] = {}
27+
_token_cache_lock = threading.Lock()
28+
1629

1730
class OidcAuthClientManager(AuthenticationClientManager):
1831
def __init__(self, auth_config: OidcClientAuthConfig):
@@ -67,7 +80,96 @@ def _read_sa_token() -> Optional[str]:
6780
return None
6881

6982
def _fetch_token_from_idp(self) -> str:
70-
"""Obtain an access token via client_credentials or ROPG flow."""
83+
"""Return a cached IdP token, or obtain a fresh one.
84+
85+
The cache stores each token's true expiry (its ``exp`` claim, falling
86+
back to the token response's ``expires_in``), and this config's
87+
``token_refresh_margin_seconds`` is applied when reading. Keeping the
88+
margin out of the stored value lets configs that share IdP credentials
89+
but set different margins share tokens while each honors its own
90+
margin. A token whose expiry is unknowable is not cached, preserving
91+
the previous per-call behavior for opaque tokens.
92+
"""
93+
cache_key = self._cache_key()
94+
with _token_cache_lock:
95+
cached = _token_cache.get(cache_key)
96+
if cached is not None:
97+
cached_token, cached_expiry = cached
98+
margin = self.auth_config.token_refresh_margin_seconds
99+
if time.time() < cached_expiry - margin:
100+
return cached_token
101+
102+
access_token, expires_in = self._request_token_from_idp()
103+
104+
expiry = self._token_expiry(access_token, expires_in)
105+
if expiry is not None:
106+
now = time.time()
107+
with _token_cache_lock:
108+
# Prune on miss. Entries are keyed by credential identity, so
109+
# the cache is bounded by the number of distinct configs, but
110+
# a long-lived process that rotates credentials would otherwise
111+
# keep every retired identity forever.
112+
for key in [k for k, (_, exp) in _token_cache.items() if exp <= now]:
113+
del _token_cache[key]
114+
_token_cache[cache_key] = (access_token, expiry)
115+
return access_token
116+
117+
def _cache_key(self) -> Tuple:
118+
"""Identity of the token request: same credentials, same token."""
119+
return (
120+
self.auth_config.auth_discovery_url,
121+
self.auth_config.client_id,
122+
self.auth_config.client_secret,
123+
self.auth_config.username,
124+
self.auth_config.password,
125+
)
126+
127+
def invalidate_token(self) -> bool:
128+
"""Drop this config's cached token so the next call refetches.
129+
130+
Returns whether an entry was actually removed.
131+
132+
Reuse means a token the IdP revokes mid-life keeps being presented
133+
until its own expiry, where fetching per call self-corrected. Callers
134+
that can observe an authentication failure should invalidate on it, so
135+
the staleness costs one rejected request rather than the remaining
136+
lifetime of the token.
137+
"""
138+
with _token_cache_lock:
139+
return _token_cache.pop(self._cache_key(), None) is not None
140+
141+
@staticmethod
142+
def _token_expiry(
143+
access_token: str, expires_in: Optional[float]
144+
) -> Optional[float]:
145+
"""Epoch expiry of *access_token*, or ``None`` when it is unknowable.
146+
147+
Prefers the token's own ``exp`` claim (authoritative); falls back to
148+
the token endpoint's ``expires_in``.
149+
150+
The refresh margin is deliberately not subtracted here. Storing one
151+
caller's deadline would let another config with a wider margin reuse
152+
the token past its own safety window, since the margin is not part of
153+
the cache key.
154+
"""
155+
exp: Optional[float] = None
156+
try:
157+
claims = jwt.decode(access_token, options={"verify_signature": False})
158+
claim = claims.get("exp")
159+
if isinstance(claim, (int, float)):
160+
exp = float(claim)
161+
except jwt.exceptions.DecodeError:
162+
pass
163+
if exp is None and isinstance(expires_in, (int, float)):
164+
exp = time.time() + float(expires_in)
165+
return exp
166+
167+
def _request_token_from_idp(self) -> Tuple[str, Optional[float]]:
168+
"""Obtain an access token via client_credentials or ROPG flow.
169+
170+
Returns the token and the token response's ``expires_in`` (seconds),
171+
when the IdP provides one.
172+
"""
71173
if self.auth_config.auth_discovery_url is None:
72174
raise ValueError(
73175
"auth_discovery_url is required for IDP token fetch "
@@ -106,13 +208,17 @@ def _fetch_token_from_idp(self) -> str:
106208
)
107209

108210
if token_response.status_code == 200:
109-
access_token = token_response.json()["access_token"]
211+
response_body = token_response.json()
212+
access_token = response_body["access_token"]
110213
if not access_token:
111214
logger.debug(
112215
f"access_token is empty for the client_id=${self.auth_config.client_id}"
113216
)
114217
raise RuntimeError("access token is empty")
115-
return access_token
218+
expires_in = response_body.get("expires_in")
219+
if not isinstance(expires_in, (int, float)):
220+
expires_in = None
221+
return access_token, expires_in
116222
else:
117223
raise RuntimeError(
118224
f"""Failed to obtain oidc access token:url=[{token_endpoint}] {token_response.status_code} - {token_response.text}"""

0 commit comments

Comments
 (0)