diff --git a/tests/tools/test_video_analyze.py b/tests/tools/test_video_analyze.py index 62987d96b208b..a8864393d5b69 100644 --- a/tests/tools/test_video_analyze.py +++ b/tests/tools/test_video_analyze.py @@ -21,6 +21,53 @@ ) +class _FakeVideoStreamResponse: + def __init__(self, chunks, *, headers=None, url="https://example.com/final.mp4"): + self._chunks = list(chunks) + self.headers = headers or {} + self.url = url + self.content_accessed = False + + @property + def content(self): + self.content_accessed = True + raise AssertionError("download should stream chunks instead of reading content") + + def raise_for_status(self): + return None + + async def aiter_bytes(self): + for chunk in self._chunks: + yield chunk + + +class _FakeVideoStreamContext: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _FakeVideoAsyncClient: + def __init__(self, response): + self.response = response + self.stream_calls = [] + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def stream(self, method, url, **kwargs): + self.stream_calls.append((method, url, kwargs)) + return _FakeVideoStreamContext(self.response) + + # --------------------------------------------------------------------------- # _detect_video_mime_type # --------------------------------------------------------------------------- @@ -309,6 +356,57 @@ async def capture_llm(**kwargs): assert content[1]["video_url"]["url"].startswith("data:video/mp4;base64,") +class TestDownloadVideo: + """Remote video download safety checks.""" + + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_download_streams_video_without_loading_full_response(self, tmp_path): + from tools.vision_tools import _download_video + + response = _FakeVideoStreamResponse( + [b"\x00\x00\x00\x18ftyp", b"video-bytes"], + headers={"content-length": "18"}, + url="https://example.com/final.mp4", + ) + + with ( + patch("tools.vision_tools.check_website_access", return_value=None), + patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls, + ): + mock_client_cls.return_value = _FakeVideoAsyncClient(response) + dest = self._run( + _download_video("https://example.com/clip.mp4", tmp_path / "clip.mp4", max_retries=1) + ) + + assert dest.read_bytes() == b"\x00\x00\x00\x18ftypvideo-bytes" + assert response.content_accessed is False + + def test_download_rejects_video_when_stream_exceeds_limit(self, tmp_path): + from tools.vision_tools import _download_video + + response = _FakeVideoStreamResponse( + [b"12345", b"67890", b"!"], + headers={}, + url="https://example.com/final.mp4", + ) + + with ( + patch("tools.vision_tools._MAX_VIDEO_BASE64_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="Video too large"), + ): + mock_client_cls.return_value = _FakeVideoAsyncClient(response) + self._run( + _download_video("https://example.com/clip.mp4", tmp_path / "clip.mp4", max_retries=1) + ) + + assert not (tmp_path / "clip.mp4").exists() + assert response.content_accessed is False + + # --------------------------------------------------------------------------- # Toolset registration # --------------------------------------------------------------------------- diff --git a/tests/tools/test_vision_tools.py b/tests/tools/test_vision_tools.py index d8977f84927fb..9ad680488e0e5 100644 --- a/tests/tools/test_vision_tools.py +++ b/tests/tools/test_vision_tools.py @@ -24,6 +24,56 @@ ) +class _FakeStreamResponse: + def __init__(self, chunks, *, headers=None, url="https://example.com/final.png"): + self._chunks = list(chunks) + self.headers = headers or {} + self.url = url + self.content_accessed = False + + @property + def content(self): + self.content_accessed = True + raise AssertionError("download should stream chunks instead of reading content") + + def raise_for_status(self): + return None + + async def aiter_bytes(self): + for chunk in self._chunks: + yield chunk + + +class _FakeStreamContext: + def __init__(self, response): + self.response = response + + async def __aenter__(self): + return self.response + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _FakeAsyncClient: + def __init__(self, response=None, stream_error=None): + self.response = response + self.stream_error = stream_error + self.stream_calls = [] + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def stream(self, method, url, **kwargs): + self.stream_calls.append((method, url, kwargs)) + if self.stream_error: + raise self.stream_error + return _FakeStreamContext(self.response) + + # --------------------------------------------------------------------------- # _validate_image_url — urlparse-based validation # --------------------------------------------------------------------------- @@ -272,10 +322,7 @@ async def test_download_failure_logs_exc_info(self, tmp_path, caplog): from tools.vision_tools import _download_image with patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls: - mock_client = AsyncMock() - 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 = _FakeAsyncClient(stream_error=ConnectionError("network down")) mock_client_cls.return_value = mock_client dest = tmp_path / "image.jpg" @@ -475,28 +522,65 @@ def fake_check(url): } raise AssertionError(f"unexpected URL checked: {url}") - class FakeResponse: - 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 + response = _FakeStreamResponse( + [b"\x89PNG\r\n\x1a\n", b"\x00" * 16], + headers={"content-length": "24"}, + url="https://blocked.test/final.png", + ) 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.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.get = AsyncMock(return_value=FakeResponse()) - mock_client_cls.return_value = mock_client + mock_client_cls.return_value = _FakeAsyncClient(response) await _download_image("https://allowed.test/cat.png", tmp_path / "cat.png", max_retries=1) assert not (tmp_path / "cat.png").exists() + assert response.content_accessed is False + + @pytest.mark.asyncio + async def test_download_streams_image_without_loading_full_response(self, tmp_path): + from tools.vision_tools import _download_image + + response = _FakeStreamResponse( + [b"\x89PNG\r\n\x1a\n", b"image-bytes"], + headers={"content-length": "19"}, + url="https://example.com/final.png", + ) + + with ( + patch("tools.vision_tools.check_website_access", return_value=None), + patch("tools.vision_tools.httpx.AsyncClient") as mock_client_cls, + ): + mock_client_cls.return_value = _FakeAsyncClient(response) + dest = await _download_image("https://example.com/cat.png", tmp_path / "cat.png", max_retries=1) + + assert dest.read_bytes() == b"\x89PNG\r\n\x1a\nimage-bytes" + assert response.content_accessed is False + + @pytest.mark.asyncio + async def test_download_rejects_image_when_stream_exceeds_limit(self, tmp_path): + from tools.vision_tools import _download_image + + response = _FakeStreamResponse( + [b"12345", b"67890", b"!"], + headers={}, + url="https://example.com/final.png", + ) + + 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_cls.return_value = _FakeAsyncClient(response) + await _download_image("https://example.com/cat.png", tmp_path / "cat.png", max_retries=1) + + assert not (tmp_path / "cat.png").exists() + assert response.content_accessed is False # --------------------------------------------------------------------------- diff --git a/tools/vision_tools.py b/tools/vision_tools.py index 912777e2e255d..1f3f78d440436 100644 --- a/tools/vision_tools.py +++ b/tools/vision_tools.py @@ -74,6 +74,55 @@ def _resolve_download_timeout() -> float: _VISION_MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024 +async def _stream_download_to_file( + client: httpx.AsyncClient, + url: str, + destination: Path, + *, + headers: dict[str, str], + max_bytes: int, + media_label: str, +) -> None: + """Stream an HTTP response to disk while enforcing a byte limit.""" + try: + async with client.stream("GET", url, headers=headers) as response: + response.raise_for_status() + + cl = response.headers.get("content-length") + if cl: + try: + content_length = int(cl) + except ValueError: + content_length = None + if content_length is not None and content_length > max_bytes: + raise ValueError( + f"{media_label} too large ({content_length} bytes, max {max_bytes})" + ) + + final_url = str(response.url) + blocked = check_website_access(final_url) + if blocked: + raise PermissionError(blocked["message"]) + + bytes_written = 0 + with 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) + except Exception: + try: + destination.unlink(missing_ok=True) + except OSError: + pass + raise + + def _validate_image_url(url: str) -> bool: """ Basic validation of image URL format. @@ -178,34 +227,17 @@ async def _ssrf_redirect_guard(response): follow_redirects=True, event_hooks={"response": [_ssrf_redirect_guard]}, ) as client: - response = await client.get( + await _stream_download_to_file( + client, image_url, + destination, 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", }, + max_bytes=_VISION_MAX_DOWNLOAD_BYTES, + 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: @@ -1118,32 +1150,17 @@ async def _ssrf_redirect_guard(response): follow_redirects=True, event_hooks={"response": [_ssrf_redirect_guard]}, ) as client: - response = await client.get( + await _stream_download_to_file( + client, video_url, + destination, 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", }, + max_bytes=_MAX_VIDEO_BASE64_BYTES, + 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: