@W-23493263 Add JWT Bearer authentication to Arrow Flight

docs/server-config.md

@@ -331,14 +331,37 @@ To authenticate a request, send the JWT as a Bearer token:
 curl -H "Authorization: Bearer <token>" 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

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 = {}
 

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,
+)

tabpy/tabpy_server/handlers/basic_auth_server_middleware_factory.py

@@ -1,48 +1,126 @@
 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
 
     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:
             return False
         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")

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