fix: enforce trusted-endpoint validation for all authentication methods

Author: komal-rani522

azure-kusto-data/azure/kusto/data/client_base.py

@@ -9,7 +9,6 @@
 from requests import Response, Session
 
 from azure.kusto.data._cloud_settings import CloudSettings
-from azure.kusto.data._token_providers import CloudInfoTokenProvider
 from .client_details import ClientDetails
 from .client_request_properties import ClientRequestProperties
 from .exceptions import KustoServiceError, KustoThrottlingError, KustoApiError
@@ -78,16 +77,20 @@ def set_proxy(self, proxy_url: str):
 
     def validate_endpoint(self):
         if not self._endpoint_validated and self._aad_helper is not None:
-            if isinstance(self._aad_helper.token_provider, CloudInfoTokenProvider):
-                endpoint = CloudSettings.get_cloud_info_for_cluster(
-                    self._kusto_cluster,
-                    self._aad_helper.token_provider._proxy_dict,
-                    self._session if isinstance(self._session, Session) else None,
-                ).login_endpoint
-                well_known_kusto_endpoints.validate_trusted_endpoint(
-                    self._kusto_cluster,
-                    endpoint,
-                )
+            # Trusted-endpoint validation must run for every authentication method. Gating it on the
+            # token provider type let token-based and callback-based flows send the Authorization
+            # header to an arbitrary host named in the connection string.
+            # The login endpoint is resolved lazily because doing so contacts the cluster itself.
+            well_known_kusto_endpoints.validate_trusted_endpoint(
+                self._kusto_cluster,
+                lambda: (
+                    CloudSettings.get_cloud_info_for_cluster(
+                        self._kusto_cluster,
+                        self._aad_helper.token_provider._proxy_dict,
+                        self._session if isinstance(self._session, Session) else None,
+                    ).login_endpoint
+                ),
+            )
             self._endpoint_validated = True
 
     @staticmethod

azure-kusto-data/azure/kusto/data/kusto_trusted_endpoints.py

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Dict
+from typing import Callable, List, Dict, Union
 from urllib.parse import urlparse
 
 from azure.kusto.data.helpers import get_string_tail_lower_case
@@ -73,18 +73,30 @@ def add_trusted_hosts(self, rules, replace):
 
         self._additional_matcher = create_fast_suffix_matcher_from_existing(rules, None if replace else self._additional_matcher)
 
-    def validate_trusted_endpoint(self, endpoint: str, login_endpoint: str):
+    def validate_trusted_endpoint(self, endpoint: str, login_endpoint: Union[str, Callable[[], str]]):
+        """Validates that the endpoint is trusted.
+
+        `login_endpoint` may be a callable, in which case it is only resolved if the built-in
+        per-cloud allow lists actually have to be consulted. Resolving it may require contacting
+        the cluster itself, so callers should pass a callable to avoid reaching out to hosts that
+        can never be trusted.
+        """
         hostname = urlparse(endpoint).hostname
         self.validate_hostname_is_trusted(hostname if hostname is not None else endpoint, login_endpoint)
 
-    def validate_hostname_is_trusted(self, hostname: str, login_endpoint: str):
+    def validate_hostname_is_trusted(self, hostname: str, login_endpoint: Union[str, Callable[[], str]]):
         if _is_local_address(hostname):
             return
         if self._override_matcher is not None:
             if self._override_matcher(hostname):
                 return
-        else:
-            matcher = self._matchers.get(login_endpoint.lower())
+        elif any(matcher.is_match(hostname) for matcher in self._matchers.values()):
+            # Only resolve the login endpoint once the hostname is known to appear in at least one
+            # cloud's allow list. Resolving it may contact the cluster, so doing it first would let
+            # an untrusted connection string drive a request to an arbitrary host before it is
+            # rejected.
+            resolved_login_endpoint = login_endpoint() if callable(login_endpoint) else login_endpoint
+            matcher = self._matchers.get(resolved_login_endpoint.lower())
             if matcher is not None and matcher.is_match(hostname):
                 return
 

azure-kusto-data/tests/test_endpoint_validation.py

