mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge 6f150655c2 into e6c4580a31
This commit is contained in:
commit
78b21c1dbf
2 changed files with 77 additions and 7 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,
|
||||
|
|
@ -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
|
||||
|
|
|
|||
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