Skip to content

fix: reject non-200 download responses #9085

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
VectorPeak wants to merge 2 commits into
base: master
Choose a base branch
from
+219 −118
Open

fix: reject non-200 download responses #9085

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
59f79cd
fix: reject non-200 download responses
VectorPeak Jun 30, 2026
403bbb6
fix: share download response handling
VectorPeak Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
230 changes: 112 additions & 118 deletions astrbot/core/utils/io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
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,
},
)
_raise_for_download_status(resp, url)
with open(path, "wb") as f:
while True:
chunk = await resp.content.read(8192)
if not chunk:
break
f.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,
},
)
await _download_response_to_file(
resp,
f,
url,
show_progress,
progress_callback,
)
except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
if not allow_insecure_ssl_fallback:
raise
Expand All @@ -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))
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,
},
)
_raise_for_download_status(resp, url)
with open(path, "wb") as f:
while True:
chunk = await resp.content.read(8192)
if not chunk:
break
f.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,
},
)
await _download_response_to_file(
resp,
f,
url,
show_progress,
progress_callback,
show_downloading_label=False,
)
if show_progress:
print()

Expand Down
107 changes: 107 additions & 0 deletions tests/unit/test_io_download_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Loading