75d1a0116e
Three review items addressed: * **Veria (Medium): SSRF via redirect.** ``fetch_validated_image_bytes`` was calling ``validate_url(url)`` once and then fetching with the default httpx client, so a 3xx to an internal IP would have been followed unvalidated. Switched to ``async_safe_get`` (the existing SSRF primitive used elsewhere in the codebase) which walks each redirect hop, re-validates, and rejects redirects to blocked networks. Default ``litellm.user_url_validation`` is True so protection is on out of the box. * **Greptile (P2): SVG can embed JS.** Removed ``image/svg+xml`` from the allowed-Content-Type set. The hardcoded response media type (``image/jpeg`` / ``image/x-icon``) means a real SVG body wouldn't render as SVG anyway in modern browsers — the allowlist entry was giving up XSS surface for no actual SVG-rendering benefit. If real SVG support is wanted later, that's a deliberate feature PR with CSP / nosniff bundled. * **Greptile (P2): cache-write OSError drops validated bytes.** When the upstream fetch succeeded but ``open(cache_path, "wb")`` raised (read-only assets dir), the bytes were discarded and the default logo was served — a silent regression for that deployment. Now serve the validated bytes inline via ``Response(...)`` as a fallback before falling back to default. Tests: - Replaced low-level mocks of ``validate_url`` with mocks of ``async_safe_get`` directly, exercising the helper's contract rather than the SSRF primitive's internals. - New ``test_rejects_svg_content_type`` confirms SVG is blocked. - ``test_get_image_cache_logic`` fixture now sets ``mock_response.is_redirect = False`` so ``async_safe_get`` doesn't treat the Mock's truthy attribute as a redirect. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
import os
|
|
import sys
|
|
from unittest import mock
|
|
|
|
# Standard path insertion
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
|
|
import pytest
|
|
import httpx
|
|
from litellm.proxy.proxy_server import app
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_image_error_handling():
|
|
"""
|
|
Test that get_image handles network errors gracefully and doesn't hang.
|
|
"""
|
|
# Set an unreachable URL
|
|
os.environ["UI_LOGO_PATH"] = "http://invalid-url-12345.com/logo.jpg"
|
|
|
|
# Clear cache
|
|
parent_dir = os.path.dirname(
|
|
os.path.dirname(
|
|
app.__file__
|
|
if hasattr(app, "__file__")
|
|
else "litellm/proxy/proxy_server.py"
|
|
)
|
|
)
|
|
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
|
|
if os.path.exists(cache_path):
|
|
os.remove(cache_path)
|
|
|
|
# Mock AsyncHTTPHandler to simulate a timeout or connection error
|
|
with mock.patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
|
|
) as mock_get:
|
|
mock_get.side_effect = httpx.ConnectError("Network is unreachable")
|
|
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
|
|
) as ac:
|
|
response = await ac.get("/get_image")
|
|
|
|
assert response.status_code == 200
|
|
assert response.headers["content-type"] == "image/jpeg"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_image_cache_logic():
|
|
"""
|
|
Test that once cached, get_image doesn't hit the network.
|
|
"""
|
|
os.environ["UI_LOGO_PATH"] = "http://example.com/logo.jpg"
|
|
|
|
# Clear cache
|
|
parent_dir = os.path.dirname(
|
|
os.path.dirname(
|
|
app.__file__
|
|
if hasattr(app, "__file__")
|
|
else "litellm/proxy/proxy_server.py"
|
|
)
|
|
)
|
|
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
|
|
if os.path.exists(cache_path):
|
|
os.remove(cache_path)
|
|
|
|
# Mock response — set headers explicitly so the Content-Type
|
|
# validation accepts the response as a legitimate image, and set
|
|
# ``is_redirect=False`` so ``async_safe_get`` doesn't try to walk
|
|
# a redirect chain.
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.content = b"fake image data"
|
|
mock_response.headers = {"content-type": "image/jpeg"}
|
|
mock_response.is_redirect = False
|
|
|
|
with mock.patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
|
|
) as mock_get:
|
|
mock_get.return_value = mock_response
|
|
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
|
|
) as ac:
|
|
# First call - should hit download logic
|
|
response1 = await ac.get("/get_image")
|
|
assert response1.status_code == 200
|
|
assert mock_get.call_count == 1
|
|
|
|
# Second call - should hit cache
|
|
response2 = await ac.get("/get_image")
|
|
assert response2.status_code == 200
|
|
# If cache works, mock_get shouldn't be called again
|
|
assert mock_get.call_count == 1
|