Fix trusted endpoint validation for token providers

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,14 @@ 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(
+            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
-                well_known_kusto_endpoints.validate_trusted_endpoint(
-                    self._kusto_cluster,
-                    endpoint,
-                )
+                ).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,23 +73,29 @@ 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]]):
         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())
-            if matcher is not None and matcher.is_match(hostname):
+            additional_matcher = self._additional_matcher
+            if additional_matcher is not None and additional_matcher.is_match(hostname):
                 return
 
-        matcher = self._additional_matcher
-        if matcher is not None and matcher.is_match(hostname):
+            if any(matcher.is_match(hostname) for matcher in self._matchers.values()):
+                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
+
+        additional_matcher = self._additional_matcher
+        if additional_matcher is not None and additional_matcher.is_match(hostname):
             return
 
         raise KustoClientInvalidConnectionStringException(

azure-kusto-data/tests/test_endpoint_validation.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
+from azure.kusto.data._cloud_settings import DEFAULT_PUBLIC_LOGIN_URL
+from azure.kusto.data.aio import KustoClient as AsyncKustoClient
+from azure.kusto.data.exceptions import KustoClientInvalidConnectionStringException
+from azure.kusto.data.kusto_trusted_endpoints import KustoTrustedEndpoints, MatchRule
+
+
+UNTRUSTED_HOST = "https://kusto.attacker.example.com"
+TRUSTED_HOST = "https://somecluster.kusto.windows.net"
+TOKEN = "a token that must not be transmitted"
+
+
+def _sync_kcsb(authentication_method, token_callback):
+    if authentication_method == "user_token":
+        return KustoConnectionStringBuilder.with_aad_user_token_authentication(UNTRUSTED_HOST, TOKEN)
+    if authentication_method == "application_token":
+        return KustoConnectionStringBuilder.with_aad_application_token_authentication(UNTRUSTED_HOST, TOKEN)
+    return KustoConnectionStringBuilder.with_token_provider(UNTRUSTED_HOST, token_callback)
+
+
+def _async_kcsb(authentication_method, token_callback):
+    if authentication_method == "user_token":
+        return KustoConnectionStringBuilder.with_aad_user_token_authentication(UNTRUSTED_HOST, TOKEN)
+    if authentication_method == "application_token":
+        return KustoConnectionStringBuilder.with_aad_application_token_authentication(UNTRUSTED_HOST, TOKEN)
+    return KustoConnectionStringBuilder.with_async_token_provider(UNTRUSTED_HOST, token_callback)
+
+
+@pytest.mark.parametrize("authentication_method", ["user_token", "application_token", "callback"])
+def test_sync_authentication_rejects_untrusted_host_before_token_use(authentication_method):
+    token_callback = Mock(return_value=TOKEN)
+    kcsb = _sync_kcsb(authentication_method, token_callback)
+
+    with (
+        patch("azure.kusto.data.client_base.CloudSettings.get_cloud_info_for_cluster") as cloud_info,
+        patch("azure.kusto.data.security._get_header_from_dict") as build_header,
+        patch("requests.get") as requests_get,
+        patch("requests.Session.get") as session_get,
+        patch("requests.Session.post") as session_post,
+    ):
+        with KustoClient(kcsb) as client:
+            with pytest.raises(KustoClientInvalidConnectionStringException):
+                client.execute_query("database", "print 1")
+
+    token_callback.assert_not_called()
+    cloud_info.assert_not_called()
+    build_header.assert_not_called()
+    requests_get.assert_not_called()
+    session_get.assert_not_called()
+    session_post.assert_not_called()
+
+
+def test_cloud_login_endpoint_is_resolved_only_for_builtin_candidate():
+    trusted_endpoints = KustoTrustedEndpoints()
+    login_endpoint = Mock(return_value=DEFAULT_PUBLIC_LOGIN_URL)
+
+    trusted_endpoints.validate_trusted_endpoint(TRUSTED_HOST, login_endpoint)
+    login_endpoint.assert_called_once_with()
+
+    login_endpoint.reset_mock()
+    with pytest.raises(KustoClientInvalidConnectionStringException):
+        trusted_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, login_endpoint)
+    login_endpoint.assert_not_called()
+
+
+def test_additional_trusted_host_does_not_resolve_cloud_login_endpoint():
+    trusted_endpoints = KustoTrustedEndpoints()
+    trusted_endpoints.add_trusted_hosts([MatchRule(".example.com", False)], False)
+    login_endpoint = Mock()
+
+    trusted_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, login_endpoint)
+
+    login_endpoint.assert_not_called()
+
+
+def test_override_policy_preserves_additional_trusted_host_fallback():
+    trusted_endpoints = KustoTrustedEndpoints()
+    trusted_endpoints.set_override_policy(lambda hostname: False)
+    trusted_endpoints.add_trusted_hosts([MatchRule(".example.com", False)], False)
+    login_endpoint = Mock()
+
+    trusted_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, login_endpoint)
+
+    login_endpoint.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("authentication_method", ["user_token", "application_token", "callback"])
+async def test_async_authentication_rejects_untrusted_host_before_token_use(authentication_method):
+    token_callback = AsyncMock(return_value=TOKEN)
+    kcsb = _async_kcsb(authentication_method, token_callback)
+
+    with (
+        patch("azure.kusto.data.client_base.CloudSettings.get_cloud_info_for_cluster") as cloud_info,
+        patch("azure.kusto.data.security._get_header_from_dict") as build_header,
+        patch("requests.get") as requests_get,
+        patch("requests.Session.get") as session_get,
+        patch("azure.kusto.data.aio.client.ClientSession.post") as session_post,
+    ):
+        async with AsyncKustoClient(kcsb) as client:
+            with pytest.raises(KustoClientInvalidConnectionStringException):
+                await client.execute_query("database", "print 1")
+
+    token_callback.assert_not_awaited()
+    cloud_info.assert_not_called()
+    build_header.assert_not_called()
+    requests_get.assert_not_called()
+    session_get.assert_not_called()
+    session_post.assert_not_called()