mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
chore(proxy): clean up request parameter validation and provider destination handling (#34189)
(cherry picked from commit 065faf6e69)
Backport adaptation for stable/1.91.x:
- litellm/proxy/auth/user_api_key_auth.py: dropped `Iterator` from the typing
import as upstream does, but kept this line's import shape rather than taking
`Protocol`, which reached staging with a feature this line does not carry and
is unused here.
- tests/code_coverage_tests/recursive_detector.py: added only this commit's own
`_iter_fallback_targets` allowlist entry; the three neighbouring entries guard
staging-prior recursive helpers absent from this line.
- tests/test_litellm/proxy/auth/test_auth_utils.py: kept all four test classes
this commit adds (TestClientsideBaseOverrideOutboundKey,
TestIsRequestBodySafeBlocksFallbackSmuggle,
TestIsRequestBodySafeRejectsUrlValuedFallback,
TestIsRequestBodySafeBlocksVertexCredentialAlias) and omitted staging-prior
tests for nvcf_function_id, use_ssl and bedrock_tags, whose subjects are not
on this line.
This commit is contained in:
yucheng-berri
Yuneng Jiang
parent
08429626f9
commit
8498dc0bb8
13 changed files with 700 additions and 104 deletions
|
|
@ -148,6 +148,15 @@ def _parse_url_destination_allowlist_entry(
|
|||
return _normalize_host(parsed.hostname), scheme, port
|
||||
|
||||
|
||||
def provider_url_destination_candidates(value: str) -> Tuple[str, ...]:
|
||||
return tuple(
|
||||
candidate
|
||||
for part in value.split(",")
|
||||
for candidate in (part.strip(), part.strip().split("/", 1)[1] if "/" in part.strip() else "")
|
||||
if candidate
|
||||
)
|
||||
|
||||
|
||||
def is_url_destination_allowed_by_host(url: str, allowed_hosts: List[str]) -> bool:
|
||||
"""Return True when a credential-bearing provider URL is admin-allowlisted.
|
||||
|
||||
|
|
|
|||
|
|
@ -322,7 +322,7 @@ class HuggingFaceEmbedding(BaseLLM):
|
|||
task = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL)
|
||||
# print_verbose(f"{model}, {task}")
|
||||
embed_url = ""
|
||||
if "https" in model:
|
||||
if model.startswith(("http://", "https://")):
|
||||
embed_url = model
|
||||
elif api_base:
|
||||
embed_url = api_base
|
||||
|
|
|
|||
|
|
@ -316,25 +316,6 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
|
|||
|
||||
return data
|
||||
|
||||
def get_api_base(self, api_base: Optional[str], model: str) -> str:
|
||||
"""
|
||||
Get the API base for the Huggingface API.
|
||||
|
||||
Do not add the chat/embedding/rerank extension here. Let the handler do this.
|
||||
"""
|
||||
if "https" in model:
|
||||
completion_url = model
|
||||
elif api_base is not None:
|
||||
completion_url = api_base
|
||||
elif "HF_API_BASE" in os.environ:
|
||||
completion_url = os.getenv("HF_API_BASE", "")
|
||||
elif "HUGGINGFACE_API_BASE" in os.environ:
|
||||
completion_url = os.getenv("HUGGINGFACE_API_BASE", "")
|
||||
else:
|
||||
completion_url = f"https://api-inference.huggingface.co/models/{model}"
|
||||
|
||||
return completion_url
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ def completion(
|
|||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
if "https" in model:
|
||||
if model.startswith(("http://", "https://")):
|
||||
completion_url = model
|
||||
elif api_base:
|
||||
completion_url = api_base
|
||||
|
|
@ -96,7 +96,7 @@ def embedding(
|
|||
encoding=None,
|
||||
):
|
||||
# Create completion URL
|
||||
if "https" in model:
|
||||
if model.startswith(("http://", "https://")):
|
||||
embeddings_url = model
|
||||
elif api_base:
|
||||
embeddings_url = f"{api_base}/v1/embeddings"
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import re
|
|||
import sys
|
||||
from functools import lru_cache
|
||||
from logging import Logger
|
||||
from typing import Any, Dict, FrozenSet, List, Mapping, Optional, Tuple, Union
|
||||
from typing import Any, Dict, FrozenSet, Iterator, List, Mapping, Optional, Tuple, Union
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
|
||||
|
|
@ -12,7 +12,12 @@ from litellm import Router, provider_list
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
|
||||
from litellm.litellm_core_utils.url_utils import (
|
||||
SSRFError,
|
||||
is_url_destination_allowed_by_host,
|
||||
provider_url_destination_candidates,
|
||||
validate_url,
|
||||
)
|
||||
from litellm.proxy._types import *
|
||||
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams
|
||||
|
|
@ -280,6 +285,7 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = (
|
|||
"deployment_url",
|
||||
# SDK-only field; also rejected outright in is_request_body_safe.
|
||||
"model_list",
|
||||
"vertex_ai_credentials",
|
||||
# Observability credentials, hosts, and project identifiers: derived
|
||||
# from the canonical ``_supported_callback_params`` allowlist so new
|
||||
# integrations are covered automatically. Sorted for stable iteration
|
||||
|
|
@ -332,6 +338,60 @@ def _check_banned_params(
|
|||
)
|
||||
|
||||
|
||||
_FALLBACK_FIELDS: tuple[str, ...] = (
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
)
|
||||
|
||||
|
||||
def _iter_fallback_field_values(request_body: Mapping[str, object]) -> Iterator[object]:
|
||||
override = request_body.get("router_settings_override")
|
||||
for source in (request_body, override):
|
||||
if isinstance(source, Mapping):
|
||||
for field in _FALLBACK_FIELDS:
|
||||
yield source.get(field)
|
||||
|
||||
|
||||
def _iter_fallback_targets(value: object, depth: int) -> Iterator[str | Mapping[str, object]]:
|
||||
if depth > 2 * litellm.ROUTER_MAX_FALLBACKS:
|
||||
raise ValueError("Rejected Request: fallback nesting exceeds the allowed validation depth.")
|
||||
if not isinstance(value, list):
|
||||
return
|
||||
for item in value:
|
||||
if isinstance(item, str):
|
||||
yield item
|
||||
elif isinstance(item, Mapping):
|
||||
values = tuple(item.values())
|
||||
if not (values and all(isinstance(v, list) for v in values)):
|
||||
yield item
|
||||
if isinstance(item.get("model"), str):
|
||||
for field in _FALLBACK_FIELDS:
|
||||
yield from _iter_fallback_targets(item.get(field), depth + 1)
|
||||
else:
|
||||
for target_list in values:
|
||||
yield from _iter_fallback_targets(target_list, depth + 1)
|
||||
|
||||
|
||||
def iter_request_fallback_targets(request_body: Mapping[str, object]) -> Iterator[str | Mapping[str, object]]:
|
||||
for value in _iter_fallback_field_values(request_body):
|
||||
yield from _iter_fallback_targets(value, 0)
|
||||
|
||||
|
||||
def _reject_url_valued_fallback_target(value: str) -> None:
|
||||
allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or []
|
||||
for candidate in provider_url_destination_candidates(value):
|
||||
if not candidate.lower().startswith(("http://", "https://")):
|
||||
continue
|
||||
if is_url_destination_allowed_by_host(candidate, allowed_hosts):
|
||||
continue
|
||||
raise ValueError(
|
||||
f"Rejected Request: URL-valued fallback destination '{value}' is not allowed. "
|
||||
"Configure custom endpoints with api_base instead, or add the destination host to "
|
||||
"`provider_url_destination_allowed_hosts` in litellm_settings."
|
||||
)
|
||||
|
||||
|
||||
def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str) -> bool:
|
||||
"""
|
||||
Check if the request body is safe.
|
||||
|
|
@ -369,6 +429,14 @@ def is_request_body_safe(request_body: dict, general_settings: dict, llm_router:
|
|||
metadata = _coerce_metadata_to_dict(request_body.get(metadata_key))
|
||||
if metadata is not None:
|
||||
_check_banned_params(metadata, general_settings, llm_router, model)
|
||||
for target in iter_request_fallback_targets(request_body):
|
||||
if isinstance(target, dict):
|
||||
_check_banned_params(target, general_settings, llm_router, model)
|
||||
target_model = target.get("model")
|
||||
if isinstance(target_model, str):
|
||||
_reject_url_valued_fallback_target(target_model)
|
||||
elif isinstance(target, str):
|
||||
_reject_url_valued_fallback_target(target)
|
||||
litellm_params = _coerce_metadata_to_dict(request_body.get("litellm_params"))
|
||||
if litellm_params is not None:
|
||||
litellm_params_metadata = _coerce_metadata_to_dict(litellm_params.get("metadata"))
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import fnmatch
|
|||
import re
|
||||
import secrets
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, Iterator, NamedTuple, List, Optional, Tuple, Union, cast
|
||||
from typing import Any, Dict, NamedTuple, List, Optional, Tuple, Union, cast
|
||||
|
||||
import fastapi
|
||||
from fastapi import HTTPException, Request, WebSocket, status
|
||||
|
|
@ -55,6 +55,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
get_model_from_request,
|
||||
get_request_route,
|
||||
get_request_route_template,
|
||||
iter_request_fallback_targets,
|
||||
normalize_request_route,
|
||||
pre_db_read_auth_checks,
|
||||
route_in_additonal_public_routes,
|
||||
|
|
@ -2694,19 +2695,11 @@ async def _enforce_key_and_fallback_model_access(
|
|||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
# Validate every fallback model name reachable by this request.
|
||||
# All three fields (``fallbacks``, ``context_window_fallbacks``,
|
||||
# ``content_policy_fallbacks``) are forwarded to the router as
|
||||
# per-request kwargs whether they appear at the top level of
|
||||
# ``request_data`` or nested under ``router_settings_override``.
|
||||
# Both surfaces must be validated against the API key's model
|
||||
# allowlist or a caller can smuggle a restricted model. VERIA-44.
|
||||
fallback_names: List[str] = []
|
||||
override_settings = request_data.get("router_settings_override")
|
||||
for _fb_key in ROUTER_FALLBACK_FIELDS:
|
||||
fallback_names.extend(iter_router_fallback_model_names(request_data.get(_fb_key)))
|
||||
if isinstance(override_settings, dict):
|
||||
fallback_names.extend(iter_router_fallback_model_names(override_settings.get(_fb_key)))
|
||||
fallback_names = tuple(
|
||||
name
|
||||
for target in iter_request_fallback_targets(request_data)
|
||||
if (name := _fallback_target_model_name(target)) is not None
|
||||
)
|
||||
|
||||
for _name in dict.fromkeys(fallback_names): # dedupe, preserve order
|
||||
await can_key_call_model(
|
||||
|
|
@ -2722,36 +2715,14 @@ async def _enforce_key_and_fallback_model_access(
|
|||
)
|
||||
|
||||
|
||||
ROUTER_FALLBACK_FIELDS: Tuple[str, ...] = (
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
)
|
||||
|
||||
|
||||
def iter_router_fallback_model_names(fallbacks: Any) -> Iterator[str]:
|
||||
"""Yield leaf model names from any of the supported fallbacks shapes.
|
||||
|
||||
Handles the simple top-level shape (``str`` or ``{"model": str}``) and
|
||||
the nested router-config shape (``[{primary: [fallback_list]}]``).
|
||||
"""
|
||||
if not isinstance(fallbacks, list):
|
||||
return
|
||||
for entry in fallbacks:
|
||||
if isinstance(entry, str):
|
||||
yield entry
|
||||
elif isinstance(entry, dict):
|
||||
if isinstance(entry.get("model"), str):
|
||||
yield entry["model"]
|
||||
continue
|
||||
for fallback_list in entry.values():
|
||||
if not isinstance(fallback_list, list):
|
||||
continue
|
||||
for m in fallback_list:
|
||||
if isinstance(m, str):
|
||||
yield m
|
||||
elif isinstance(m, dict) and isinstance(m.get("model"), str):
|
||||
yield m["model"]
|
||||
def _fallback_target_model_name(target: object) -> str | None:
|
||||
if isinstance(target, str):
|
||||
return target
|
||||
if isinstance(target, dict):
|
||||
model = target.get("model")
|
||||
if isinstance(model, str):
|
||||
return model
|
||||
return None
|
||||
|
||||
|
||||
async def _run_post_custom_auth_checks(
|
||||
|
|
|
|||
|
|
@ -19,7 +19,10 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
|||
iter_client_callback_metadata_dicts,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host
|
||||
from litellm.litellm_core_utils.url_utils import (
|
||||
is_url_destination_allowed_by_host,
|
||||
provider_url_destination_candidates,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
AddTeamCallback,
|
||||
CommonProxyErrors,
|
||||
|
|
@ -225,23 +228,26 @@ def _reject_url_valued_destinations(data: Dict[str, Any]) -> None:
|
|||
allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or []
|
||||
for field in _URL_DESTINATION_REQUEST_FIELDS:
|
||||
value = data.get(field)
|
||||
if not isinstance(value, str) or not value.startswith(("http://", "https://")):
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
if is_url_destination_allowed_by_host(value, allowed_hosts):
|
||||
continue
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "invalid_request",
|
||||
"param": field,
|
||||
"message": (
|
||||
f"URL-valued '{field}' is not allowed. Configure custom "
|
||||
"endpoints with api_base instead, or add the destination "
|
||||
"host to `provider_url_destination_allowed_hosts` in "
|
||||
"litellm_settings."
|
||||
),
|
||||
},
|
||||
)
|
||||
for candidate in provider_url_destination_candidates(value):
|
||||
if not candidate.lower().startswith(("http://", "https://")):
|
||||
continue
|
||||
if is_url_destination_allowed_by_host(candidate, allowed_hosts):
|
||||
continue
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "invalid_request",
|
||||
"param": field,
|
||||
"message": (
|
||||
f"URL-valued '{field}' is not allowed. Configure custom "
|
||||
"endpoints with api_base instead, or add the destination "
|
||||
"host to `provider_url_destination_allowed_hosts` in "
|
||||
"litellm_settings."
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _strip_untrusted_request_header_controls(
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ IGNORE_FUNCTIONS = [
|
|||
"resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input).
|
||||
"sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth.
|
||||
"_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap.
|
||||
"_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -121,6 +121,20 @@ class TestHuggingFaceEmbedding:
|
|||
assert response.usage.prompt_tokens > 0
|
||||
assert response.usage.total_tokens == response.usage.prompt_tokens
|
||||
|
||||
def test_model_name_with_https_substring_uses_api_base(self):
|
||||
api_base = "https://legit.example/embed"
|
||||
|
||||
litellm.embedding(
|
||||
model="huggingface/my-https-endpoint",
|
||||
input=["hello world"],
|
||||
input_type="embed",
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
self.mock_http.assert_called_once()
|
||||
called_url = self.mock_http.call_args[0][0]
|
||||
assert called_url == api_base
|
||||
|
||||
def test_embedding_with_sentence_similarity_task(self):
|
||||
"""Test embedding when task type is sentence-similarity (requires 2+ sentences)"""
|
||||
|
||||
|
|
|
|||
55
tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py
Normal file
55
tests/test_litellm/llms/oobabooga/chat/test_oobabooga.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
import litellm
|
||||
|
||||
MOCK_COMPLETION_RESPONSE = {
|
||||
"choices": [{"message": {"role": "assistant", "content": "hi there"}}],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
|
||||
}
|
||||
|
||||
|
||||
def _mock_post_response():
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = "ok"
|
||||
mock_response.json.return_value = MOCK_COMPLETION_RESPONSE
|
||||
return mock_response
|
||||
|
||||
|
||||
def test_model_name_with_https_substring_uses_api_base():
|
||||
api_base = "https://legit.example"
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post"
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_post_response()
|
||||
|
||||
litellm.completion(
|
||||
model="oobabooga/my-https-model",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
called_url = mock_post.call_args[0][0]
|
||||
assert called_url == f"{api_base}/v1/chat/completions"
|
||||
|
||||
|
||||
def test_url_valued_model_still_targets_that_url():
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post"
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_post_response()
|
||||
|
||||
litellm.completion(
|
||||
model="oobabooga/https://sdk-user.example",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
called_url = mock_post.call_args[0][0]
|
||||
assert called_url == "https://sdk-user.example/v1/chat/completions"
|
||||
|
|
@ -1512,7 +1512,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
|||
}
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params=dict(admin_params),
|
||||
request_kwargs={"base_url": "https://attacker.example"},
|
||||
request_kwargs={"base_url": "https://attacker.example", "api_key": "sk-caller"},
|
||||
)
|
||||
assert "aws_access_key_id" not in out
|
||||
assert "aws_secret_access_key" not in out
|
||||
|
|
@ -1540,6 +1540,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
|||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"api_key": "sk-caller",
|
||||
"organization": "org-attacker",
|
||||
"extra_body": {"attacker": "value"},
|
||||
},
|
||||
|
|
@ -1563,6 +1564,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
|||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"api_key": "sk-caller",
|
||||
"organization": "",
|
||||
"extra_body": "",
|
||||
},
|
||||
|
|
@ -1590,6 +1592,310 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
|||
assert out["api_version"] == "2026-04-01"
|
||||
assert out["api_base"] == "https://admin.upstream/v1"
|
||||
|
||||
def test_client_api_key_used_when_supplied_with_base_override(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-admin-secret",
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"api_key": "sk-client-byok",
|
||||
},
|
||||
)
|
||||
assert out["api_key"] == "sk-client-byok"
|
||||
assert "sk-admin-secret" not in str(out)
|
||||
|
||||
|
||||
_OPENAI_CHAT_RESPONSE = {
|
||||
"id": "chatcmpl-x",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
class TestClientsideBaseOverrideOutboundKey:
|
||||
"""Drive a completion through the router and assert on the outbound request
|
||||
when the caller overrides ``api_base``."""
|
||||
|
||||
def _router(self):
|
||||
from litellm import Router
|
||||
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4",
|
||||
"api_key": "sk-SERVER-CONFIG",
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _ambient_server_key(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-SERVER-ENV")
|
||||
monkeypatch.setattr(litellm, "api_key", None, raising=False)
|
||||
|
||||
def test_caller_key_override_sends_caller_key_never_server_key(self):
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
with respx.mock:
|
||||
route = respx.post("https://caller.example/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json=_OPENAI_CHAT_RESPONSE)
|
||||
)
|
||||
self._router().completion(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base="https://caller.example/v1",
|
||||
api_key="sk-CALLER",
|
||||
)
|
||||
authorization = route.calls.last.request.headers.get("authorization")
|
||||
assert authorization == "Bearer sk-CALLER"
|
||||
assert "SERVER" not in (authorization or "")
|
||||
|
||||
|
||||
def _rounds_deep_api_base_payload(rounds, field):
|
||||
"""Build a fallbacks payload with ``api_base`` on a target nested ``rounds``
|
||||
fallback-rounds deep, each round wrapped in its own grouping dict."""
|
||||
node = {"model": "leaf", "api_base": "https://attacker.example"}
|
||||
for i in range(rounds):
|
||||
node = {"model": f"m{i}", field: [{"grp": [node]}]}
|
||||
return {"model": "gpt-4", field: [{"grp": [node]}]}
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksFallbackSmuggle:
|
||||
"""``is_request_body_safe`` runs the banned-param check on every dict target
|
||||
inside the fallback lists."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _disable_url_validation(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fallback_key",
|
||||
["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"],
|
||||
)
|
||||
def test_api_base_smuggled_via_nested_fallback_is_rejected(self, fallback_key):
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
fallback_key: [
|
||||
{
|
||||
"gpt-4": [
|
||||
{"model": "evil", "api_base": "https://attacker.example"},
|
||||
]
|
||||
}
|
||||
],
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_string_only_fallbacks_are_accepted(self):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"fallbacks": [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}],
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_benign_dict_fallback_entry_is_accepted(self):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"fallbacks": [{"gpt-4": [{"model": "gpt-3.5-turbo"}]}],
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_smuggled_fallback_allowed_under_proxy_wide_opt_in(self):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"fallbacks": [
|
||||
{"gpt-4": [{"model": "byok", "api_base": "https://my-byok.example"}]}
|
||||
],
|
||||
},
|
||||
general_settings={"allow_client_side_credentials": True},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fallback_field",
|
||||
["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"],
|
||||
)
|
||||
@pytest.mark.parametrize("surface", ["top_level", "router_settings_override"])
|
||||
def test_deeply_nested_api_base_smuggle_rejected_on_both_surfaces(self, fallback_field, surface):
|
||||
nested = [
|
||||
{
|
||||
"always-fail": [
|
||||
{
|
||||
"model": "x",
|
||||
fallback_field: [
|
||||
{"x": [{"model": "deepseek-chat", "api_base": "http://attacker"}]}
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
request_body = {"model": "gpt-4"}
|
||||
if surface == "top_level":
|
||||
request_body[fallback_field] = nested
|
||||
else:
|
||||
request_body["router_settings_override"] = {fallback_field: nested}
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
is_request_body_safe(
|
||||
request_body=request_body,
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_router_settings_override_single_level_api_base_rejected(self):
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"router_settings_override": {
|
||||
"fallbacks": [{"gpt-4": [{"model": "x", "api_base": "http://attacker"}]}]
|
||||
},
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_model_less_config_dict_api_base_rejected(self):
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"fallbacks": [{"gpt-4": [{"api_base": "http://attacker"}]}],
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_nested_api_base_caught_across_router_fallback_rounds(self):
|
||||
"""An ``api_base`` target nested ``ROUTER_MAX_FALLBACKS - 1`` rounds deep
|
||||
is still reached and rejected."""
|
||||
import litellm
|
||||
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
is_request_body_safe(
|
||||
request_body=_rounds_deep_api_base_payload(litellm.ROUTER_MAX_FALLBACKS - 1, "fallbacks"),
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_grouping_only_deep_chain_is_rejected_at_depth_limit(self):
|
||||
"""A deep grouping-only chain (``{"g": [{"g": [...]}]}``) is rejected at the
|
||||
validation-depth limit rather than accepted or raising RecursionError."""
|
||||
node: object = ["safe-model"]
|
||||
for _ in range(5000):
|
||||
node = [{"grp": node}]
|
||||
with pytest.raises(ValueError, match="depth"):
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", "fallbacks": node},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_pathologically_deep_model_nesting_is_rejected(self):
|
||||
with pytest.raises(ValueError, match="depth"):
|
||||
is_request_body_safe(
|
||||
request_body=_rounds_deep_api_base_payload(5000, "fallbacks"),
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
|
||||
class TestIsRequestBodySafeRejectsUrlValuedFallback:
|
||||
@pytest.mark.parametrize("fallback_field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"])
|
||||
def test_url_valued_string_fallback_is_rejected(self, fallback_field):
|
||||
with pytest.raises(ValueError, match="URL-valued fallback"):
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
fallback_field: [{"gpt-4": ["huggingface/http://attacker.example/path"]}],
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("fallback_field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"])
|
||||
def test_url_valued_dict_model_fallback_is_rejected(self, fallback_field):
|
||||
with pytest.raises(ValueError, match="URL-valued fallback"):
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
fallback_field: [{"gpt-4": [{"model": "huggingface/http://attacker.example/path"}]}],
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
def test_ordinary_string_fallback_is_allowed(self):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", "fallbacks": [{"gpt-4": ["gpt-4-backup"]}]},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_ordinary_dict_model_fallback_is_allowed(self):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", "fallbacks": [{"gpt-4": [{"model": "gpt-4-backup"}]}]},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksEndpointTargetingFields:
|
||||
"""
|
||||
|
|
@ -1712,6 +2018,46 @@ class TestIsRequestBodySafeBlocksBedrockProjectOverride:
|
|||
)
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksVertexCredentialAlias:
|
||||
@pytest.mark.parametrize("field", ["vertex_ai_credentials"])
|
||||
def test_field_in_request_body_is_rejected(self, field):
|
||||
with pytest.raises(ValueError, match=field):
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", field: "attacker-supplied"},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("field", ["vertex_ai_credentials"])
|
||||
def test_admin_opt_in_proxy_wide_allows(self, field):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", field: "byok-supplied"},
|
||||
general_settings={"allow_client_side_credentials": True},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_legitimate_request_body_param_still_allowed(self):
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 128,
|
||||
"user": "end-user-123",
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
# ── is_request_body_safe nested-config recursion (VERIA-6) ────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -11,12 +11,22 @@ from unittest.mock import AsyncMock, patch
|
|||
import pytest
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import iter_request_fallback_targets
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_enforce_key_and_fallback_model_access,
|
||||
iter_router_fallback_model_names,
|
||||
_fallback_target_model_name,
|
||||
)
|
||||
|
||||
|
||||
def _fallback_model_names(fallbacks):
|
||||
"""Model names the auth check validates for a top-level ``fallbacks`` value."""
|
||||
return [
|
||||
name
|
||||
for target in iter_request_fallback_targets({"fallbacks": fallbacks})
|
||||
if (name := _fallback_target_model_name(target)) is not None
|
||||
]
|
||||
|
||||
|
||||
def _key_with_models(models: List[str]) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed",
|
||||
|
|
@ -26,37 +36,40 @@ def _key_with_models(models: List[str]) -> UserAPIKeyAuth:
|
|||
)
|
||||
|
||||
|
||||
# ── iter_router_fallback_model_names ─────────────────────────────────────────
|
||||
# ── fallback model-name extraction ───────────────────────────────────────────
|
||||
|
||||
|
||||
def testiter_router_fallback_model_names_router_config_shape():
|
||||
def test_fallback_model_names_router_config_shape():
|
||||
"""Router-config shape: ``[{primary: [fallback_list]}]``."""
|
||||
assert list(
|
||||
iter_router_fallback_model_names(
|
||||
[{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}]
|
||||
)
|
||||
assert _fallback_model_names(
|
||||
[{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}]
|
||||
) == ["gpt-4", "claude-3", "o1"]
|
||||
|
||||
|
||||
def testiter_router_fallback_model_names_simple_string_shape():
|
||||
def test_fallback_model_names_simple_string_shape():
|
||||
"""Simple top-level shape: list of strings."""
|
||||
assert list(iter_router_fallback_model_names(["gpt-4", "claude-3"])) == [
|
||||
assert _fallback_model_names(["gpt-4", "claude-3"]) == ["gpt-4", "claude-3"]
|
||||
|
||||
|
||||
def test_fallback_model_names_client_side_shape():
|
||||
"""ClientSideFallbackModel shape: ``[{"model": "..."}]``."""
|
||||
assert _fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}]) == [
|
||||
"gpt-4",
|
||||
"claude-3",
|
||||
]
|
||||
|
||||
|
||||
def testiter_router_fallback_model_names_client_side_shape():
|
||||
"""ClientSideFallbackModel shape: ``[{"model": "..."}]``."""
|
||||
assert list(
|
||||
iter_router_fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}])
|
||||
) == ["gpt-4", "claude-3"]
|
||||
def test_fallback_model_names_nested_deployment_fallbacks():
|
||||
"""A deployment target's own nested fallback field is unrolled too."""
|
||||
assert _fallback_model_names(
|
||||
[{"primary": [{"model": "gpt-4", "fallbacks": [{"gpt-4": ["deepseek-chat"]}]}]}]
|
||||
) == ["gpt-4", "deepseek-chat"]
|
||||
|
||||
|
||||
def testiter_router_fallback_model_names_empty_or_none():
|
||||
assert list(iter_router_fallback_model_names(None)) == []
|
||||
assert list(iter_router_fallback_model_names([])) == []
|
||||
assert list(iter_router_fallback_model_names("not a list")) == []
|
||||
def test_fallback_model_names_empty_or_none():
|
||||
assert _fallback_model_names(None) == []
|
||||
assert _fallback_model_names([]) == []
|
||||
assert _fallback_model_names("not a list") == []
|
||||
|
||||
|
||||
# ── _enforce_key_and_fallback_model_access ────────────────────────────────────
|
||||
|
|
@ -200,6 +213,98 @@ async def test_top_level_fallback_fields_validated(fallback_field):
|
|||
assert "top-level-smuggled" in seen
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_deployment_fallback_inner_model_validated():
|
||||
"""A model name nested several fallback rounds deep, inside a deployment
|
||||
target's own ``fallbacks``, is extracted and passed to can_key_call_model."""
|
||||
valid_token = _key_with_models(["gpt-3.5-turbo"])
|
||||
request_data = {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"fallbacks": [
|
||||
{
|
||||
"gpt-3.5-turbo": [
|
||||
{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"fallbacks": [{"gpt-3.5-turbo": ["deep-smuggled-model"]}],
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
seen: List[str] = []
|
||||
|
||||
async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router):
|
||||
seen.append(model)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
|
||||
side_effect=fake_can_key_call_model,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model",
|
||||
new=AsyncMock(),
|
||||
),
|
||||
):
|
||||
await _enforce_key_and_fallback_model_access(
|
||||
valid_token=valid_token,
|
||||
request_data=request_data,
|
||||
route="/v1/chat/completions",
|
||||
request=None,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
assert "deep-smuggled-model" in seen
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_less_fallback_dict_is_skipped_never_passed_as_none():
|
||||
"""A fallback target dict without a ``model`` key is skipped, never passed
|
||||
as ``None`` into can_key_call_model / is_valid_fallback_model."""
|
||||
valid_token = _key_with_models(["gpt-3.5-turbo"])
|
||||
request_data = {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"fallbacks": [
|
||||
{
|
||||
"gpt-3.5-turbo": [
|
||||
{"model": "real-fallback"},
|
||||
{"api_base": "http://attacker"},
|
||||
"string-fallback",
|
||||
]
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
seen: List[str] = []
|
||||
|
||||
async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router):
|
||||
seen.append(model)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
|
||||
side_effect=fake_can_key_call_model,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model",
|
||||
new=AsyncMock(),
|
||||
),
|
||||
):
|
||||
await _enforce_key_and_fallback_model_access(
|
||||
valid_token=valid_token,
|
||||
request_data=request_data,
|
||||
route="/v1/chat/completions",
|
||||
request=None,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
assert None not in seen
|
||||
assert seen == ["gpt-3.5-turbo", "real-fallback", "string-fallback"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_override_without_fallbacks_does_not_break_auth():
|
||||
"""``router_settings_override`` set without any fallback fields is a
|
||||
|
|
|
|||
|
|
@ -39,6 +39,46 @@ class TestRejectUrlValuedDestinations:
|
|||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail["param"] == "model"
|
||||
|
||||
def test_provider_prefixed_url_rejected(self):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_reject_url_valued_destinations(
|
||||
{"model": "huggingface/https://attacker.example/v1"}
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail["param"] == "model"
|
||||
|
||||
def test_comma_batch_smuggled_url_rejected(self):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_reject_url_valued_destinations(
|
||||
{"model": "gpt-4,huggingface/https://attacker.example/v1"}
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail["param"] == "model"
|
||||
|
||||
def test_provider_prefixed_uppercase_scheme_url_rejected(self):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_reject_url_valued_destinations(
|
||||
{"model": "huggingface/HTTPS://evil.example/v1"}
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail["param"] == "model"
|
||||
|
||||
def test_provider_prefixed_plain_model_passes(self):
|
||||
_reject_url_valued_destinations({"model": "huggingface/BAAI/bge-small-en"})
|
||||
|
||||
def test_comma_batch_plain_models_pass(self):
|
||||
_reject_url_valued_destinations({"model": "gpt-4,huggingface/BAAI/bge-small-en"})
|
||||
|
||||
def test_provider_prefixed_url_respects_allowlist(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"provider_url_destination_allowed_hosts",
|
||||
["trusted.example"],
|
||||
)
|
||||
_reject_url_valued_destinations(
|
||||
{"model": "huggingface/https://trusted.example/v1"}
|
||||
)
|
||||
|
||||
def test_url_valued_file_id_rejected(self):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_reject_url_valued_destinations(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue