fix(guardrails): only send an Entra token to an Azure endpoint

A Microsoft Entra token is scoped to every Cognitive Services resource
the identity can reach, not to the one resource a key would unlock, so a
typo, a stale DNS name, or a hostile api_base receives far more than a
leaked resource key would give away.

Check the destination before minting: the Entra path now requires HTTPS
and a recognised Azure Content Safety host, and refuses anything else
with an error naming api_key as the way to reach another endpoint. The
check runs at startup and again per request, so clearing api_key on a
running guardrail cannot quietly promote its host

The api_key path is untouched. A resource key only unlocks the resource
it belongs to, so gateways and test doubles stay reachable with one

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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)

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 == []