fix(vision): stream image and video downloads with chunk-by-chunk size cap
_download_image() and _download_video() both used client.get() + response.content, buffering the entire media body into memory before checking the size cap. A server that omits Content-Length could send an arbitrarily large payload, causing OOM. Extract _stream_download_to_file() shared helper: streams via client.stream() + aiter_bytes(), writes chunks to a temp file, enforces the running byte count against the cap after each chunk, and atomically replaces onto the destination on success. Cleans up the temp file on failure. Uses utils.atomic_replace() for cross-device/symlink safety. Malformed Content-Length values are now caught and ignored instead of crashing with ValueError; the streaming cap is the authoritative guard. Approach adapted from PR #10440 by @WuKongAI-CMU (closed as stale — 14923 commits behind, reverted 32 commits of vision_tools.py evolution including SSRF-safe client, retry classification, and lazy imports). Closes #10440
This commit is contained in:
kshitij
parent
23dce021a5
commit
b7eb97a835
|
|
@ -179,10 +179,10 @@ class TestErrorLoggingExcInfo:
|
|||
from tools.vision_tools import _download_image
|
||||
|
||||
with patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls:
|
||||
mock_client = AsyncMock()
|
||||
mock_client = MagicMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client.get = AsyncMock(side_effect=ConnectionError("network down"))
|
||||
mock_client.stream.side_effect = ConnectionError("network down")
|
||||
mock_client_cls.return_value = mock_client
|
||||
|
||||
dest = tmp_path / "image.jpg"
|
||||
|
|
@ -365,29 +365,131 @@ class TestVisionSafetyGuards:
|
|||
}
|
||||
raise AssertionError(f"unexpected URL checked: {url}")
|
||||
|
||||
class FakeResponse:
|
||||
class _FakeStreamResponse:
|
||||
url = "https://blocked.test/final.png"
|
||||
headers = {"content-length": "24"}
|
||||
content = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
async def aiter_bytes(self):
|
||||
yield b"\x89PNG\r\n\x1a\n" + b"\x00" * 16
|
||||
|
||||
class _FakeAsyncStream:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.response
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
with (
|
||||
patch("tools.vision_tools.check_website_access", side_effect=fake_check),
|
||||
patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls,
|
||||
pytest.raises(PermissionError, match="Blocked by website policy"),
|
||||
):
|
||||
mock_client = AsyncMock()
|
||||
mock_client = MagicMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client.get = AsyncMock(return_value=FakeResponse())
|
||||
mock_client.stream.return_value = _FakeAsyncStream(_FakeStreamResponse())
|
||||
mock_client_cls.return_value = mock_client
|
||||
|
||||
await _download_image("https://allowed.test/cat.png", tmp_path / "cat.png", max_retries=1)
|
||||
await _download_image(
|
||||
"https://allowed.test/cat.png", tmp_path / "cat.png", max_retries=1
|
||||
)
|
||||
|
||||
assert not (tmp_path / "cat.png").exists()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_enforces_size_cap_while_streaming(self, tmp_path):
|
||||
"""Streaming download rejects oversize payloads chunk-by-chunk and cleans up."""
|
||||
from tools.vision_tools import _download_image
|
||||
|
||||
class _FakeStreamResponse:
|
||||
url = "https://example.com/big.png"
|
||||
headers = {}
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
async def aiter_bytes(self):
|
||||
yield b"12345"
|
||||
yield b"678901"
|
||||
|
||||
class _FakeAsyncStream:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.response
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
with (
|
||||
patch("tools.vision_tools._VISION_MAX_DOWNLOAD_BYTES", 10),
|
||||
patch("tools.vision_tools.check_website_access", return_value=None),
|
||||
patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls,
|
||||
pytest.raises(ValueError, match="Image too large"),
|
||||
):
|
||||
mock_client = MagicMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client.stream.return_value = _FakeAsyncStream(_FakeStreamResponse())
|
||||
mock_client_cls.return_value = mock_client
|
||||
|
||||
await _download_image(
|
||||
"https://example.com/big.png", tmp_path / "big.png", max_retries=1
|
||||
)
|
||||
|
||||
assert not (tmp_path / "big.png").exists()
|
||||
assert not list(tmp_path.glob(".big.png.*.tmp"))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_ignores_malformed_content_length(self, tmp_path):
|
||||
"""Malformed Content-Length is ignored; streaming size cap still works."""
|
||||
from tools.vision_tools import _download_image
|
||||
|
||||
body = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16
|
||||
dest = tmp_path / "cat.png"
|
||||
|
||||
class _FakeStreamResponse:
|
||||
url = "https://example.com/cat.png"
|
||||
headers = {"content-length": "not-a-number"}
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
async def aiter_bytes(self):
|
||||
yield body[:4]
|
||||
yield body[4:]
|
||||
|
||||
class _FakeAsyncStream:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.response
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
with (
|
||||
patch("tools.vision_tools.check_website_access", return_value=None),
|
||||
patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls,
|
||||
):
|
||||
mock_client = MagicMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client.stream.return_value = _FakeAsyncStream(_FakeStreamResponse())
|
||||
mock_client_cls.return_value = mock_client
|
||||
|
||||
await _download_image("https://example.com/cat.png", dest, max_retries=1)
|
||||
|
||||
assert dest.read_bytes() == body
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# check_vision_requirements
|
||||
|
|
@ -702,15 +804,18 @@ class TestDownloadRetryClassification:
|
|||
)
|
||||
|
||||
def _make_client_raising_status(self, status_code):
|
||||
"""AsyncClient whose response.raise_for_status() raises HTTPStatusError."""
|
||||
"""AsyncClient whose stream response.raise_for_status() raises HTTPStatusError."""
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock(
|
||||
side_effect=self._status_error(status_code)
|
||||
)
|
||||
mock_client = AsyncMock()
|
||||
mock_stream = MagicMock()
|
||||
mock_stream.__aenter__ = AsyncMock(return_value=mock_response)
|
||||
mock_stream.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client = MagicMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
mock_client.stream = MagicMock(return_value=mock_stream)
|
||||
return mock_client
|
||||
|
||||
def test_is_retryable_classification(self):
|
||||
|
|
@ -746,7 +851,7 @@ class TestDownloadRetryClassification:
|
|||
"https://example.com/flaky.jpg", tmp_path / "y.jpg", max_retries=3
|
||||
)
|
||||
# All three attempts used, two backoff sleeps between them.
|
||||
assert mock_client.get.await_count == 3
|
||||
assert mock_client.stream.call_count == 3
|
||||
assert mock_sleep.await_count == 2
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -403,6 +403,77 @@ def _is_retryable_download_error(error: Exception) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
async def _stream_download_to_file(
|
||||
client,
|
||||
url: str,
|
||||
destination: Path,
|
||||
max_bytes: int,
|
||||
*,
|
||||
headers: dict,
|
||||
media_label: str = "Image",
|
||||
) -> Path:
|
||||
"""Stream an HTTP download to *destination* via a temp file with a running size cap.
|
||||
|
||||
Uses ``client.stream("GET", ...)`` so the response body is never fully
|
||||
buffered in memory — chunks are written to a temp file and the running
|
||||
byte count is checked against *max_bytes* after each chunk. On success
|
||||
the temp file is atomically replaced onto *destination*; on failure the
|
||||
temp file is deleted.
|
||||
|
||||
A ``Content-Length`` header, when present and parseable, is used for an
|
||||
early rejection before any bytes are streamed, but the streaming cap is
|
||||
the authoritative guard (servers can omit or lie about the header).
|
||||
"""
|
||||
from utils import atomic_replace
|
||||
|
||||
async with client.stream("GET", url, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
# Early rejection via Content-Length when present and valid.
|
||||
cl = response.headers.get("content-length")
|
||||
if cl:
|
||||
try:
|
||||
declared_size = int(cl)
|
||||
except ValueError:
|
||||
declared_size = None
|
||||
if declared_size is not None and declared_size > max_bytes:
|
||||
raise ValueError(
|
||||
f"{media_label} too large ({declared_size} bytes, max {max_bytes})"
|
||||
)
|
||||
|
||||
final_url = str(response.url)
|
||||
blocked = check_website_access(final_url)
|
||||
if blocked:
|
||||
raise PermissionError(blocked["message"])
|
||||
|
||||
tmp_destination = destination.with_name(
|
||||
f".{destination.name}.{uuid.uuid4().hex}.tmp"
|
||||
)
|
||||
bytes_written = 0
|
||||
try:
|
||||
with tmp_destination.open("wb") as f:
|
||||
async for chunk in response.aiter_bytes():
|
||||
if not chunk:
|
||||
continue
|
||||
bytes_written += len(chunk)
|
||||
if bytes_written > max_bytes:
|
||||
raise ValueError(
|
||||
f"{media_label} too large ({bytes_written} bytes, max {max_bytes})"
|
||||
)
|
||||
f.write(chunk)
|
||||
atomic_replace(tmp_destination, destination)
|
||||
except Exception:
|
||||
try:
|
||||
tmp_destination.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
logger.debug(
|
||||
"Could not delete partial download: %s", tmp_destination, exc_info=True
|
||||
)
|
||||
raise
|
||||
|
||||
return destination
|
||||
|
||||
|
||||
async def _download_image(image_url: str, destination: Path, max_retries: int = 3) -> Path:
|
||||
"""
|
||||
Download an image from a URL to a local destination (async) with retry logic.
|
||||
|
|
@ -447,43 +518,28 @@ async def _download_image(image_url: str, destination: Path, max_retries: int =
|
|||
|
||||
from tools.url_safety import create_ssrf_safe_async_client
|
||||
|
||||
# Download the image with appropriate headers using async httpx
|
||||
# Enable follow_redirects to handle image CDNs that redirect (e.g., Imgur, Picsum)
|
||||
# Download the image with appropriate headers using async httpx.
|
||||
# Enable follow_redirects to handle image CDNs that redirect (e.g., Imgur, Picsum).
|
||||
# SSRF: the client validates DNS at TCP connect time; event_hooks
|
||||
# validate each redirect target against private IP ranges.
|
||||
# Streaming: body is written chunk-by-chunk to a temp file so the
|
||||
# size cap bounds memory, not just disk.
|
||||
async with create_ssrf_safe_async_client(
|
||||
timeout=_VISION_DOWNLOAD_TIMEOUT,
|
||||
follow_redirects=True,
|
||||
event_hooks={"response": [_ssrf_redirect_guard]},
|
||||
) as client:
|
||||
response = await client.get(
|
||||
await _stream_download_to_file(
|
||||
client,
|
||||
image_url,
|
||||
destination,
|
||||
_VISION_MAX_DOWNLOAD_BYTES,
|
||||
headers={
|
||||
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
|
||||
"Accept": "image/*,*/*;q=0.8",
|
||||
},
|
||||
media_label="Image",
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
# Reject overly large images early via Content-Length header.
|
||||
cl = response.headers.get("content-length")
|
||||
if cl and int(cl) > _VISION_MAX_DOWNLOAD_BYTES:
|
||||
raise ValueError(
|
||||
f"Image too large ({int(cl)} bytes, max {_VISION_MAX_DOWNLOAD_BYTES})"
|
||||
)
|
||||
|
||||
final_url = str(response.url)
|
||||
blocked = check_website_access(final_url)
|
||||
if blocked:
|
||||
raise PermissionError(blocked["message"])
|
||||
|
||||
# Save the image content (double-check actual size)
|
||||
body = response.content
|
||||
if len(body) > _VISION_MAX_DOWNLOAD_BYTES:
|
||||
raise ValueError(
|
||||
f"Image too large ({len(body)} bytes, max {_VISION_MAX_DOWNLOAD_BYTES})"
|
||||
)
|
||||
destination.write_bytes(body)
|
||||
|
||||
return destination
|
||||
except Exception as e:
|
||||
|
|
@ -1607,32 +1663,17 @@ async def _download_video(video_url: str, destination: Path, max_retries: int =
|
|||
follow_redirects=True,
|
||||
event_hooks={"response": [_ssrf_redirect_guard]},
|
||||
) as client:
|
||||
response = await client.get(
|
||||
await _stream_download_to_file(
|
||||
client,
|
||||
video_url,
|
||||
destination,
|
||||
_MAX_VIDEO_BASE64_BYTES,
|
||||
headers={
|
||||
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
|
||||
"Accept": "video/*,*/*;q=0.8",
|
||||
},
|
||||
media_label="Video",
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
cl = response.headers.get("content-length")
|
||||
if cl and int(cl) > _MAX_VIDEO_BASE64_BYTES:
|
||||
raise ValueError(
|
||||
f"Video too large ({int(cl)} bytes, max {_MAX_VIDEO_BASE64_BYTES})"
|
||||
)
|
||||
|
||||
final_url = str(response.url)
|
||||
blocked = check_website_access(final_url)
|
||||
if blocked:
|
||||
raise PermissionError(blocked["message"])
|
||||
|
||||
body = response.content
|
||||
if len(body) > _MAX_VIDEO_BASE64_BYTES:
|
||||
raise ValueError(
|
||||
f"Video too large ({len(body)} bytes, max {_MAX_VIDEO_BASE64_BYTES})"
|
||||
)
|
||||
destination.write_bytes(body)
|
||||
|
||||
return destination
|
||||
except Exception as e:
|
||||
|
|
|
|||
Loading…
Reference in New Issue