diff --git a/src/databricks/sql/auth/oauth.py b/src/databricks/sql/auth/oauth.py index 231a18905..3fa82650b 100644 --- a/src/databricks/sql/auth/oauth.py +++ b/src/databricks/sql/auth/oauth.py @@ -1,4 +1,5 @@ import base64 +import errno import hashlib import json import logging @@ -62,18 +63,43 @@ def refresh(self) -> Token: class OAuthManager: + # Default maximum time (in seconds) to wait for the browser OAuth redirect + # callback before giving up. Without a bound, the local callback server + # would block forever in a headless environment (e.g. a notebook/job with + # no browser), making the connection appear to hang indefinitely. See issue + # #458. This is generous for interactive logins (slow MFA, IdP re-auth, SSO + # redirects) and is a fixed ceiling for connector end users: there is no + # public ``connect()``/``Connection`` parameter to change it. The + # ``redirect_callback_timeout_seconds`` argument below is an internal-only + # override for callers constructing ``OAuthManager`` directly (private API); + # it is not plumbed through ``DatabricksOAuthProvider`` or the public + # connection kwargs. It accepts a ``float`` so callers (and tests) may pass + # fractional-second timeouts; this class default is a whole number of + # seconds. + REDIRECT_CALLBACK_TIMEOUT_SECONDS = 60 * 5 + def __init__( self, port_range: List[int], client_id: str, idp_endpoint: OAuthEndpointCollection, http_client, + redirect_callback_timeout_seconds: Optional[float] = None, ): self.port_range = port_range self.client_id = client_id self.redirect_port = None self.idp_endpoint = idp_endpoint self.http_client = http_client + # Fall back to the class default when not overridden. Kept as a + # lowercase instance attribute (distinct from the ALL_CAPS class + # constant that provides the default) so the rest of the flow reads a + # single, clearly per-instance source of truth. + self._redirect_callback_timeout_seconds = ( + redirect_callback_timeout_seconds + if redirect_callback_timeout_seconds is not None + else self.REDIRECT_CALLBACK_TIMEOUT_SECONDS + ) @staticmethod def __token_urlsafe(nbytes=32): @@ -128,11 +154,40 @@ def __get_challenge(): def __get_authorization_code(self, client, auth_url, scope, state, challenge): handler = OAuthHttpSingleRequestHandler("Databricks Sql Connector") + # Bound the read of the accepted connection too (not just the accept + # wait below). StreamRequestHandler.setup() applies this to the accepted + # socket via settimeout(), so a client that connects but never completes + # the request line can no longer block indefinitely. + handler.timeout = self._redirect_callback_timeout_seconds last_error = None + callback_timed_out = False for port in self.port_range: try: with HTTPServer(("", port), handler) as httpd: + # Bound how long we wait for the browser redirect callback so + # that a headless environment (no browser to complete the + # flow) fails with a clear error instead of hanging forever. + # NOTE: httpd.timeout bounds only the wait to ACCEPT an + # incoming connection (it is the select() timeout used by + # handle_request). The subsequent read of the HTTP request + # line is bounded separately by the handler.timeout set + # above (applied to the accepted socket by + # StreamRequestHandler.setup()), so a client that connects + # but never completes the request can no longer block. The + # headless "no connection ever arrives" case (issue #458) is + # covered by this accept-wait timeout. + httpd.timeout = self._redirect_callback_timeout_seconds + + # HTTPServer.handle_request() returns normally (via + # handle_timeout()) when the wait elapses without a + # connection, so record that case to distinguish it from a + # received-but-empty callback below. + def _on_timeout(): + nonlocal callback_timed_out + callback_timed_out = True + + httpd.handle_timeout = _on_timeout redirect_url = OAuthManager.__get_redirect_url(port) auth_req_uri, _, _ = client.prepare_authorization_request( authorization_url=auth_url, @@ -152,11 +207,61 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge): self.redirect_port = port break except OSError as e: - if e.errno == 48: + # Record every bind-time OSError so the fall-through + # `raise last_error` below always re-raises a real exception. + # If we only recorded the port-in-use case, an unexpected error + # on every port (e.g. EACCES, EADDRNOTAVAIL) would leave + # last_error == None and turn `raise last_error` into + # `raise None`, masking the real failure with a confusing + # TypeError. + last_error = e + # errno.EADDRINUSE resolves to the platform's value (48 on + # macOS, 98 on Linux), so a port-in-use bind failure is + # recognized cross-platform and downgraded to an info log while + # we try the next port. + if e.errno == errno.EADDRINUSE: logger.info(f"Port {port} is in use") - last_error = e + else: + logger.warning(f"Unexpected error binding port {port}: {e}") + except webbrowser.Error as e: + # No browser could be launched at all (webbrowser.get() found no + # runnable browser). This is a strong, reliable headless signal + # (e.g. a notebook/job with no display), so fail fast with a + # clear error instead of waiting out the full redirect-callback + # timeout. We do NOT fail fast merely on a falsy open_new return + # value: that is unreliable across platforms (it can be falsy + # even when a browser did open) and would break working + # interactive logins. Note this fast-fail only covers the case + # where webbrowser.get() raises; the more common headless case + # (no browser registered, so open_new() returns False without + # raising) is NOT caught here and instead falls through to the + # bounded redirect-callback timeout below, which is the real + # safeguard against blocking forever. See issue #458. Placed + # before the generic + # ``except Exception`` so this RuntimeError propagates instead of + # being swallowed and retried against the remaining ports. + msg = ( + "Could not launch a web browser to complete the U2M OAuth " + f"login flow ({e}). This looks like a headless environment " + "(e.g. a notebook or job with no browser), where browser-" + "based OAuth cannot complete. See issue #458." + ) + logger.error(msg) + raise RuntimeError(msg) from e except Exception as e: + # A non-OSError/non-webbrowser.Error raised here is not + # port-related (e.g. an OAuth2Error from + # prepare_authorization_request, or an unexpected programming + # bug). Binding the same call to a different local port cannot + # help, so retrying it against every remaining port would only + # multiply latency and log noise and then re-raise the *last* + # attempt's error, obscuring the original failure. Fail fast + # instead, preserving the original exception and traceback via a + # bare `raise`. This also sidesteps the `raise None` hazard that + # the OSError branch guards against, since we never fall through + # to `raise last_error` for this case. logger.error("unexpected error: %s", e) + raise if self.redirect_port is None: logger.error( f"Tried all the ports {self.port_range} for oauth redirect, but can't find free port" @@ -164,7 +269,18 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge): raise last_error if not handler.request_path: - msg = f"No path parameters were returned to the callback at {redirect_url}" + if callback_timed_out: + msg = ( + f"Timed out after {self._redirect_callback_timeout_seconds} " + f"seconds waiting for the OAuth redirect callback at " + f"{redirect_url}. The login flow was not completed in time — " + "either the interactive login was not finished within the " + "timeout, or this is a headless environment (e.g. a notebook " + "or job with no browser) where no browser can complete the " + "flow. See issue #458." + ) + else: + msg = f"No path parameters were returned to the callback at {redirect_url}" logger.error(msg) raise RuntimeError(msg) # This is a kludge because the parsing library expects https callbacks diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index d1b941208..927d9d150 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -207,6 +207,243 @@ def test_get_python_sql_connector_default_auth(self, mock__initial_get_token): self.assertEqual(auth_provider.external_provider._client_id, PYSQL_OAUTH_CLIENT_ID) + @patch("databricks.sql.auth.oauth.webbrowser.open_new") + @patch.object(OAuthManager, "_OAuthManager__fetch_well_known_config") + def test_get_tokens_does_not_hang_when_no_callback_received( + self, mock_fetch_config, mock_open_new + ): + """When the U2M browser OAuth callback never arrives (e.g. a headless + notebook/job with a null token), the local redirect server must not + block forever. It should time out and surface a clear error rather than + hang. See issue #458.""" + mock_fetch_config.return_value = { + "authorization_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/authorize", + "token_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/token", + } + # Do not actually launch a browser during the test. + mock_open_new.return_value = True + + # Bind to an OS-assigned ephemeral port (0) rather than a fixed port so + # the test does not depend on a specific port being free. A fixed port + # that happens to be occupied would fail to bind and take the + # can't-find-free-port branch instead of the timeout path we exercise. + # Keep the test fast: shorten the callback wait via the constructor + # keyword (the actual new surface this PR introduces) rather than + # mutating the private attribute after the fact. This also exercises the + # Optional[int] fall-back logic. The production default is minutes; the + # bug is that WITHOUT any bound the wait is infinite. + oauth_manager = OAuthManager( + port_range=[0], + client_id="mock-id", + idp_endpoint=InHouseOAuthEndpointCollection(), + http_client=MagicMock(), + # Sub-second so this real-wallclock wait stays negligible in the + # fast/mocked unit suite. The wait IS the timeout, so there is no + # setup racing against it that a shorter bound could make flaky. + redirect_callback_timeout_seconds=0.5, + ) + + # No callback is ever delivered to the redirect server. Run the flow in + # a daemon thread and join with a wall-clock bound so a regression to + # the infinite-block behaviour fails this test loudly instead of hanging + # the whole suite. + import threading + + result = {} + + def run(): + try: + oauth_manager.get_tokens( + hostname="foo.cloud.databricks.com", scope="offline_access sql" + ) + result["outcome"] = "returned" + except BaseException as e: # noqa: BLE001 + result["outcome"] = "raised" + result["error"] = e + + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(timeout=30) + + self.assertFalse( + worker.is_alive(), + "OAuth callback server blocked indefinitely waiting for a callback " + "that never arrives (issue #458)", + ) + self.assertEqual(result.get("outcome"), "raised") + self.assertIsInstance(result.get("error"), RuntimeError) + # The headless timeout path must surface a clear, timeout-specific error + # (not the generic received-but-empty callback message). See issue #458. + error_message = str(result.get("error")) + self.assertIn("Timed out", error_message) + self.assertIn( + str(oauth_manager._redirect_callback_timeout_seconds), error_message + ) + + @patch("databricks.sql.auth.oauth.webbrowser.open_new") + @patch.object(OAuthManager, "_OAuthManager__fetch_well_known_config") + def test_get_tokens_does_not_hang_when_callback_connects_but_never_completes( + self, mock_fetch_config, mock_open_new + ): + """A client that connects to the local redirect server but never sends a + complete HTTP request line must not block forever either. This is the + separate defense the PR adds via ``handler.timeout``: setup() applies it + to the accepted socket, so the stalled request-line read raises + socket.timeout, which socketserver swallows via handle_error. Because no + callback path is ever recorded (request_path stays None) and this is not + the accept-wait timeout (callback_timed_out stays False), the flow raises + the generic "No path parameters" error rather than the "Timed out" one. + See oauth.py __get_authorization_code.""" + import socket + import threading + from urllib.parse import urlparse, parse_qs + + mock_fetch_config.return_value = { + "authorization_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/authorize", + "token_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/token", + } + + stalled_sockets = [] + + def connect_but_stall(auth_req_uri): + # Instead of launching a browser, open a bare TCP connection to the + # local redirect server (parsed out of the authorization request's + # redirect_uri) and never send a request line, so the accepted + # socket's read stalls until handler.timeout fires. + redirect_uri = parse_qs(urlparse(auth_req_uri).query)["redirect_uri"][0] + port = urlparse(redirect_uri).port + stalled_sockets.append(socket.create_connection(("localhost", port))) + return True + + mock_open_new.side_effect = connect_but_stall + + def run_flow(port): + # Target a specific free port so the connecting client above knows + # where to reach the redirect server; unlike the accept-wait test, + # port_range=[0] can't be used here because the source builds the + # redirect URL from the requested port, so it would point the client + # at localhost:0. + oauth_manager = OAuthManager( + port_range=[port], + client_id="mock-id", + idp_endpoint=InHouseOAuthEndpointCollection(), + http_client=MagicMock(), + # Sub-second: bounds the stalled request-line read. Keeps the + # per-attempt wait negligible so even the rare TOCTOU retry loop + # stays fast in the mocked unit suite. + redirect_callback_timeout_seconds=0.5, + ) + + result = {} + + def run(): + try: + oauth_manager.get_tokens( + hostname="foo.cloud.databricks.com", scope="offline_access sql" + ) + result["outcome"] = "returned" + except BaseException as e: # noqa: BLE001 + result["outcome"] = "raised" + result["error"] = e + + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(timeout=30) + return worker, result + + try: + # Pre-selecting a fixed port has a small TOCTOU window: between + # closing the probe socket and the HTTPServer bind inside + # __get_authorization_code, another process can claim the port. That + # would take the can't-bind branch instead of the stalled-request- + # line path this test exercises, surfacing a bind error rather than + # the RuntimeError asserted below. Re-probe a fresh port and retry + # on that rare race so the test stays deterministic in CI. + worker, result = None, {} + for _ in range(5): + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + probe.bind(("localhost", 0)) + free_port = probe.getsockname()[1] + probe.close() + + worker, result = run_flow(free_port) + if isinstance(result.get("error"), RuntimeError): + break + + self.assertFalse( + worker.is_alive(), + "OAuth callback server blocked indefinitely on a client that " + "connected but never completed the request line", + ) + self.assertEqual(result.get("outcome"), "raised") + self.assertIsInstance(result.get("error"), RuntimeError) + # A stalled request-line read is NOT the accept-wait timeout: + # request_path stays None and callback_timed_out stays False, so the + # generic message is raised, not the "Timed out" one. + error_message = str(result.get("error")) + self.assertIn("No path parameters", error_message) + self.assertNotIn("Timed out", error_message) + finally: + for sock in stalled_sockets: + sock.close() + + @patch("databricks.sql.auth.oauth.webbrowser.open_new") + @patch.object(OAuthManager, "_OAuthManager__fetch_well_known_config") + def test_get_tokens_fails_fast_when_no_browser_available( + self, mock_fetch_config, mock_open_new + ): + """When no browser can be launched at all (webbrowser.open_new raises + webbrowser.Error), the flow must fail fast with a clear headless error + rather than waiting out the full redirect-callback timeout. See issue + #458.""" + import webbrowser + + mock_fetch_config.return_value = { + "authorization_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/authorize", + "token_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/token", + } + mock_open_new.side_effect = webbrowser.Error("no browser available") + + # A large timeout: if the fail-fast path regressed, joining below would + # hit the wall-clock bound and the test would fail loudly instead of + # silently waiting minutes. + oauth_manager = OAuthManager( + port_range=[0], + client_id="mock-id", + idp_endpoint=InHouseOAuthEndpointCollection(), + http_client=MagicMock(), + redirect_callback_timeout_seconds=600, + ) + + import threading + + result = {} + + def run(): + try: + oauth_manager.get_tokens( + hostname="foo.cloud.databricks.com", scope="offline_access sql" + ) + result["outcome"] = "returned" + except BaseException as e: # noqa: BLE001 + result["outcome"] = "raised" + result["error"] = e + + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(timeout=30) + + self.assertFalse( + worker.is_alive(), + "OAuth flow did not fail fast when no browser was available", + ) + self.assertEqual(result.get("outcome"), "raised") + self.assertIsInstance(result.get("error"), RuntimeError) + error_message = str(result.get("error")) + self.assertIn("browser", error_message.lower()) + # The fail-fast path must not surface the accept-wait timeout message. + self.assertNotIn("Timed out", error_message) + class TestClientCredentialsTokenSource: @pytest.fixture