Stream remote media downloads with size limits

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
 # ---------------------------------------------------------------------------

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
 
 
 # ---------------------------------------------------------------------------

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: