From 9af89d4e2eaeedd83488ead65245c96f73117132 Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Fri, 11 Sep 2026 12:37:54 +0800 Subject: [PATCH] feat(updater): prefer CERNET PyPI source packages for releases --- astrbot/core/updater.py | 331 ++++++++++++++++++++++++++++++------ tests/test_updater_pypi.py | 297 ++++++++++++++++++++++++++++++++ tests/test_updater_socks.py | 3 + 3 files changed, 576 insertions(+), 55 deletions(-) create mode 100644 tests/test_updater_pypi.py diff --git a/astrbot/core/updater.py b/astrbot/core/updater.py index 8d56ca98a4..cb62fb0398 100644 --- a/astrbot/core/updater.py +++ b/astrbot/core/updater.py @@ -1,12 +1,20 @@ import asyncio +import hashlib import os import shutil +import tarfile import tempfile import zipfile from collections.abc import Awaitable, Callable from dataclasses import dataclass -from pathlib import Path +from email.parser import BytesParser +from html.parser import HTMLParser +from pathlib import Path, PurePosixPath from typing import Literal +from urllib.parse import parse_qs, unquote, urldefrag, urljoin, urlsplit + +from packaging.utils import InvalidSdistFilename, parse_sdist_filename +from packaging.version import Version from astrbot.core import logger from astrbot.core.config.default import VERSION @@ -59,6 +67,20 @@ class UpdateProgress: UpdateProgressCallback = Callable[[UpdateProgress], Awaitable[None]] +class _SourceDistributionLinks(HTMLParser): + """Collect non-yanked download links from a PyPI Simple HTML index.""" + + def __init__(self) -> None: + super().__init__() + self.links: list[str] = [] + + def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: + attributes = dict(attrs) + href = attributes.get("href") + if tag == "a" and href and "data-yanked" not in attributes: + self.links.append(href) + + class AstrBotUpdater(_RepoZipUpdater): """Expose the complete, high-level AstrBot Core update operations.""" @@ -201,6 +223,12 @@ async def core_progress(payload: dict) -> None: payload, ) + if os.environ.get("ASTRBOT_CLI") or os.environ.get("ASTRBOT_LAUNCHER"): + raise RuntimeError( + "You are running AstrBot via CLI; use pip or uv tool upgrade " + "to update AstrBot." + ) + target_version = version target_release = None if not target_version or target_version == "latest": @@ -211,20 +239,6 @@ async def core_progress(payload: dict) -> None: target_version = target_release["tag_name"] if self._compare_version(VERSION, target_version) >= 0: raise RuntimeError("AstrBot is already up to date.") - elif target_version.startswith("v"): - releases = await self._fetch_release_info(self._release_api) - target_release = next( - ( - release - for release in releases - if release["tag_name"] == target_version - ), - None, - ) - if target_release is None: - raise RuntimeError( - f"No update package was found for version {target_version}." - ) update_temp_parent = Path(get_astrbot_temp_path()) / "updates" if update_temp_parent.is_symlink(): @@ -240,47 +254,84 @@ async def core_progress(payload: dict) -> None: dashboard_zip_path = update_temp_dir / "dashboard.zip" core_zip_path = update_temp_dir / "core.zip" - await emit_progress( - "dashboard", - "running", - "正在下载 WebUI...", - 0, - ) - await _download_package( - path=str(dashboard_zip_path), - version=target_version, - proxy=proxy, - progress_callback=dashboard_progress, - extract=False, - allow_insecure_ssl_fallback=False, - ) - await emit_progress( - "dashboard", - "done", - "WebUI 下载完成。", - 45, - ) + mirror_prepared = False + if target_version.startswith("v"): + await emit_progress( + "dashboard", + "running", + "Checking the PyPI mirror for the update...", + 0, + ) + mirror_prepared = await self._download_pypi_package( + target_version, + core_zip_path, + dashboard_zip_path, + progress_callback=dashboard_progress, + ) + if mirror_prepared: + await emit_progress( + "dashboard", "done", "Bundled WebUI prepared from PyPI.", 45 + ) + await emit_progress( + "core", "done", "AstrBot source prepared from PyPI.", 90 + ) + else: + if target_version.startswith("v") and target_release is None: + releases = await self._fetch_release_info(self._release_api) + target_release = next( + ( + release + for release in releases + if release["tag_name"] == target_version + ), + None, + ) + if target_release is None: + raise RuntimeError( + f"No update package was found for version {target_version}." + ) + + await emit_progress( + "dashboard", + "running", + "正在下载 WebUI...", + 0, + ) + await _download_package( + path=str(dashboard_zip_path), + version=target_version, + proxy=proxy, + progress_callback=dashboard_progress, + extract=False, + allow_insecure_ssl_fallback=False, + ) + await emit_progress( + "dashboard", + "done", + "WebUI 下载完成。", + 45, + ) - await emit_progress( - "core", - "running", - "正在下载 AstrBot 项目代码...", - 45, - ) - await self._download_core_package( - latest=False, - version=target_version, - proxy=proxy, - path=core_zip_path, - progress_callback=core_progress, - release_data=target_release, - ) - await emit_progress( - "core", - "done", - "项目代码下载完成。", - 90, - ) + await emit_progress( + "core", + "running", + "正在下载 AstrBot 项目代码...", + 45, + ) + await self._download_core_package( + latest=False, + version=target_version, + proxy=proxy, + path=core_zip_path, + progress_callback=core_progress, + release_data=target_release, + ) + await emit_progress( + "core", + "done", + "项目代码下载完成。", + 90, + ) await emit_progress( "verify", @@ -324,6 +375,176 @@ def verify_packages() -> None: 92, ) + async def _download_pypi_package( + self, + version: str, + core_zip_path: Path, + dashboard_zip_path: Path, + progress_callback=None, + ) -> bool: + """Prepare a matching source distribution and its bundled Dashboard. + + Args: + version: Target release tag. + core_zip_path: Temporary destination for the existing Core apply flow. + dashboard_zip_path: Temporary destination for the Dashboard apply flow. + progress_callback: Download progress observer. + + Returns: + Whether both packages are validated and ready. Mirror failures return + False so the caller can use the existing update sources. + """ + index_url = "https://mirrors.cernet.edu.cn/pypi/web/simple/astrbot/" + source_path = core_zip_path.with_suffix(".tar.gz") + try: + target = Version(version) + async with self._create_httpx_client(timeout=10.0) as client: + response = await client.get(index_url) + response.raise_for_status() + parser = _SourceDistributionLinks() + parser.feed(response.text) + for link in parser.links: + url = urljoin(str(response.url), link) + parsed = urlsplit(url) + filename = PurePosixPath(unquote(parsed.path)).name + if parsed.scheme != "https" or not filename.endswith(".tar.gz"): + continue + try: + name, candidate = parse_sdist_filename(filename) + except InvalidSdistFilename: + continue + if name != "astrbot" or candidate != target: + continue + digest = parse_qs(parsed.fragment).get("sha256", [""])[0] + if len(digest) != 64 or any( + char not in "0123456789abcdef" for char in digest.lower() + ): + raise ValueError("PyPI source link has no valid SHA-256 digest") + logger.info("Downloading AstrBot source and WebUI from %s", url) + await self._download_file( + urldefrag(url)[0], + str(source_path), + timeout=60.0, + progress_callback=progress_callback, + ) + await asyncio.to_thread( + self._prepare_pypi_package, + source_path, + digest, + version, + core_zip_path, + dashboard_zip_path, + ) + return True + logger.info("AstrBot %s is not available in the PyPI mirror.", version) + except Exception as exc: + logger.warning( + "PyPI update package failed: %s. Falling back to the existing " + "Core and Dashboard download sources.", + exc, + ) + finally: + source_path.unlink(missing_ok=True) + core_zip_path.unlink(missing_ok=True) + dashboard_zip_path.unlink(missing_ok=True) + return False + + @staticmethod + def _prepare_pypi_package( + source_path: Path, + expected_digest: str, + version: str, + core_zip_path: Path, + dashboard_zip_path: Path, + ) -> None: + """Validate an sdist and stage ZIPs for the existing update application. + + Args: + source_path: Downloaded tar.gz source distribution. + expected_digest: SHA-256 advertised by the mirror index. + version: Target release tag. + core_zip_path: Prepared Core ZIP destination. + dashboard_zip_path: Prepared Dashboard ZIP destination. + + Raises: + ValueError: If the archive hash, paths, metadata, or assets are invalid. + """ + digest = hashlib.sha256() + with source_path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + if digest.hexdigest() != expected_digest.lower(): + raise ValueError("PyPI source archive SHA-256 mismatch") + + with tempfile.TemporaryDirectory( + prefix="pypi-source-", dir=source_path.parent + ) as staging_name: + staging = Path(staging_name) + with tarfile.open(source_path, "r:gz") as archive: + members = archive.getmembers() + roots: set[str] = set() + seen: set[PurePosixPath] = set() + for member in members: + path = PurePosixPath(member.name) + if ( + not path.parts + or path.is_absolute() + or ".." in path.parts + or "\\" in member.name + or ":" in member.name + or not (member.isfile() or member.isdir()) + or path in seen + ): + raise ValueError(f"Unsafe PyPI archive member: {member.name}") + roots.add(path.parts[0]) + seen.add(path) + if len(roots) != 1: + raise ValueError("PyPI source archive must have one root directory") + root_name = roots.pop() + name, source_version = parse_sdist_filename(f"{root_name}.tar.gz") + if name != "astrbot" or source_version != Version(version): + raise ValueError("PyPI source directory does not match the release") + + # Copy only regular files; never follow archive links or execute code. + for member in members: + destination = staging.joinpath(*PurePosixPath(member.name).parts) + if member.isdir(): + destination.mkdir(parents=True, exist_ok=True) + continue + destination.parent.mkdir(parents=True, exist_ok=True) + with ( + archive.extractfile(member) as source, + destination.open("wb") as output, + ): + shutil.copyfileobj(source, output) + + root = staging / root_name + for required in ( + "PKG-INFO", + "main.py", + "pyproject.toml", + "requirements.txt", + "astrbot/__init__.py", + ): + if not (root / required).is_file(): + raise ValueError(f"PyPI source archive is missing {required}") + with (root / "PKG-INFO").open("rb") as metadata_file: + metadata = BytesParser().parse(metadata_file, headersonly=True) + if str(metadata.get("Name", "")).lower() != "astrbot" or Version( + str(metadata.get("Version", "")) + ) != Version(version): + raise ValueError("PyPI source metadata does not match the release") + dashboard = root / "astrbot" / "dashboard" / "dist" + if not _is_dist_compatible(dashboard, version): + raise ValueError("PyPI source has an incomplete or mismatched WebUI") + + shutil.make_archive( + str(core_zip_path.with_suffix("")), "zip", staging, root_name + ) + shutil.make_archive( + str(dashboard_zip_path.with_suffix("")), "zip", dashboard.parent, "dist" + ) + async def ensure_dashboard(self) -> Path: """Ensure acceptable Dashboard assets exist for the active runtime. diff --git a/tests/test_updater_pypi.py b/tests/test_updater_pypi.py new file mode 100644 index 0000000000..85b751416f --- /dev/null +++ b/tests/test_updater_pypi.py @@ -0,0 +1,297 @@ +"""Regression coverage for mirror-first source and Dashboard updates.""" + +import hashlib +import io +import tarfile +import zipfile +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import httpx +import pytest + +from astrbot.core import updater as updater_module +from astrbot.core.updater import AstrBotUpdater + + +def _sdist(files=None, extra=None, version="99.0.0"): + """Build a complete source archive with optional corruptions. + + Args: + files: Relative file overrides; None values remove a file. + extra: Optional tar member for archive safety tests. + version: Package version shared by metadata and bundled assets. + + Returns: + Compressed source archive bytes. + """ + contents = { + "PKG-INFO": f"Metadata-Version: 2.4\nName: AstrBot\nVersion: {version}\n\n", + "main.py": "SOURCE = 'mirror'\n", + "pyproject.toml": f'[project]\nname="AstrBot"\nversion="{version}"\n', + "requirements.txt": "httpx\n", + "astrbot/__init__.py": f'__version__ = "{version}"\n', + "astrbot/dashboard/dist/index.html": '', + "astrbot/dashboard/dist/assets/app.js": "// mirror dashboard\n", + "astrbot/dashboard/dist/assets/version": f"v{version}", + } + contents.update(files or {}) + stream = io.BytesIO() + with tarfile.open(fileobj=stream, mode="w:gz") as archive: + for name, content in contents.items(): + if content is None: + continue + data = content.encode() + member = tarfile.TarInfo(f"astrbot-{version}/{name}") + member.size = len(data) + archive.addfile(member, io.BytesIO(data)) + if extra: + archive.addfile(extra, io.BytesIO(b"")) + return stream.getvalue() + + +@pytest.fixture +def mirror_update(monkeypatch, tmp_path): + """Isolate all update writes and mirror requests from the running installation.""" + monkeypatch.delenv("ASTRBOT_CLI", raising=False) + monkeypatch.delenv("ASTRBOT_LAUNCHER", raising=False) + install = tmp_path / "install" + install.mkdir() + (install / "main.py").write_text("old core") + data = tmp_path / "data" + (data / "dist").mkdir(parents=True) + (data / "dist/index.html").write_text("old dashboard") + updater = AstrBotUpdater() + updater._main_path = str(install) + monkeypatch.setattr( + updater_module, "get_astrbot_temp_path", lambda: str(tmp_path / "temp") + ) + monkeypatch.setattr(updater_module, "get_astrbot_data_path", lambda: str(data)) + state = SimpleNamespace( + updater=updater, + install=install, + data=data, + temp=tmp_path / "temp", + payload=_sdist(), + index=None, + index_status=200, + download_status=200, + network_error=None, + requests=[], + fallback=[], + events=[], + ) + + def handle(request): + state.requests.append(str(request.url)) + if request.url.host == "mirrors.cernet.edu.cn": + return httpx.Response( + 302, + headers={"Location": "https://mirror.example/pypi/web/simple/astrbot/"}, + ) + if request.url.path.endswith("/simple/astrbot/"): + if state.network_error == "index": + raise httpx.ReadTimeout("mirror timed out", request=request) + index = state.index + if index is None: + digest = hashlib.sha256(state.payload).hexdigest() + index = f'source' + return httpx.Response(state.index_status, text=index) + assert request.url.host == "mirror.example" + assert request.url.path == "/pypi/web/packages/astrbot-99.0.0.tar.gz" + assert not request.url.fragment + if state.network_error == "download": + raise httpx.ReadTimeout("download timed out", request=request) + return httpx.Response(state.download_status, content=state.payload) + + monkeypatch.setattr( + updater, + "_create_httpx_client", + lambda timeout=30: httpx.AsyncClient( + transport=httpx.MockTransport(handle), + follow_redirects=True, + timeout=timeout, + ), + ) + state.releases = AsyncMock( + return_value=[ + { + "tag_name": "v99.0.0", + "zipball_url": "https://github.example/core.zip", + } + ] + ) + monkeypatch.setattr(updater, "_fetch_release_info", state.releases) + + async def legacy_dashboard(*, path, **kwargs): + state.fallback.append(("dashboard", kwargs)) + assert (install / "main.py").read_text() == "old core" + assert (data / "dist/index.html").read_text() == "old dashboard" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr( + "dist/index.html", '' + ) + archive.writestr("dist/assets/fallback.js", "// fallback") + archive.writestr("dist/assets/version", "v99.0.0") + + async def legacy_core(*, path, **kwargs): + state.fallback.append(("core", kwargs)) + assert (install / "main.py").read_text() == "old core" + assert (data / "dist/index.html").read_text() == "old dashboard" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("AstrBot-release/main.py", "fallback core") + return path + + monkeypatch.setattr(updater_module, "_download_package", legacy_dashboard) + monkeypatch.setattr(updater, "_download_core_package", legacy_core) + return state + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", ["v99.0.0", None, "latest"]) +async def test_mirror_applies_matching_source_and_bundled_dashboard( + mirror_update, version +): + state = mirror_update + + async def progress(event): + state.events.append(event) + + await state.updater.update( + version, proxy="https://github-proxy.example", progress_callback=progress + ) + assert (state.install / "main.py").read_text() == "SOURCE = 'mirror'\n" + assert (state.data / "dist/assets/app.js").read_text() == "// mirror dashboard\n" + assert (state.install / "astrbot/dashboard/dist/assets/app.js").read_bytes() == ( + state.data / "dist/assets/app.js" + ).read_bytes() + assert state.fallback == [] + assert len(state.requests) == 3 + assert state.releases.await_count == (0 if version == "v99.0.0" else 1) + assert [(event.stage, event.status) for event in state.events][-1] == ( + "apply", + "done", + ) + assert {event.stage for event in state.events} == { + "dashboard", + "core", + "verify", + "apply", + } + assert not list((state.temp / "updates").iterdir()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", + [ + "missing-version", + "index-error", + "index-timeout", + "download-timeout", + "download-error", + "hash-mismatch", + "missing-hash", + "corrupt-tar", + "missing-dashboard", + "incomplete-dashboard", + "wrong-dashboard-version", + "missing-core", + "wrong-metadata", + "traversal", + "symlink", + "hardlink", + "absolute", + "windows-path", + "duplicate", + "multiple-roots", + "yanked", + "wheel-only", + ], +) +async def test_mirror_failure_uses_existing_downloads_before_applying( + mirror_update, failure +): + state = mirror_update + if failure == "missing-version": + state.index = 'other version' + elif failure in ("index-timeout", "download-timeout"): + state.network_error = failure.split("-")[0] + elif failure == "index-error": + state.index_status = 503 + elif failure == "download-error": + state.download_status = 404 + elif failure in ("hash-mismatch", "missing-hash"): + fragment = "#sha256=" + "0" * 64 if failure == "hash-mismatch" else "" + state.index = ( + f'source' + ) + elif failure == "corrupt-tar": + state.payload = b"not an archive" + elif failure == "missing-dashboard": + state.payload = _sdist({"astrbot/dashboard/dist/index.html": None}) + elif failure == "incomplete-dashboard": + state.payload = _sdist({"astrbot/dashboard/dist/assets/app.js": None}) + elif failure == "wrong-dashboard-version": + state.payload = _sdist({"astrbot/dashboard/dist/assets/version": "v98.0.0"}) + elif failure == "missing-core": + state.payload = _sdist({"main.py": None}) + elif failure == "wrong-metadata": + state.payload = _sdist({"PKG-INFO": "Name: AstrBot\nVersion: 98.0.0\n"}) + elif failure in ( + "traversal", + "absolute", + "windows-path", + "multiple-roots", + "duplicate", + ): + path = { + "traversal": "astrbot-99.0.0/../outside", + "absolute": "/outside", + "windows-path": "astrbot-99.0.0/C:\\outside", + "multiple-roots": "other/outside", + "duplicate": "astrbot-99.0.0/main.py", + }[failure] + state.payload = _sdist(extra=tarfile.TarInfo(path)) + elif failure in ("symlink", "hardlink"): + member = tarfile.TarInfo("astrbot-99.0.0/link") + member.type = tarfile.SYMTYPE if failure == "symlink" else tarfile.LNKTYPE + member.linkname = "../../outside" + state.payload = _sdist(extra=member) + elif failure == "yanked": + state.index = 'yanked' + elif failure == "wheel-only": + state.index = 'wheel' + await state.updater.update("v99.0.0", proxy="https://github-proxy.example") + assert [kind for kind, _ in state.fallback] == ["dashboard", "core"] + assert all( + kwargs["proxy"] == "https://github-proxy.example" + for _, kwargs in state.fallback + ) + assert (state.install / "main.py").read_text() == "fallback core" + assert (state.data / "dist/assets/fallback.js").is_file() + assert not (state.install / "astrbot").exists() + assert not list((state.temp / "updates").iterdir()) + state.releases.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_commit_updates_skip_pypi(mirror_update): + state = mirror_update + revision = "a" * 40 + await state.updater.update(revision) + assert state.requests == [] + state.releases.assert_not_awaited() + assert all(kwargs["version"] == revision for _, kwargs in state.fallback) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("environment", ["ASTRBOT_CLI", "ASTRBOT_LAUNCHER"]) +async def test_managed_installs_cannot_bypass_update_restriction( + mirror_update, monkeypatch, environment +): + monkeypatch.setenv(environment, "1") + with pytest.raises(RuntimeError, match="pip or uv tool upgrade"): + await mirror_update.updater.update("v99.0.0") + assert mirror_update.requests == [] + assert mirror_update.fallback == [] diff --git a/tests/test_updater_socks.py b/tests/test_updater_socks.py index 15bc1f168c..fc1250b01e 100644 --- a/tests/test_updater_socks.py +++ b/tests/test_updater_socks.py @@ -4,6 +4,7 @@ from dataclasses import dataclass, field from pathlib import Path from types import SimpleNamespace +from unittest.mock import AsyncMock from urllib.parse import urlparse import certifi @@ -521,6 +522,7 @@ async def test_astrbot_updater_prepares_both_packages_before_applying( tmp_path: Path, ) -> None: updater = AstrBotUpdater() + monkeypatch.setattr(updater, "_download_pypi_package", AsyncMock(return_value=False)) calls: list[str] = [] progress_events: list[UpdateProgress] = [] fetch_count = 0 @@ -650,6 +652,7 @@ async def test_astrbot_updater_does_not_apply_unverified_packages( tmp_path: Path, ) -> None: updater = AstrBotUpdater() + monkeypatch.setattr(updater, "_download_pypi_package", AsyncMock(return_value=False)) calls: list[str] = [] async def fake_fetch_release_info(_url: str):