diff --git a/docs/server-config.md b/docs/server-config.md index b223baa3..27e870c9 100755 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -331,14 +331,37 @@ 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-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 (`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. 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. 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. + +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 @@ -350,6 +373,26 @@ 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 or when +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 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..c7dbd084 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 @@ -119,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" @@ -129,11 +137,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 +626,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..0b7573a2 100644 --- a/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py +++ b/tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py @@ -1,11 +1,31 @@ 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.flight_headers import ( + get_flight_authorization_header, +) 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 +# 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 + +# Avoid handing a freshly authenticated client a token that is about to +# 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): def __init__(self, token): self.token = token @@ -13,10 +33,16 @@ 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. 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() def is_valid_user(self, username, password): if username not in self.creds: @@ -24,25 +50,77 @@ 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): + 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): + with self._tokens_lock: + now = time.monotonic() + self._evict_expired_tokens(now) + + active_tokens = self._tokens_by_username.get(username, []) + if active_tokens: + 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, []) + 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): + self.tokens.pop(token, None) + active_tokens = self._tokens_by_username.get(username) + if active_tokens is None: + return + active_tokens.remove(token) + if not active_tokens: + 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]: + username, _ = self.tokens[token] + self._remove_token(token, username) + 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") + username = username.lower() 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/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_auth.py b/tabpy/tabpy_server/handlers/jwt_auth.py index 8569a0d5..7fea7023 100644 --- a/tabpy/tabpy_server/handlers/jwt_auth.py +++ b/tabpy/tabpy_server/handlers/jwt_auth.py @@ -1,39 +1,55 @@ -import logging +import threading import time import jwt from jwt import PyJWKClient -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 # 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 +# 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 + +# 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 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. _jwks_state_lock only covers these +# 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 -> (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 = {} + # 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 @@ -53,27 +69,180 @@ 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 + with _jwks_state_lock: + client = _jwks_clients.get(jwks_uri) + if client is None: + client = PyJWKClient( + jwks_uri, + cache_jwk_set=True, + lifespan=JWKS_CACHE_LIFESPAN_SECONDS, + 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) + return ( + last is not None + and 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 _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): + 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: + 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}" ) - _jwks_clients[jwks_uri] = client - return client + 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 -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 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'JWKS endpoint "{jwks_uri}" failed recently; not retrying yet' + f"Unable to find a signing key that matches: {kid!r}" + ) + + fetch_lock = _fetch_lock_for(jwks_uri) + 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}"' ) try: - return jwks_client.get_signing_keys(refresh=refresh) - except Exception: - _jwks_last_fetch_failure[jwks_uri] = now - raise + # 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}" + ) + + 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() + 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): @@ -86,26 +255,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) + 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 - now = time.monotonic() - last_failure = _jwks_last_failed_refresh.get(jwks_uri, 0) - if now - last_failure < JWKS_MIN_REFRESH_INTERVAL_SECONDS: - 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) - signing_key = PyJWKClient.match_kid(signing_keys, kid) - if signing_key is None: - _jwks_last_failed_refresh[jwks_uri] = now - 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): @@ -155,11 +310,11 @@ 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 + 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). - logger.log(logging.ERROR, f"Unexpected error resolving JWT signing key: {str(ex)}") raise JwtValidationError("Unable to resolve JWT signing key") from ex try: @@ -172,8 +327,11 @@ 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. + raise JwtValidationError("JWT validation failed") from ex if required_scopes: try: @@ -183,7 +341,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 @@ -191,6 +348,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..220b15c1 --- /dev/null +++ b/tabpy/tabpy_server/handlers/jwt_server_middleware_factory.py @@ -0,0 +1,84 @@ +import logging + +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 + +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.WARNING, 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/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 b3420159..18e4f4f4 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,18 @@ 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_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() + + 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..577fc0c7 --- /dev/null +++ b/tests/unit/server_tests/test_basic_auth_server_middleware_factory.py @@ -0,0 +1,204 @@ +import base64 +import time +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, 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): + 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_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( + 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_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.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.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_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"), + "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: + 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_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]) diff --git a/tests/unit/server_tests/test_jwt_auth.py b/tests/unit/server_tests/test_jwt_auth.py index a9da4075..002662bf 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,15 +165,20 @@ 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.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): @@ -237,16 +241,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 +286,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( @@ -316,19 +314,47 @@ 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 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( @@ -341,6 +367,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 @@ -374,6 +421,394 @@ 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 + 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() + 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() + + 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)) + 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_fetch_only_blocks_refresh_for_bounded_time(self): + """ + 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 + + 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() + 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)) + started = time.monotonic() + 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() + 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_refresh_and_cooldown_publication_are_atomic(self): + """ + 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 + + 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") + 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 ( + threading.current_thread() is first_thread + and fetch_returned.is_set() + and not first_matching.is_set() + ): + first_matching.set() + self.assertTrue(release_first.wait(5)) + return original_match_kid(signing_keys, kid) + + 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 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) + + 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.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) + + # 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_concurrent_requests_for_new_kid_share_one_refresh(self): + """ + 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 + + 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 + ) + + 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 = [] + results = [] + errors = [] + 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 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 worker(): + try: + barrier.wait(5) + claims = validate_jwt( + rotated_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, + ), patch.object( + 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: + thread.start() + for thread in threads: + thread.join(5) + + self.assertTrue(all(not thread.is_alive() for thread in threads)) + self.assertEqual(errors, []) + 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): """ _get_jwks_client must return the same PyJWKClient instance for @@ -382,21 +817,19 @@ 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) - 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__": 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)) 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",