diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 3989cc9934e..c6adc05c191 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -3,6 +3,7 @@ import re from collections.abc import Callable from functools import cache from typing import TYPE_CHECKING, Any, Final +from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -28,6 +29,28 @@ AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000 AZURE_CONTENT_SAFETY_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +AZURE_CONTENT_SAFETY_ENTRA_HOST_SUFFIXES: Final = ( + ".cognitiveservices.azure.com", + ".cognitiveservices.azure.us", + ".cognitiveservices.azure.cn", + ".services.ai.azure.com", +) + + +def _assert_entra_destination_is_azure(api_base: str) -> None: + """An Entra token is scoped to every Cognitive Services resource the identity can reach, + not to one resource, so it is only ever sent to an Azure endpoint over TLS. Entra also + requires the resource's custom subdomain, so any other host is not a valid destination.""" + parsed: Final = urlparse(api_base) + host: Final = (parsed.hostname or "").lower() + if parsed.scheme == "https" and host.endswith(AZURE_CONTENT_SAFETY_ENTRA_HOST_SUFFIXES): + return + raise ValueError( + f"Azure Content Safety: refusing to send a Microsoft Entra token to api_base {api_base!r}. " + "Entra authentication needs the resource's HTTPS custom subdomain endpoint, for example " + "https://your-resource.cognitiveservices.azure.com. Set api_key to reach any other host" + ) + @cache def _default_entra_token_provider() -> Callable[[], str]: @@ -66,6 +89,8 @@ class AzureGuardrailBase: self.api_key = api_key self.api_base = api_base self.api_version: str = kwargs.get("api_version") or "2024-09-01" + if not api_key: + _assert_entra_destination_is_azure(api_base) self._entra_token_provider: Final = entra_token_provider or ( None if api_key else _default_entra_token_provider() ) @@ -79,6 +104,7 @@ class AzureGuardrailBase: if self.api_key: return ("Ocp-Apim-Subscription-Key", self.api_key) + _assert_entra_destination_is_azure(self.api_base) minter: Final = self._entra_token_provider or _default_entra_token_provider() try: token: Final = await asyncio.to_thread(minter) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py index 3179be458e0..bb0a1ad4c9a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py @@ -186,3 +186,77 @@ async def test_api_key_guardrail_never_reaches_azure_identity(api_base, capturin await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" + + +@pytest.mark.parametrize( + "bad_api_base", + [ + "http://contoso.cognitiveservices.azure.com", + "https://contoso.cognitiveservices.azure.com.attacker.example", + "https://attacker.example", + "https://australiaeast.api.cognitive.microsoft.com", + ], +) +def test_keyless_guardrail_refuses_a_non_azure_destination(bad_api_base): + """The Entra token covers every Cognitive Services resource the identity can reach, so a + typo or a hijacked host would receive far more than one resource's key would give away.""" + with pytest.raises(ValueError, match="refusing to send a Microsoft Entra token") as exc_info: + AzureContentSafetyPromptShieldGuardrail(guardrail_name="azure-guard", api_base=bad_api_base) + + assert "api_key" in str(exc_info.value) + + +@pytest.mark.parametrize( + "good_api_base", + [ + "https://contoso.cognitiveservices.azure.com", + "https://contoso.privatelink.cognitiveservices.azure.com", + "https://contoso.services.ai.azure.com", + "https://contoso.cognitiveservices.azure.us", + ], +) +def test_keyless_guardrail_accepts_azure_content_safety_endpoints(good_api_base): + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=good_api_base, + entra_token_provider=lambda: "entra-token", + ) + + assert guardrail.api_base == good_api_base + + +@pytest.mark.asyncio +async def test_api_key_guardrail_may_use_any_destination(capturing_handler): + """A key is scoped to one resource, so gateways and test doubles stay reachable with one.""" + handler, sent = capturing_handler + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base="https://gateway.internal.example", + api_key="secret-key", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" + + +@pytest.mark.asyncio +async def test_clearing_api_key_cannot_send_a_token_to_a_non_azure_destination(capturing_handler): + """A guardrail admitted on its api_key must not start minting tokens for that same host.""" + handler, sent = capturing_handler + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base="https://gateway.internal.example", + api_key="secret-key", + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + guardrail.update_in_memory_litellm_params({"api_key": None}) + + with pytest.raises(ValueError, match="refusing to send a Microsoft Entra token"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent == []