From 001253c47865c5409fedecfdb09efe36d6837178 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sat, 2 May 2026 02:17:46 +0000 Subject: [PATCH] chore(proxy): move URL-valued model/file_id guard from SDK to proxy The previous per-provider guards in HuggingFace, Oobabooga, and Gemini files lived in the SDK layer, breaking SDK callers who legitimately pass URL-valued model identifiers. Move the check to the proxy boundary in add_litellm_data_to_request so SDK users keep working while proxy users default-deny URL-valued model and file_id, with admin opt-in via litellm.provider_url_destination_allowed_hosts. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/llms/gemini/files/transformation.py | 38 +---- .../llms/huggingface/chat/transformation.py | 13 +- litellm/llms/huggingface/common_utils.py | 30 ---- litellm/llms/huggingface/embedding/handler.py | 28 +--- .../huggingface/embedding/transformation.py | 12 +- litellm/llms/oobabooga/chat/oobabooga.py | 12 +- litellm/llms/oobabooga/common_utils.py | 30 ---- litellm/proxy/litellm_pre_call_utils.py | 43 +++++- .../files/test_gemini_files_transformation.py | 77 ---------- .../test_huggingface_model_url_guard.py | 134 ----------------- .../test_oobabooga_model_url_guard.py | 111 -------------- .../test_provider_url_destination_guard.py | 139 ++++++++++++++++++ 12 files changed, 201 insertions(+), 466 deletions(-) delete mode 100644 tests/test_litellm/llms/huggingface/test_huggingface_model_url_guard.py delete mode 100644 tests/test_litellm/llms/oobabooga/test_oobabooga_model_url_guard.py create mode 100644 tests/test_litellm/proxy/test_provider_url_destination_guard.py diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 00822da43e..63a383ebd3 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -12,11 +12,8 @@ import httpx from openai.types.file_deleted import FileDeleted from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import ( - encode_url_path_segment, - is_url_destination_allowed_by_host, -) from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, @@ -228,21 +225,20 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not api_key: raise ValueError("api_key is required") + file_part = self._normalize_gemini_file_id(file_id) + api_base = ( self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" ) api_base = api_base.rstrip("/") - file_part = self._normalize_gemini_file_id(file_id, api_base=api_base) url = f"{api_base}/v1beta/{file_part}" # API key is passed via x-goog-api-key header (set in validate_environment) return url, {} - def _normalize_gemini_file_id( - self, file_id: str, api_base: Optional[str] = None - ) -> str: + def _normalize_gemini_file_id(self, file_id: str) -> str: """ Normalize file identifier into `files/{id}` form. @@ -253,14 +249,6 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ if file_id.startswith(("http://", "https://")): parsed = urlparse(file_id) - if ( - parsed.username is not None - or parsed.password is not None - or not self._is_allowed_gemini_file_url( - file_url=file_id, api_base=api_base - ) - ): - raise ValueError("Invalid Gemini file URL") path = parsed.path.lstrip("/") files_index = path.find("files/") if files_index != -1: @@ -280,20 +268,6 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): return f"files/{encoded_file_id}" - @staticmethod - def _is_allowed_gemini_file_url( - file_url: str, api_base: Optional[str] = None - ) -> bool: - import litellm - - allowed_hosts: List[str] = [] - if api_base: - allowed_hosts.append(api_base) - allowed_hosts.extend( - getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] - ) - return is_url_destination_allowed_by_host(file_url, allowed_hosts) - def transform_retrieve_file_response( self, raw_response: httpx.Response, @@ -368,8 +342,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not api_key: raise ValueError("api_key is required") - api_base = api_base.rstrip("/") - file_name = self._normalize_gemini_file_id(file_id, api_base=api_base) + # Normalize and encode the file name before interpolating it into the URL. + file_name = self._normalize_gemini_file_id(file_id) # Construct the delete URL url = f"{api_base}/v1beta/{file_name}" diff --git a/litellm/llms/huggingface/chat/transformation.py b/litellm/llms/huggingface/chat/transformation.py index d138ade51e..557aa48550 100644 --- a/litellm/llms/huggingface/chat/transformation.py +++ b/litellm/llms/huggingface/chat/transformation.py @@ -16,12 +16,7 @@ else: from litellm.llms.base_llm.chat.transformation import BaseLLMException from ...openai.chat.gpt_transformation import OpenAIGPTConfig -from ..common_utils import ( - HuggingFaceError, - _fetch_inference_provider_mapping, - is_url_model_destination, - validate_huggingface_model_identifier, -) +from ..common_utils import HuggingFaceError, _fetch_inference_provider_mapping logger = logging.getLogger(__name__) @@ -81,8 +76,7 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): Do not add the chat/embedding/rerank extension here. Let the handler do this. """ - validate_huggingface_model_identifier(model) - if is_url_model_destination(model): + if model.startswith(("http://", "https://")): base_url = model elif base_url is None: base_url = os.getenv("HF_API_BASE") or os.getenv("HUGGINGFACE_API_BASE", "") @@ -101,7 +95,6 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): Get the complete URL for the API call. For provider-specific routing through huggingface """ - validate_huggingface_model_identifier(model) # Check if api_base is provided if api_base is not None: complete_url = api_base @@ -110,7 +103,7 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): complete_url = str(os.getenv("HF_API_BASE")) or str( os.getenv("HUGGINGFACE_API_BASE") ) - elif is_url_model_destination(model): + elif model.startswith(("http://", "https://")): complete_url = model complete_url = _build_chat_completion_url(complete_url) # Default construction with provider diff --git a/litellm/llms/huggingface/common_utils.py b/litellm/llms/huggingface/common_utils.py index 9a8c4895c9..9ab4367c9b 100644 --- a/litellm/llms/huggingface/common_utils.py +++ b/litellm/llms/huggingface/common_utils.py @@ -4,41 +4,11 @@ from typing import Literal, Optional, Union import httpx -from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host from litellm.llms.base_llm.chat.transformation import BaseLLMException HF_HUB_URL = "https://huggingface.co" -def is_url_model_destination(model: str) -> bool: - return model.startswith(("http://", "https://")) - - -def _is_url_model_destination_allowed(model: str) -> bool: - import litellm - - allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] - return is_url_destination_allowed_by_host(model, allowed_hosts) - - -def validate_huggingface_model_identifier(model: str) -> None: - """Reject URL-valued model identifiers before provider credentials are added.""" - if "://" not in model: - return - if is_url_model_destination(model) and _is_url_model_destination_allowed(model): - return - raise HuggingFaceError( - status_code=400, - message=( - "Invalid Hugging Face model identifier. Configure custom endpoints with " - "api_base or HF_API_BASE/HUGGINGFACE_API_BASE instead of passing a URL " - "as the model. To keep legacy URL-valued models for trusted endpoints, " - "add the destination host or origin to " - "`provider_url_destination_allowed_hosts` in litellm_settings." - ), - ) - - class HuggingFaceError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 2f7232a726..226f6b2eba 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -14,11 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.types.utils import EmbeddingResponse from ...base import BaseLLM -from ..common_utils import ( - HuggingFaceError, - is_url_model_destination, - validate_huggingface_model_identifier, -) +from ..common_utils import HuggingFaceError from .transformation import HuggingFaceEmbeddingConfig config = HuggingFaceEmbeddingConfig() @@ -158,8 +154,6 @@ class HuggingFaceEmbedding(BaseLLM): embed_url: str, ) -> dict: data: Dict = {} - validate_huggingface_model_identifier(model) - model_uses_url_destination = is_url_model_destination(model) ## TRANSFORMATION ## if "sentence-transformers" in model: @@ -175,12 +169,8 @@ class HuggingFaceEmbedding(BaseLLM): task_type = optional_params.pop("input_type", None) if call_type == "sync": - hf_task = ( - task_type - if model_uses_url_destination - else get_hf_task_embedding_for_model( - model=model, task_type=task_type, api_base=HF_HUB_URL - ) + hf_task = get_hf_task_embedding_for_model( + model=model, task_type=task_type, api_base=HF_HUB_URL ) elif call_type == "async": return self._async_transform_input( @@ -344,7 +334,6 @@ class HuggingFaceEmbedding(BaseLLM): headers={}, ) -> EmbeddingResponse: super().embedding() - validate_huggingface_model_identifier(model) headers = config.validate_environment( api_key=api_key, headers=headers, @@ -354,17 +343,12 @@ class HuggingFaceEmbedding(BaseLLM): litellm_params=litellm_params, ) task_type = optional_params.get("input_type", None) - model_uses_url_destination = is_url_model_destination(model) - task = ( - task_type - if model_uses_url_destination - else get_hf_task_embedding_for_model( - model=model, task_type=task_type, api_base=HF_HUB_URL - ) + 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 model_uses_url_destination: + if "https" in model: embed_url = model elif api_base: embed_url = api_base diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index ce8319e82d..88d42cfcdc 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -21,14 +21,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import token_counter -from ..common_utils import ( - HuggingFaceError, - hf_task_list, - hf_tasks, - is_url_model_destination, - output_parser, - validate_huggingface_model_identifier, -) +from ..common_utils import HuggingFaceError, hf_task_list, hf_tasks, output_parser if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -343,8 +336,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): Do not add the chat/embedding/rerank extension here. Let the handler do this. """ - validate_huggingface_model_identifier(model) - if is_url_model_destination(model): + if "https" in model: completion_url = model elif api_base is not None: completion_url = api_base diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 090ee11d1a..5eb68a03d4 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -5,11 +5,7 @@ import litellm from litellm.llms.custom_httpx.http_handler import _get_httpx_client from litellm.utils import EmbeddingResponse, ModelResponse, Usage -from ..common_utils import ( - OobaboogaError, - is_url_model_destination, - validate_oobabooga_model_identifier, -) +from ..common_utils import OobaboogaError from .transformation import OobaboogaConfig oobabooga_config = OobaboogaConfig() @@ -30,7 +26,6 @@ def completion( logger_fn=None, default_max_tokens_to_sample=None, ): - validate_oobabooga_model_identifier(model) headers = oobabooga_config.validate_environment( api_key=api_key, headers={}, @@ -39,7 +34,7 @@ def completion( optional_params=optional_params, litellm_params=litellm_params, ) - if is_url_model_destination(model): + if "https" in model: completion_url = model elif api_base: completion_url = api_base @@ -101,8 +96,7 @@ def embedding( encoding=None, ): # Create completion URL - validate_oobabooga_model_identifier(model) - if is_url_model_destination(model): + if "https" in model: embeddings_url = model elif api_base: embeddings_url = f"{api_base}/v1/embeddings" diff --git a/litellm/llms/oobabooga/common_utils.py b/litellm/llms/oobabooga/common_utils.py index 6c863ea555..82f8cda951 100644 --- a/litellm/llms/oobabooga/common_utils.py +++ b/litellm/llms/oobabooga/common_utils.py @@ -2,7 +2,6 @@ from typing import Optional, Union import httpx -from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -14,32 +13,3 @@ class OobaboogaError(BaseLLMException): headers: Optional[Union[dict, httpx.Headers]] = None, ): super().__init__(status_code=status_code, message=message, headers=headers) - - -def is_url_model_destination(model: str) -> bool: - return model.startswith(("http://", "https://")) - - -def _is_url_model_destination_allowed(model: str) -> bool: - import litellm - - allowed_hosts = getattr(litellm, "provider_url_destination_allowed_hosts", []) or [] - return is_url_destination_allowed_by_host(model, allowed_hosts) - - -def validate_oobabooga_model_identifier(model: str) -> None: - """Oobabooga endpoints must be configured with api_base, not model URLs.""" - if "://" not in model: - return - if is_url_model_destination(model) and _is_url_model_destination_allowed(model): - return - raise OobaboogaError( - status_code=400, - message=( - "Invalid Oobabooga model identifier. Configure the endpoint with " - "api_base instead of passing a URL as the model. To keep legacy " - "URL-valued models for trusted endpoints, add the destination host " - "or origin to `provider_url_destination_allowed_hosts` in " - "litellm_settings." - ), - ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3077efe116..52bdb95fad 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -5,7 +5,7 @@ import time from collections import OrderedDict from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union -from fastapi import Request +from fastapi import HTTPException, Request from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers @@ -14,6 +14,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.litellm_core_utils.credential_accessor import CredentialAccessor 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.proxy._types import ( AddTeamCallback, CommonProxyErrors, @@ -154,6 +155,45 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY = ( "allow_client_message_redaction_opt_out" ) +# Request fields whose value, when URL-valued, becomes the outbound destination +# for a provider call. Letting a proxy caller pin the destination is an SSRF +# primitive (HuggingFace/Oobabooga `model`, Gemini files `file_id`); guard +# them centrally so SDK users keep working but proxy users default-deny. +_URL_DESTINATION_REQUEST_FIELDS = ("model", "file_id") + + +def _reject_url_valued_destinations(data: Dict[str, Any]) -> None: + """Reject URL-valued ``model``/``file_id`` unless admin-allowlisted. + + Some providers (HuggingFace, Oobabooga, Gemini files) accept a URL in the + identifier field and use it as the outbound destination. On the proxy that + is an SSRF primitive — a low-privilege caller can point traffic at any + host the proxy can reach, including internal services. Reject here at the + proxy boundary so SDK users (who legitimately pass URL-valued identifiers) + are unaffected, while admins can opt specific hosts back in via + ``litellm.provider_url_destination_allowed_hosts``. + """ + 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://")): + 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." + ), + }, + ) + def _strip_untrusted_request_header_controls( headers: Any, @@ -1105,6 +1145,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915 if _allow_client_mock_response and _internal_key in _CLIENT_MOCK_CONTROL_FIELDS: continue data.pop(_internal_key, None) + _reject_url_valued_destinations(data) # Strip spoofable auth metadata from user-supplied metadata dict _user_metadata = data.get("metadata") if isinstance(_user_metadata, dict): diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py index 33fe5d24d7..a2f9572468 100644 --- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py +++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py @@ -7,7 +7,6 @@ from unittest.mock import Mock, patch import httpx import pytest -import litellm from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler from litellm.types.llms.openai import OpenAIFileObject @@ -93,54 +92,6 @@ class TestGoogleAIStudioFilesTransformation: assert "key=" not in url assert params == {} - def test_transform_retrieve_file_request_rejects_untrusted_full_url(self): - file_id = "https://attacker.example/v1beta/files/test123" - litellm_params = {"api_key": "test-api-key"} - - with pytest.raises(ValueError, match="Invalid Gemini file URL"): - self.handler.transform_retrieve_file_request( - file_id=file_id, - optional_params={}, - litellm_params=litellm_params, - ) - - def test_transform_retrieve_file_request_allows_full_url_matching_api_base(self): - file_id = "https://custom-gemini.example/v1beta/files/test123" - litellm_params = { - "api_key": "test-api-key", - "api_base": "https://custom-gemini.example", - } - - url, params = self.handler.transform_retrieve_file_request( - file_id=file_id, - optional_params={}, - litellm_params=litellm_params, - ) - - assert url == "https://custom-gemini.example/v1beta/files/test123" - assert params == {} - - def test_transform_retrieve_file_request_allows_full_url_when_host_allowlisted( - self, - monkeypatch, - ): - monkeypatch.setattr( - litellm, - "provider_url_destination_allowed_hosts", - ["trusted-gemini.example"], - ) - file_id = "https://trusted-gemini.example/v1beta/files/test123" - litellm_params = {"api_key": "test-api-key"} - - url, params = self.handler.transform_retrieve_file_request( - file_id=file_id, - optional_params={}, - litellm_params=litellm_params, - ) - - assert url == "https://generativelanguage.googleapis.com/v1beta/files/test123" - assert params == {} - def test_transform_retrieve_file_request_encodes_file_id_path_segment(self): file_id = "files/../../models/gemini-pro?x=1#frag" litellm_params = {"api_key": "test-api-key"} @@ -393,34 +344,6 @@ class TestGoogleAIStudioFilesTransformation: assert "generativelanguage.googleapis.com" in url assert params == {} - def test_transform_delete_file_request_rejects_untrusted_full_url(self): - litellm_params = { - "api_key": "test-api-key", - "api_base": "https://generativelanguage.googleapis.com", - } - - with pytest.raises(ValueError, match="Invalid Gemini file URL"): - self.handler.transform_delete_file_request( - file_id="https://attacker.example/v1beta/files/test123", - optional_params={}, - litellm_params=litellm_params, - ) - - def test_transform_delete_file_request_allows_full_url_matching_api_base(self): - litellm_params = { - "api_key": "test-api-key", - "api_base": "https://custom-gemini.example", - } - - url, params = self.handler.transform_delete_file_request( - file_id="https://custom-gemini.example/v1beta/files/test123", - optional_params={}, - litellm_params=litellm_params, - ) - - assert url == "https://custom-gemini.example/v1beta/files/test123" - assert params == {} - def test_transform_delete_file_request_encodes_file_id_path_segment(self): file_id = "files/../../models/gemini-pro?x=1#frag" litellm_params = { diff --git a/tests/test_litellm/llms/huggingface/test_huggingface_model_url_guard.py b/tests/test_litellm/llms/huggingface/test_huggingface_model_url_guard.py deleted file mode 100644 index 27e3b855de..0000000000 --- a/tests/test_litellm/llms/huggingface/test_huggingface_model_url_guard.py +++ /dev/null @@ -1,134 +0,0 @@ -from unittest.mock import MagicMock, patch - -import pytest - -import litellm -from litellm.llms.huggingface.chat.transformation import HuggingFaceChatConfig -from litellm.llms.huggingface.common_utils import HuggingFaceError -from litellm.llms.huggingface.embedding.handler import HuggingFaceEmbedding -from litellm.llms.huggingface.embedding.transformation import ( - HuggingFaceEmbeddingConfig, -) - - -def test_huggingface_chat_rejects_url_valued_model(): - config = HuggingFaceChatConfig() - - with pytest.raises(HuggingFaceError) as exc_info: - config.get_complete_url( - api_base=None, - api_key="hf-secret", - model="https://attacker.example/v1", - optional_params={}, - litellm_params={}, - ) - - assert exc_info.value.status_code == 400 - - -def test_huggingface_chat_allows_legacy_url_model_when_host_is_allowlisted( - monkeypatch, -): - monkeypatch.setattr( - litellm, "provider_url_destination_allowed_hosts", ["trusted.example"] - ) - config = HuggingFaceChatConfig() - - complete_url = config.get_complete_url( - api_base=None, - api_key="hf-secret", - model="https://trusted.example", - optional_params={}, - litellm_params={}, - ) - - assert complete_url == "https://trusted.example/v1/chat/completions" - - -def test_huggingface_chat_rejects_url_model_when_host_is_not_allowlisted( - monkeypatch, -): - monkeypatch.setattr( - litellm, "provider_url_destination_allowed_hosts", ["trusted.example"] - ) - config = HuggingFaceChatConfig() - - with pytest.raises(HuggingFaceError) as exc_info: - config.get_complete_url( - api_base=None, - api_key="hf-secret", - model="https://other.example", - optional_params={}, - litellm_params={}, - ) - - assert exc_info.value.status_code == 400 - - -def test_huggingface_chat_keeps_explicit_api_base_for_custom_endpoints(): - config = HuggingFaceChatConfig() - - complete_url = config.get_complete_url( - api_base="https://admin-configured.example", - api_key="hf-secret", - model="huggingface/mistral", - optional_params={}, - litellm_params={}, - ) - - assert complete_url == "https://admin-configured.example/v1/chat/completions" - - -def test_huggingface_embedding_config_rejects_url_valued_model(): - config = HuggingFaceEmbeddingConfig() - - with pytest.raises(HuggingFaceError) as exc_info: - config.get_api_base( - api_base=None, - model="prefixhttps://attacker.example/embeddings", - ) - - assert exc_info.value.status_code == 400 - - -def test_huggingface_embedding_config_allows_legacy_url_model_when_origin_is_allowlisted( - monkeypatch, -): - monkeypatch.setattr( - litellm, - "provider_url_destination_allowed_hosts", - ["https://trusted.example"], - ) - config = HuggingFaceEmbeddingConfig() - - api_base = config.get_api_base( - api_base=None, - model="https://trusted.example/embeddings", - ) - - assert api_base == "https://trusted.example/embeddings" - - -def test_huggingface_embedding_handler_rejects_before_task_lookup(): - handler = HuggingFaceEmbedding() - logging_obj = MagicMock() - encoding = MagicMock() - encoding.encode.return_value = [] - - with patch( - "litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model" - ) as mock_task_lookup: - with pytest.raises(HuggingFaceError) as exc_info: - handler.embedding( - model="https://attacker.example/embeddings", - input=["hello"], - model_response=MagicMock(), - optional_params={}, - litellm_params={}, - logging_obj=logging_obj, - encoding=encoding, - api_key="hf-secret", - ) - - assert exc_info.value.status_code == 400 - mock_task_lookup.assert_not_called() diff --git a/tests/test_litellm/llms/oobabooga/test_oobabooga_model_url_guard.py b/tests/test_litellm/llms/oobabooga/test_oobabooga_model_url_guard.py deleted file mode 100644 index 5ed0cdca03..0000000000 --- a/tests/test_litellm/llms/oobabooga/test_oobabooga_model_url_guard.py +++ /dev/null @@ -1,111 +0,0 @@ -from unittest.mock import MagicMock, patch - -import pytest - -import litellm -from litellm.llms.oobabooga.chat.oobabooga import completion, embedding -from litellm.llms.oobabooga.common_utils import OobaboogaError - - -def test_oobabooga_completion_rejects_url_valued_model_before_request(): - with patch("litellm.llms.oobabooga.chat.oobabooga._get_httpx_client") as mock_get: - with pytest.raises(OobaboogaError) as exc_info: - completion( - model="https://attacker.example/v1", - messages=[], - api_base="https://admin-configured.example", - model_response=MagicMock(), - print_verbose=MagicMock(), - encoding=MagicMock(), - api_key="ooba-secret", - logging_obj=MagicMock(), - optional_params={}, - litellm_params={}, - ) - - assert exc_info.value.status_code == 400 - mock_get.assert_not_called() - - -def test_oobabooga_completion_allows_legacy_url_model_when_host_is_allowlisted( - monkeypatch, -): - monkeypatch.setattr( - litellm, "provider_url_destination_allowed_hosts", ["trusted.example"] - ) - response = MagicMock() - client = MagicMock() - client.post.return_value = response - - with patch( - "litellm.llms.oobabooga.chat.oobabooga._get_httpx_client", - return_value=client, - ): - with patch( - "litellm.llms.oobabooga.chat.oobabooga.oobabooga_config.transform_response", - return_value="ok", - ): - result = completion( - model="https://trusted.example", - messages=[], - api_base=None, - model_response=MagicMock(), - print_verbose=MagicMock(), - encoding=MagicMock(), - api_key="ooba-secret", - logging_obj=MagicMock(), - optional_params={}, - litellm_params={}, - ) - - assert result == "ok" - client.post.assert_called_once() - assert ( - client.post.call_args.args[0] == "https://trusted.example/v1/chat/completions" - ) - - -def test_oobabooga_completion_rejects_url_model_when_host_is_not_allowlisted( - monkeypatch, -): - monkeypatch.setattr( - litellm, "provider_url_destination_allowed_hosts", ["trusted.example"] - ) - - with patch("litellm.llms.oobabooga.chat.oobabooga._get_httpx_client") as mock_get: - with pytest.raises(OobaboogaError) as exc_info: - completion( - model="https://other.example", - messages=[], - api_base=None, - model_response=MagicMock(), - print_verbose=MagicMock(), - encoding=MagicMock(), - api_key="ooba-secret", - logging_obj=MagicMock(), - optional_params={}, - litellm_params={}, - ) - - assert exc_info.value.status_code == 400 - mock_get.assert_not_called() - - -def test_oobabooga_embedding_rejects_url_valued_model_before_request(): - with patch( - "litellm.llms.oobabooga.chat.oobabooga.litellm.module_level_client.post" - ) as mock_post: - with pytest.raises(OobaboogaError) as exc_info: - embedding( - model="prefixhttps://attacker.example/embeddings", - input=["hello"], - model_response=MagicMock(), - api_key="ooba-secret", - api_base="https://admin-configured.example", - logging_obj=MagicMock(), - optional_params={}, - encoding=MagicMock(), - ) - - assert exc_info.value.status_code == 400 - mock_post.assert_not_called() diff --git a/tests/test_litellm/proxy/test_provider_url_destination_guard.py b/tests/test_litellm/proxy/test_provider_url_destination_guard.py new file mode 100644 index 0000000000..51cd76105d --- /dev/null +++ b/tests/test_litellm/proxy/test_provider_url_destination_guard.py @@ -0,0 +1,139 @@ +"""Proxy-level guard against URL-valued ``model`` / ``file_id`` request fields. + +Some providers (HuggingFace, Oobabooga, Gemini files) accept a URL in the +identifier field and use it as the outbound destination. On the proxy that is +an SSRF primitive — guarded centrally in ``litellm_pre_call_utils`` so SDK +users keep working but proxy users default-deny. +""" + +import os +import sys +from unittest.mock import MagicMock + +import pytest +from fastapi import HTTPException, Request + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import ( + _reject_url_valued_destinations, + add_litellm_data_to_request, +) + +sys.path.insert(0, os.path.abspath("../../..")) + + +class TestRejectUrlValuedDestinations: + def test_plain_model_passes(self): + _reject_url_valued_destinations({"model": "gpt-4"}) + + def test_plain_file_id_passes(self): + _reject_url_valued_destinations({"file_id": "files/abc123"}) + + def test_no_destination_field_passes(self): + _reject_url_valued_destinations({"messages": [{"role": "user"}]}) + + def test_url_valued_model_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations({"model": "https://attacker.example/v1"}) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model" + + def test_url_valued_file_id_rejected(self): + with pytest.raises(HTTPException) as exc_info: + _reject_url_valued_destinations( + {"file_id": "https://attacker.example/v1beta/files/abc"} + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "file_id" + + def test_http_scheme_also_rejected(self): + with pytest.raises(HTTPException): + _reject_url_valued_destinations({"model": "http://10.0.0.1:8080/v1"}) + + def test_non_string_value_ignored(self): + # Defensive: malformed inputs (list, dict, None) shouldn't crash here; + # downstream Pydantic validation handles the type error. + _reject_url_valued_destinations({"model": None}) + _reject_url_valued_destinations({"model": 42}) + _reject_url_valued_destinations({"file_id": ["a", "b"]}) + + def test_allowlisted_host_passes(self, monkeypatch): + monkeypatch.setattr( + litellm, + "provider_url_destination_allowed_hosts", + ["trusted.example"], + ) + _reject_url_valued_destinations({"model": "https://trusted.example/v1"}) + + def test_allowlisted_origin_rejects_mismatched_scheme(self, monkeypatch): + monkeypatch.setattr( + litellm, + "provider_url_destination_allowed_hosts", + ["https://trusted.example"], + ) + _reject_url_valued_destinations({"model": "https://trusted.example/v1"}) + with pytest.raises(HTTPException): + _reject_url_valued_destinations({"model": "http://trusted.example/v1"}) + + def test_allowlisted_host_port_rejects_other_ports(self, monkeypatch): + monkeypatch.setattr( + litellm, + "provider_url_destination_allowed_hosts", + ["trusted.example:8443"], + ) + _reject_url_valued_destinations({"model": "https://trusted.example:8443/v1"}) + with pytest.raises(HTTPException): + _reject_url_valued_destinations({"model": "https://trusted.example/v1"}) + + def test_userinfo_in_url_rejected_even_when_host_allowlisted(self, monkeypatch): + # Embedded credentials in URL are an exfil channel — must never pass. + monkeypatch.setattr( + litellm, + "provider_url_destination_allowed_hosts", + ["trusted.example"], + ) + with pytest.raises(HTTPException): + _reject_url_valued_destinations( + {"model": "https://user:pass@trusted.example/v1"} + ) + + +def _make_request_mock() -> Request: + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + return request_mock + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_rejects_url_valued_model(): + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + data = {"model": "https://attacker.example/v1", "messages": []} + + with pytest.raises(HTTPException) as exc_info: + await add_litellm_data_to_request( + data=data, + request=_make_request_mock(), + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["param"] == "model"