This commit is contained in:
Sash 2026-08-27 19:35:34 -05:00 committed by GitHub
commit 78b21c1dbf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 77 additions and 7 deletions

View file

@ -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

View file

@ -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/"
)