From e6abe4e29938319c0df63d0b6389e3061d6c5102 Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Mon, 17 Aug 2026 13:27:50 -0700 Subject: [PATCH 01/10] Add JWT Bearer authentication to Arrow Flight. --- docs/server-config.md | 35 ++- tabpy/tabpy_server/app/app.py | 56 ++-- tabpy/tabpy_server/handlers/__init__.py | 3 + .../basic_auth_server_middleware_factory.py | 81 ++++-- tabpy/tabpy_server/handlers/jwt_auth.py | 66 +++-- .../handlers/jwt_server_middleware_factory.py | 82 ++++++ tabpy/tabpy_server/handlers/util.py | 25 ++ tests/unit/server_tests/jwt_test_helpers.py | 16 +- ...st_basic_auth_server_middleware_factory.py | 125 +++++++++ tests/unit/server_tests/test_config.py | 61 ++++- tests/unit/server_tests/test_jwt_auth.py | 77 ++++-- .../test_jwt_server_middleware_factory.py | 245 ++++++++++++++++++ 12 files changed, 771 insertions(+), 101 deletions(-) create mode 100644 tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py create mode 100644 tests/unit/server_tests/test_basic_auth_server_middleware_factory.py create mode 100644 tests/unit/server_tests/test_jwt_server_middleware_factory.py diff --git a/docs/server-config.md b/docs/server-config.md index b223baa3..09983a6a 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -331,14 +331,31 @@ To authenticate a request, send the JWT as a Bearer token: curl -H "Authorization: Bearer " http://localhost:9004/info ``` +The same Bearer token is accepted on the Arrow Flight (gRPC) path when Arrow +is enabled. Failed Flight authentication is rejected with gRPC +`UNAUTHENTICATED` rather than HTTP 401. Using Basic on Flight while OAuth is +enabled also requires `TABPY_PWD_FILE`. When `TABPY_OAUTH_ENABLED` is false, +Flight continues to use basic-auth middleware unchanged. + When both basic access authentication and OAuth are enabled, TabPy picks the method based on the scheme of the `Authorization` header sent by the client (`Basic` or `Bearer`), so both can be used against the same server. -JWKS lookups are cached, but because TabPy serves requests on a single -thread, a slow or unresponsive IdP during a cache-cold fetch (startup, or a -key rotation) will briefly stall all concurrent requests, not just the one -that triggered the fetch. +With `TABPY_TRANSFER_PROTOCOL = http`, Flight uses `grpc+tcp` and the Bearer +token is sent in cleartext, the same as HTTP Basic/Bearer on an unencrypted +port. + +JWKS lookups are cached. On the HTTP path a cache-cold fetch runs on TabPy's +single IO-loop thread. Arrow Flight auth runs on the gRPC thread pool. JWKS +client creation and fetches are serialized with a shared lock so concurrent +Flight calls cannot bypass the refresh rate limit. A JWT check that can't +take that lock within one second is rejected rather than left waiting for +the in-flight fetch, so a slow or unresponsive identity provider can't stall +concurrent requests for the full fetch timeout. + +A request carrying two conflicting `Authorization` values is rejected, +because which one wins would otherwise decide the caller's identity. +Repeating the same value is accepted. ### Endpoint Security @@ -350,6 +367,16 @@ TabPy can be configured to enable Arrow Flight. This will cause a Flight server to start up alongside the HTTP server and will allow for handling incoming streamed data in the Arrow columnar format. +When authentication is enabled, Flight accepts the same credentials as HTTP. +See [Authentication](#authentication). Failed Flight auth is rejected with +gRPC `UNAUTHENTICATED` rather than HTTP 401. + +After a successful Basic call, Flight also returns an opaque session token +in the response `authorization` header. A client may send that token back as +a Bearer credential instead of repeating its Basic credentials. The token is +server-issued, is not a JWT, and expires an hour after it is issued, at +which point the client authenticates with Basic again. + **As of May 2023, the Arrow Flight feature can only be used by compatible versions of Tableau Prep. The Arrow Flight feature is not used by Tableau Desktop, Tableau Server, or Tableau Cloud, regardless of the diff --git a/tabpy/tabpy_server/app/app.py b/tabpy/tabpy_server/app/app.py index d8a2cc06..715947b1 100644 --- a/tabpy/tabpy_server/app/app.py +++ b/tabpy/tabpy_server/app/app.py @@ -20,7 +20,12 @@ from tabpy.tabpy import __version__ from tabpy.tabpy_server.app.app_parameters import ConfigParameters, SettingsParameters from tabpy.tabpy_server.app.util import parse_pwd_file -from tabpy.tabpy_server.handlers.basic_auth_server_middleware_factory import BasicAuthServerMiddlewareFactory +from tabpy.tabpy_server.handlers.basic_auth_server_middleware_factory import ( + BasicAuthServerMiddlewareFactory, +) +from tabpy.tabpy_server.handlers.jwt_server_middleware_factory import ( + JwtAuthServerMiddlewareFactory, +) from tabpy.tabpy_server.handlers.no_op_auth_handler import NoOpAuthHandler from tabpy.tabpy_server.management.state import TabPyState from tabpy.tabpy_server.management.util import _get_state_from_file @@ -129,11 +134,32 @@ def _get_arrow_server(self, config): location = "{}://{}:{}".format(scheme, host, port) auth_middleware = None - if "authentication" in config[SettingsParameters.ApiVersions]["v1"]["features"]: - _, creds = parse_pwd_file(config[ConfigParameters.TABPY_PWD_FILE]) - auth_middleware = { - "basic": BasicAuthServerMiddlewareFactory(creds) - } + features = config[SettingsParameters.ApiVersions]["v1"]["features"] + if "authentication" in features: + basic_factory = None + if ConfigParameters.TABPY_PWD_FILE in config: + basic_factory = BasicAuthServerMiddlewareFactory(self.credentials) + + # pyarrow invokes every registered middleware factory on each + # call. Installing Basic as a sibling key would reject valid + # Bearer tokens, so Basic is delegated through the JWT factory + # instead of registered alongside it. + if config.get(SettingsParameters.OAuthEnabled): + auth_middleware = { + "jwt": JwtAuthServerMiddlewareFactory( + issuer=config[SettingsParameters.OAuthIssuer], + jwks_uri=config[SettingsParameters.OAuthJwksUri], + audience=config[SettingsParameters.OAuthAudience], + required_scopes=config.get( + SettingsParameters.OAuthRequiredScopes + ), + basic_factory=basic_factory, + ) + } + elif basic_factory is not None: + auth_middleware = { + "basic": basic_factory + } server = pa.FlightServer(host, location, tls_certificates=tls_certificates, @@ -597,24 +623,6 @@ def _validate_oauth_settings(self): logger.critical(msg) raise RuntimeError(msg) - # Arrow Flight's auth middleware currently only supports basic auth - # (see _get_arrow_server): whenever any auth method is enabled it - # unconditionally reads TABPY_PWD_FILE, so OAuth-only + Arrow would - # otherwise crash at startup looking for a pwd file that was never - # configured. Revisit this check if Arrow Flight ever adds its own - # OAuth/JWT middleware option. - if ( - self.settings[SettingsParameters.ArrowEnabled] - and ConfigParameters.TABPY_PWD_FILE not in self.settings - ): - msg = ( - f"{ConfigParameters.TABPY_ARROW_ENABLE} requires " - f"{ConfigParameters.TABPY_PWD_FILE} to be set: Arrow Flight does not " - "support OAuth authentication" - ) - logger.critical(msg) - raise RuntimeError(msg) - def _get_features(self): features = {} diff --git a/tabpy/tabpy_server/handlers/__init__.py b/tabpy/tabpy_server/handlers/__init__.py index d909800e..2ef6e8c5 100644 --- a/tabpy/tabpy_server/handlers/__init__.py +++ b/tabpy/tabpy_server/handlers/__init__.py @@ -15,3 +15,6 @@ from tabpy.tabpy_server.handlers.basic_auth_server_middleware_factory import ( BasicAuthServerMiddlewareFactory, ) +from tabpy.tabpy_server.handlers.jwt_server_middleware_factory import ( + JwtAuthServerMiddlewareFactory, +) diff --git a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py index b80b4a91..2c61d381 100644 --- a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py @@ -1,10 +1,29 @@ import base64 +import binascii import secrets +import threading +import time from pyarrow.flight import ServerMiddlewareFactory, ServerMiddleware from pyarrow.flight import FlightUnauthenticatedError -from tabpy.tabpy_server.handlers.util import hash_password +from tabpy.tabpy_server.handlers.util import ( + get_flight_authorization_header, + hash_password, +) + +# A successful Basic call mints an opaque token and hands it back to the +# client via sending_headers(). The client may replay it as a Bearer +# credential, so it is a real credential and gets an expiry. Once it +# lapses the client re-authenticates with Basic, which it can always do: +# that is how it obtained the token. +FLIGHT_TOKEN_TTL_SECONDS = 3600 + +# Expiry alone doesn't bound the store, since a client that ignores the +# token still mints one per call. Cap it and drop the soonest-to-expire +# entries once full. +MAX_FLIGHT_TOKENS = 1024 + class BasicAuthServerMiddleware(ServerMiddleware): def __init__(self, token): @@ -13,10 +32,17 @@ def __init__(self, token): def sending_headers(self): return {"authorization": f"Bearer {self.token}"} + class BasicAuthServerMiddlewareFactory(ServerMiddlewareFactory): def __init__(self, creds): self.creds = creds + # token -> (username, monotonic expiry). Read from the gRPC thread + # pool on every call; every mutation goes through _issue_token + # under _tokens_lock, so readers never see the store being + # resized and concurrent calls can't grow it past + # MAX_FLIGHT_TOKENS. self.tokens = {} + self._tokens_lock = threading.Lock() def is_valid_user(self, username, password): if username not in self.creds: @@ -24,25 +50,50 @@ def is_valid_user(self, username, password): hashed_pwd = hash_password(username, password) return self.creds[username].lower() == hashed_pwd.lower() + def is_valid_token(self, token): + entry = self.tokens.get(token) + return entry is not None and time.monotonic() < entry[1] + + def _issue_token(self, username): + token = secrets.token_urlsafe(32) + with self._tokens_lock: + self._evict_tokens() + self.tokens[token] = (username, time.monotonic() + FLIGHT_TOKEN_TTL_SECONDS) + return token + + def _evict_tokens(self): + now = time.monotonic() + for token in [t for t, (_, expiry) in self.tokens.items() if expiry <= now]: + self.tokens.pop(token, None) + while len(self.tokens) >= MAX_FLIGHT_TOKENS: + oldest = min(self.tokens, key=lambda t: self.tokens[t][1]) + self.tokens.pop(oldest, None) + def start_call(self, info, headers): - auth_header = None - for header in headers: - if header.lower() == "authorization": - auth_header = headers[header][0] - break + auth_header = get_flight_authorization_header(headers) if not auth_header: raise FlightUnauthenticatedError("No credentials supplied") - - auth_type, _, value = auth_header.partition(" ") - if auth_type == "Basic": - decoded = base64.b64decode(value).decode("utf-8") - username, _, password = decoded.partition(":") + parts = auth_header.split(" ") + if len(parts) != 2: + raise FlightUnauthenticatedError("No credentials supplied") + auth_type, value = parts + + if auth_type.lower() == "basic": + try: + decoded = base64.b64decode(value, validate=True).decode("utf-8") + except (binascii.Error, UnicodeDecodeError): + raise FlightUnauthenticatedError("Invalid credentials") from None + + username, separator, password = decoded.partition(":") + if not separator or not username: + raise FlightUnauthenticatedError("Invalid credentials") if not self.is_valid_user(username, password): raise FlightUnauthenticatedError("Invalid credentials") - token = secrets.token_urlsafe(32) - self.tokens[token] = username - return BasicAuthServerMiddleware(token) - + return BasicAuthServerMiddleware(self._issue_token(username)) + + if auth_type.lower() == "bearer" and self.is_valid_token(value): + return BasicAuthServerMiddleware(value) + raise FlightUnauthenticatedError("No credentials supplied") diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 8569a0d5..ee92e4da 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -1,4 +1,5 @@ import logging +import threading import time import jwt @@ -6,32 +7,36 @@ logger = logging.getLogger(__name__) -# PyJWKClient fetches JWKS synchronously (via requests), and that call runs -# directly on TabPy's single Tornado IO-loop thread -- a slow/unresponsive -# IdP therefore stalls every concurrent request on the server, not just the -# one that triggered the fetch, for up to this many seconds. Caching (see -# _jwks_clients) and the refresh rate limit below bound how often this can -# happen, but a cold start or a legitimate key rotation still pays this -# cost. Moving the fetch to a thread pool (e.g. via IOLoop.run_in_executor) -# would remove the stall entirely, at the cost of making the auth path -# asynchronous; not done here. +# PyJWKClient fetches JWKS synchronously (via requests). On the HTTP path +# that call runs on TabPy's single Tornado IO-loop thread, so a +# slow/unresponsive IdP stalls every concurrent HTTP request for up to this +# many seconds. Caching (see _jwks_clients) bounds how often this happens. +# A cold start or legitimate key rotation still pays this cost. JWKS_FETCH_TIMEOUT_SECONDS = 10 +# Bounds how long a JWT check waits for _jwks_lock. Arrow Flight auth runs +# on the gRPC thread pool and HTTP auth runs on the Tornado IO loop, and +# both share the lock, so without a bound an unauthenticated Flight caller +# could force a slow JWKS fetch and stall every HTTP request behind it for +# up to JWKS_FETCH_TIMEOUT_SECONDS. Callers that time out waiting fail +# closed instead; a responsive IdP resolves well inside this window. +JWKS_LOCK_WAIT_SECONDS = 1 + # An unauthenticated caller can force a fresh JWKS fetch just by sending a -# made-up `kid` (read from the token header pre-signature-check), and that -# fetch blocks the single IO-loop thread. This bounds, per jwks_uri, how -# often a *failed* forced refresh (kid still not found) can refetch again. -# Only failures set the cooldown -- a successful refresh (e.g. a genuine key -# rotation) must not be penalized. Also used to rate-limit retrying a JWKS -# endpoint that just failed to fetch at all (network error, timeout, -# malformed response), so a down/unreachable IdP can't be hammered with a -# fresh blocking fetch on every single request. +# made-up `kid` (read from the token header pre-signature-check). This +# bounds, per jwks_uri, how often a *failed* forced refresh (kid still not +# found) can refetch again. Only failures set the cooldown -- a successful +# refresh (e.g. a genuine key rotation) must not be penalized. Also used +# to rate-limit retrying a JWKS endpoint that just failed to fetch at all +# (network error, timeout, malformed response), so a down/unreachable IdP +# can't be hammered with a fresh blocking fetch on every single request. JWKS_MIN_REFRESH_INTERVAL_SECONDS = 30 # One PyJWKClient per JWKS URI, reused so its JWK Set cache actually avoids -# per-request fetches. Process-global and not lock-protected: safe only -# because TabPy runs a single app instance per process on a single IO-loop -# thread. Would need a lock (or per-app scoping) if that ever changes. +# per-request fetches. Process-global; mutations are serialized by +# _jwks_lock because HTTP (IO-loop) and Arrow Flight (gRPC thread pool) +# share these dicts. +_jwks_lock = threading.Lock() _jwks_clients = {} # jwks_uri -> monotonic timestamp of the last failed JWKS fetch (network @@ -152,8 +157,15 @@ def validate_jwt( raise JwtValidationError("Missing JWT") try: - jwks_client = _get_jwks_client(jwks_uri) - signing_key = _get_signing_key(jwks_client, jwks_uri, token) + if not _jwks_lock.acquire(timeout=JWKS_LOCK_WAIT_SECONDS): + raise jwt.exceptions.PyJWKClientError( + f'Timed out waiting on an in-flight JWKS fetch for "{jwks_uri}"' + ) + try: + jwks_client = _get_jwks_client(jwks_uri) + signing_key = _get_signing_key(jwks_client, jwks_uri, token) + finally: + _jwks_lock.release() except (jwt.exceptions.PyJWKClientError, jwt.exceptions.InvalidTokenError) as ex: logger.log(logging.ERROR, f"Unable to resolve JWT signing key: {str(ex)}") raise JwtValidationError("Unable to resolve JWT signing key") from ex @@ -174,6 +186,11 @@ def validate_jwt( except jwt.exceptions.InvalidTokenError as ex: logger.log(logging.ERROR, f"JWT validation failed: {str(ex)}") raise JwtValidationError(f"JWT validation failed: {str(ex)}") from ex + except Exception as ex: + # Must still fail closed as a 401 / UNAUTHENTICATED, not a 500 + # or an ArrowInvalid traceback to an unauthenticated caller. + logger.log(logging.ERROR, f"Unexpected error decoding JWT: {str(ex)}") + raise JwtValidationError("JWT validation failed") from ex if required_scopes: try: @@ -191,6 +208,9 @@ def validate_jwt( def _check_scopes(claims: dict, required_scopes: str) -> None: granted = set(claims.get("scope", "").split()) - missing = [s for s in (s.strip() for s in required_scopes.split(",")) if s and s not in granted] + missing = [ + s for s in (s.strip() for s in required_scopes.split(",")) + if s and s not in granted + ] if missing: raise JwtValidationError(f"JWT missing required scope(s): {', '.join(missing)}") diff --git a/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py new file mode 100644 index 00000000..d307641f --- /dev/null +++ b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py @@ -0,0 +1,82 @@ +import logging + +from pyarrow.flight import FlightUnauthenticatedError +from pyarrow.flight import ServerMiddleware, ServerMiddlewareFactory + +from tabpy.tabpy_server.handlers.jwt_auth import JwtValidationError, validate_jwt +from tabpy.tabpy_server.handlers.util import get_flight_authorization_header + +logger = logging.getLogger(__name__) + + +class JwtAuthServerMiddleware(ServerMiddleware): + def __init__(self, claims): + self.claims = claims + + +class JwtAuthServerMiddlewareFactory(ServerMiddlewareFactory): + """ + Arrow Flight middleware that validates JWT Bearer tokens via + jwt_auth.validate_jwt. When a Basic-auth factory is provided, Basic + credentials are delegated to it so both methods can coexist on Flight. + """ + + def __init__( + self, + issuer, + jwks_uri, + audience, + required_scopes=None, + basic_factory=None, + ): + self.issuer = issuer + self.jwks_uri = jwks_uri + self.audience = audience + self.required_scopes = required_scopes + self.basic_factory = basic_factory + + def start_call(self, info, headers): + auth_header = get_flight_authorization_header(headers) + if not auth_header: + raise FlightUnauthenticatedError("No credentials supplied") + + # Match HTTP: exactly two Authorization parts (scheme + value). + parts = auth_header.split(" ") + if len(parts) != 2: + raise FlightUnauthenticatedError("No credentials supplied") + auth_type, value = parts + + if auth_type.lower() == "bearer": + if ( + self.basic_factory is not None + and self.basic_factory.is_valid_token(value) + ): + return self.basic_factory.start_call(info, headers) + + try: + claims = validate_jwt( + value, + issuer=self.issuer, + jwks_uri=self.jwks_uri, + audience=self.audience, + required_scopes=self.required_scopes, + ) + except JwtValidationError as ex: + logger.log( + logging.ERROR, f"Flight JWT authentication failed: {ex}" + ) + raise FlightUnauthenticatedError("Invalid credentials") from ex + except Exception as ex: + # Must surface as UNAUTHENTICATED, not an ArrowInvalid + # carrying a traceback back to an unauthenticated caller. + logger.log( + logging.ERROR, + f"Unexpected error validating Flight JWT: {ex}", + ) + raise FlightUnauthenticatedError("Invalid credentials") from ex + return JwtAuthServerMiddleware(claims) + + if auth_type.lower() == "basic" and self.basic_factory is not None: + return self.basic_factory.start_call(info, headers) + + raise FlightUnauthenticatedError("No credentials supplied") diff --git a/tabpy/tabpy_server/handlers/util.py b/tabpy/tabpy_server/handlers/util.py index c9fc0e43..66f171c3 100644 --- a/tabpy/tabpy_server/handlers/util.py +++ b/tabpy/tabpy_server/handlers/util.py @@ -8,6 +8,7 @@ class AuthErrorStates(Enum): NotAuthorized = auto() NotRequired = auto() + def hash_password(username, pwd): """ Hashes password using PKDBF2 method: @@ -36,3 +37,27 @@ def hash_password(username, pwd): hash_name="sha512", password=pwd.encode(), salt=salt.encode(), iterations=10000 ) return binascii.hexlify(hash).decode() + + +def get_flight_authorization_header(headers): + """ + Returns the Authorization value from pyarrow Flight headers. + + Header names are matched case-insensitively. A missing or empty value + returns None so callers fail closed, as does a request carrying two + different Authorization values: which one wins would otherwise decide + the identity. Repeating the same value is harmless and accepted, + since some clients set the header through more than one mechanism. + """ + values = set() + for header, header_values in headers.items(): + if header.lower() == "authorization": + if isinstance(header_values, (list, tuple)): + values.update(header_values) + else: + values.add(header_values) + + if len(values) != 1: + return None + value = values.pop() + return value or None diff --git a/tests/unit/server_tests/jwt_test_helpers.py b/tests/unit/server_tests/jwt_test_helpers.py index b3420159..1f5dd20f 100644 --- a/tests/unit/server_tests/jwt_test_helpers.py +++ b/tests/unit/server_tests/jwt_test_helpers.py @@ -1,7 +1,8 @@ """ -Shared JWT test fixtures for test_jwt_auth.py and test_oauth_handler.py, -so the two suites can't silently drift apart on how tokens/signing keys -are faked. +Shared JWT test fixtures for test_jwt_auth.py, test_oauth_handler.py, +and test_jwt_server_middleware_factory.py, so the suites can't silently +drift apart on how tokens/signing keys are faked or how JWKS module +state is reset. """ import datetime from unittest.mock import patch @@ -13,6 +14,15 @@ JWKS_URI = "https://idp.example.com/.well-known/jwks.json" +def reset_jwks_state(): + """Clears process-global JWKS caches used by jwt_auth.""" + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + jwt_auth_module._jwks_clients.clear() + jwt_auth_module._jwks_last_failed_refresh.clear() + jwt_auth_module._jwks_last_fetch_failure.clear() + + def make_token(private_key, claims_override=None, headers=None): now = datetime.datetime.now(datetime.timezone.utc) claims = { diff --git a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py new file mode 100644 index 00000000..194634bf --- /dev/null +++ b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py @@ -0,0 +1,125 @@ +import base64 +import unittest +from unittest.mock import patch + +from pyarrow.flight import FlightUnauthenticatedError + +from tabpy.tabpy_server.handlers import basic_auth_server_middleware_factory as mod +from tabpy.tabpy_server.handlers.basic_auth_server_middleware_factory import ( + BasicAuthServerMiddleware, + BasicAuthServerMiddlewareFactory, +) +from tabpy.tabpy_server.handlers.util import hash_password + + +class TestBasicAuthServerMiddlewareFactory(unittest.TestCase): + def setUp(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + self.factory = BasicAuthServerMiddlewareFactory(creds) + + def _headers(self, authorization=None): + if authorization is None: + return {} + return {"authorization": [authorization]} + + def _basic_header(self, username, password): + encoded = base64.b64encode(f"{username}:{password}".encode()).decode() + return f"Basic {encoded}" + + def _authenticate(self): + return self.factory.start_call( + None, self._headers(self._basic_header("user1", "P@ssw0rd")) + ) + + def test_valid_basic_is_accepted(self): + middleware = self._authenticate() + self.assertIsInstance(middleware, BasicAuthServerMiddleware) + self.assertEqual( + middleware.sending_headers(), + {"authorization": f"Bearer {middleware.token}"}, + ) + + def test_issued_token_records_the_username(self): + middleware = self._authenticate() + username, _ = self.factory.tokens[middleware.token] + self.assertEqual(username, "user1") + + def test_invalid_password_is_rejected(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call( + None, self._headers(self._basic_header("user1", "wrong")) + ) + + def test_missing_credentials_are_rejected(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, {}) + + def test_issued_token_is_accepted_on_the_next_call(self): + handshake = self._authenticate() + subsequent_call = self.factory.start_call( + None, self._headers(f"Bearer {handshake.token}") + ) + self.assertEqual(subsequent_call.token, handshake.token) + self.assertEqual(len(self.factory.tokens), 1) + + def test_unknown_bearer_token_is_rejected(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers("Bearer not-a-real-token")) + + def test_jwt_shaped_bearer_is_rejected_without_oauth(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call( + None, self._headers("Bearer header.payload.signature") + ) + + def test_expired_token_is_rejected(self): + with patch.object(mod, "FLIGHT_TOKEN_TTL_SECONDS", -1): + handshake = self._authenticate() + self.assertFalse(self.factory.is_valid_token(handshake.token)) + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers(f"Bearer {handshake.token}")) + + def test_expired_tokens_are_evicted_on_the_next_mint(self): + with patch.object(mod, "FLIGHT_TOKEN_TTL_SECONDS", -1): + expired = self._authenticate() + self._authenticate() + self.assertNotIn(expired.token, self.factory.tokens) + + def test_token_store_is_bounded(self): + with patch.object(mod, "MAX_FLIGHT_TOKENS", 4): + first = self._authenticate() + for _ in range(20): + latest = self._authenticate() + + self.assertLessEqual(len(self.factory.tokens), 4) + self.assertNotIn(first.token, self.factory.tokens) + self.assertTrue(self.factory.is_valid_token(latest.token)) + + def test_invalid_base64_is_unauthenticated(self): + with self.assertRaises(FlightUnauthenticatedError) as err: + self.factory.start_call(None, self._headers("Basic !!!not-base64!!!")) + self.assertIn("Invalid credentials", str(err.exception)) + + def test_invalid_utf8_is_unauthenticated(self): + encoded = base64.b64encode(b"\xff\xfe").decode() + with self.assertRaises(FlightUnauthenticatedError) as err: + self.factory.start_call(None, self._headers(f"Basic {encoded}")) + self.assertIn("Invalid credentials", str(err.exception)) + + def test_missing_password_separator_is_unauthenticated(self): + encoded = base64.b64encode(b"user1").decode() + with self.assertRaises(FlightUnauthenticatedError) as err: + self.factory.start_call(None, self._headers(f"Basic {encoded}")) + self.assertIn("Invalid credentials", str(err.exception)) + + def test_conflicting_authorization_values_are_rejected(self): + headers = { + "authorization": [self._basic_header("user1", "P@ssw0rd")], + "Authorization": [self._basic_header("user1", "wrong")], + } + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, headers) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/server_tests/test_config.py b/tests/unit/server_tests/test_config.py index ed37e641..823319ee 100644 --- a/tests/unit/server_tests/test_config.py +++ b/tests/unit/server_tests/test_config.py @@ -582,12 +582,11 @@ def test_oauth_enabled_with_unresolvable_jwks_uri_raises(self): "tabpy.tabpy_server.app.app.socket.getaddrinfo", return_value=PUBLIC_JWKS_ADDRINFO, ) - def test_oauth_only_with_arrow_enabled_raises(self, mock_getaddrinfo): + def test_oauth_only_with_arrow_enabled_succeeds(self, mock_getaddrinfo): """ - Arrow Flight's auth middleware only supports basic auth (see - TabPyApp._get_arrow_server), so OAuth-only + Arrow must be - rejected at startup rather than crashing the Arrow thread with a - KeyError looking for a password file that was never configured. + Arrow Flight JWT middleware authenticates OAuth-only deployments, + so OAuth-only + Arrow is a valid startup configuration and must + not KeyError looking for a password file. """ self.fp.write( "[TabPy]\n" @@ -599,9 +598,22 @@ def test_oauth_only_with_arrow_enabled_raises(self, mock_getaddrinfo): ) self.fp.close() - with self.assertRaises(RuntimeError) as err: - TabPyApp(self.fp.name) - self.assertIn("TABPY_ARROW_ENABLE", err.exception.args[0]) + app = TabPyApp(self.fp.name) + self.assertTrue(app.settings["oauth_enabled"]) + self.assertTrue(app.settings["arrow_enabled"]) + with patch("tabpy.tabpy_server.app.app.pa.FlightServer") as mock_fs: + app._get_arrow_server(app.settings) + middleware = mock_fs.call_args.kwargs["middleware"] + jwt_mw = middleware["jwt"] + self.assertIn("jwt", middleware) + self.assertNotIn("basic", middleware) + self.assertEqual(jwt_mw.issuer, "https://idp.example.com/") + self.assertEqual( + jwt_mw.jwks_uri, "https://idp.example.com/.well-known/jwks.json" + ) + self.assertEqual(jwt_mw.audience, "tabpy") + self.assertIsNone(jwt_mw.required_scopes) + self.assertIsNone(jwt_mw.basic_factory) @patch( "tabpy.tabpy_server.app.app.socket.getaddrinfo", @@ -620,11 +632,44 @@ def test_oauth_and_basic_auth_with_arrow_enabled_succeeds(self, mock_getaddrinfo "TABPY_OAUTH_ISSUER = https://idp.example.com/\n" "TABPY_OAUTH_JWKS_URI = https://idp.example.com/.well-known/jwks.json\n" "TABPY_OAUTH_AUDIENCE = tabpy\n" + "TABPY_OAUTH_REQUIRED_SCOPES = tabpy:query\n" ) self.fp.close() app = TabPyApp(self.fp.name) self.assertTrue(app.settings["oauth_enabled"]) + with patch("tabpy.tabpy_server.app.app.pa.FlightServer") as mock_fs: + app._get_arrow_server(app.settings) + middleware = mock_fs.call_args.kwargs["middleware"] + jwt_mw = middleware["jwt"] + self.assertIn("jwt", middleware) + self.assertNotIn("basic", middleware) + self.assertEqual(jwt_mw.issuer, "https://idp.example.com/") + self.assertEqual( + jwt_mw.jwks_uri, "https://idp.example.com/.well-known/jwks.json" + ) + self.assertEqual(jwt_mw.audience, "tabpy") + self.assertEqual(jwt_mw.required_scopes, "tabpy:query") + self.assertIsNotNone(jwt_mw.basic_factory) + + def test_basic_auth_with_arrow_installs_basic_middleware(self): + pwd_file = os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname(__file__))), + "integration", "resources", "pwdfile.txt", + ) + self.fp.write( + "[TabPy]\n" + f"TABPY_PWD_FILE = {pwd_file}\n" + "TABPY_ARROW_ENABLE = true\n" + ) + self.fp.close() + + app = TabPyApp(self.fp.name) + with patch("tabpy.tabpy_server.app.app.pa.FlightServer") as mock_fs: + app._get_arrow_server(app.settings) + middleware = mock_fs.call_args.kwargs["middleware"] + self.assertIn("basic", middleware) + self.assertNotIn("jwt", middleware) @patch( "tabpy.tabpy_server.app.app.socket.getaddrinfo", diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index a9da4075..cbe2bdaa 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -3,6 +3,8 @@ import hmac import hashlib import json +import threading +import time import unittest from unittest.mock import patch @@ -17,6 +19,7 @@ JWKS_URI, make_token, patched_jwks_client, + reset_jwks_state, ) @@ -30,11 +33,7 @@ def setUpClass(cls): cls.private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) def setUp(self): - import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - - jwt_auth_module._jwks_clients.clear() - jwt_auth_module._jwks_last_failed_refresh.clear() - jwt_auth_module._jwks_last_fetch_failure.clear() + reset_jwks_state() def _make_token(self, claims_override=None, headers=None): return make_token( @@ -166,13 +165,14 @@ def test_unexpected_jwks_error_is_rejected_not_raised(self): def test_jwks_client_is_created_with_bounded_timeout(self): """ - TabPy serves requests on a single IO-loop thread, so the JWKS HTTP - client must not be allowed to hang indefinitely on a slow/unreachable - IdP -- that would stall the entire server, not just one request. + TabPy's HTTP path serves requests on a single IO-loop thread, so + the JWKS HTTP client must not be allowed to hang indefinitely on a + slow/unreachable IdP -- that would stall concurrent HTTP requests. + Arrow Flight auth shares the same client on the gRPC thread pool. """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - jwt_auth_module._jwks_clients.clear() + reset_jwks_state() client = jwt_auth_module._get_jwks_client(JWKS_URI) self.assertEqual(client.timeout, jwt_auth_module.JWKS_FETCH_TIMEOUT_SECONDS) self.assertLess(jwt_auth_module.JWKS_FETCH_TIMEOUT_SECONDS, 30) @@ -237,16 +237,13 @@ def test_repeated_unknown_kid_does_not_force_unbounded_jwks_refresh(self): (read from the unverified header, before any signature check). Repeatedly retrying an unknown `kid` -- even a different one each time -- must not each force a fresh JWKS fetch, since that fetch is - a blocking network call on TabPy's single IO-loop thread. The + a blocking network call (HTTP IO-loop or a Flight gRPC thread). The cooldown is keyed only by jwks_uri (not by kid, since kid is attacker-controlled pre-signature-check): only one forced refresh per jwks_uri is allowed within JWKS_MIN_REFRESH_INTERVAL_SECONDS, no matter how many distinct bogus kids are tried. """ - import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - - jwt_auth_module._jwks_clients.clear() - jwt_auth_module._jwks_last_failed_refresh.clear() + reset_jwks_state() token1 = self._make_token(headers={"kid": "unknown-kid-1"}) token2 = self._make_token(headers={"kid": "unknown-kid-2"}) @@ -285,10 +282,7 @@ def test_unknown_kid_refresh_cooldown_blocks_a_different_kid(self): varying the kid, which is a worse (unauthenticated DoS) outcome than briefly delaying visibility of a rotated key. """ - import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - - jwt_auth_module._jwks_clients.clear() - jwt_auth_module._jwks_last_failed_refresh.clear() + reset_jwks_state() bogus_token = self._make_token(headers={"kid": "unknown-kid"}) with patch( @@ -320,15 +314,12 @@ def test_failed_jwks_fetch_is_rate_limited(self): """ A down/unreachable IdP must not be hammered with a fresh blocking fetch (see JWKS_FETCH_TIMEOUT_SECONDS) on every single request -- - that fetch runs directly on TabPy's single IO-loop thread. After + that fetch blocks the HTTP IO-loop or a Flight gRPC thread. After one failed fetch, subsequent requests within JWKS_MIN_REFRESH_INTERVAL_SECONDS must fail fast without calling get_signing_keys again. """ - import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - - jwt_auth_module._jwks_clients.clear() - jwt_auth_module._jwks_last_fetch_failure.clear() + reset_jwks_state() token = self._make_token() with patch( @@ -374,6 +365,44 @@ def test_hs256_substitution_using_public_key_is_rejected(self): with self.assertRaises(JwtValidationError): validate_jwt(forged_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) + def test_waiting_on_an_in_flight_jwks_fetch_fails_fast(self): + """ + _jwks_lock is held across the blocking JWKS fetch, and Arrow Flight + (gRPC thread pool) shares it with HTTP (single IO-loop thread). A + caller that can't take the lock promptly must fail closed instead + of blocking for up to JWKS_FETCH_TIMEOUT_SECONDS, which would let + an unauthenticated Flight caller stall every concurrent HTTP + request behind the fetch it triggered. + """ + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + token = self._make_token() + acquired = threading.Event() + release = threading.Event() + + def hold_lock(): + with jwt_auth_module._jwks_lock: + acquired.set() + release.wait(10) + + holder = threading.Thread(target=hold_lock) + holder.start() + try: + self.assertTrue(acquired.wait(5)) + with patch.object(jwt_auth_module, "JWKS_LOCK_WAIT_SECONDS", 0.05): + with self._patched_jwks_client(): + started = time.monotonic() + with self.assertRaises(JwtValidationError): + validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + elapsed = time.monotonic() - started + finally: + release.set() + holder.join() + + self.assertLess(elapsed, 1) + def test_jwks_client_is_reused_for_same_uri(self): """ _get_jwks_client must return the same PyJWKClient instance for @@ -382,7 +411,7 @@ def test_jwks_client_is_reused_for_same_uri(self): """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - jwt_auth_module._jwks_clients.clear() + reset_jwks_state() first = jwt_auth_module._get_jwks_client(JWKS_URI) second = jwt_auth_module._get_jwks_client(JWKS_URI) self.assertIs(first, second) diff --git a/tests/unit/server_tests/test_jwt_server_middleware_factory.py b/tests/unit/server_tests/test_jwt_server_middleware_factory.py new file mode 100644 index 00000000..7722520f --- /dev/null +++ b/tests/unit/server_tests/test_jwt_server_middleware_factory.py @@ -0,0 +1,245 @@ +import base64 +import datetime +import unittest +from unittest.mock import patch + +from cryptography.hazmat.primitives.asymmetric import rsa +from pyarrow.flight import FlightUnauthenticatedError + +from tabpy.tabpy_server.handlers.basic_auth_server_middleware_factory import ( + BasicAuthServerMiddleware, + BasicAuthServerMiddlewareFactory, +) +from tabpy.tabpy_server.handlers.jwt_server_middleware_factory import ( + JwtAuthServerMiddleware, + JwtAuthServerMiddlewareFactory, +) +from tabpy.tabpy_server.handlers.util import hash_password +from tests.unit.server_tests.jwt_test_helpers import ( + AUDIENCE, + ISSUER, + JWKS_URI, + make_token, + patched_jwks_client, + reset_jwks_state, +) + + +class TestJwtServerMiddlewareFactory(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + + def setUp(self): + reset_jwks_state() + self.factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + + def _make_token(self, claims_override=None): + return make_token(self.private_key, claims_override=claims_override) + + def _headers(self, authorization=None): + if authorization is None: + return {} + return {"authorization": [authorization]} + + def _basic_header(self, username, password): + encoded = base64.b64encode(f"{username}:{password}".encode()).decode() + return f"Basic {encoded}" + + def test_valid_token_is_accepted(self): + token = self._make_token() + with patched_jwks_client(self.private_key): + middleware = self.factory.start_call( + None, self._headers(f"Bearer {token}") + ) + self.assertIsInstance(middleware, JwtAuthServerMiddleware) + self.assertEqual(middleware.claims["sub"], "user1") + + def test_expired_token_is_rejected(self): + past = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( + hours=1 + ) + token = self._make_token( + {"iat": past - datetime.timedelta(minutes=5), "exp": past} + ) + with patched_jwks_client(self.private_key): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers(f"Bearer {token}")) + + def test_wrong_issuer_is_rejected(self): + token = self._make_token({"iss": "https://wrong-idp.example.com/"}) + with patched_jwks_client(self.private_key): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers(f"Bearer {token}")) + + def test_wrong_audience_is_rejected(self): + token = self._make_token({"aud": "wrong-audience"}) + with patched_jwks_client(self.private_key): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers(f"Bearer {token}")) + + def test_missing_token_is_rejected(self): + with self.assertRaises(FlightUnauthenticatedError) as err: + self.factory.start_call(None, {}) + self.assertIn("No credentials supplied", str(err.exception)) + + def test_empty_bearer_is_rejected(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers("Bearer")) + + def test_bearer_with_extra_whitespace_is_rejected(self): + token = self._make_token() + with patched_jwks_client(self.private_key): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call( + None, self._headers(f"Bearer {token}") + ) + + def test_unknown_scheme_is_rejected(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call(None, self._headers("Token abc")) + + def test_repeated_identical_authorization_value_is_accepted(self): + token = self._make_token() + headers = {"authorization": [f"Bearer {token}", f"Bearer {token}"]} + with patched_jwks_client(self.private_key): + middleware = self.factory.start_call(None, headers) + self.assertIsInstance(middleware, JwtAuthServerMiddleware) + + def test_case_varied_identical_authorization_headers_are_accepted(self): + token = self._make_token() + headers = { + "Authorization": [f"Bearer {token}"], + "authorization": [f"Bearer {token}"], + } + with patched_jwks_client(self.private_key): + middleware = self.factory.start_call(None, headers) + self.assertIsInstance(middleware, JwtAuthServerMiddleware) + + def test_conflicting_authorization_values_are_rejected(self): + token = self._make_token() + headers = { + "authorization": [f"Bearer {token}"], + "Authorization": [self._basic_header("user1", "P@ssw0rd")], + } + with patched_jwks_client(self.private_key): + with self.assertRaises(FlightUnauthenticatedError) as err: + self.factory.start_call(None, headers) + self.assertIn("No credentials supplied", str(err.exception)) + + def test_unexpected_validation_error_is_unauthenticated(self): + token = self._make_token() + with patch( + "tabpy.tabpy_server.handlers.jwt_server_middleware_factory.validate_jwt", + side_effect=RuntimeError("boom"), + ): + with self.assertRaises(FlightUnauthenticatedError) as err: + self.factory.start_call(None, self._headers(f"Bearer {token}")) + self.assertIn("Invalid credentials", str(err.exception)) + + def test_basic_is_rejected_when_no_basic_factory(self): + with self.assertRaises(FlightUnauthenticatedError): + self.factory.start_call( + None, self._headers(self._basic_header("user1", "P@ssw0rd")) + ) + + def test_valid_basic_is_accepted_when_basic_factory_provided(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=BasicAuthServerMiddlewareFactory(creds), + ) + middleware = factory.start_call( + None, self._headers(self._basic_header("user1", "P@ssw0rd")) + ) + self.assertIsInstance(middleware, BasicAuthServerMiddleware) + + def test_basic_issued_bearer_token_is_accepted_on_next_call(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + basic_factory = BasicAuthServerMiddlewareFactory(creds) + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=basic_factory, + ) + + handshake = factory.start_call( + None, self._headers(self._basic_header("user1", "P@ssw0rd")) + ) + subsequent_call = factory.start_call( + None, self._headers(f"Bearer {handshake.token}") + ) + + self.assertIsInstance(subsequent_call, BasicAuthServerMiddleware) + self.assertEqual(subsequent_call.token, handshake.token) + + def test_lowercase_basic_is_accepted_when_basic_factory_provided(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=BasicAuthServerMiddlewareFactory(creds), + ) + encoded = base64.b64encode(b"user1:P@ssw0rd").decode() + middleware = factory.start_call( + None, self._headers(f"basic {encoded}") + ) + self.assertIsInstance(middleware, BasicAuthServerMiddleware) + + def test_invalid_basic_is_rejected_when_basic_factory_provided(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=BasicAuthServerMiddlewareFactory(creds), + ) + with self.assertRaises(FlightUnauthenticatedError): + factory.start_call( + None, + self._headers(self._basic_header("user1", "wrong_password")), + ) + + def test_invalid_base64_basic_is_unauthenticated(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=BasicAuthServerMiddlewareFactory(creds), + ) + with self.assertRaises(FlightUnauthenticatedError) as err: + factory.start_call(None, self._headers("Basic !!!not-base64!!!")) + self.assertIn("Invalid credentials", str(err.exception)) + + def test_invalid_utf8_basic_is_unauthenticated(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=BasicAuthServerMiddlewareFactory(creds), + ) + encoded = base64.b64encode(b"\xff\xfe").decode() + with self.assertRaises(FlightUnauthenticatedError) as err: + factory.start_call(None, self._headers(f"Basic {encoded}")) + self.assertIn("Invalid credentials", str(err.exception)) + + def test_basic_without_password_separator_is_unauthenticated(self): + creds = {"user1": hash_password("user1", "P@ssw0rd")} + factory = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + basic_factory=BasicAuthServerMiddlewareFactory(creds), + ) + encoded = base64.b64encode(b"user1").decode() + with self.assertRaises(FlightUnauthenticatedError) as err: + factory.start_call(None, self._headers(f"Basic {encoded}")) + self.assertIn("Invalid credentials", str(err.exception)) From 8475d2bbaa7f673c7537ea3caecdcd839522051a Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Mon, 17 Aug 2026 14:15:40 -0700 Subject: [PATCH 02/10] Send JSON Content-Type in max-request-size tests Without it Tornado treats the 2MB body as form-urlencoded and newer Python parse_qsl rejects it with 400 before TabPy's 413 check runs. --- tests/unit/server_tests/test_evaluation_plane_handler.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/server_tests/test_evaluation_plane_handler.py b/tests/unit/server_tests/test_evaluation_plane_handler.py index e6aba894..ceadb462 100755 --- a/tests/unit/server_tests/test_evaluation_plane_handler.py +++ b/tests/unit/server_tests/test_evaluation_plane_handler.py @@ -497,7 +497,8 @@ def test_evaluation_payload_exceeds_max_request_size(self): response = self.fetch( "/evaluate", method="POST", - body=self.create_large_payload() + body=self.create_large_payload(), + headers={"Content-Type": "application/json"}, ) self.assertEqual(413, response.code) @@ -506,7 +507,8 @@ def test_evaluation_max_request_size_not_applied(self): response = self.fetch( "/evaluate", method="POST", - body=self.create_large_payload() + body=self.create_large_payload(), + headers={"Content-Type": "application/json"}, ) self.assertEqual(200, response.code) self.assertEqual(1, json.loads(response.body)[0]) From d4da0cf532e945f812b28d97f124d84bc1268e13 Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Mon, 17 Aug 2026 17:06:20 -0700 Subject: [PATCH 03/10] Move Flight Authorization helper out of password util Keep util.py unchanged so the pre-existing PBKDF2 hash stays out of this diff. Flight header parsing now lives in flight_headers.py. --- .../basic_auth_server_middleware_factory.py | 4 +-- tabpy/tabpy_server/handlers/flight_headers.py | 22 ++++++++++++++++ .../handlers/jwt_server_middleware_factory.py | 4 ++- tabpy/tabpy_server/handlers/util.py | 25 ------------------- 4 files changed, 27 insertions(+), 28 deletions(-) create mode 100644 tabpy/tabpy_server/handlers/flight_headers.py diff --git a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py index 2c61d381..2cd919c2 100644 --- a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py @@ -7,10 +7,10 @@ from pyarrow.flight import ServerMiddlewareFactory, ServerMiddleware from pyarrow.flight import FlightUnauthenticatedError -from tabpy.tabpy_server.handlers.util import ( +from tabpy.tabpy_server.handlers.flight_headers import ( get_flight_authorization_header, - hash_password, ) +from tabpy.tabpy_server.handlers.util import hash_password # A successful Basic call mints an opaque token and hands it back to the # client via sending_headers(). The client may replay it as a Bearer diff --git a/tabpy/tabpy_server/handlers/flight_headers.py b/tabpy/tabpy_server/handlers/flight_headers.py new file mode 100644 index 00000000..eb2ae8ab --- /dev/null +++ b/tabpy/tabpy_server/handlers/flight_headers.py @@ -0,0 +1,22 @@ +def get_flight_authorization_header(headers): + """ + Returns the Authorization value from pyarrow Flight headers. + + Header names are matched case-insensitively. A missing or empty value + returns None so callers fail closed, as does a request carrying two + different Authorization values: which one wins would otherwise decide + the identity. Repeating the same value is harmless and accepted, + since some clients set the header through more than one mechanism. + """ + values = set() + for header, header_values in headers.items(): + if header.lower() == "authorization": + if isinstance(header_values, (list, tuple)): + values.update(header_values) + else: + values.add(header_values) + + if len(values) != 1: + return None + value = values.pop() + return value or None diff --git a/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py index d307641f..9a3a3953 100644 --- a/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py @@ -3,8 +3,10 @@ from pyarrow.flight import FlightUnauthenticatedError from pyarrow.flight import ServerMiddleware, ServerMiddlewareFactory +from tabpy.tabpy_server.handlers.flight_headers import ( + get_flight_authorization_header, +) from tabpy.tabpy_server.handlers.jwt_auth import JwtValidationError, validate_jwt -from tabpy.tabpy_server.handlers.util import get_flight_authorization_header logger = logging.getLogger(__name__) diff --git a/tabpy/tabpy_server/handlers/util.py b/tabpy/tabpy_server/handlers/util.py index 66f171c3..c9fc0e43 100644 --- a/tabpy/tabpy_server/handlers/util.py +++ b/tabpy/tabpy_server/handlers/util.py @@ -8,7 +8,6 @@ class AuthErrorStates(Enum): NotAuthorized = auto() NotRequired = auto() - def hash_password(username, pwd): """ Hashes password using PKDBF2 method: @@ -37,27 +36,3 @@ def hash_password(username, pwd): hash_name="sha512", password=pwd.encode(), salt=salt.encode(), iterations=10000 ) return binascii.hexlify(hash).decode() - - -def get_flight_authorization_header(headers): - """ - Returns the Authorization value from pyarrow Flight headers. - - Header names are matched case-insensitively. A missing or empty value - returns None so callers fail closed, as does a request carrying two - different Authorization values: which one wins would otherwise decide - the identity. Repeating the same value is harmless and accepted, - since some clients set the header through more than one mechanism. - """ - values = set() - for header, header_values in headers.items(): - if header.lower() == "authorization": - if isinstance(header_values, (list, tuple)): - values.update(header_values) - else: - values.add(header_values) - - if len(values) != 1: - return None - value = values.pop() - return value or None From 14bb7752e518873892730f50def7ca5bf410989a Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Tue, 18 Aug 2026 15:05:07 -0700 Subject: [PATCH 04/10] Stop JWKS refreshes from blocking cached JWT validation Hold a per-URI fetch lock only around forced JWKS refreshes so a forged unknown kid cannot reject valid HTTP or Flight tokens. Flight still follows TABPY_TRANSFER_PROTOCOL, including http, matching the existing HTTP Basic/Bearer cleartext behavior. --- docs/server-config.md | 15 ++- tabpy/tabpy_server/app/app.py | 3 + tabpy/tabpy_server/handlers/jwt_auth.py | 110 +++++++++++++------- tests/unit/server_tests/jwt_test_helpers.py | 1 + tests/unit/server_tests/test_jwt_auth.py | 70 +++++++++---- 5 files changed, 134 insertions(+), 65 deletions(-) diff --git a/docs/server-config.md b/docs/server-config.md index 09983a6a..df6133f1 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -342,16 +342,15 @@ method based on the scheme of the `Authorization` header sent by the client (`Basic` or `Bearer`), so both can be used against the same server. With `TABPY_TRANSFER_PROTOCOL = http`, Flight uses `grpc+tcp` and the Bearer -token is sent in cleartext, the same as HTTP Basic/Bearer on an unencrypted -port. +token is sent in cleartext. That matches HTTP Basic/Bearer on the same +setting; TabPy does not require HTTPS for Flight auth alone. JWKS lookups are cached. On the HTTP path a cache-cold fetch runs on TabPy's -single IO-loop thread. Arrow Flight auth runs on the gRPC thread pool. JWKS -client creation and fetches are serialized with a shared lock so concurrent -Flight calls cannot bypass the refresh rate limit. A JWT check that can't -take that lock within one second is rejected rather than left waiting for -the in-flight fetch, so a slow or unresponsive identity provider can't stall -concurrent requests for the full fetch timeout. +single IO-loop thread. Arrow Flight auth runs on the gRPC thread pool. +Forced JWKS refreshes are serialized per JWKS URI so concurrent Flight +calls cannot bypass the refresh rate limit. A cached signing-key lookup +does not wait on that refresh, so a forged unknown `kid` cannot reject +unrelated valid tokens. A request carrying two conflicting `Authorization` values is rejected, because which one wins would otherwise decide the caller's identity. diff --git a/tabpy/tabpy_server/app/app.py b/tabpy/tabpy_server/app/app.py index 715947b1..c7dbd084 100644 --- a/tabpy/tabpy_server/app/app.py +++ b/tabpy/tabpy_server/app/app.py @@ -124,6 +124,9 @@ def _get_tls_certificates(self, config): def _get_arrow_server(self, config): verify_client = None tls_certificates = None + # Same transport as the HTTP server: http -> grpc+tcp, https -> + # grpc+tls. Cleartext Flight is allowed because HTTP Basic/Bearer + # already travel in the clear when TABPY_TRANSFER_PROTOCOL=http. scheme = "grpc+tcp" if config[SettingsParameters.TransferProtocol] == "https": scheme = "grpc+tls" diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index ee92e4da..2d6f880d 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -14,14 +14,6 @@ # A cold start or legitimate key rotation still pays this cost. JWKS_FETCH_TIMEOUT_SECONDS = 10 -# Bounds how long a JWT check waits for _jwks_lock. Arrow Flight auth runs -# on the gRPC thread pool and HTTP auth runs on the Tornado IO loop, and -# both share the lock, so without a bound an unauthenticated Flight caller -# could force a slow JWKS fetch and stall every HTTP request behind it for -# up to JWKS_FETCH_TIMEOUT_SECONDS. Callers that time out waiting fail -# closed instead; a responsive IdP resolves well inside this window. -JWKS_LOCK_WAIT_SECONDS = 1 - # An unauthenticated caller can force a fresh JWKS fetch just by sending a # made-up `kid` (read from the token header pre-signature-check). This # bounds, per jwks_uri, how often a *failed* forced refresh (kid still not @@ -33,10 +25,12 @@ JWKS_MIN_REFRESH_INTERVAL_SECONDS = 30 # One PyJWKClient per JWKS URI, reused so its JWK Set cache actually avoids -# per-request fetches. Process-global; mutations are serialized by -# _jwks_lock because HTTP (IO-loop) and Arrow Flight (gRPC thread pool) -# share these dicts. -_jwks_lock = threading.Lock() +# per-request fetches. Process-global. _jwks_state_lock only covers these +# dicts. A per-URI fetch lock serializes refreshes for that IdP so two +# Flight threads cannot bypass the cooldown. Cached kid lookups do not +# take that fetch lock, so a refresh cannot reject unrelated valid tokens. +_jwks_state_lock = threading.Lock() +_jwks_fetch_locks = {} _jwks_clients = {} # jwks_uri -> monotonic timestamp of the last failed JWKS fetch (network @@ -58,27 +52,74 @@ def _get_jwks_client(jwks_uri: str) -> PyJWKClient: - client = _jwks_clients.get(jwks_uri) - if client is None: - client = PyJWKClient( - jwks_uri, cache_jwk_set=True, timeout=JWKS_FETCH_TIMEOUT_SECONDS - ) - _jwks_clients[jwks_uri] = client - return client + with _jwks_state_lock: + client = _jwks_clients.get(jwks_uri) + if client is None: + client = PyJWKClient( + jwks_uri, cache_jwk_set=True, timeout=JWKS_FETCH_TIMEOUT_SECONDS + ) + _jwks_clients[jwks_uri] = client + return client + + +def _fetch_lock_for(jwks_uri: str) -> threading.Lock: + with _jwks_state_lock: + lock = _jwks_fetch_locks.get(jwks_uri) + if lock is None: + lock = threading.Lock() + _jwks_fetch_locks[jwks_uri] = lock + return lock + + +def _recently_recorded(store: dict, jwks_uri: str) -> bool: + with _jwks_state_lock: + last = store.get(jwks_uri, 0) + return time.monotonic() - last < JWKS_MIN_REFRESH_INTERVAL_SECONDS + + +def _record_now(store: dict, jwks_uri: str) -> None: + with _jwks_state_lock: + store[jwks_uri] = time.monotonic() def _fetch_signing_keys(jwks_client: PyJWKClient, jwks_uri: str, refresh: bool): - now = time.monotonic() - last_failure = _jwks_last_fetch_failure.get(jwks_uri, 0) - if now - last_failure < JWKS_MIN_REFRESH_INTERVAL_SECONDS: + if not refresh: + # Cache lookup. Must not wait on another thread's refresh, or a + # forged unknown kid can reject valid tokens on HTTP and Flight. + # A cold cache may still hit the network here; that is the first + # request to this IdP, not an attacker-forced refresh. + if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): + raise jwt.exceptions.PyJWKClientError( + f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' + ) + try: + return jwks_client.get_signing_keys(refresh=False) + except Exception: + _record_now(_jwks_last_fetch_failure, jwks_uri) + raise + + if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): raise jwt.exceptions.PyJWKClientError( f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' ) + + fetch_lock = _fetch_lock_for(jwks_uri) + if not fetch_lock.acquire(blocking=False): + raise jwt.exceptions.PyJWKClientError( + f'JWKS refresh already in flight for "{jwks_uri}"' + ) try: - return jwks_client.get_signing_keys(refresh=refresh) - except Exception: - _jwks_last_fetch_failure[jwks_uri] = now - raise + if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): + raise jwt.exceptions.PyJWKClientError( + f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' + ) + try: + return jwks_client.get_signing_keys(refresh=True) + except Exception: + _record_now(_jwks_last_fetch_failure, jwks_uri) + raise + finally: + fetch_lock.release() def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): @@ -96,9 +137,7 @@ def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): if signing_key is not None: return signing_key - now = time.monotonic() - last_failure = _jwks_last_failed_refresh.get(jwks_uri, 0) - if now - last_failure < JWKS_MIN_REFRESH_INTERVAL_SECONDS: + if _recently_recorded(_jwks_last_failed_refresh, jwks_uri): raise jwt.exceptions.PyJWKClientError( f"Unable to find a signing key that matches: {kid!r}" ) @@ -106,7 +145,7 @@ def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=True) signing_key = PyJWKClient.match_kid(signing_keys, kid) if signing_key is None: - _jwks_last_failed_refresh[jwks_uri] = now + _record_now(_jwks_last_failed_refresh, jwks_uri) raise jwt.exceptions.PyJWKClientError( f"Unable to find a signing key that matches: {kid!r}" ) @@ -157,15 +196,8 @@ def validate_jwt( raise JwtValidationError("Missing JWT") try: - if not _jwks_lock.acquire(timeout=JWKS_LOCK_WAIT_SECONDS): - raise jwt.exceptions.PyJWKClientError( - f'Timed out waiting on an in-flight JWKS fetch for "{jwks_uri}"' - ) - try: - jwks_client = _get_jwks_client(jwks_uri) - signing_key = _get_signing_key(jwks_client, jwks_uri, token) - finally: - _jwks_lock.release() + jwks_client = _get_jwks_client(jwks_uri) + signing_key = _get_signing_key(jwks_client, jwks_uri, token) except (jwt.exceptions.PyJWKClientError, jwt.exceptions.InvalidTokenError) as ex: logger.log(logging.ERROR, f"Unable to resolve JWT signing key: {str(ex)}") raise JwtValidationError("Unable to resolve JWT signing key") from ex diff --git a/tests/unit/server_tests/jwt_test_helpers.py b/tests/unit/server_tests/jwt_test_helpers.py index 1f5dd20f..ddf85f55 100644 --- a/tests/unit/server_tests/jwt_test_helpers.py +++ b/tests/unit/server_tests/jwt_test_helpers.py @@ -19,6 +19,7 @@ def reset_jwks_state(): import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module jwt_auth_module._jwks_clients.clear() + jwt_auth_module._jwks_fetch_locks.clear() jwt_auth_module._jwks_last_failed_refresh.clear() jwt_auth_module._jwks_last_fetch_failure.clear() diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index cbe2bdaa..47f1791f 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -365,38 +365,72 @@ def test_hs256_substitution_using_public_key_is_rejected(self): with self.assertRaises(JwtValidationError): validate_jwt(forged_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) - def test_waiting_on_an_in_flight_jwks_fetch_fails_fast(self): + def test_cached_key_lookup_does_not_wait_on_a_refresh(self): """ - _jwks_lock is held across the blocking JWKS fetch, and Arrow Flight - (gRPC thread pool) shares it with HTTP (single IO-loop thread). A - caller that can't take the lock promptly must fail closed instead - of blocking for up to JWKS_FETCH_TIMEOUT_SECONDS, which would let - an unauthenticated Flight caller stall every concurrent HTTP - request behind the fetch it triggered. + A forged unknown kid can start a per-URI JWKS refresh. Valid tokens + that already have a cached key must still validate while that + refresh is in flight, instead of being rejected behind the lock. """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module token = self._make_token() + fetch_lock = jwt_auth_module._fetch_lock_for(JWKS_URI) acquired = threading.Event() release = threading.Event() - def hold_lock(): - with jwt_auth_module._jwks_lock: + def hold_refresh(): + with fetch_lock: acquired.set() release.wait(10) - holder = threading.Thread(target=hold_lock) + holder = threading.Thread(target=hold_refresh) holder.start() try: self.assertTrue(acquired.wait(5)) - with patch.object(jwt_auth_module, "JWKS_LOCK_WAIT_SECONDS", 0.05): - with self._patched_jwks_client(): - started = time.monotonic() - with self.assertRaises(JwtValidationError): - validate_jwt( - token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE - ) - elapsed = time.monotonic() - started + with self._patched_jwks_client(): + started = time.monotonic() + claims = validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + elapsed = time.monotonic() - started + finally: + release.set() + holder.join() + + self.assertEqual(claims["sub"], "user1") + self.assertLess(elapsed, 1) + + def test_in_flight_refresh_does_not_block_another_refresh_attempt(self): + """ + A second forced refresh for the same jwks_uri must fail immediately + rather than wait for the in-flight fetch timeout. + """ + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + token = self._make_token(headers={"kid": "unknown-kid"}) + fetch_lock = jwt_auth_module._fetch_lock_for(JWKS_URI) + acquired = threading.Event() + release = threading.Event() + + def hold_refresh(): + with fetch_lock: + acquired.set() + release.wait(10) + + holder = threading.Thread(target=hold_refresh) + holder.start() + try: + self.assertTrue(acquired.wait(5)) + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", + return_value=[], + ): + started = time.monotonic() + with self.assertRaises(JwtValidationError): + validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + elapsed = time.monotonic() - started finally: release.set() holder.join() From e30de8fe4caf8d9b6d365d8f2061a04ec9d94b75 Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Tue, 18 Aug 2026 15:13:31 -0700 Subject: [PATCH 05/10] Single-flight cold JWKS fetches and recheck unknown-kid cooldown --- tabpy/tabpy_server/handlers/jwt_auth.py | 104 ++++++--- tests/unit/server_tests/jwt_test_helpers.py | 1 + tests/unit/server_tests/test_jwt_auth.py | 230 ++++++++++++++++++-- 3 files changed, 293 insertions(+), 42 deletions(-) diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 2d6f880d..47cecba1 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -24,15 +24,23 @@ # can't be hammered with a fresh blocking fetch on every single request. JWKS_MIN_REFRESH_INTERVAL_SECONDS = 30 +# Matches PyJWKClient's default JWK Set cache lifespan. After this, the +# next lookup treats the snapshot as expired and single-flights one fetch. +JWKS_CACHE_LIFESPAN_SECONDS = 300 + # One PyJWKClient per JWKS URI, reused so its JWK Set cache actually avoids # per-request fetches. Process-global. _jwks_state_lock only covers these -# dicts. A per-URI fetch lock serializes refreshes for that IdP so two -# Flight threads cannot bypass the cooldown. Cached kid lookups do not -# take that fetch lock, so a refresh cannot reject unrelated valid tokens. +# dicts. A per-URI fetch lock serializes cache-miss fetches and forced +# refreshes for that IdP. Warm (unexpired) kid lookups do not take that +# fetch lock, so a refresh cannot reject unrelated valid tokens. _jwks_state_lock = threading.Lock() _jwks_fetch_locks = {} _jwks_clients = {} +# jwks_uri -> (signing_keys, monotonic timestamp). Lets a warm cache +# return keys without waiting on another thread's in-flight fetch. +_jwks_cached_keys = {} + # jwks_uri -> monotonic timestamp of the last failed JWKS fetch (network # error, timeout, malformed response -- not a kid mismatch). Rate-limits # retrying a broken/unreachable IdP, independent of which kid was @@ -82,26 +90,69 @@ def _record_now(store: dict, jwks_uri: str) -> None: store[jwks_uri] = time.monotonic() -def _fetch_signing_keys(jwks_client: PyJWKClient, jwks_uri: str, refresh: bool): +def _read_cached_keys(jwks_uri: str): + with _jwks_state_lock: + entry = _jwks_cached_keys.get(jwks_uri) + if entry is None: + return None + keys, fetched_at = entry + if time.monotonic() - fetched_at >= JWKS_CACHE_LIFESPAN_SECONDS: + return None + return keys + + +def _write_cached_keys(jwks_uri: str, keys) -> None: + if not keys: + return + with _jwks_state_lock: + _jwks_cached_keys[jwks_uri] = (list(keys), time.monotonic()) + + +def _fetch_failed_recently_error(jwks_uri: str) -> jwt.exceptions.PyJWKClientError: + return jwt.exceptions.PyJWKClientError( + f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' + ) + + +def _call_get_signing_keys(jwks_client: PyJWKClient, jwks_uri: str, refresh: bool): + if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): + raise _fetch_failed_recently_error(jwks_uri) + try: + keys = jwks_client.get_signing_keys(refresh=refresh) + except Exception: + _record_now(_jwks_last_fetch_failure, jwks_uri) + raise + _write_cached_keys(jwks_uri, keys) + return keys + + +def _fetch_signing_keys( + jwks_client: PyJWKClient, jwks_uri: str, refresh: bool, kid=None +): if not refresh: - # Cache lookup. Must not wait on another thread's refresh, or a - # forged unknown kid can reject valid tokens on HTTP and Flight. - # A cold cache may still hit the network here; that is the first - # request to this IdP, not an attacker-forced refresh. + cached = _read_cached_keys(jwks_uri) + if cached is not None: + # Warm cache. Must not wait on another thread's refresh, or a + # forged unknown kid can reject valid tokens on HTTP and Flight. + return cached + + # Cold or expired snapshot: single-flight one outbound fetch so + # concurrent Flight threads cannot each pay JWKS_FETCH_TIMEOUT. if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): - raise jwt.exceptions.PyJWKClientError( - f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' - ) + raise _fetch_failed_recently_error(jwks_uri) + + fetch_lock = _fetch_lock_for(jwks_uri) + fetch_lock.acquire() try: - return jwks_client.get_signing_keys(refresh=False) - except Exception: - _record_now(_jwks_last_fetch_failure, jwks_uri) - raise + cached = _read_cached_keys(jwks_uri) + if cached is not None: + return cached + return _call_get_signing_keys(jwks_client, jwks_uri, refresh=False) + finally: + fetch_lock.release() if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): - raise jwt.exceptions.PyJWKClientError( - f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' - ) + raise _fetch_failed_recently_error(jwks_uri) fetch_lock = _fetch_lock_for(jwks_uri) if not fetch_lock.acquire(blocking=False): @@ -109,15 +160,16 @@ def _fetch_signing_keys(jwks_client: PyJWKClient, jwks_uri: str, refresh: bool): f'JWKS refresh already in flight for "{jwks_uri}"' ) try: + # Recheck both cooldowns under the lock. The unknown-kid cooldown + # is set by a sibling thread after this caller already passed the + # unlocked check in _get_signing_key. if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): + raise _fetch_failed_recently_error(jwks_uri) + if _recently_recorded(_jwks_last_failed_refresh, jwks_uri): raise jwt.exceptions.PyJWKClientError( - f'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' + f"Unable to find a signing key that matches: {kid!r}" ) - try: - return jwks_client.get_signing_keys(refresh=True) - except Exception: - _record_now(_jwks_last_fetch_failure, jwks_uri) - raise + return _call_get_signing_keys(jwks_client, jwks_uri, refresh=True) finally: fetch_lock.release() @@ -132,7 +184,7 @@ def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): header = jwt.get_unverified_header(token) kid = header.get("kid") - signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=False) + signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=False, kid=kid) signing_key = PyJWKClient.match_kid(signing_keys, kid) if signing_key is not None: return signing_key @@ -142,7 +194,7 @@ def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): f"Unable to find a signing key that matches: {kid!r}" ) - signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=True) + signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=True, kid=kid) signing_key = PyJWKClient.match_kid(signing_keys, kid) if signing_key is None: _record_now(_jwks_last_failed_refresh, jwks_uri) diff --git a/tests/unit/server_tests/jwt_test_helpers.py b/tests/unit/server_tests/jwt_test_helpers.py index ddf85f55..acdf5c87 100644 --- a/tests/unit/server_tests/jwt_test_helpers.py +++ b/tests/unit/server_tests/jwt_test_helpers.py @@ -20,6 +20,7 @@ def reset_jwks_state(): jwt_auth_module._jwks_clients.clear() jwt_auth_module._jwks_fetch_locks.clear() + jwt_auth_module._jwks_cached_keys.clear() jwt_auth_module._jwks_last_failed_refresh.clear() jwt_auth_module._jwks_last_fetch_failure.clear() diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index 47f1791f..19bb8d98 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -365,6 +365,13 @@ def test_hs256_substitution_using_public_key_is_rejected(self): with self.assertRaises(JwtValidationError): validate_jwt(forged_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) + def _signing_key(self, kid=None): + signing_key = type("SigningKey", (), {})() + signing_key.key = self.private_key.public_key() + signing_key.algorithm_name = "RS256" + signing_key.key_id = kid + return signing_key + def test_cached_key_lookup_does_not_wait_on_a_refresh(self): """ A forged unknown kid can start a per-URI JWKS refresh. Valid tokens @@ -374,6 +381,9 @@ def test_cached_key_lookup_does_not_wait_on_a_refresh(self): import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module token = self._make_token() + with self._patched_jwks_client(): + validate_jwt(token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) + fetch_lock = jwt_auth_module._fetch_lock_for(JWKS_URI) acquired = threading.Event() release = threading.Event() @@ -387,12 +397,11 @@ def hold_refresh(): holder.start() try: self.assertTrue(acquired.wait(5)) - with self._patched_jwks_client(): - started = time.monotonic() - claims = validate_jwt( - token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE - ) - elapsed = time.monotonic() - started + started = time.monotonic() + claims = validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + elapsed = time.monotonic() - started finally: release.set() holder.join() @@ -407,6 +416,12 @@ def test_in_flight_refresh_does_not_block_another_refresh_attempt(self): """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + warm_token = self._make_token() + with self._patched_jwks_client(): + validate_jwt( + warm_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + token = self._make_token(headers={"kid": "unknown-kid"}) fetch_lock = jwt_auth_module._fetch_lock_for(JWKS_URI) acquired = threading.Event() @@ -421,22 +436,205 @@ def hold_refresh(): holder.start() try: self.assertTrue(acquired.wait(5)) - with patch( - "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", - return_value=[], - ): - started = time.monotonic() - with self.assertRaises(JwtValidationError): - validate_jwt( - token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE - ) - elapsed = time.monotonic() - started + started = time.monotonic() + with self.assertRaises(JwtValidationError): + validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + elapsed = time.monotonic() - started finally: release.set() holder.join() self.assertLess(elapsed, 1) + def test_concurrent_cold_cache_performs_one_jwks_fetch(self): + """ + Concurrent cache-cold validations for the same jwks_uri must + single-flight one outbound JWKS fetch, then reuse that snapshot. + """ + token = self._make_token() + signing_key = self._signing_key() + started = threading.Event() + release = threading.Event() + call_count = [] + count_lock = threading.Lock() + + def fake_get_signing_keys(refresh=False): + with count_lock: + call_count.append(refresh) + started.set() + self.assertTrue(release.wait(5)) + return [signing_key] + + workers = 8 + barrier = threading.Barrier(workers) + results = [] + errors = [] + + def worker(): + try: + barrier.wait(5) + claims = validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + results.append(claims) + except Exception as ex: + errors.append(ex) + + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", + side_effect=fake_get_signing_keys, + ): + threads = [threading.Thread(target=worker) for _ in range(workers)] + for thread in threads: + thread.start() + self.assertTrue(started.wait(5)) + time.sleep(0.1) + release.set() + for thread in threads: + thread.join(5) + + self.assertEqual(errors, []) + self.assertEqual(len(results), workers) + self.assertEqual(call_count, [False]) + + def test_concurrent_expired_cache_performs_one_jwks_fetch(self): + """ + Concurrent lookups against an expired snapshot must also + single-flight one outbound JWKS fetch per URI. + """ + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + token = self._make_token() + with self._patched_jwks_client(): + validate_jwt(token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) + + keys, _fetched_at = jwt_auth_module._jwks_cached_keys[JWKS_URI] + jwt_auth_module._jwks_cached_keys[JWKS_URI] = ( + keys, + time.monotonic() - jwt_auth_module.JWKS_CACHE_LIFESPAN_SECONDS - 1, + ) + + signing_key = self._signing_key() + started = threading.Event() + release = threading.Event() + call_count = [] + count_lock = threading.Lock() + + def fake_get_signing_keys(refresh=False): + with count_lock: + call_count.append(refresh) + started.set() + self.assertTrue(release.wait(5)) + return [signing_key] + + workers = 8 + barrier = threading.Barrier(workers) + results = [] + errors = [] + + def worker(): + try: + barrier.wait(5) + claims = validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + results.append(claims) + except Exception as ex: + errors.append(ex) + + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", + side_effect=fake_get_signing_keys, + ): + threads = [threading.Thread(target=worker) for _ in range(workers)] + for thread in threads: + thread.start() + self.assertTrue(started.wait(5)) + time.sleep(0.1) + release.set() + for thread in threads: + thread.join(5) + + self.assertEqual(errors, []) + self.assertEqual(len(results), workers) + self.assertEqual(call_count, [False]) + + def test_unknown_kid_cooldown_is_rechecked_under_fetch_lock(self): + """ + Two unknown-kid callers can both pass the unlocked cooldown check. + The loser of the fetch lock must recheck _jwks_last_failed_refresh + and skip a second outbound refresh. + """ + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + warm_token = self._make_token() + with self._patched_jwks_client(): + validate_jwt( + warm_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + + unknown = self._make_token(headers={"kid": "unknown-kid"}) + signing_key = self._signing_key(kid="the-real-kid") + a_passed_outer = threading.Event() + b_finished = threading.Event() + thread_a_holder = [] + original_recently = jwt_auth_module._recently_recorded + refresh_true_calls = [] + + def gated_recently(store, uri): + result = original_recently(store, uri) + if ( + store is jwt_auth_module._jwks_last_failed_refresh + and thread_a_holder + and threading.current_thread() is thread_a_holder[0] + and not a_passed_outer.is_set() + ): + a_passed_outer.set() + self.assertFalse(result) + self.assertTrue(b_finished.wait(5)) + return result + + def fake_get_signing_keys(refresh=False): + if refresh: + refresh_true_calls.append(1) + return [signing_key] + + def run_a(): + with self.assertRaises(JwtValidationError): + validate_jwt( + unknown, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + + def run_b(): + self.assertTrue(a_passed_outer.wait(5)) + with self.assertRaises(JwtValidationError): + validate_jwt( + unknown, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + b_finished.set() + + thread_a = threading.Thread(target=run_a) + thread_a_holder.append(thread_a) + thread_b = threading.Thread(target=run_b) + + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", + side_effect=fake_get_signing_keys, + ), patch( + "tabpy.tabpy_server.handlers.jwt_auth._recently_recorded", + side_effect=gated_recently, + ): + thread_a.start() + thread_b.start() + thread_a.join(10) + thread_b.join(10) + + self.assertFalse(thread_a.is_alive()) + self.assertFalse(thread_b.is_alive()) + self.assertEqual(refresh_true_calls, [1]) + def test_jwks_client_is_reused_for_same_uri(self): """ _get_jwks_client must return the same PyJWKClient instance for From cdf7e34a195381de8360672297652f24d2ba5dd7 Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Tue, 18 Aug 2026 16:52:41 -0700 Subject: [PATCH 06/10] Make JWKS refresh outcomes atomic Keep unknown-kid cooldown publication under refresh ownership and preserve opaque Flight sessions without cross-user capacity eviction. --- docs/server-config.md | 21 +- .../basic_auth_server_middleware_factory.py | 61 +++-- tabpy/tabpy_server/handlers/jwt_auth.py | 85 +++---- ...st_basic_auth_server_middleware_factory.py | 60 ++++- tests/unit/server_tests/test_jwt_auth.py | 219 ++++++++++++++---- tests/unit/server_tests/test_oauth_handler.py | 8 + 6 files changed, 329 insertions(+), 125 deletions(-) diff --git a/docs/server-config.md b/docs/server-config.md index df6133f1..6067dc29 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -347,10 +347,14 @@ setting; TabPy does not require HTTPS for Flight auth alone. JWKS lookups are cached. On the HTTP path a cache-cold fetch runs on TabPy's single IO-loop thread. Arrow Flight auth runs on the gRPC thread pool. -Forced JWKS refreshes are serialized per JWKS URI so concurrent Flight -calls cannot bypass the refresh rate limit. A cached signing-key lookup -does not wait on that refresh, so a forged unknown `kid` cannot reject -unrelated valid tokens. +Cold and expired-cache fetches are single-flighted per JWKS URI. Forced +unknown-`kid` refreshes are mutually excluded; a concurrent refresh attempt +fails authentication rather than waiting on the network request. A cached +signing-key lookup does not wait on that refresh, so unrelated valid tokens +continue to work. After a refresh still cannot find a requested `kid`, TabPy +rate-limits another forced refresh for 30 seconds. This also means a legitimate +new key may be rejected for up to 30 seconds after a bogus unknown-`kid` +request. A request carrying two conflicting `Authorization` values is rejected, because which one wins would otherwise decide the caller's identity. @@ -373,8 +377,13 @@ gRPC `UNAUTHENTICATED` rather than HTTP 401. After a successful Basic call, Flight also returns an opaque session token in the response `authorization` header. A client may send that token back as a Bearer credential instead of repeating its Basic credentials. The token is -server-issued, is not a JWT, and expires an hour after it is issued, at -which point the client authenticates with Basic again. +server-issued, is not a JWT, and expires an hour after it is issued or when +the server restarts. TabPy retains one active opaque token per username, so +repeated Basic calls do not grow the token store or invalidate another user's +unexpired token. A repeated Basic call normally returns the same token without +extending its original expiry. During its final five minutes, Basic +authentication rotates it to a fresh one-hour token. After expiry, the client +authenticates with Basic again. **As of May 2023, the Arrow Flight feature can only be used by compatible versions of Tableau Prep. The Arrow Flight feature is not used by Tableau diff --git a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py index 2cd919c2..624a5d44 100644 --- a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py @@ -19,10 +19,10 @@ # that is how it obtained the token. FLIGHT_TOKEN_TTL_SECONDS = 3600 -# Expiry alone doesn't bound the store, since a client that ignores the -# token still mints one per call. Cap it and drop the soonest-to-expire -# entries once full. -MAX_FLIGHT_TOKENS = 1024 +# Avoid handing a freshly authenticated client a token that is about to +# expire. Rotation still happens under the per-factory lock and retains +# only one active token for the username. +FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS = 300 class BasicAuthServerMiddleware(ServerMiddleware): @@ -37,11 +37,11 @@ class BasicAuthServerMiddlewareFactory(ServerMiddlewareFactory): def __init__(self, creds): self.creds = creds # token -> (username, monotonic expiry). Read from the gRPC thread - # pool on every call; every mutation goes through _issue_token - # under _tokens_lock, so readers never see the store being - # resized and concurrent calls can't grow it past - # MAX_FLIGHT_TOKENS. + # pool on every call. One active token is retained per username, + # bounding request-driven growth without evicting another user's + # unexpired credential. self.tokens = {} + self._tokens_by_username = {} self._tokens_lock = threading.Lock() def is_valid_user(self, username, password): @@ -51,23 +51,42 @@ def is_valid_user(self, username, password): return self.creds[username].lower() == hashed_pwd.lower() def is_valid_token(self, token): - entry = self.tokens.get(token) - return entry is not None and time.monotonic() < entry[1] + with self._tokens_lock: + entry = self.tokens.get(token) + if entry is None: + return False + username, expiry = entry + if time.monotonic() >= expiry: + self._remove_token(token, username) + return False + return True def _issue_token(self, username): - token = secrets.token_urlsafe(32) with self._tokens_lock: - self._evict_tokens() - self.tokens[token] = (username, time.monotonic() + FLIGHT_TOKEN_TTL_SECONDS) - return token - - def _evict_tokens(self): - now = time.monotonic() + now = time.monotonic() + self._evict_expired_tokens(now) + + existing = self._tokens_by_username.get(username) + if existing is not None: + _, expiry = self.tokens[existing] + if expiry - now > FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS: + return existing + self._remove_token(existing, username) + + token = secrets.token_urlsafe(32) + self.tokens[token] = (username, now + FLIGHT_TOKEN_TTL_SECONDS) + self._tokens_by_username[username] = token + return token + + def _remove_token(self, token, username): + self.tokens.pop(token, None) + if self._tokens_by_username.get(username) == token: + self._tokens_by_username.pop(username, None) + + def _evict_expired_tokens(self, now): for token in [t for t, (_, expiry) in self.tokens.items() if expiry <= now]: - self.tokens.pop(token, None) - while len(self.tokens) >= MAX_FLIGHT_TOKENS: - oldest = min(self.tokens, key=lambda t: self.tokens[t][1]) - self.tokens.pop(oldest, None) + username, _ = self.tokens[token] + self._remove_token(token, username) def start_call(self, info, headers): auth_header = get_flight_authorization_header(headers) diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 47cecba1..3695b1df 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -81,8 +81,11 @@ def _fetch_lock_for(jwks_uri: str) -> threading.Lock: def _recently_recorded(store: dict, jwks_uri: str) -> bool: with _jwks_state_lock: - last = store.get(jwks_uri, 0) - return time.monotonic() - last < JWKS_MIN_REFRESH_INTERVAL_SECONDS + last = store.get(jwks_uri) + return ( + last is not None + and time.monotonic() - last < JWKS_MIN_REFRESH_INTERVAL_SECONDS + ) def _record_now(store: dict, jwks_uri: str) -> None: @@ -126,33 +129,38 @@ def _call_get_signing_keys(jwks_client: PyJWKClient, jwks_uri: str, refresh: boo return keys -def _fetch_signing_keys( - jwks_client: PyJWKClient, jwks_uri: str, refresh: bool, kid=None -): - if not refresh: +def _fetch_signing_keys(jwks_client: PyJWKClient, jwks_uri: str): + cached = _read_cached_keys(jwks_uri) + if cached is not None: + # Warm cache. Must not wait on another thread's refresh, or a + # forged unknown kid can reject valid tokens on HTTP and Flight. + return cached + + # Cold or expired snapshot: single-flight one outbound fetch so + # concurrent Flight threads cannot each pay JWKS_FETCH_TIMEOUT. + if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): + raise _fetch_failed_recently_error(jwks_uri) + + fetch_lock = _fetch_lock_for(jwks_uri) + fetch_lock.acquire() + try: cached = _read_cached_keys(jwks_uri) if cached is not None: - # Warm cache. Must not wait on another thread's refresh, or a - # forged unknown kid can reject valid tokens on HTTP and Flight. return cached + return _call_get_signing_keys(jwks_client, jwks_uri, refresh=False) + finally: + fetch_lock.release() - # Cold or expired snapshot: single-flight one outbound fetch so - # concurrent Flight threads cannot each pay JWKS_FETCH_TIMEOUT. - if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): - raise _fetch_failed_recently_error(jwks_uri) - - fetch_lock = _fetch_lock_for(jwks_uri) - fetch_lock.acquire() - try: - cached = _read_cached_keys(jwks_uri) - if cached is not None: - return cached - return _call_get_signing_keys(jwks_client, jwks_uri, refresh=False) - finally: - fetch_lock.release() +def _refresh_signing_key( + jwks_client: PyJWKClient, jwks_uri: str, kid +): if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): raise _fetch_failed_recently_error(jwks_uri) + if _recently_recorded(_jwks_last_failed_refresh, jwks_uri): + raise jwt.exceptions.PyJWKClientError( + f"Unable to find a signing key that matches: {kid!r}" + ) fetch_lock = _fetch_lock_for(jwks_uri) if not fetch_lock.acquire(blocking=False): @@ -160,16 +168,25 @@ def _fetch_signing_keys( f'JWKS refresh already in flight for "{jwks_uri}"' ) try: - # Recheck both cooldowns under the lock. The unknown-kid cooldown - # is set by a sibling thread after this caller already passed the - # unlocked check in _get_signing_key. + # Refresh ownership includes key matching and cooldown publication. + # No sibling can fetch between a mismatch and recording its cooldown. if _recently_recorded(_jwks_last_fetch_failure, jwks_uri): raise _fetch_failed_recently_error(jwks_uri) if _recently_recorded(_jwks_last_failed_refresh, jwks_uri): raise jwt.exceptions.PyJWKClientError( f"Unable to find a signing key that matches: {kid!r}" ) - return _call_get_signing_keys(jwks_client, jwks_uri, refresh=True) + + signing_keys = _call_get_signing_keys( + jwks_client, jwks_uri, refresh=True + ) + signing_key = PyJWKClient.match_kid(signing_keys, kid) + if signing_key is None: + _record_now(_jwks_last_failed_refresh, jwks_uri) + raise jwt.exceptions.PyJWKClientError( + f"Unable to find a signing key that matches: {kid!r}" + ) + return signing_key finally: fetch_lock.release() @@ -184,24 +201,12 @@ def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): header = jwt.get_unverified_header(token) kid = header.get("kid") - signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=False, kid=kid) + signing_keys = _fetch_signing_keys(jwks_client, jwks_uri) signing_key = PyJWKClient.match_kid(signing_keys, kid) if signing_key is not None: return signing_key - if _recently_recorded(_jwks_last_failed_refresh, jwks_uri): - raise jwt.exceptions.PyJWKClientError( - f"Unable to find a signing key that matches: {kid!r}" - ) - - signing_keys = _fetch_signing_keys(jwks_client, jwks_uri, refresh=True, kid=kid) - signing_key = PyJWKClient.match_kid(signing_keys, kid) - if signing_key is None: - _record_now(_jwks_last_failed_refresh, jwks_uri) - raise jwt.exceptions.PyJWKClientError( - f"Unable to find a signing key that matches: {kid!r}" - ) - return signing_key + return _refresh_signing_key(jwks_client, jwks_uri, kid) class JwtValidationError(Exception): diff --git a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py index 194634bf..d0f5c508 100644 --- a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py +++ b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py @@ -1,4 +1,5 @@ import base64 +import time import unittest from unittest.mock import patch @@ -26,9 +27,10 @@ def _basic_header(self, username, password): encoded = base64.b64encode(f"{username}:{password}".encode()).decode() return f"Basic {encoded}" - def _authenticate(self): - return self.factory.start_call( - None, self._headers(self._basic_header("user1", "P@ssw0rd")) + def _authenticate(self, username="user1", password="P@ssw0rd", factory=None): + target = factory or self.factory + return target.start_call( + None, self._headers(self._basic_header(username, password)) ) def test_valid_basic_is_accepted(self): @@ -85,15 +87,51 @@ def test_expired_tokens_are_evicted_on_the_next_mint(self): self._authenticate() self.assertNotIn(expired.token, self.factory.tokens) - def test_token_store_is_bounded(self): - with patch.object(mod, "MAX_FLIGHT_TOKENS", 4): - first = self._authenticate() - for _ in range(20): - latest = self._authenticate() + def test_repeated_basic_auth_reuses_one_unexpired_token_per_user(self): + first = self._authenticate() + first_expiry = self.factory.tokens[first.token][1] + for _ in range(20): + latest = self._authenticate() - self.assertLessEqual(len(self.factory.tokens), 4) - self.assertNotIn(first.token, self.factory.tokens) - self.assertTrue(self.factory.is_valid_token(latest.token)) + self.assertEqual(latest.token, first.token) + self.assertEqual(self.factory.tokens[first.token][1], first_expiry) + self.assertEqual(len(self.factory.tokens), 1) + self.assertTrue(self.factory.is_valid_token(first.token)) + + def test_basic_auth_rotates_a_token_near_expiry(self): + first = self._authenticate() + username, _ = self.factory.tokens[first.token] + self.factory.tokens[first.token] = ( + username, + time.monotonic() + mod.FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS - 1, + ) + + renewed = self._authenticate() + + self.assertNotEqual(renewed.token, first.token) + self.assertFalse(self.factory.is_valid_token(first.token)) + self.assertTrue(self.factory.is_valid_token(renewed.token)) + remaining = self.factory.tokens[renewed.token][1] - time.monotonic() + self.assertGreater(remaining, mod.FLIGHT_TOKEN_TTL_SECONDS - 1) + + def test_other_users_cannot_evict_an_unexpired_token(self): + creds = { + "user1": hash_password("user1", "P@ssw0rd"), + "user2": hash_password("user2", "OtherP@ssw0rd"), + } + factory = BasicAuthServerMiddlewareFactory(creds) + first = self._authenticate(factory=factory) + + for _ in range(20): + second = self._authenticate( + username="user2", + password="OtherP@ssw0rd", + factory=factory, + ) + + self.assertEqual(len(factory.tokens), 2) + self.assertTrue(factory.is_valid_token(first.token)) + self.assertTrue(factory.is_valid_token(second.token)) def test_invalid_base64_is_unauthenticated(self): with self.assertRaises(FlightUnauthenticatedError) as err: diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index 19bb8d98..91fd2ddf 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -332,6 +332,27 @@ def test_failed_jwks_fetch_is_rate_limited(self): validate_jwt(token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) self.assertEqual(mock_get_signing_keys.call_count, 1) + def test_empty_cooldown_store_is_not_recent_just_after_host_boot(self): + """ + time.monotonic() is host-uptime based. An absent entry must not use + zero as a timestamp or JWT auth fails during the first cooldown + interval after a host boot. + """ + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + store = {} + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.time.monotonic", + return_value=12, + ): + self.assertFalse( + jwt_auth_module._recently_recorded(store, JWKS_URI) + ) + jwt_auth_module._record_now(store, JWKS_URI) + self.assertTrue( + jwt_auth_module._recently_recorded(store, JWKS_URI) + ) + def test_hs256_substitution_using_public_key_is_rejected(self): """ Guards against the classic RS256->HS256 confusion attack: an @@ -561,79 +582,183 @@ def worker(): self.assertEqual(len(results), workers) self.assertEqual(call_count, [False]) - def test_unknown_kid_cooldown_is_rechecked_under_fetch_lock(self): + def test_unknown_kid_refresh_and_cooldown_publication_are_atomic(self): """ - Two unknown-kid callers can both pass the unlocked cooldown check. - The loser of the fetch lock must recheck _jwks_last_failed_refresh - and skip a second outbound refresh. + Pause the first caller after its refresh returns but before the + unknown-kid result is published. A second caller must not fetch in + that window, and the cooldown must be recorded before lock release. """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module - warm_token = self._make_token() - with self._patched_jwks_client(): + cached_token = self._make_token(headers={"kid": "cached-kid"}) + with self._patched_jwks_client(kid="cached-kid"): validate_jwt( - warm_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + cached_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE ) unknown = self._make_token(headers={"kid": "unknown-kid"}) - signing_key = self._signing_key(kid="the-real-kid") - a_passed_outer = threading.Event() - b_finished = threading.Event() - thread_a_holder = [] - original_recently = jwt_auth_module._recently_recorded - refresh_true_calls = [] - - def gated_recently(store, uri): - result = original_recently(store, uri) + refreshed_key = self._signing_key(kid="the-real-kid") + fetch_returned = threading.Event() + first_matching = threading.Event() + release_first = threading.Event() + second_finished = threading.Event() + refresh_calls = [] + outcomes = [] + errors = [] + record_lock_states = [] + fetch_lock = jwt_auth_module._fetch_lock_for(JWKS_URI) + original_match_kid = jwt_auth_module.PyJWKClient.match_kid + original_record_now = jwt_auth_module._record_now + + def fake_get_signing_keys(refresh=False): + refresh_calls.append(refresh) + fetch_returned.set() + return [refreshed_key] + + def gated_match_kid(signing_keys, kid): if ( - store is jwt_auth_module._jwks_last_failed_refresh - and thread_a_holder - and threading.current_thread() is thread_a_holder[0] - and not a_passed_outer.is_set() + threading.current_thread() is first_thread + and fetch_returned.is_set() + and not first_matching.is_set() ): - a_passed_outer.set() - self.assertFalse(result) - self.assertTrue(b_finished.wait(5)) - return result + first_matching.set() + self.assertTrue(release_first.wait(5)) + return original_match_kid(signing_keys, kid) - def fake_get_signing_keys(refresh=False): - if refresh: - refresh_true_calls.append(1) - return [signing_key] + def checked_record_now(store, uri): + if store is jwt_auth_module._jwks_last_failed_refresh: + record_lock_states.append(fetch_lock.locked()) + return original_record_now(store, uri) - def run_a(): - with self.assertRaises(JwtValidationError): + def validate_unknown(label): + try: validate_jwt( unknown, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE ) + except JwtValidationError: + outcomes.append(label) + except Exception as ex: + errors.append(ex) - def run_b(): - self.assertTrue(a_passed_outer.wait(5)) - with self.assertRaises(JwtValidationError): + first_thread = threading.Thread( + target=validate_unknown, args=("first",) + ) + + def run_second(): + validate_unknown("second") + second_finished.set() + + second_thread = threading.Thread(target=run_second) + + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", + side_effect=fake_get_signing_keys, + ), patch.object( + jwt_auth_module.PyJWKClient, + "match_kid", + side_effect=gated_match_kid, + ), patch( + "tabpy.tabpy_server.handlers.jwt_auth._record_now", + side_effect=checked_record_now, + ): + first_thread.start() + self.assertTrue(first_matching.wait(5)) + + second_thread.start() + self.assertTrue(second_finished.wait(5)) + self.assertEqual(refresh_calls, [True]) + + release_first.set() + first_thread.join(5) + second_thread.join(5) + + # Once the first caller releases the lock, a later unknown-kid + # request must observe the published cooldown without fetching. + validate_unknown("third") + + self.assertFalse(first_thread.is_alive()) + self.assertFalse(second_thread.is_alive()) + self.assertEqual(errors, []) + self.assertCountEqual(outcomes, ["first", "second", "third"]) + self.assertEqual(refresh_calls, [True]) + self.assertEqual(record_lock_states, [True]) + + def test_unknown_kid_cooldown_is_rechecked_after_lock_handoff(self): + """ + Pause one caller after its unlocked cooldown check but before lock + acquisition. Another caller publishes the cooldown and releases; + the paused caller must recheck under the lock without refetching. + """ + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + cached_token = self._make_token(headers={"kid": "cached-kid"}) + with self._patched_jwks_client(kid="cached-kid"): + validate_jwt( + cached_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + + unknown = self._make_token(headers={"kid": "unknown-kid"}) + refreshed_key = self._signing_key(kid="the-real-kid") + waiting_before_acquire = threading.Event() + release_waiting = threading.Event() + refresh_calls = [] + outcomes = [] + errors = [] + original_fetch_lock_for = jwt_auth_module._fetch_lock_for + + def fake_get_signing_keys(refresh=False): + refresh_calls.append(refresh) + return [refreshed_key] + + def gated_fetch_lock_for(uri): + lock = original_fetch_lock_for(uri) + if ( + threading.current_thread() is waiting_thread + and not waiting_before_acquire.is_set() + ): + waiting_before_acquire.set() + self.assertTrue(release_waiting.wait(5)) + return lock + + def validate_unknown(label): + try: validate_jwt( unknown, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE ) - b_finished.set() + except JwtValidationError: + outcomes.append(label) + except Exception as ex: + errors.append(ex) - thread_a = threading.Thread(target=run_a) - thread_a_holder.append(thread_a) - thread_b = threading.Thread(target=run_b) + waiting_thread = threading.Thread( + target=validate_unknown, args=("waiting",) + ) + owner_thread = threading.Thread( + target=validate_unknown, args=("owner",) + ) with patch( "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", side_effect=fake_get_signing_keys, ), patch( - "tabpy.tabpy_server.handlers.jwt_auth._recently_recorded", - side_effect=gated_recently, + "tabpy.tabpy_server.handlers.jwt_auth._fetch_lock_for", + side_effect=gated_fetch_lock_for, ): - thread_a.start() - thread_b.start() - thread_a.join(10) - thread_b.join(10) - - self.assertFalse(thread_a.is_alive()) - self.assertFalse(thread_b.is_alive()) - self.assertEqual(refresh_true_calls, [1]) + waiting_thread.start() + self.assertTrue(waiting_before_acquire.wait(5)) + + owner_thread.start() + owner_thread.join(5) + self.assertFalse(owner_thread.is_alive()) + self.assertEqual(refresh_calls, [True]) + + release_waiting.set() + waiting_thread.join(5) + + self.assertFalse(waiting_thread.is_alive()) + self.assertEqual(errors, []) + self.assertCountEqual(outcomes, ["owner", "waiting"]) + self.assertEqual(refresh_calls, [True]) def test_jwks_client_is_reused_for_same_uri(self): """ diff --git a/tests/unit/server_tests/test_oauth_handler.py b/tests/unit/server_tests/test_oauth_handler.py index 6bbc0460..19145efe 100644 --- a/tests/unit/server_tests/test_oauth_handler.py +++ b/tests/unit/server_tests/test_oauth_handler.py @@ -17,6 +17,7 @@ JWKS_URI, make_token, patched_jwks_client, + reset_jwks_state, ) # The fake "idp.example.com" JWKS host used by these tests doesn't @@ -29,6 +30,13 @@ class BaseTestOAuthHandler(AsyncHTTPTestCase): + def setUp(self): + # Each class generates a different signing key but intentionally + # reuses the same fake JWKS URI. Do not leak a cached key between + # otherwise-independent HTTP auth tests. + reset_jwks_state() + super().setUp() + def get_app(self): with patch( "tabpy.tabpy_server.app.app.socket.getaddrinfo", From 7a81ae9b651c1cdace12c4d8d5089ab97787ece1 Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Tue, 18 Aug 2026 17:04:00 -0700 Subject: [PATCH 07/10] Pin the JWKS cache lifetime Keep TabPy's parsed-key snapshot synchronized with PyJWT's cache and clarify the remaining cold-fetch availability tradeoff. --- docs/server-config.md | 6 ++++-- tabpy/tabpy_server/handlers/jwt_auth.py | 5 ++++- tests/unit/server_tests/test_jwt_auth.py | 4 ++++ 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/docs/server-config.md b/docs/server-config.md index 6067dc29..3379282c 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -335,7 +335,7 @@ The same Bearer token is accepted on the Arrow Flight (gRPC) path when Arrow is enabled. Failed Flight authentication is rejected with gRPC `UNAUTHENTICATED` rather than HTTP 401. Using Basic on Flight while OAuth is enabled also requires `TABPY_PWD_FILE`. When `TABPY_OAUTH_ENABLED` is false, -Flight continues to use basic-auth middleware unchanged. +Flight continues to use Basic-only middleware rather than JWT middleware. When both basic access authentication and OAuth are enabled, TabPy picks the method based on the scheme of the `Authorization` header sent by the client @@ -346,7 +346,9 @@ token is sent in cleartext. That matches HTTP Basic/Bearer on the same setting; TabPy does not require HTTPS for Flight auth alone. JWKS lookups are cached. On the HTTP path a cache-cold fetch runs on TabPy's -single IO-loop thread. Arrow Flight auth runs on the gRPC thread pool. +single IO-loop thread, so waiting on the IdP or an in-flight Flight fetch +briefly stalls other HTTP requests. Arrow Flight auth runs on the gRPC thread +pool. Cold and expired-cache fetches are single-flighted per JWKS URI. Forced unknown-`kid` refreshes are mutually excluded; a concurrent refresh attempt fails authentication rather than waiting on the network request. A cached diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 3695b1df..0a92c809 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -64,7 +64,10 @@ def _get_jwks_client(jwks_uri: str) -> PyJWKClient: client = _jwks_clients.get(jwks_uri) if client is None: client = PyJWKClient( - jwks_uri, cache_jwk_set=True, timeout=JWKS_FETCH_TIMEOUT_SECONDS + jwks_uri, + cache_jwk_set=True, + lifespan=JWKS_CACHE_LIFESPAN_SECONDS, + timeout=JWKS_FETCH_TIMEOUT_SECONDS, ) _jwks_clients[jwks_uri] = client return client diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index 91fd2ddf..9e6823b9 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -175,6 +175,10 @@ def test_jwks_client_is_created_with_bounded_timeout(self): reset_jwks_state() client = jwt_auth_module._get_jwks_client(JWKS_URI) self.assertEqual(client.timeout, jwt_auth_module.JWKS_FETCH_TIMEOUT_SECONDS) + self.assertEqual( + client.jwk_set_cache.lifespan, + jwt_auth_module.JWKS_CACHE_LIFESPAN_SECONDS, + ) self.assertLess(jwt_auth_module.JWKS_FETCH_TIMEOUT_SECONDS, 30) def test_algorithm_is_pinned_to_jwks_key_not_token_header(self): From 5a1eedf3270681a00322f04d006b0a1e1a997681 Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Wed, 19 Aug 2026 13:02:37 -0700 Subject: [PATCH 08/10] Harden Flight authentication concurrency and token lifecycle --- docs/server-config.md | 37 ++-- .../basic_auth_server_middleware_factory.py | 33 ++-- tabpy/tabpy_server/handlers/jwt_auth.py | 69 ++++++-- .../handlers/jwt_server_middleware_factory.py | 2 +- .../integration/test_arrow_server_jwt_auth.py | 103 +++++++++++ tests/unit/server_tests/jwt_test_helpers.py | 1 + ...st_basic_auth_server_middleware_factory.py | 26 ++- tests/unit/server_tests/test_jwt_auth.py | 163 +++++++++++------- 8 files changed, 335 insertions(+), 99 deletions(-) create mode 100644 tests/integration/test_arrow_server_jwt_auth.py diff --git a/docs/server-config.md b/docs/server-config.md index 3379282c..f32f6aae 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -346,21 +346,23 @@ token is sent in cleartext. That matches HTTP Basic/Bearer on the same setting; TabPy does not require HTTPS for Flight auth alone. JWKS lookups are cached. On the HTTP path a cache-cold fetch runs on TabPy's -single IO-loop thread, so waiting on the IdP or an in-flight Flight fetch -briefly stalls other HTTP requests. Arrow Flight auth runs on the gRPC thread -pool. -Cold and expired-cache fetches are single-flighted per JWKS URI. Forced -unknown-`kid` refreshes are mutually excluded; a concurrent refresh attempt -fails authentication rather than waiting on the network request. A cached +single IO-loop thread, so waiting on the IdP briefly stalls other HTTP +requests. Arrow Flight auth runs on the gRPC thread pool. Cold and +expired-cache fetches are single-flighted per JWKS URI. A caller waits at most +one second for another in-flight fetch before authentication fails. + +Forced unknown-`kid` refreshes are mutually excluded. Concurrent requests for +the same `kid` wait up to one second for the refresh result; requests for a +different unknown `kid` fail without starting another fetch. A cached signing-key lookup does not wait on that refresh, so unrelated valid tokens continue to work. After a refresh still cannot find a requested `kid`, TabPy rate-limits another forced refresh for 30 seconds. This also means a legitimate new key may be rejected for up to 30 seconds after a bogus unknown-`kid` request. -A request carrying two conflicting `Authorization` values is rejected, -because which one wins would otherwise decide the caller's identity. -Repeating the same value is accepted. +On Flight, a request carrying two conflicting `Authorization` values is +rejected, because which one wins would otherwise decide the caller's identity. +Flight accepts repeated copies of the same value. ### Endpoint Security @@ -380,12 +382,17 @@ After a successful Basic call, Flight also returns an opaque session token in the response `authorization` header. A client may send that token back as a Bearer credential instead of repeating its Basic credentials. The token is server-issued, is not a JWT, and expires an hour after it is issued or when -the server restarts. TabPy retains one active opaque token per username, so -repeated Basic calls do not grow the token store or invalidate another user's -unexpired token. A repeated Basic call normally returns the same token without -extending its original expiry. During its final five minutes, Basic -authentication rotates it to a fresh one-hour token. After expiry, the client -authenticates with Basic again. +the server restarts. TabPy normally retains one active opaque token per +username. A repeated Basic call returns the same token without extending its +original expiry. During its final five minutes, Basic authentication rotates +it to a fresh one-hour token while the prior token remains valid until its +original expiry. This overlap is bounded to two tokens per username. + +Opaque tokens are not rechecked against the password file after issuance. +Changing a password or removing a user therefore does not revoke an existing +token before its one-hour expiry. Restart TabPy to revoke all issued opaque +tokens immediately, and use `grpc+tls` so Basic and Bearer credentials are +encrypted in transit. After expiry, the client authenticates with Basic again. **As of May 2023, the Arrow Flight feature can only be used by compatible versions of Tableau Prep. The Arrow Flight feature is not used by Tableau diff --git a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py index 624a5d44..d10a43c7 100644 --- a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py @@ -20,9 +20,10 @@ FLIGHT_TOKEN_TTL_SECONDS = 3600 # Avoid handing a freshly authenticated client a token that is about to -# expire. Rotation still happens under the per-factory lock and retains -# only one active token for the username. +# expire. Rotation keeps the prior token valid until its original expiry, +# so another client using the same username is not disconnected. FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS = 300 +MAX_ACTIVE_FLIGHT_TOKENS_PER_USER = 2 class BasicAuthServerMiddleware(ServerMiddleware): @@ -37,9 +38,8 @@ class BasicAuthServerMiddlewareFactory(ServerMiddlewareFactory): def __init__(self, creds): self.creds = creds # token -> (username, monotonic expiry). Read from the gRPC thread - # pool on every call. One active token is retained per username, - # bounding request-driven growth without evicting another user's - # unexpired credential. + # pool on every call. Normally one token is retained per username; + # rotation briefly permits the old and new tokens to overlap. self.tokens = {} self._tokens_by_username = {} self._tokens_lock = threading.Lock() @@ -66,21 +66,33 @@ def _issue_token(self, username): now = time.monotonic() self._evict_expired_tokens(now) - existing = self._tokens_by_username.get(username) - if existing is not None: + active_tokens = self._tokens_by_username.get(username, set()) + if active_tokens: + existing = max( + active_tokens, key=lambda active: self.tokens[active][1] + ) _, expiry = self.tokens[existing] if expiry - now > FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS: return existing - self._remove_token(existing, username) token = secrets.token_urlsafe(32) self.tokens[token] = (username, now + FLIGHT_TOKEN_TTL_SECONDS) - self._tokens_by_username[username] = token + active_tokens = self._tokens_by_username.setdefault(username, set()) + active_tokens.add(token) + while len(active_tokens) > MAX_ACTIVE_FLIGHT_TOKENS_PER_USER: + oldest = min( + active_tokens, key=lambda active: self.tokens[active][1] + ) + self._remove_token(oldest, username) return token def _remove_token(self, token, username): self.tokens.pop(token, None) - if self._tokens_by_username.get(username) == token: + active_tokens = self._tokens_by_username.get(username) + if active_tokens is None: + return + active_tokens.discard(token) + if not active_tokens: self._tokens_by_username.pop(username, None) def _evict_expired_tokens(self, now): @@ -108,6 +120,7 @@ def start_call(self, info, headers): username, separator, password = decoded.partition(":") if not separator or not username: raise FlightUnauthenticatedError("Invalid credentials") + username = username.lower() if not self.is_valid_user(username, password): raise FlightUnauthenticatedError("Invalid credentials") return BasicAuthServerMiddleware(self._issue_token(username)) diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 0a92c809..4f9ec7f5 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -1,12 +1,9 @@ -import logging import threading import time import jwt from jwt import PyJWKClient -logger = logging.getLogger(__name__) - # PyJWKClient fetches JWKS synchronously (via requests). On the HTTP path # that call runs on TabPy's single Tornado IO-loop thread, so a # slow/unresponsive IdP stalls every concurrent HTTP request for up to this @@ -28,6 +25,12 @@ # next lookup treats the snapshot as expired and single-flights one fetch. JWKS_CACHE_LIFESPAN_SECONDS = 300 +# Concurrent requests for the same newly published kid may briefly wait for +# the refresh owner instead of failing spuriously. Keep the wait short and +# bounded so this synchronous helper cannot tie up an HTTP or Flight worker +# for the full IdP fetch timeout. +JWKS_REFRESH_WAIT_SECONDS = 1 + # One PyJWKClient per JWKS URI, reused so its JWK Set cache actually avoids # per-request fetches. Process-global. _jwks_state_lock only covers these # dicts. A per-URI fetch lock serializes cache-miss fetches and forced @@ -37,6 +40,12 @@ _jwks_fetch_locks = {} _jwks_clients = {} +# jwks_uri -> (kid, completion event). Reserving refresh ownership under +# _jwks_state_lock closes the race between taking the fetch lock and +# publishing which kid is being refreshed. Same-kid callers may wait for +# the owner; other unknown kids fail fast and cannot amplify JWKS traffic. +_jwks_in_flight_refreshes = {} + # jwks_uri -> (signing_keys, monotonic timestamp). Lets a warm cache # return keys without waiting on another thread's in-flight fetch. _jwks_cached_keys = {} @@ -145,7 +154,10 @@ def _fetch_signing_keys(jwks_client: PyJWKClient, jwks_uri: str): raise _fetch_failed_recently_error(jwks_uri) fetch_lock = _fetch_lock_for(jwks_uri) - fetch_lock.acquire() + if not fetch_lock.acquire(timeout=JWKS_REFRESH_WAIT_SECONDS): + raise jwt.exceptions.PyJWKClientError( + f'JWKS fetch already in flight for "{jwks_uri}"' + ) try: cached = _read_cached_keys(jwks_uri) if cached is not None: @@ -165,8 +177,45 @@ def _refresh_signing_key( f"Unable to find a signing key that matches: {kid!r}" ) + with _jwks_state_lock: + in_flight = _jwks_in_flight_refreshes.get(jwks_uri) + if in_flight is None: + refresh_done = threading.Event() + _jwks_in_flight_refreshes[jwks_uri] = (kid, refresh_done) + owns_refresh = True + else: + in_flight_kid, refresh_done = in_flight + owns_refresh = False + + if not owns_refresh: + if in_flight_kid != kid: + raise jwt.exceptions.PyJWKClientError( + f'JWKS refresh already in flight for "{jwks_uri}"' + ) + if not refresh_done.wait(JWKS_REFRESH_WAIT_SECONDS): + raise jwt.exceptions.PyJWKClientError( + f'JWKS refresh timed out for "{jwks_uri}"' + ) + signing_keys = _read_cached_keys(jwks_uri) + signing_key = ( + PyJWKClient.match_kid(signing_keys, kid) + if signing_keys is not None + else None + ) + if signing_key is not None: + return signing_key + raise jwt.exceptions.PyJWKClientError( + f"Unable to find a signing key that matches: {kid!r}" + ) + fetch_lock = _fetch_lock_for(jwks_uri) - if not fetch_lock.acquire(blocking=False): + acquired = fetch_lock.acquire(timeout=JWKS_REFRESH_WAIT_SECONDS) + if not acquired: + with _jwks_state_lock: + current = _jwks_in_flight_refreshes.get(jwks_uri) + if current is not None and current[1] is refresh_done: + _jwks_in_flight_refreshes.pop(jwks_uri) + refresh_done.set() raise jwt.exceptions.PyJWKClientError( f'JWKS refresh already in flight for "{jwks_uri}"' ) @@ -192,6 +241,11 @@ def _refresh_signing_key( return signing_key finally: fetch_lock.release() + with _jwks_state_lock: + current = _jwks_in_flight_refreshes.get(jwks_uri) + if current is not None and current[1] is refresh_done: + _jwks_in_flight_refreshes.pop(jwks_uri) + refresh_done.set() def _get_signing_key(jwks_client: PyJWKClient, jwks_uri: str, token: str): @@ -259,11 +313,9 @@ def validate_jwt( jwks_client = _get_jwks_client(jwks_uri) signing_key = _get_signing_key(jwks_client, jwks_uri, token) except (jwt.exceptions.PyJWKClientError, jwt.exceptions.InvalidTokenError) as ex: - logger.log(logging.ERROR, f"Unable to resolve JWT signing key: {str(ex)}") raise JwtValidationError("Unable to resolve JWT signing key") from ex except Exception as ex: # Must still surface as a 401, not a 500 (e.g. malformed JWKS response). - logger.log(logging.ERROR, f"Unexpected error resolving JWT signing key: {str(ex)}") raise JwtValidationError("Unable to resolve JWT signing key") from ex try: @@ -276,12 +328,10 @@ def validate_jwt( options={"require": ["exp", "iat"]}, ) except jwt.exceptions.InvalidTokenError as ex: - logger.log(logging.ERROR, f"JWT validation failed: {str(ex)}") raise JwtValidationError(f"JWT validation failed: {str(ex)}") from ex except Exception as ex: # Must still fail closed as a 401 / UNAUTHENTICATED, not a 500 # or an ArrowInvalid traceback to an unauthenticated caller. - logger.log(logging.ERROR, f"Unexpected error decoding JWT: {str(ex)}") raise JwtValidationError("JWT validation failed") from ex if required_scopes: @@ -292,7 +342,6 @@ def validate_jwt( except Exception as ex: # Must still fail closed as a 401 (e.g. a `scope` claim that # isn't a space-separated string), not an uncaught 500. - logger.log(logging.ERROR, f"Unable to evaluate JWT scopes: {str(ex)}") raise JwtValidationError("Unable to evaluate JWT scopes") from ex return claims diff --git a/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py index 9a3a3953..220b15c1 100644 --- a/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py @@ -65,7 +65,7 @@ def start_call(self, info, headers): ) except JwtValidationError as ex: logger.log( - logging.ERROR, f"Flight JWT authentication failed: {ex}" + logging.WARNING, f"Flight JWT authentication failed: {ex}" ) raise FlightUnauthenticatedError("Invalid credentials") from ex except Exception as ex: diff --git a/tests/integration/test_arrow_server_jwt_auth.py b/tests/integration/test_arrow_server_jwt_auth.py new file mode 100644 index 00000000..5c09ebaf --- /dev/null +++ b/tests/integration/test_arrow_server_jwt_auth.py @@ -0,0 +1,103 @@ +import contextlib +import threading +import unittest + +from cryptography.hazmat.primitives.asymmetric import rsa +import pyarrow.flight + +from tabpy.tabpy_server.app.arrow_server import FlightServer +from tabpy.tabpy_server.handlers.jwt_server_middleware_factory import ( + JwtAuthServerMiddlewareFactory, +) +from tabpy.tabpy_server.handlers.no_op_auth_handler import NoOpAuthHandler +from tests.unit.server_tests.jwt_test_helpers import ( + AUDIENCE, + ISSUER, + JWKS_URI, + make_token, + patched_jwks_client, + reset_jwks_state, +) + + +class TestArrowServerJwtAuth(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.private_key = rsa.generate_private_key( + public_exponent=65537, key_size=2048 + ) + + def setUp(self): + reset_jwks_state() + + @contextlib.contextmanager + def _server_and_client(self, required_scopes=None): + middleware = JwtAuthServerMiddlewareFactory( + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + required_scopes=required_scopes, + ) + server = FlightServer( + host="localhost", + location="grpc+tcp://localhost:0", + auth_handler=NoOpAuthHandler(), + middleware={"jwt": middleware}, + ) + server_thread = threading.Thread(target=server.serve, daemon=True) + server_thread.start() + client = pyarrow.flight.FlightClient( + f"grpc+tcp://localhost:{server.port}" + ) + try: + yield client + finally: + server.shutdown() + server_thread.join(5) + + def _options(self, token): + return pyarrow.flight.FlightCallOptions( + headers=[(b"authorization", f"Bearer {token}".encode())] + ) + + def test_valid_jwt_authenticates_through_flight_transport(self): + token = make_token(self.private_key) + + with self._server_and_client() as client, patched_jwks_client( + self.private_key + ): + actions = list(client.list_actions(options=self._options(token))) + + self.assertTrue(actions) + + def test_invalid_jwt_is_rejected_by_flight_transport(self): + invalid_token = make_token( + self.private_key, claims_override={"aud": "wrong-audience"} + ) + + with self._server_and_client() as client, patched_jwks_client( + self.private_key + ): + with self.assertRaises(pyarrow.flight.FlightUnauthenticatedError): + list(client.list_actions(options=self._options(invalid_token))) + + def test_required_scope_is_enforced_through_flight_transport(self): + allowed = make_token( + self.private_key, claims_override={"scope": "read execute"} + ) + denied = make_token( + self.private_key, claims_override={"scope": "read"} + ) + + with self._server_and_client( + required_scopes="execute" + ) as client, patched_jwks_client(self.private_key): + self.assertTrue( + list(client.list_actions(options=self._options(allowed))) + ) + with self.assertRaises(pyarrow.flight.FlightUnauthenticatedError): + list(client.list_actions(options=self._options(denied))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/server_tests/jwt_test_helpers.py b/tests/unit/server_tests/jwt_test_helpers.py index acdf5c87..18e4f4f4 100644 --- a/tests/unit/server_tests/jwt_test_helpers.py +++ b/tests/unit/server_tests/jwt_test_helpers.py @@ -20,6 +20,7 @@ def reset_jwks_state(): jwt_auth_module._jwks_clients.clear() jwt_auth_module._jwks_fetch_locks.clear() + jwt_auth_module._jwks_in_flight_refreshes.clear() jwt_auth_module._jwks_cached_keys.clear() jwt_auth_module._jwks_last_failed_refresh.clear() jwt_auth_module._jwks_last_fetch_failure.clear() diff --git a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py index d0f5c508..eebf29b3 100644 --- a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py +++ b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py @@ -46,6 +46,11 @@ def test_issued_token_records_the_username(self): username, _ = self.factory.tokens[middleware.token] self.assertEqual(username, "user1") + def test_basic_username_is_case_insensitive(self): + middleware = self._authenticate(username="UsEr1") + username, _ = self.factory.tokens[middleware.token] + self.assertEqual(username, "user1") + def test_invalid_password_is_rejected(self): with self.assertRaises(FlightUnauthenticatedError): self.factory.start_call( @@ -109,11 +114,30 @@ def test_basic_auth_rotates_a_token_near_expiry(self): renewed = self._authenticate() self.assertNotEqual(renewed.token, first.token) - self.assertFalse(self.factory.is_valid_token(first.token)) + self.assertTrue(self.factory.is_valid_token(first.token)) self.assertTrue(self.factory.is_valid_token(renewed.token)) + self.assertEqual(len(self.factory.tokens), 2) remaining = self.factory.tokens[renewed.token][1] - time.monotonic() self.assertGreater(remaining, mod.FLIGHT_TOKEN_TTL_SECONDS - 1) + def test_token_overlap_is_bounded_per_username(self): + latest = self._authenticate() + issued_tokens = [latest.token] + + for _ in range(3): + username, _ = self.factory.tokens[latest.token] + self.factory.tokens[latest.token] = ( + username, + time.monotonic() + mod.FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS - 1, + ) + latest = self._authenticate() + issued_tokens.append(latest.token) + + self.assertEqual(len(self.factory.tokens), 2) + self.assertNotIn(issued_tokens[0], self.factory.tokens) + self.assertTrue(self.factory.is_valid_token(issued_tokens[-2])) + self.assertTrue(self.factory.is_valid_token(issued_tokens[-1])) + def test_other_users_cannot_evict_an_unexpired_token(self): creds = { "user1": hash_password("user1", "P@ssw0rd"), diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index 9e6823b9..3578b252 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -314,6 +314,37 @@ def test_unknown_kid_refresh_cooldown_blocks_a_different_kid(self): # to the cooldown armed by the bogus kid moments earlier. self.assertEqual(mock_get_signing_keys.call_count, 1) + def test_unknown_kid_refresh_recovers_after_cooldown_expires(self): + import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module + + cached_token = self._make_token(headers={"kid": "cached-kid"}) + with self._patched_jwks_client(kid="cached-kid"): + validate_jwt( + cached_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + + jwt_auth_module._jwks_last_failed_refresh[JWKS_URI] = ( + time.monotonic() + - jwt_auth_module.JWKS_MIN_REFRESH_INTERVAL_SECONDS + - 1 + ) + rotated_token = self._make_token(headers={"kid": "rotated-kid"}) + rotated_signing_key = self._signing_key(kid="rotated-kid") + + with patch( + "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", + return_value=[rotated_signing_key], + ) as mock_get_signing_keys: + claims = validate_jwt( + rotated_token, + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + ) + + self.assertEqual(claims["sub"], "user1") + mock_get_signing_keys.assert_called_once_with(refresh=True) + def test_failed_jwks_fetch_is_rate_limited(self): """ A down/unreachable IdP must not be hammered with a fresh blocking @@ -434,10 +465,10 @@ def hold_refresh(): self.assertEqual(claims["sub"], "user1") self.assertLess(elapsed, 1) - def test_in_flight_refresh_does_not_block_another_refresh_attempt(self): + def test_in_flight_fetch_only_blocks_refresh_for_bounded_time(self): """ - A second forced refresh for the same jwks_uri must fail immediately - rather than wait for the in-flight fetch timeout. + A forced refresh may wait briefly for the per-URI fetch owner, but + must not wait for the full outbound JWKS timeout. """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module @@ -462,10 +493,16 @@ def hold_refresh(): try: self.assertTrue(acquired.wait(5)) started = time.monotonic() - with self.assertRaises(JwtValidationError): - validate_jwt( - token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE - ) + with patch.object( + jwt_auth_module, "JWKS_REFRESH_WAIT_SECONDS", 0.05 + ): + with self.assertRaises(JwtValidationError): + validate_jwt( + token, + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, + ) elapsed = time.monotonic() - started finally: release.set() @@ -669,10 +706,11 @@ def run_second(): self.assertTrue(first_matching.wait(5)) second_thread.start() - self.assertTrue(second_finished.wait(5)) + self.assertFalse(second_finished.wait(0.05)) self.assertEqual(refresh_calls, [True]) release_first.set() + self.assertTrue(second_finished.wait(5)) first_thread.join(5) second_thread.join(5) @@ -687,11 +725,10 @@ def run_second(): self.assertEqual(refresh_calls, [True]) self.assertEqual(record_lock_states, [True]) - def test_unknown_kid_cooldown_is_rechecked_after_lock_handoff(self): + def test_concurrent_requests_for_new_kid_share_one_refresh(self): """ - Pause one caller after its unlocked cooldown check but before lock - acquisition. Another caller publishes the cooldown and releases; - the paused caller must recheck under the lock without refetching. + Requests for the same newly rotated kid wait briefly for one + refresh owner, then all validate against the published snapshot. """ import tabpy.tabpy_server.handlers.jwt_auth as jwt_auth_module @@ -701,67 +738,71 @@ def test_unknown_kid_cooldown_is_rechecked_after_lock_handoff(self): cached_token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE ) - unknown = self._make_token(headers={"kid": "unknown-kid"}) - refreshed_key = self._signing_key(kid="the-real-kid") - waiting_before_acquire = threading.Event() - release_waiting = threading.Event() + rotated_private_key = rsa.generate_private_key( + public_exponent=65537, key_size=2048 + ) + rotated_token = make_token( + rotated_private_key, headers={"kid": "rotated-kid"} + ) + refreshed_key = type("SigningKey", (), {})() + refreshed_key.key = rotated_private_key.public_key() + refreshed_key.algorithm_name = "RS256" + refreshed_key.key_id = "rotated-kid" + + workers = 8 + barrier = threading.Barrier(workers) + all_refresh_callers_entered = threading.Event() + refresh_callers = 0 + refresh_callers_lock = threading.Lock() refresh_calls = [] - outcomes = [] + results = [] errors = [] - original_fetch_lock_for = jwt_auth_module._fetch_lock_for + original_refresh = jwt_auth_module._refresh_signing_key def fake_get_signing_keys(refresh=False): refresh_calls.append(refresh) + self.assertTrue(all_refresh_callers_entered.wait(5)) return [refreshed_key] - def gated_fetch_lock_for(uri): - lock = original_fetch_lock_for(uri) - if ( - threading.current_thread() is waiting_thread - and not waiting_before_acquire.is_set() - ): - waiting_before_acquire.set() - self.assertTrue(release_waiting.wait(5)) - return lock + def counted_refresh(*args): + nonlocal refresh_callers + with refresh_callers_lock: + refresh_callers += 1 + if refresh_callers == workers: + all_refresh_callers_entered.set() + return original_refresh(*args) - def validate_unknown(label): + def worker(): try: - validate_jwt( - unknown, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + barrier.wait(5) + claims = validate_jwt( + rotated_token, + issuer=ISSUER, + jwks_uri=JWKS_URI, + audience=AUDIENCE, ) - except JwtValidationError: - outcomes.append(label) + results.append(claims) except Exception as ex: errors.append(ex) - waiting_thread = threading.Thread( - target=validate_unknown, args=("waiting",) - ) - owner_thread = threading.Thread( - target=validate_unknown, args=("owner",) - ) - with patch( "tabpy.tabpy_server.handlers.jwt_auth.PyJWKClient.get_signing_keys", side_effect=fake_get_signing_keys, - ), patch( - "tabpy.tabpy_server.handlers.jwt_auth._fetch_lock_for", - side_effect=gated_fetch_lock_for, + ), patch.object( + jwt_auth_module, + "_refresh_signing_key", + side_effect=counted_refresh, ): - waiting_thread.start() - self.assertTrue(waiting_before_acquire.wait(5)) - - owner_thread.start() - owner_thread.join(5) - self.assertFalse(owner_thread.is_alive()) - self.assertEqual(refresh_calls, [True]) - - release_waiting.set() - waiting_thread.join(5) + threads = [threading.Thread(target=worker) for _ in range(workers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(5) - self.assertFalse(waiting_thread.is_alive()) + self.assertTrue(all(not thread.is_alive() for thread in threads)) self.assertEqual(errors, []) - self.assertCountEqual(outcomes, ["owner", "waiting"]) + self.assertEqual(len(results), workers) + self.assertTrue(all(claims["sub"] == "user1" for claims in results)) self.assertEqual(refresh_calls, [True]) def test_jwks_client_is_reused_for_same_uri(self): @@ -777,16 +818,14 @@ def test_jwks_client_is_reused_for_same_uri(self): second = jwt_auth_module._get_jwks_client(JWKS_URI) self.assertIs(first, second) - def test_validation_failure_does_not_log_raw_token(self): + def test_validation_failure_does_not_expose_raw_token(self): token = self._make_token({"iss": "https://wrong-idp.example.com/"}) with self._patched_jwks_client(): - with self.assertLogs( - "tabpy.tabpy_server.handlers.jwt_auth", level="ERROR" - ) as log_ctx: - with self.assertRaises(JwtValidationError): - validate_jwt(token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE) - logged_text = " ".join(log_ctx.output) - self.assertNotIn(token, logged_text) + with self.assertRaises(JwtValidationError) as error: + validate_jwt( + token, issuer=ISSUER, jwks_uri=JWKS_URI, audience=AUDIENCE + ) + self.assertNotIn(token, str(error.exception)) if __name__ == "__main__": From 36505a995990fe59717caf6b247e153d8b84dabb Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Wed, 19 Aug 2026 13:36:17 -0700 Subject: [PATCH 09/10] Preserve JWKS single-flight availability --- docs/server-config.md | 7 +++---- tabpy/tabpy_server/handlers/jwt_auth.py | 9 ++++----- tests/unit/server_tests/test_jwt_auth.py | 4 ++++ 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/server-config.md b/docs/server-config.md index f32f6aae..27e870c9 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -346,10 +346,9 @@ token is sent in cleartext. That matches HTTP Basic/Bearer on the same setting; TabPy does not require HTTPS for Flight auth alone. JWKS lookups are cached. On the HTTP path a cache-cold fetch runs on TabPy's -single IO-loop thread, so waiting on the IdP briefly stalls other HTTP -requests. Arrow Flight auth runs on the gRPC thread pool. Cold and -expired-cache fetches are single-flighted per JWKS URI. A caller waits at most -one second for another in-flight fetch before authentication fails. +single IO-loop thread, so waiting on the IdP or an in-flight Flight fetch +briefly stalls other HTTP requests. Arrow Flight auth runs on the gRPC thread +pool. Cold and expired-cache fetches are single-flighted per JWKS URI. Forced unknown-`kid` refreshes are mutually excluded. Concurrent requests for the same `kid` wait up to one second for the refresh result; requests for a diff --git a/tabpy/tabpy_server/handlers/jwt_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 4f9ec7f5..7fea7023 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -154,10 +154,7 @@ def _fetch_signing_keys(jwks_client: PyJWKClient, jwks_uri: str): raise _fetch_failed_recently_error(jwks_uri) fetch_lock = _fetch_lock_for(jwks_uri) - if not fetch_lock.acquire(timeout=JWKS_REFRESH_WAIT_SECONDS): - raise jwt.exceptions.PyJWKClientError( - f'JWKS fetch already in flight for "{jwks_uri}"' - ) + fetch_lock.acquire() try: cached = _read_cached_keys(jwks_uri) if cached is not None: @@ -313,7 +310,9 @@ def validate_jwt( jwks_client = _get_jwks_client(jwks_uri) signing_key = _get_signing_key(jwks_client, jwks_uri, token) except (jwt.exceptions.PyJWKClientError, jwt.exceptions.InvalidTokenError) as ex: - raise JwtValidationError("Unable to resolve JWT signing key") from ex + raise JwtValidationError( + f"Unable to resolve JWT signing key: {str(ex)}" + ) from ex except Exception as ex: # Must still surface as a 401, not a 500 (e.g. malformed JWKS response). raise JwtValidationError("Unable to resolve JWT signing key") from ex diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index 3578b252..002662bf 100644 --- a/tests/unit/server_tests/test_jwt_auth.py +++ b/tests/unit/server_tests/test_jwt_auth.py @@ -792,6 +792,10 @@ def worker(): jwt_auth_module, "_refresh_signing_key", side_effect=counted_refresh, + ), patch.object( + jwt_auth_module, + "JWKS_REFRESH_WAIT_SECONDS", + 5, ): threads = [threading.Thread(target=worker) for _ in range(workers)] for thread in threads: From c931871c0b0035a816609a2e6ae3ddb4eb1de80a Mon Sep 17 00:00:00 2001 From: jakeichikawasalesforce Date: Wed, 19 Aug 2026 14:29:34 -0700 Subject: [PATCH 10/10] Make Flight token overlap eviction deterministic --- .../basic_auth_server_middleware_factory.py | 19 +++++++------------ ...st_basic_auth_server_middleware_factory.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py index d10a43c7..0b7573a2 100644 --- a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py @@ -66,24 +66,19 @@ def _issue_token(self, username): now = time.monotonic() self._evict_expired_tokens(now) - active_tokens = self._tokens_by_username.get(username, set()) + active_tokens = self._tokens_by_username.get(username, []) if active_tokens: - existing = max( - active_tokens, key=lambda active: self.tokens[active][1] - ) + existing = active_tokens[-1] _, expiry = self.tokens[existing] if expiry - now > FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS: return existing token = secrets.token_urlsafe(32) self.tokens[token] = (username, now + FLIGHT_TOKEN_TTL_SECONDS) - active_tokens = self._tokens_by_username.setdefault(username, set()) - active_tokens.add(token) - while len(active_tokens) > MAX_ACTIVE_FLIGHT_TOKENS_PER_USER: - oldest = min( - active_tokens, key=lambda active: self.tokens[active][1] - ) - self._remove_token(oldest, username) + active_tokens = self._tokens_by_username.setdefault(username, []) + active_tokens.append(token) + if len(active_tokens) > MAX_ACTIVE_FLIGHT_TOKENS_PER_USER: + self._remove_token(active_tokens[0], username) return token def _remove_token(self, token, username): @@ -91,7 +86,7 @@ def _remove_token(self, token, username): active_tokens = self._tokens_by_username.get(username) if active_tokens is None: return - active_tokens.discard(token) + active_tokens.remove(token) if not active_tokens: self._tokens_by_username.pop(username, None) diff --git a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py index eebf29b3..577fc0c7 100644 --- a/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py +++ b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py @@ -138,6 +138,23 @@ def test_token_overlap_is_bounded_per_username(self): self.assertTrue(self.factory.is_valid_token(issued_tokens[-2])) self.assertTrue(self.factory.is_valid_token(issued_tokens[-1])) + def test_tied_expiries_evict_oldest_issued_token(self): + first = self._authenticate() + username, _ = self.factory.tokens[first.token] + near_expiry = ( + time.monotonic() + mod.FLIGHT_TOKEN_RENEWAL_WINDOW_SECONDS - 1 + ) + self.factory.tokens[first.token] = (username, near_expiry) + second = self._authenticate() + self.factory.tokens[first.token] = (username, near_expiry) + self.factory.tokens[second.token] = (username, near_expiry) + + third = self._authenticate() + + self.assertNotIn(first.token, self.factory.tokens) + self.assertTrue(self.factory.is_valid_token(second.token)) + self.assertTrue(self.factory.is_valid_token(third.token)) + def test_other_users_cannot_evict_an_unexpired_token(self): creds = { "user1": hash_password("user1", "P@ssw0rd"),