mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(gigachat): fetch chat image_url via safe_get
User-controlled multimodal image URLs were downloaded with bare client.get, allowing SSRF to private/metadata targets. Route through safe_get/async_safe_get and fail closed on SSRFError. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
6a919aec6a
commit
5b963cf5e0
2 changed files with 80 additions and 4 deletions
|
|
@ -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,
|
||||
|
|
@ -52,9 +53,10 @@ 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."""
|
||||
"""Download image from URL synchronously with SSRF guards."""
|
||||
client: Final = _get_httpx_client(params={"ssl_verify": False})
|
||||
response: Final = client.get(url)
|
||||
# Chat image_url is user/LLM controlled; use the shared redirect-aware guard.
|
||||
response: Final = safe_get(client, url)
|
||||
response.raise_for_status()
|
||||
|
||||
content_type: Final = response.headers.get("content-type", "image/jpeg")
|
||||
|
|
@ -64,12 +66,12 @@ 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."""
|
||||
"""Download image from URL asynchronously with SSRF guards."""
|
||||
client: Final = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.GIGACHAT,
|
||||
params={"ssl_verify": False},
|
||||
)
|
||||
response: Final = await client.get(url)
|
||||
response: Final = await async_safe_get(client, url)
|
||||
response.raise_for_status()
|
||||
|
||||
content_type: Final = response.headers.get("content-type", "image/jpeg")
|
||||
|
|
@ -138,6 +140,9 @@ def upload_file_sync(
|
|||
|
||||
return file_id
|
||||
|
||||
except SSRFError:
|
||||
# Fail closed: do not treat blocked URLs as a soft upload miss.
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.error("Error uploading file to GigaChat: %s", e)
|
||||
return None
|
||||
|
|
@ -206,6 +211,9 @@ async def upload_file_async(
|
|||
|
||||
return file_id
|
||||
|
||||
except SSRFError:
|
||||
# Fail closed: do not treat blocked URLs as a soft upload miss.
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.error("Error uploading file to GigaChat: %s", e)
|
||||
return None
|
||||
|
|
|
|||
68
tests/test_litellm/llms/gigachat/test_file_handler_ssrf.py
Normal file
68
tests/test_litellm/llms/gigachat/test_file_handler_ssrf.py
Normal 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/"
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue