fix: reject non-200 download responses (#9085)

* fix: reject non-200 download responses

* fix: share download response handling

astrbot/core/utils/io.py

@@ -178,6 +178,101 @@ async def _emit_download_progress(progress_callback, payload: dict) -> None:
         await result
 
 
+class DownloadFileHTTPError(RuntimeError):
+    """Raised when a file download returns an unsuccessful HTTP status."""
+
+
+def _raise_for_download_status(resp, url: str) -> None:
+    if resp.status == 200:
+        return
+    logger.error(
+        "Failed to download file from %s. HTTP status code: %s",
+        _safe_url_for_log(url),
+        resp.status,
+    )
+    raise DownloadFileHTTPError(
+        "Failed to download file from "
+        f"{_safe_url_for_log(url)}. HTTP status code: {resp.status}"
+    )
+
+
+async def _download_response_to_file(
+    resp,
+    file_obj,
+    url: str,
+    show_progress: bool,
+    progress_callback,
+    show_downloading_label: bool = True,
+) -> None:
+    """Write a successful download response to a local file.
+
+    Args:
+        resp: aiohttp response object to read from.
+        file_obj: Open writable binary file object.
+        url: Source URL used for progress events and sanitized errors.
+        show_progress: Whether to print progress to stdout.
+        progress_callback: Optional callback for progress payloads.
+        show_downloading_label: Whether to use the standard download heading.
+
+    """
+
+    total_size = int(resp.headers.get("content-length", 0))
+    downloaded_size = 0
+    start_time = time.time()
+    if show_progress:
+        if show_downloading_label:
+            print(
+                f"Downloading: {_safe_url_for_log(url)} | "
+                f"Size: {total_size / 1024:.2f} KB"
+            )
+        else:
+            print(f"Size: {total_size / 1024:.2f} KB | URL: {_safe_url_for_log(url)}")
+    await _emit_download_progress(
+        progress_callback,
+        {
+            "url": url,
+            "downloaded": 0,
+            "total": total_size,
+            "percent": 0,
+            "speed": 0,
+        },
+    )
+    while True:
+        chunk = await resp.content.read(8192)
+        if not chunk:
+            break
+        file_obj.write(chunk)
+        downloaded_size += len(chunk)
+        elapsed_time = time.time() - start_time if time.time() - start_time > 0 else 1
+        speed = downloaded_size / 1024 / elapsed_time  # KB/s
+        percent = downloaded_size / total_size if total_size > 0 else 0
+        await _emit_download_progress(
+            progress_callback,
+            {
+                "url": url,
+                "downloaded": downloaded_size,
+                "total": total_size,
+                "percent": percent,
+                "speed": speed,
+            },
+        )
+        if show_progress:
+            print(
+                f"\rProgress: {percent:.2%} Speed: {speed:.2f} KB/s",
+                end="",
+            )
+    await _emit_download_progress(
+        progress_callback,
+        {
+            "url": url,
+            "downloaded": downloaded_size,
+            "total": total_size,
+            "percent": 1,
+            "speed": 0,
+        },
+    )
+
+
 async def download_file(
     url: str,
     path: str,
@@ -209,69 +304,15 @@ async def download_file(
             connector=connector,
         ) as session:
             async with session.get(url, timeout=1800) as resp:
-                if resp.status != 200:
+                _raise_for_download_status(resp, url)
-                    logger.error(
-                        "Failed to download file from %s. HTTP status code: %s",
-                        _safe_url_for_log(url),
-                        resp.status,
-                    )
-                total_size = int(resp.headers.get("content-length", 0))
-                downloaded_size = 0
-                start_time = time.time()
-                if show_progress:
-                    print(
-                        f"Downloading: {_safe_url_for_log(url)} | "
-                        f"Size: {total_size / 1024:.2f} KB"
-                    )
-                await _emit_download_progress(
-                    progress_callback,
-                    {
-                        "url": url,
-                        "downloaded": 0,
-                        "total": total_size,
-                        "percent": 0,
-                        "speed": 0,
-                    },
-                )
                 with open(path, "wb") as f:
-                    while True:
+                    await _download_response_to_file(
-                        chunk = await resp.content.read(8192)
+                        resp,
-                        if not chunk:
+                        f,
-                            break
+                        url,
-                        f.write(chunk)
+                        show_progress,
-                        downloaded_size += len(chunk)
+                        progress_callback,
-                        elapsed_time = (
+                    )
-                            time.time() - start_time
-                            if time.time() - start_time > 0
-                            else 1
-                        )
-                        speed = downloaded_size / 1024 / elapsed_time  # KB/s
-                        percent = downloaded_size / total_size if total_size > 0 else 0
-                        await _emit_download_progress(
-                            progress_callback,
-                            {
-                                "url": url,
-                                "downloaded": downloaded_size,
-                                "total": total_size,
-                                "percent": percent,
-                                "speed": speed,
-                            },
-                        )
-                        if show_progress:
-                            print(
-                                f"\rProgress: {percent:.2%} Speed: {speed:.2f} KB/s",
-                                end="",
-                            )
-                await _emit_download_progress(
-                    progress_callback,
-                    {
-                        "url": url,
-                        "downloaded": downloaded_size,
-                        "total": total_size,
-                        "percent": 1,
-                        "speed": 0,
-                    },
-                )
     except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
         if not allow_insecure_ssl_fallback:
             raise
@@ -291,63 +332,16 @@ async def download_file(
         ssl_context.verify_mode = ssl.CERT_NONE
         async with aiohttp.ClientSession() as session:
             async with session.get(url, ssl=ssl_context, timeout=120) as resp:
-                total_size = int(resp.headers.get("content-length", 0))
+                _raise_for_download_status(resp, url)
-                downloaded_size = 0
-                start_time = time.time()
-                if show_progress:
-                    print(
-                        f"Size: {total_size / 1024:.2f} KB | "
-                        f"URL: {_safe_url_for_log(url)}"
-                    )
-                await _emit_download_progress(
-                    progress_callback,
-                    {
-                        "url": url,
-                        "downloaded": 0,
-                        "total": total_size,
-                        "percent": 0,
-                        "speed": 0,
-                    },
-                )
                 with open(path, "wb") as f:
-                    while True:
+                    await _download_response_to_file(
-                        chunk = await resp.content.read(8192)
+                        resp,
-                        if not chunk:
+                        f,
-                            break
+                        url,
-                        f.write(chunk)
+                        show_progress,
-                        downloaded_size += len(chunk)
+                        progress_callback,
-                        elapsed_time = (
+                        show_downloading_label=False,
-                            time.time() - start_time
+                    )
-                            if time.time() - start_time > 0
-                            else 1
-                        )
-                        speed = downloaded_size / 1024 / elapsed_time  # KB/s
-                        percent = downloaded_size / total_size if total_size > 0 else 0
-                        await _emit_download_progress(
-                            progress_callback,
-                            {
-                                "url": url,
-                                "downloaded": downloaded_size,
-                                "total": total_size,
-                                "percent": percent,
-                                "speed": speed,
-                            },
-                        )
-                        if show_progress:
-                            print(
-                                f"\rProgress: {percent:.2%} Speed: {speed:.2f} KB/s",
-                                end="",
-                            )
-                await _emit_download_progress(
-                    progress_callback,
-                    {
-                        "url": url,
-                        "downloaded": downloaded_size,
-                        "total": total_size,
-                        "percent": 1,
-                        "speed": 0,
-                    },
-                )
     if show_progress:
         print()
 

tests/unit/test_io_download_file.py

@@ -0,0 +1,107 @@
+import pytest
+
+from astrbot.core.utils import io
+
+
+class _FakeContent:
+    def __init__(self, chunks: list[bytes]):
+        self._chunks = chunks
+
+    async def read(self, _size: int) -> bytes:
+        if self._chunks:
+            return self._chunks.pop(0)
+        return b""
+
+
+class _FakeResponse:
+    def __init__(self, *, status: int, chunks: list[bytes]):
+        self.status = status
+        self.headers = {"content-length": str(sum(len(chunk) for chunk in chunks))}
+        self.content = _FakeContent(chunks)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+
+class _FakeSession:
+    def __init__(self, response: _FakeResponse | Exception):
+        self._response = response
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    def get(self, *_args, **_kwargs):
+        if isinstance(self._response, Exception):
+            raise self._response
+        return self._response
+
+
+def _patch_download_session(monkeypatch, response: _FakeResponse):
+    _patch_download_sessions(monkeypatch, [response])
+
+
+def _patch_download_sessions(monkeypatch, responses: list[_FakeResponse | Exception]):
+    monkeypatch.setattr(io.aiohttp, "TCPConnector", lambda **_kwargs: object())
+    monkeypatch.setattr(
+        io.aiohttp,
+        "ClientSession",
+        lambda **_kwargs: _FakeSession(responses.pop(0)),
+    )
+
+
+@pytest.mark.asyncio
+async def test_download_file_rejects_non_200_response(monkeypatch, tmp_path):
+    target_path = tmp_path / "missing.bin"
+    _patch_download_session(
+        monkeypatch,
+        _FakeResponse(status=404, chunks=[b"not found"]),
+    )
+
+    with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 404"):
+        await io.download_file("https://example.test/missing", str(target_path))
+
+    assert not target_path.exists()
+
+
+@pytest.mark.asyncio
+async def test_download_file_rejects_non_200_response_after_ssl_fallback(
+    monkeypatch,
+    tmp_path,
+):
+    class FakeSSLError(Exception):
+        pass
+
+    target_path = tmp_path / "missing.bin"
+    _patch_download_sessions(
+        monkeypatch,
+        [
+            FakeSSLError(),
+            _FakeResponse(status=404, chunks=[b"not found"]),
+        ],
+    )
+    monkeypatch.setattr(io.aiohttp, "ClientConnectorSSLError", FakeSSLError)
+    monkeypatch.setattr(io.aiohttp, "ClientConnectorCertificateError", FakeSSLError)
+
+    with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 404"):
+        await io.download_file("https://example.test/missing", str(target_path))
+
+    assert not target_path.exists()
+
+
+@pytest.mark.asyncio
+async def test_download_file_writes_successful_response(monkeypatch, tmp_path):
+    target_path = tmp_path / "ok.bin"
+    _patch_download_session(
+        monkeypatch,
+        _FakeResponse(status=200, chunks=[b"hello", b" world"]),
+    )
+
+    await io.download_file("https://example.test/ok.bin", str(target_path))
+
+    assert target_path.read_bytes() == b"hello world"