Visitar URL original
fix: Prune expired cache entries and drop a token the server rejects · feast-dev/feast@2a73cb7 · GitHub
Skip to content

Commit 2a73cb7

Browse files
larrysingleton007larrysingleton007
committed
fix: Prune expired cache entries and drop a token the server rejects
Two gaps ntkathole raised on the token cache. Expired entries were never evicted. Keys are credential identities, so the cache is bounded by distinct configs, but a long-lived process that rotates credentials would retain every retired identity. Inserting on a miss now first drops entries whose stored expiry has passed. Reuse also meant a token the IdP revoked mid-life kept 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), and wires it into the gRPC interceptor, which already inspects the response: an UNAUTHENTICATED status drops the cached token so the next call refetches, bounding the staleness to the request that was rejected. Deliberately does not retry the rejected call. All four interceptor methods share that path and a stream's request_iterator may already be consumed, so a blanket retry could replay a partially-sent stream. Retrying belongs per-transport rather than in the shared handler. 328 permissions tests pass. Signed-off-by: Larry Singleton <166439969+larrysingleton007@users.noreply.github.com>
1 parent ce3ffd4 commit 2a73cb7

5 files changed

Lines changed: 126 additions & 12 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/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: 32 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -90,13 +90,7 @@ def _fetch_token_from_idp(self) -> str:
9090
margin. A token whose expiry is unknowable is not cached, preserving
9191
the previous per-call behavior for opaque tokens.
9292
"""
93-
cache_key = (
94-
self.auth_config.auth_discovery_url,
95-
self.auth_config.client_id,
96-
self.auth_config.client_secret,
97-
self.auth_config.username,
98-
self.auth_config.password,
99-
)
93+
cache_key = self._cache_key()
10094
with _token_cache_lock:
10195
cached = _token_cache.get(cache_key)
10296
if cached is not None:
@@ -109,10 +103,41 @@ def _fetch_token_from_idp(self) -> str:
109103

110104
expiry = self._token_expiry(access_token, expires_in)
111105
if expiry is not None:
106+
now = time.time()
112107
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]
113114
_token_cache[cache_key] = (access_token, expiry)
114115
return access_token
115116

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+
116141
@staticmethod
117142
def _token_expiry(
118143
access_token: str, expires_in: Optional[float]

‎sdk/python/tests/unit/permissions/test_oidc_auth_client.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,10 @@
1212
OidcClientAuthConfig,
1313
)
1414
from feast.permissions.client import oidc_authentication_client_manager
15-
from feast.permissions.client.client_auth_token import get_auth_token
15+
from feast.permissions.client.client_auth_token import (
16+
get_auth_token,
17+
invalidate_auth_token,
18+
)
1619
from feast.permissions.client.http_auth_requests_wrapper import (
1720
AuthenticatedRequestsSession,
1821
get_http_auth_requests_session,
@@ -266,3 +269,47 @@ def test_shared_credentials_each_honor_their_own_margin(mock_post, mock_discover
266269
# of remaining life clears its 30s margin.
267270
assert get_auth_token(narrow_margin) == wide_token
268271
assert mock_post.call_count == 2
272+
273+
274+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
275+
@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post")
276+
def test_invalidate_auth_token_forces_a_refetch(mock_post, mock_discovery):
277+
"""A token the IdP revokes mid-life still looks valid to the client until
278+
its own expiry. Invalidating drops it so the next call refetches, bounding
279+
the staleness to the request that was rejected."""
280+
mock_discovery.return_value = _DISCOVERY
281+
first = _jwt_expiring_in(3600, subject="first")
282+
second = _jwt_expiring_in(3600, subject="second")
283+
mock_post.side_effect = [_token_response(first), _token_response(second)]
284+
285+
config = _get_dummy_oidc_auth_type()
286+
assert get_auth_token(config) == first
287+
assert get_auth_token(config) == first
288+
assert mock_post.call_count == 1
289+
290+
assert invalidate_auth_token(config) is True
291+
assert get_auth_token(config) == second
292+
assert mock_post.call_count == 2
293+
294+
# Nothing cached for this config any more, so a second invalidate is a no-op.
295+
invalidate_auth_token(config)
296+
assert invalidate_auth_token(config) is False
297+
298+
299+
@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data")
300+
@patch("feast.permissions.client.oidc_authentication_client_manager.requests.post")
301+
def test_expired_entries_are_pruned_on_miss(mock_post, mock_discovery):
302+
"""Entries are keyed by credential identity, so a long-lived process that
303+
rotates credentials would otherwise retain every retired identity."""
304+
mock_discovery.return_value = _DISCOVERY
305+
cache = oidc_authentication_client_manager._token_cache
306+
stale_key = ("retired-idp", "retired-client", "secret", None, None)
307+
cache[stale_key] = ("stale-token", time.time() - 1)
308+
live_key = ("live-idp", "live-client", "secret", None, None)
309+
cache[live_key] = ("live-token", time.time() + 3600)
310+
311+
mock_post.return_value = _token_response(_jwt_expiring_in(3600))
312+
get_auth_token(_get_dummy_oidc_auth_type())
313+
314+
assert stale_key not in cache, "expired entry should be pruned on insert"
315+
assert live_key in cache, "unexpired entries must survive the prune"

0 commit comments

Comments
 (0)