diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 401d7bb9f48..f21296ef26b 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -29,6 +29,8 @@ from litellm.types.utils import LlmProviders from ..common_utils import GeminiModelInfo +_GEMINI_FILES_HOST = "generativelanguage.googleapis.com" + class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def __init__(self): @@ -248,6 +250,13 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ if file_id.startswith(("http://", "https://")): parsed = urlparse(file_id) + if ( + parsed.scheme != "https" + or parsed.hostname != _GEMINI_FILES_HOST + or parsed.username is not None + or parsed.password is not None + ): + raise ValueError("Invalid Gemini file URL") path = parsed.path.lstrip("/") files_index = path.find("files/") if files_index != -1: @@ -260,9 +269,22 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): normalized_file_id = normalized_file_id.strip("/") if not normalized_file_id.startswith("files/"): normalized_file_id = f"files/{normalized_file_id}" + self._validate_gemini_file_name(normalized_file_id) return normalized_file_id + @staticmethod + def _validate_gemini_file_name(file_name: str) -> None: + parts = file_name.split("/") + if ( + len(parts) != 2 + or parts[0] != "files" + or not parts[1] + or parts[1] in {".", ".."} + or any(char in parts[1] for char in ("\\", "?", "#")) + ): + raise ValueError("Invalid Gemini file name") + def transform_retrieve_file_response( self, raw_response: httpx.Response, @@ -337,13 +359,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not api_key: raise ValueError("api_key is required") - # Extract file name from URI if full URI is provided - # file_id could be "files/abc123" or "https://generativelanguage.googleapis.com/v1beta/files/abc123" - if file_id.startswith("http"): - # Extract the file path from full URI - file_name = file_id.split("/v1beta/")[-1] - else: - file_name = file_id if file_id.startswith("files/") else f"files/{file_id}" + 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 557aa48550b..1591ce0e888 100644 --- a/litellm/llms/huggingface/chat/transformation.py +++ b/litellm/llms/huggingface/chat/transformation.py @@ -16,7 +16,11 @@ 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 +from ..common_utils import ( + HuggingFaceError, + _fetch_inference_provider_mapping, + validate_huggingface_model_identifier, +) logger = logging.getLogger(__name__) @@ -76,9 +80,8 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): Do not add the chat/embedding/rerank extension here. Let the handler do this. """ - if model.startswith(("http://", "https://")): - base_url = model - elif base_url is None: + validate_huggingface_model_identifier(model) + if base_url is None: base_url = os.getenv("HF_API_BASE") or os.getenv("HUGGINGFACE_API_BASE", "") return base_url @@ -95,6 +98,7 @@ 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 @@ -103,9 +107,6 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): complete_url = str(os.getenv("HF_API_BASE")) or str( os.getenv("HUGGINGFACE_API_BASE") ) - elif model.startswith(("http://", "https://")): - complete_url = model - complete_url = _build_chat_completion_url(complete_url) # Default construction with provider else: # Parse provider and model diff --git a/litellm/llms/huggingface/common_utils.py b/litellm/llms/huggingface/common_utils.py index 9ab4367c9b3..2d5f583873a 100644 --- a/litellm/llms/huggingface/common_utils.py +++ b/litellm/llms/huggingface/common_utils.py @@ -9,6 +9,20 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException HF_HUB_URL = "https://huggingface.co" +def validate_huggingface_model_identifier(model: str) -> None: + """Reject URL-valued model identifiers before provider credentials are added.""" + if "://" not in 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." + ), + ) + + class HuggingFaceError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 226f6b2ebad..4e3ac0ec0ba 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -14,7 +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 +from ..common_utils import HuggingFaceError, validate_huggingface_model_identifier from .transformation import HuggingFaceEmbeddingConfig config = HuggingFaceEmbeddingConfig() @@ -154,6 +154,7 @@ class HuggingFaceEmbedding(BaseLLM): embed_url: str, ) -> dict: data: Dict = {} + validate_huggingface_model_identifier(model) ## TRANSFORMATION ## if "sentence-transformers" in model: @@ -334,6 +335,7 @@ class HuggingFaceEmbedding(BaseLLM): headers={}, ) -> EmbeddingResponse: super().embedding() + validate_huggingface_model_identifier(model) headers = config.validate_environment( api_key=api_key, headers=headers, @@ -348,9 +350,7 @@ class HuggingFaceEmbedding(BaseLLM): ) # print_verbose(f"{model}, {task}") embed_url = "" - if "https" in model: - embed_url = model - elif api_base: + if api_base: embed_url = api_base elif "HF_API_BASE" in os.environ: embed_url = os.getenv("HF_API_BASE", "") diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 88d42cfcdcc..3a17b0a114f 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -21,7 +21,13 @@ 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, output_parser +from ..common_utils import ( + HuggingFaceError, + hf_task_list, + hf_tasks, + output_parser, + validate_huggingface_model_identifier, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -336,9 +342,8 @@ class HuggingFaceEmbeddingConfig(BaseConfig): 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: + validate_huggingface_model_identifier(model) + if api_base is not None: completion_url = api_base elif "HF_API_BASE" in os.environ: completion_url = os.getenv("HF_API_BASE", "") diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 5eb68a03d4b..bcd3584efbe 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -5,7 +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 +from ..common_utils import OobaboogaError, validate_oobabooga_model_identifier from .transformation import OobaboogaConfig oobabooga_config = OobaboogaConfig() @@ -26,6 +26,7 @@ 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={}, @@ -34,9 +35,7 @@ def completion( optional_params=optional_params, litellm_params=litellm_params, ) - if "https" in model: - completion_url = model - elif api_base: + if api_base: completion_url = api_base else: raise OobaboogaError( @@ -96,9 +95,8 @@ def embedding( encoding=None, ): # Create completion URL - if "https" in model: - embeddings_url = model - elif api_base: + validate_oobabooga_model_identifier(model) + if api_base: embeddings_url = f"{api_base}/v1/embeddings" else: raise OobaboogaError( diff --git a/litellm/llms/oobabooga/common_utils.py b/litellm/llms/oobabooga/common_utils.py index 82f8cda9511..69f2cb519b7 100644 --- a/litellm/llms/oobabooga/common_utils.py +++ b/litellm/llms/oobabooga/common_utils.py @@ -13,3 +13,16 @@ class OobaboogaError(BaseLLMException): headers: Optional[Union[dict, httpx.Headers]] = None, ): super().__init__(status_code=status_code, message=message, headers=headers) + + +def validate_oobabooga_model_identifier(model: str) -> None: + """Oobabooga endpoints must be configured with api_base, not model URLs.""" + if "://" not in 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." + ), + ) 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 2431c9a9c4f..41006fb916c 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 @@ -2,7 +2,6 @@ Test Google AI Studio (Gemini) files transformation functionality """ -import os from unittest.mock import Mock, patch import httpx @@ -93,6 +92,27 @@ 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_rejects_traversal_name(self): + litellm_params = {"api_key": "test-api-key"} + + with pytest.raises(ValueError, match="Invalid Gemini file name"): + self.handler.transform_retrieve_file_request( + file_id="files/../secrets", + optional_params={}, + litellm_params=litellm_params, + ) + @patch.dict("os.environ", {}, clear=True) @patch("litellm.llms.gemini.common_utils.get_secret_str", return_value=None) def test_transform_retrieve_file_request_missing_api_key(self, mock_get_secret): @@ -322,3 +342,16 @@ class TestGoogleAIStudioFilesTransformation: assert file_id in url 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, + ) 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 new file mode 100644 index 00000000000..c0bb56c1717 --- /dev/null +++ b/tests/test_litellm/llms/huggingface/test_huggingface_model_url_guard.py @@ -0,0 +1,76 @@ +from unittest.mock import MagicMock, patch + +import pytest + +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_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_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 new file mode 100644 index 00000000000..dab02637e50 --- /dev/null +++ b/tests/test_litellm/llms/oobabooga/test_oobabooga_model_url_guard.py @@ -0,0 +1,46 @@ +from unittest.mock import MagicMock, patch + +import pytest + +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_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()