@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License
+"""Trusted-endpoint validation must be enforced for every authentication method.
+
+Token-based and callback-based authentication used to bypass validation entirely, which let a
+connection string send its Authorization header to an arbitrary host.
+"""
+
+import asyncio
+from unittest.mock import patch
+
+import pytest
+
+from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
+from azure.kusto.data.exceptions import KustoClientInvalidConnectionStringException
+from azure.kusto.data.kusto_trusted_endpoints import MatchRule, well_known_kusto_endpoints
+
+UNTRUSTED_HOST = "https://kusto.attacker.example.com"
+TRUSTED_HOST = "https://somecluster.kusto.windows.net"
+TOKEN = "a token that must never leave the machine"
+
+
+def _bypassing_auth_kcsbs(cluster: str):
+    """Connection strings whose token providers do not derive from CloudInfoTokenProvider."""
+    return {
+        "user_token": KustoConnectionStringBuilder.with_aad_user_token_authentication(cluster, TOKEN),
+        "application_token": KustoConnectionStringBuilder.with_aad_application_token_authentication(cluster, TOKEN),
+        "token_provider": KustoConnectionStringBuilder.with_token_provider(cluster, lambda: TOKEN),
+        "async_token_provider": KustoConnectionStringBuilder.with_async_token_provider(cluster, lambda: asyncio.sleep(0, result=TOKEN)),
+    }
+
+
+@pytest.fixture(params=["user_token", "application_token", "token_provider", "async_token_provider"])
+def bypassing_auth_name(request):
+    return request.param
+
+
+class TestEndpointValidation:
+    def test_untrusted_host_is_rejected_for_token_based_auth(self, bypassing_auth_name):
+        kcsb = _bypassing_auth_kcsbs(UNTRUSTED_HOST)[bypassing_auth_name]
+        with KustoClient(kcsb) as client:
+            with pytest.raises(KustoClientInvalidConnectionStringException):
+                client.execute_query("PythonTest", "Deft")
+
+    def test_untrusted_host_is_rejected_before_any_network_call(self, bypassing_auth_name):
+        """Validation must not contact the untrusted host, otherwise it is an SSRF primitive."""
+        kcsb = _bypassing_auth_kcsbs(UNTRUSTED_HOST)[bypassing_auth_name]
+        with patch("requests.get") as mock_get, patch("requests.Session.get") as mock_session_get, patch("requests.Session.post") as mock_post:
+            with KustoClient(kcsb) as client:
+                with pytest.raises(KustoClientInvalidConnectionStringException):
+                    client.execute_query("PythonTest", "Deft")
+            assert not mock_get.called
+            assert not mock_session_get.called
+            assert not mock_post.called
+
+    def test_explicitly_trusted_host_needs_no_cloud_metadata(self):
+        """Hosts trusted via add_trusted_hosts must not require the metadata endpoint."""
+        try:
+            well_known_kusto_endpoints.add_trusted_hosts([MatchRule("kusto.attacker.example.com", True)], False)
+            resolved = []
+
+            def resolver():
+                resolved.append(True)
+                return "https://login.microsoftonline.com"
+
+            well_known_kusto_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, resolver)
+            assert not resolved
+        finally:
+            well_known_kusto_endpoints.add_trusted_hosts(None, True)
+
+    def test_login_endpoint_resolved_only_for_allow_listed_hosts(self):
+        resolved = []
+
+        def resolver():
+            resolved.append(True)
+            return "https://login.microsoftonline.com"
+
+        well_known_kusto_endpoints.validate_trusted_endpoint(TRUSTED_HOST, resolver)
+        assert resolved
+
+        resolved.clear()
+        with pytest.raises(KustoClientInvalidConnectionStringException):
+            well_known_kusto_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, resolver)
+        assert not resolved
+
+    def test_plain_string_login_endpoint_still_supported(self):
+        well_known_kusto_endpoints.validate_trusted_endpoint(TRUSTED_HOST, "https://login.microsoftonline.com")
+        with pytest.raises(KustoClientInvalidConnectionStringException):
+            well_known_kusto_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, "https://login.microsoftonline.com")