Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
115 changes: 66 additions & 49 deletions astrbot/core/utils/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import shutil
import socket
import ssl
import tempfile
import time
import uuid
from pathlib import Path
Expand Down Expand Up @@ -280,7 +281,10 @@ async def download_file(
progress_callback=None,
allow_insecure_ssl_fallback: bool = True,
) -> None:
"""Download a remote file to a local path.
"""Download a remote file and replace the destination on success.

Keep the destination intact until the response and progress callbacks have
completed. Failed or cancelled downloads leave no partial file behind.

Args:
url: Remote URL to download.
Expand All @@ -294,54 +298,67 @@ async def download_file(
None.
"""

try:
ssl_context = ssl.create_default_context(
cafile=certifi.where(),
) # 使用 certifi 提供的 CA 证书
connector = aiohttp.TCPConnector(ssl=ssl_context)
async with aiohttp.ClientSession(
trust_env=True,
connector=connector,
) as session:
async with session.get(url, timeout=1800) as resp:
_raise_for_download_status(resp, url)
with open(path, "wb") as f:
await _download_response_to_file(
resp,
f,
url,
show_progress,
progress_callback,
)
except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
if not allow_insecure_ssl_fallback:
raise
# 关闭SSL验证(仅在证书验证失败时作为fallback)
logger.warning(
f"SSL certificate verification failed for {_safe_url_for_log(url)}. "
"Falling back to unverified connection (CERT_NONE). "
)
logger.warning(
f"SSL certificate verification failed for {_safe_url_for_log(url)}. "
"Falling back to unverified connection (CERT_NONE). "
"This is insecure and exposes the application to man-in-the-middle attacks. "
"Please investigate certificate issues with the remote server."
)
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
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:
_raise_for_download_status(resp, url)
with open(path, "wb") as f:
await _download_response_to_file(
resp,
f,
url,
show_progress,
progress_callback,
show_downloading_label=False,
)
# Resolve symlinks as open(path, "wb") does, and stage on the same filesystem
# so replacement is atomic. Each call owns its own temporary directory.
target_path = Path(path).resolve()
with tempfile.TemporaryDirectory(
prefix=".astrbot-download-", dir=target_path.parent
) as staging_dir:
Comment on lines +304 to +306

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium · Writable file cannot be updated

When the destination is writable but its parent directory does not permit creating entries, TemporaryDirectory raises PermissionError before the download because it must create a sibling entry, so the existing file is not updated.

Use a staging approach that does not require creating an entry in the destination's parent directory.

Prompt for AI agents
In `astrbot/core/utils/io.py` at lines 304-306:

**Writable file cannot be updated**

When the destination is writable but its parent directory does not permit creating entries, `TemporaryDirectory` raises `PermissionError` before the download because it must create a sibling entry, so the existing file is not updated.

Use a staging approach that does not require creating an entry in the destination's parent directory.

staging_path = Path(staging_dir) / "download.part"
try:
ssl_context = ssl.create_default_context(
cafile=certifi.where(),
) # Use certifi's CA certificates.
connector = aiohttp.TCPConnector(ssl=ssl_context)
async with aiohttp.ClientSession(
trust_env=True,
connector=connector,
) as session:
async with session.get(url, timeout=1800) as resp:
_raise_for_download_status(resp, url)
with staging_path.open("wb") as f:
await _download_response_to_file(
resp,
f,
url,
show_progress,
progress_callback,
)
except (
aiohttp.ClientConnectorSSLError,
aiohttp.ClientConnectorCertificateError,
):
if not allow_insecure_ssl_fallback:
raise
# Retry without verification only after certificate failures.
logger.warning(
f"SSL certificate verification failed for {_safe_url_for_log(url)}. "
"Falling back to unverified connection (CERT_NONE). "
)
logger.warning(
f"SSL certificate verification failed for {_safe_url_for_log(url)}. "
"Falling back to unverified connection (CERT_NONE). "
"This is insecure and exposes the application to man-in-the-middle attacks. "
"Please investigate certificate issues with the remote server."
)
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
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:
_raise_for_download_status(resp, url)
with staging_path.open("wb") as f:
await _download_response_to_file(
resp,
f,
url,
show_progress,
progress_callback,
show_downloading_label=False,
)
if target_path.exists():
shutil.copymode(target_path, staging_path)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 High · Downloaded bytes gain elevated privileges

When a privileged downloader replaces a destination with setuid or setgid bits, shutil.copymode copies those bits to the staging file, and staging_path.replace publishes the downloaded contents with them, creating a setuid/setgid executable from untrusted bytes.

Clear setuid and setgid bits from the staging file before publishing it.

Prompt for AI agents
In `astrbot/core/utils/io.py` at line 360:

**Downloaded bytes gain elevated privileges**

When a privileged downloader replaces a destination with setuid or setgid bits, `shutil.copymode` copies those bits to the staging file, and `staging_path.replace` publishes the downloaded contents with them, creating a setuid/setgid executable from untrusted bytes.

Clear setuid and setgid bits from the staging file before publishing it.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium · Existing access metadata is lost

When the destination relies on ownership, ACLs, or extended attributes beyond its permission bits, shutil.copymode copies only mode bits, and staging_path.replace installs a new inode without the destination's other access metadata, changing who can access the completed download.

Preserve the destination's ownership, ACLs, and extended attributes when publishing the downloaded file.

Also at astrbot/core/utils/io.py:361-362.

Prompt for AI agents
In `astrbot/core/utils/io.py` at line 360:

**Existing access metadata is lost**

When the destination relies on ownership, ACLs, or extended attributes beyond its permission bits, `shutil.copymode` copies only mode bits, and `staging_path.replace` installs a new inode without the destination's other access metadata, changing who can access the completed download.

Preserve the destination's ownership, ACLs, and extended attributes when publishing the downloaded file.

Also at `astrbot/core/utils/io.py:361-362`.

staging_path.replace(target_path)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium · Read-only destinations become replaceable

When the destination is non-writable but its parent directory is writable, download_file writes only to the staging file, then replaces the destination using the parent directory's permissions. On POSIX, a caller that cannot write an existing file can still overwrite it when the parent directory is writable, bypassing file-level write protection.

Check and enforce the existing destination's write authorization before publishing, rejecting the download when the caller could not write the destination directly.

Prompt for AI agents
In `astrbot/core/utils/io.py` at line 361:

**Read-only destinations become replaceable**

When the destination is non-writable but its parent directory is writable, `download_file` writes only to the staging file, then replaces the destination using the parent directory's permissions. On POSIX, a caller that cannot write an existing file can still overwrite it when the parent directory is writable, bypassing file-level write protection.

Check and enforce the existing destination's write authorization before publishing, rejecting the download when the caller could not write the destination directly.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium · Older download replaces newer result

When two successful downloads to the same destination finish in the reverse order from their start, the later-finishing call replaces the destination regardless of which download was started most recently, so an older response can overwrite a newer completed file and callers see stale content.

Serialize publication per destination or reject publication from calls superseded by a newer request.

Prompt for AI agents
In `astrbot/core/utils/io.py` at line 361:

**Older download replaces newer result**

When two successful downloads to the same destination finish in the reverse order from their start, the later-finishing call replaces the destination regardless of which download was started most recently, so an older response can overwrite a newer completed file and callers see stale content.

Serialize publication per destination or reject publication from calls superseded by a newer request.

if show_progress:
print()

Expand Down
231 changes: 227 additions & 4 deletions tests/unit/test_io_download_file.py
Original file line number Diff line number Diff line change
@@ -1,22 +1,32 @@
import asyncio
import stat

import pytest

from astrbot.core.utils import io


class _FakeContent:
def __init__(self, chunks: list[bytes]):
def __init__(self, chunks: list[bytes | BaseException]):
self._chunks = chunks

async def read(self, _size: int) -> bytes:
if self._chunks:
return self._chunks.pop(0)
chunk = self._chunks.pop(0)
if isinstance(chunk, BaseException):
raise chunk
return chunk
return b""


class _FakeResponse:
def __init__(self, *, status: int, chunks: list[bytes]):
def __init__(self, *, status: int, chunks: list[bytes | BaseException]):
self.status = status
self.headers = {"content-length": str(sum(len(chunk) for chunk in chunks))}
self.headers = {
"content-length": str(
sum(len(chunk) for chunk in chunks if isinstance(chunk, bytes))
)
}
self.content = _FakeContent(chunks)

async def __aenter__(self):
Expand Down Expand Up @@ -105,3 +115,216 @@ 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.fixture(params=[False, True], ids=["verified", "ssl-fallback"])
def patch_response(request, monkeypatch):
"""Install a response on either of the downloader's TLS paths.

Args:
request: Pytest request selecting the TLS path.
monkeypatch: Pytest monkeypatch fixture.

Returns:
A callable that installs a fake response for a download.
"""

class FakeSSLError(Exception):
pass

monkeypatch.setattr(io.aiohttp, "ClientConnectorSSLError", FakeSSLError)
monkeypatch.setattr(io.aiohttp, "ClientConnectorCertificateError", FakeSSLError)

def install(response):
responses = [FakeSSLError(), response] if request.param else [response]
_patch_download_sessions(monkeypatch, responses)

return install


@pytest.mark.asyncio
@pytest.mark.parametrize("existing", [False, True], ids=["new", "overwrite"])
@pytest.mark.parametrize(
"error_type",
[io.aiohttp.ClientPayloadError, TimeoutError, asyncio.CancelledError],
)
async def test_failed_download_preserves_destination_and_cleans_staging(
patch_response, tmp_path, existing, error_type
):
target = tmp_path / "download.bin"
if existing:
target.write_bytes(b"previous complete download")
patch_response(
_FakeResponse(status=200, chunks=[b"partial", error_type("interrupted")])
)

with pytest.raises(error_type, match="interrupted"):
await io.download_file("https://example.test/file", str(target))

if existing:
assert target.read_bytes() == b"previous complete download"
else:
assert not target.exists()
assert list(tmp_path.iterdir()) == ([target] if existing else [])


@pytest.mark.asyncio
@pytest.mark.parametrize("existing", [False, True], ids=["new", "overwrite"])
async def test_download_publishes_only_after_all_chunks_are_received(
patch_response, tmp_path, existing
):
target = tmp_path / "download.bin"
if existing:
target.write_bytes(b"old")
patch_response(_FakeResponse(status=200, chunks=[b"hello", b" world"]))
progress = []

async def on_progress(payload):
progress.append(payload)
if existing:
assert target.read_bytes() == b"old"
else:
assert not target.exists()

await io.download_file(
"https://example.test/file", str(target), progress_callback=on_progress
)

assert target.read_bytes() == b"hello world"
assert list(tmp_path.iterdir()) == [target]
assert progress[0]["downloaded"] == 0
assert progress[-1]["downloaded"] == 11
assert progress[-1]["percent"] == 1


@pytest.mark.asyncio
async def test_download_preserves_existing_permissions(patch_response, tmp_path):
target = tmp_path / "download.bin"
target.write_bytes(b"old")
target.chmod(0o600)
original_mode = stat.S_IMODE(target.stat().st_mode)
patch_response(_FakeResponse(status=200, chunks=[b"new"]))

await io.download_file("https://example.test/file", str(target))

assert target.read_bytes() == b"new"
assert stat.S_IMODE(target.stat().st_mode) == original_mode


@pytest.mark.asyncio
async def test_download_preserves_destination_symlink(patch_response, require_symlink):
target = require_symlink / "target.bin"
target.write_bytes(b"old")
link = require_symlink / "download.bin"
link.symlink_to(target)
patch_response(_FakeResponse(status=200, chunks=[b"new"]))

await io.download_file("https://example.test/file", str(link))

assert link.is_symlink()
assert target.read_bytes() == b"new"
assert set(require_symlink.iterdir()) == {target, link}


@pytest.mark.asyncio
@pytest.mark.parametrize("cancelled", [False, True], ids=["disconnect", "cancel"])
async def test_interrupted_http_download_cannot_clobber_concurrent_success(
tmp_path, monkeypatch, cancelled
):
monkeypatch.setenv("NO_PROXY", "127.0.0.1")
target = tmp_path / "download.bin"
target.write_bytes(b"old")
partial_received = asyncio.Event()
release_connection = asyncio.Event()
handlers = []

async def handle_connection(reader, writer):
handlers.append(asyncio.current_task())
try:
request = await reader.readuntil(b"\r\n\r\n")
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n")
if request.startswith(b"GET /interrupted "):
writer.write(b"part")
await writer.drain()
await release_connection.wait()
else:
writer.write(b"complete!")
await writer.drain()
finally:
writer.close()
await writer.wait_closed()

def on_progress(payload):
if payload["downloaded"]:
partial_received.set()

server = await asyncio.start_server(handle_connection, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
async with server:
pending = asyncio.create_task(
io.download_file(
f"http://127.0.0.1:{port}/interrupted",
str(target),
progress_callback=on_progress,
)
)
try:
await asyncio.wait_for(partial_received.wait(), timeout=10)
assert target.read_bytes() == b"old"
await io.download_file(f"http://127.0.0.1:{port}/complete", str(target))
assert target.read_bytes() == b"complete!"
if cancelled:
pending.cancel()
else:
release_connection.set()
error_type = (
asyncio.CancelledError if cancelled else io.aiohttp.ClientPayloadError
)
with pytest.raises(error_type):
await asyncio.wait_for(pending, timeout=10)
finally:
release_connection.set()
pending.cancel()
await asyncio.gather(pending, return_exceptions=True)
await asyncio.gather(*handlers)

assert target.read_bytes() == b"complete!"
assert list(tmp_path.iterdir()) == [target]


@pytest.mark.asyncio
async def test_callback_failure_does_not_publish_download(patch_response, tmp_path):
target = tmp_path / "download.bin"
target.write_bytes(b"old")
patch_response(_FakeResponse(status=200, chunks=[b"new"]))

def on_progress(payload):
if payload["downloaded"]:
raise RuntimeError("progress callback failed")

with pytest.raises(RuntimeError, match="progress callback failed"):
await io.download_file(
"https://example.test/file", str(target), progress_callback=on_progress
)

assert target.read_bytes() == b"old"
assert list(tmp_path.iterdir()) == [target]


@pytest.mark.asyncio
async def test_failed_publication_preserves_destination(
patch_response, tmp_path, monkeypatch
):
target = tmp_path / "download.bin"
target.write_bytes(b"old")
patch_response(_FakeResponse(status=200, chunks=[b"new"]))

def fail_replace(self, destination):
raise PermissionError("destination is busy")

monkeypatch.setattr(io.Path, "replace", fail_replace)
with pytest.raises(PermissionError, match="destination is busy"):
await io.download_file("https://example.test/file", str(target))

assert target.read_bytes() == b"old"
assert list(tmp_path.iterdir()) == [target]