diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 4cbde551fa2..e6b83c9d6d1 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -12,6 +12,7 @@ import uuid from typing import Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -53,8 +54,8 @@ def _parse_data_url(data_url: str) -> tuple[bytes, str, str] | None: def _download_image_sync(url: str) -> tuple[bytes, str, str]: """Download image from URL synchronously.""" - client: Final = _get_httpx_client(params={"ssl_verify": False}) - response: Final = client.get(url) + client: Final = _get_httpx_client() + response: Final = safe_get(client, url) response.raise_for_status() content_type: Final = response.headers.get("content-type", "image/jpeg") @@ -65,11 +66,8 @@ def _download_image_sync(url: str) -> tuple[bytes, str, str]: async def _download_image_async(url: str) -> tuple[bytes, str, str]: """Download image from URL asynchronously.""" - client: Final = get_async_httpx_client( - llm_provider=LlmProviders.GIGACHAT, - params={"ssl_verify": False}, - ) - response: Final = await client.get(url) + client: Final = get_async_httpx_client(llm_provider=LlmProviders.GIGACHAT) + response: Final = await async_safe_get(client, url) response.raise_for_status() content_type: Final = response.headers.get("content-type", "image/jpeg") @@ -138,6 +136,8 @@ def upload_file_sync( return file_id + except SSRFError: + raise except Exception as e: verbose_logger.error("Error uploading file to GigaChat: %s", e) return None @@ -206,6 +206,8 @@ async def upload_file_async( return file_id + except SSRFError: + raise except Exception as e: verbose_logger.error("Error uploading file to GigaChat: %s", e) return None diff --git a/tests/test_litellm/llms/gigachat/test_file_handler_ssrf.py b/tests/test_litellm/llms/gigachat/test_file_handler_ssrf.py new file mode 100644 index 00000000000..9dccea4869b --- /dev/null +++ b/tests/test_litellm/llms/gigachat/test_file_handler_ssrf.py @@ -0,0 +1,68 @@ +"""SSRF guards for GigaChat multimodal image_url downloads.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.llms.gigachat import file_handler + + +def test_download_image_sync_blocks_private_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + mock_client = MagicMock() + + with patch.object(file_handler, "_get_httpx_client", return_value=mock_client): + with pytest.raises(SSRFError): + file_handler._download_image_sync("http://127.0.0.1/secret.png") + + mock_client.get.assert_not_called() + + +def test_upload_file_sync_propagates_ssrf_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + file_handler._file_cache.clear() + + with patch.object( + file_handler, + "_download_image_sync", + side_effect=SSRFError("blocked"), + ): + with pytest.raises(SSRFError, match="blocked"): + file_handler.upload_file_sync("http://169.254.169.254/latest/meta-data/") + + +@pytest.mark.asyncio +async def test_download_image_async_blocks_private_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch.object(file_handler, "get_async_httpx_client", return_value=mock_client): + with pytest.raises(SSRFError): + await file_handler._download_image_async("http://10.0.0.5/img.png") + + mock_client.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_upload_file_async_propagates_ssrf_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + file_handler._file_cache.clear() + + with patch.object( + file_handler, + "_download_image_async", + side_effect=SSRFError("blocked"), + ): + with pytest.raises(SSRFError, match="blocked"): + await file_handler.upload_file_async( + "http://169.254.169.254/latest/meta-data/" + )