Strip custom credential headers on cross-host redirects

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
 

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"