mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 09:40:30 +08:00
fix: reject non-200 download responses (#9085)
* fix: reject non-200 download responses * fix: share download response handling
This commit is contained in:
VectorPeak
@@ -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:
|
||||
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
|
||||
@@ -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()
|
||||
|
||||
|
||||
107
tests/unit/test_io_download_file.py
Normal file
107
tests/unit/test_io_download_file.py
Normal file
@@ -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"
|
||||
Reference in New Issue
Block a user