From d0e1dfcd172ed4589aded6e9e5df6a140e5291b9 Mon Sep 17 00:00:00 2001 From: VectorPeak <73048950+VectorPeak@users.noreply.github.com> Date: Sat, 18 Jul 2026 21:23:02 +0800 Subject: [PATCH 1/2] fix: reject non-200 image download responses Co-authored-by: chatgpt-codex-connector[bot] <199175422+chatgpt-codex-connector[bot]@users.noreply.github.com> --- astrbot/core/utils/io.py | 34 +++++------- tests/unit/test_io_download_file.py | 83 +++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 20 deletions(-) diff --git a/astrbot/core/utils/io.py b/astrbot/core/utils/io.py index 41cc87639a..88bdd1cf33 100644 --- a/astrbot/core/utils/io.py +++ b/astrbot/core/utils/io.py @@ -113,6 +113,16 @@ def save_temp_img(img: Image.Image | bytes) -> str: return p +async def _save_image_response(resp, 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 +141,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,18 +159,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 + return await _save_image_response(resp, url, path) except Exception as e: raise e diff --git a/tests/unit/test_io_download_file.py b/tests/unit/test_io_download_file.py index 5cea2f0e50..d7ddc86488 100644 --- a/tests/unit/test_io_download_file.py +++ b/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,77 @@ 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.asyncio +async def test_download_image_by_url_rejects_non_200_post_response( + monkeypatch, + tmp_path, +): + 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=True, + post_data={"json": False}, + path=str(target_path), + ) + + assert not target_path.exists() + + +@pytest.mark.asyncio +async def test_download_image_by_url_rejects_non_200_response_after_ssl_fallback( + monkeypatch, + tmp_path, +): + 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=True, + post_data={"json": False}, + path=str(target_path), + ) + + assert not target_path.exists() + + +@pytest.mark.asyncio +async def test_download_image_by_url_writes_successful_post_response( + monkeypatch, + tmp_path, +): + 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=True, + post_data={"json": False}, + path=str(target_path), + ) + + assert result == str(target_path) + assert target_path.read_bytes() == b"image bytes" From 21bbb31166ee15ae71cd41594332529c4c2efae0 Mon Sep 17 00:00:00 2001 From: VectorPeak <73048950+VectorPeak@users.noreply.github.com> Date: Sat, 18 Jul 2026 21:36:58 +0800 Subject: [PATCH 2/2] test: cover image downloads for get and post Co-authored-by: chatgpt-codex-connector[bot] <199175422+chatgpt-codex-connector[bot]@users.noreply.github.com> --- astrbot/core/utils/io.py | 8 +++++--- tests/unit/test_io_download_file.py | 22 ++++++++++++++-------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/astrbot/core/utils/io.py b/astrbot/core/utils/io.py index 88bdd1cf33..a2f7946e8d 100644 --- a/astrbot/core/utils/io.py +++ b/astrbot/core/utils/io.py @@ -113,7 +113,11 @@ def save_temp_img(img: Image.Image | bytes) -> str: return p -async def _save_image_response(resp, url: str, path: str | None) -> str: +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: @@ -163,8 +167,6 @@ async def download_image_by_url( else: async with session.get(url, ssl=ssl_context) as resp: return await _save_image_response(resp, url, path) - except Exception as e: - raise e async def _emit_download_progress(progress_callback, payload: dict) -> None: diff --git a/tests/unit/test_io_download_file.py b/tests/unit/test_io_download_file.py index d7ddc86488..71aa125e7b 100644 --- a/tests/unit/test_io_download_file.py +++ b/tests/unit/test_io_download_file.py @@ -116,10 +116,12 @@ async def test_download_file_writes_successful_response(monkeypatch, tmp_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_post_response( +async def test_download_image_by_url_rejects_non_200_response( monkeypatch, tmp_path, + post, ): target_path = tmp_path / "image.jpg" _patch_download_session( @@ -130,18 +132,20 @@ async def test_download_image_by_url_rejects_non_200_post_response( with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 404"): await io.download_image_by_url( "https://example.test/generate", - post=True, - post_data={"json": False}, + 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 @@ -160,18 +164,20 @@ class FakeSSLError(Exception): with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 500"): await io.download_image_by_url( "https://example.test/generate", - post=True, - post_data={"json": False}, + 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_post_response( +async def test_download_image_by_url_writes_successful_response( monkeypatch, tmp_path, + post, ): target_path = tmp_path / "image.jpg" _patch_download_session( @@ -181,8 +187,8 @@ async def test_download_image_by_url_writes_successful_post_response( result = await io.download_image_by_url( "https://example.test/generate", - post=True, - post_data={"json": False}, + post=post, + post_data={"json": False} if post else None, path=str(target_path), )