diff --git a/databricks/sdk/_base_client.py b/databricks/sdk/_base_client.py index 13ddb9bec..871e2c01f 100644 --- a/databricks/sdk/_base_client.py +++ b/databricks/sdk/_base_client.py @@ -38,6 +38,29 @@ def _fix_host_if_needed(host: Optional[str]) -> Optional[str]: return urllib.parse.urlunparse((o.scheme, netloc, path, o.params, o.query, o.fragment)) +class _CredentialStrippingSession(requests.Session): + """A requests session that also strips Databricks credential headers on a cross-host redirect. + + requests.Session.rebuild_auth() already removes the standard sensitive headers + (Authorization, Cookie, etc.) when a redirect changes host. The Databricks cloud-provider + token headers are custom, so requests does not know to strip them. Extend rebuild_auth to + drop them on the same host-change condition, so a cross-host redirect cannot forward the + caller's Azure/GCP token to a different host. + """ + + _sensitive_cross_host_headers = ( + "X-Databricks-Azure-SP-Management-Token", + "X-Databricks-GCP-SA-Access-Token", + "X-Databricks-Azure-Workspace-Resource-Id", + ) + + def rebuild_auth(self, prepared_request, response): + super().rebuild_auth(prepared_request, response) + if self.should_strip_auth(response.request.url, prepared_request.url): + for header in self._sensitive_cross_host_headers: + prepared_request.headers.pop(header, None) + + class _BaseClient: def __init__( self, @@ -80,7 +103,7 @@ def __init__( self._user_agent_base = user_agent_base or useragent.to_string() self._header_factory = header_factory self._clock = clock or RealClock() - self._session = requests.Session() + self._session = _CredentialStrippingSession() self._session.auth = self._authenticate self._streaming_buffer_size = streaming_buffer_size diff --git a/tests/test_base_client.py b/tests/test_base_client.py index 1b8beed05..57d859e91 100644 --- a/tests/test_base_client.py +++ b/tests/test_base_client.py @@ -561,3 +561,56 @@ def do(): do() assert session._received_requests == [test_case._expected_result for _ in range(expected_attempts_made)] + + +def _prep(url, headers): + req = PreparedRequest() + req.prepare(method="GET", url=url, headers=headers) + return req + + +def test_credential_stripping_session_drops_cloud_tokens_on_cross_host_redirect(): + from databricks.sdk._base_client import _CredentialStrippingSession + + session = _CredentialStrippingSession() + original = _prep("https://tenant.cloud.databricks.com/api/2.0/preview/scim/v2/Me", {}) + response = Response() + response.request = original + + redirected = _prep( + "https://attacker.example/api/2.0/preview/scim/v2/Me", + { + "Authorization": "Bearer pat-token", + "X-Databricks-Azure-SP-Management-Token": "azure-arm-token", + "X-Databricks-GCP-SA-Access-Token": "gcp-sa-token", + }, + ) + session.rebuild_auth(redirected, response) + + # requests strips Authorization on host change; we additionally strip the custom cloud headers. + assert "Authorization" not in redirected.headers + assert "X-Databricks-Azure-SP-Management-Token" not in redirected.headers + assert "X-Databricks-GCP-SA-Access-Token" not in redirected.headers + + +def test_credential_stripping_session_keeps_cloud_tokens_on_same_host_redirect(): + from databricks.sdk._base_client import _CredentialStrippingSession + + session = _CredentialStrippingSession() + original = _prep("https://tenant.cloud.databricks.com/first", {}) + response = Response() + response.request = original + + redirected = _prep( + "https://tenant.cloud.databricks.com/second", + { + "Authorization": "Bearer pat-token", + "X-Databricks-Azure-SP-Management-Token": "azure-arm-token", + "X-Databricks-GCP-SA-Access-Token": "gcp-sa-token", + }, + ) + session.rebuild_auth(redirected, response) + + assert redirected.headers["Authorization"] == "Bearer pat-token" + assert redirected.headers["X-Databricks-Azure-SP-Management-Token"] == "azure-arm-token" + assert redirected.headers["X-Databricks-GCP-SA-Access-Token"] == "gcp-sa-token"