fix: reject non-200 image download responses

astrbot/core/utils/io.py

@@ -113,6 +113,20 @@ def save_temp_img(img: Image.Image | bytes) -> str:
     return p
 
 
+async def _save_image_response(
+    resp: aiohttp.ClientResponse,
+    url: str,
+    path: str | None,
+) -> str:
+    _raise_for_download_status(resp, url)
+    content = await resp.read()
+    if not path:
+        return save_temp_img(content)
+    with open(path, "wb") as f:
+        f.write(content)
+    return path
+
+
 async def download_image_by_url(
     url: str,
     post: bool = False,
@@ -131,18 +145,10 @@ async def download_image_by_url(
         ) as session:
             if post:
                 async with session.post(url, json=post_data) as resp:
-                    if not path:
-                        return save_temp_img(await resp.read())
-                    with open(path, "wb") as f:
-                        f.write(await resp.read())
-                    return path
+                    return await _save_image_response(resp, url, path)
             else:
                 async with session.get(url) as resp:
-                    if not path:
-                        return save_temp_img(await resp.read())
-                    with open(path, "wb") as f:
-                        f.write(await resp.read())
-                    return path
+                    return await _save_image_response(resp, url, path)
     except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
         # 关闭SSL验证（仅在证书验证失败时作为fallback）
         logger.warning(
@@ -157,20 +163,10 @@ async def download_image_by_url(
         async with aiohttp.ClientSession() as session:
             if post:
                 async with session.post(url, json=post_data, ssl=ssl_context) as resp:
-                    if not path:
-                        return save_temp_img(await resp.read())
-                    with open(path, "wb") as f:
-                        f.write(await resp.read())
-                    return path
+                    return await _save_image_response(resp, url, path)
             else:
                 async with session.get(url, ssl=ssl_context) as resp:
-                    if not path:
-                        return save_temp_img(await resp.read())
-                    with open(path, "wb") as f:
-                        f.write(await resp.read())
-                    return path
-    except Exception as e:
-        raise e
+                    return await _save_image_response(resp, url, path)
 
 
 async def _emit_download_progress(progress_callback, payload: dict) -> None:

tests/unit/test_io_download_file.py

@@ -17,8 +17,12 @@ 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._body = b"".join(chunks)
         self.content = _FakeContent(chunks)
 
+    async def read(self) -> bytes:
+        return self._body
+
     async def __aenter__(self):
         return self
 
@@ -41,6 +45,11 @@ def get(self, *_args, **_kwargs):
             raise self._response
         return self._response
 
+    def post(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])
@@ -105,3 +114,83 @@ async def test_download_file_writes_successful_response(monkeypatch, tmp_path):
     await io.download_file("https://example.test/ok.bin", str(target_path))
 
     assert target_path.read_bytes() == b"hello world"
+
+
+@pytest.mark.parametrize("post", [False, True])
+@pytest.mark.asyncio
+async def test_download_image_by_url_rejects_non_200_response(
+    monkeypatch,
+    tmp_path,
+    post,
+):
+    target_path = tmp_path / "image.jpg"
+    _patch_download_session(
+        monkeypatch,
+        _FakeResponse(status=404, chunks=[b'{"error":"not found"}']),
+    )
+
+    with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 404"):
+        await io.download_image_by_url(
+            "https://example.test/generate",
+            post=post,
+            post_data={"json": False} if post else None,
+            path=str(target_path),
+        )
+
+    assert not target_path.exists()
+
+
+@pytest.mark.parametrize("post", [False, True])
+@pytest.mark.asyncio
+async def test_download_image_by_url_rejects_non_200_response_after_ssl_fallback(
+    monkeypatch,
+    tmp_path,
+    post,
+):
+    class FakeSSLError(Exception):
+        pass
+
+    target_path = tmp_path / "image.jpg"
+    _patch_download_sessions(
+        monkeypatch,
+        [
+            FakeSSLError(),
+            _FakeResponse(status=500, chunks=[b"server error"]),
+        ],
+    )
+    monkeypatch.setattr(io.aiohttp, "ClientConnectorSSLError", FakeSSLError)
+    monkeypatch.setattr(io.aiohttp, "ClientConnectorCertificateError", FakeSSLError)
+
+    with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 500"):
+        await io.download_image_by_url(
+            "https://example.test/generate",
+            post=post,
+            post_data={"json": False} if post else None,
+            path=str(target_path),
+        )
+
+    assert not target_path.exists()
+
+
+@pytest.mark.parametrize("post", [False, True])
+@pytest.mark.asyncio
+async def test_download_image_by_url_writes_successful_response(
+    monkeypatch,
+    tmp_path,
+    post,
+):
+    target_path = tmp_path / "image.jpg"
+    _patch_download_session(
+        monkeypatch,
+        _FakeResponse(status=200, chunks=[b"image bytes"]),
+    )
+
+    result = await io.download_image_by_url(
+        "https://example.test/generate",
+        post=post,
+        post_data={"json": False} if post else None,
+        path=str(target_path),
+    )
+
+    assert result == str(target_path)
+    assert target_path.read_bytes() == b"image bytes"