diff --git a/astrbot/core/star/updator.py b/astrbot/core/star/updator.py index 1a0c5fc260..2ae0039db0 100644 --- a/astrbot/core/star/updator.py +++ b/astrbot/core/star/updator.py @@ -11,8 +11,8 @@ class PluginUpdator(RepoZipUpdator): - def __init__(self, repo_mirror: str = "") -> None: - super().__init__(repo_mirror) + def __init__(self, repo_mirror: str = "", verify: str | bool | None = None) -> None: + super().__init__(repo_mirror, verify=verify) self.plugin_store_path = get_astrbot_plugin_path() def get_plugin_store_path(self) -> str: diff --git a/astrbot/core/updator.py b/astrbot/core/updator.py index df2cfb82cd..73c837b889 100644 --- a/astrbot/core/updator.py +++ b/astrbot/core/updator.py @@ -7,7 +7,6 @@ from astrbot.core import logger from astrbot.core.config.default import VERSION from astrbot.core.utils.astrbot_path import get_astrbot_path -from astrbot.core.utils.io import download_file from .zip_updator import ReleaseInfo, RepoZipUpdator @@ -18,8 +17,8 @@ class AstrBotUpdator(RepoZipUpdator): 功能包括检查更新、下载更新文件、解压缩更新文件等 """ - def __init__(self, repo_mirror: str = "") -> None: - super().__init__(repo_mirror) + def __init__(self, repo_mirror: str = "", verify: str | bool | None = None) -> None: + super().__init__(repo_mirror, verify=verify) self.MAIN_PATH = get_astrbot_path() self.ASTRBOT_RELEASE_API = "https://api.soulter.top/releases" @@ -176,7 +175,7 @@ async def update(self, reboot=False, latest=True, version=None, proxy="") -> Non file_url = f"{proxy}/{file_url}" try: - await download_file(file_url, "temp.zip") + await self._download_file(file_url, "temp.zip") logger.info("下载 AstrBot Core 更新文件完成,正在执行解压...") self.unzip_file("temp.zip", self.MAIN_PATH) except BaseException as e: diff --git a/astrbot/core/zip_updator.py b/astrbot/core/zip_updator.py index 6cea6b38d5..116b55c57d 100644 --- a/astrbot/core/zip_updator.py +++ b/astrbot/core/zip_updator.py @@ -1,15 +1,15 @@ import os import re import shutil -import ssl import zipfile +from pathlib import Path from typing import NoReturn -import aiohttp import certifi +import httpx from astrbot.core import logger -from astrbot.core.utils.io import download_file, on_error +from astrbot.core.utils.io import on_error from astrbot.core.utils.version_comparator import VersionComparator @@ -33,36 +33,53 @@ def __str__(self) -> str: class RepoZipUpdator: - def __init__(self, repo_mirror: str = "") -> None: + def __init__(self, repo_mirror: str = "", verify: str | bool | None = None) -> None: self.repo_mirror = repo_mirror self.rm_on_error = on_error + self.httpx_verify = certifi.where() if verify is None else verify + + def _create_httpx_client(self, timeout: float = 30.0) -> httpx.AsyncClient: + return httpx.AsyncClient( + follow_redirects=True, + timeout=timeout, + trust_env=True, + verify=self.httpx_verify, + ) + + @staticmethod + def _truncate_response_body(body: str, max_len: int = 1000) -> str: + if len(body) <= max_len: + return body + return body[:max_len] + "...[truncated]" + + async def _download_file( + self, url: str, path: str, timeout: float = 1800.0 + ) -> None: + target_path = Path(path) + target_path.parent.mkdir(parents=True, exist_ok=True) + + try: + async with self._create_httpx_client(timeout=timeout) as client: + async with client.stream("GET", url) as response: + response.raise_for_status() + with target_path.open("wb") as file: + async for chunk in response.aiter_bytes(8192): + file.write(chunk) + except Exception as e: + logger.error(f"下载文件失败: {url} -> {target_path}, 错误: {e}") + if self.rm_on_error and target_path.exists(): + target_path.unlink() + raise async def fetch_release_info(self, url: str, latest: bool = True) -> list: """请求版本信息。 返回一个列表,每个元素是一个字典,包含版本号、发布时间、更新内容、commit hash等信息。 """ try: - ssl_context = ssl.create_default_context( - cafile=certifi.where(), - ) # 新增:创建基于 certifi 的 SSL 上下文 - connector = aiohttp.TCPConnector( - ssl=ssl_context, - ) # 新增:使用 TCPConnector 指定 SSL 上下文 - async with ( - aiohttp.ClientSession( - trust_env=True, - connector=connector, - ) as session, - session.get(url) as response, - ): - # 检查 HTTP 状态码 - if response.status != 200: - text = await response.text() - logger.error( - f"请求 {url} 失败,状态码: {response.status}, 内容: {text}", - ) - raise Exception(f"请求失败,状态码: {response.status}") - result = await response.json() + async with self._create_httpx_client() as client: + response = await client.get(url) + response.raise_for_status() + result = response.json() if not result: return [] # if latest: @@ -80,9 +97,17 @@ async def fetch_release_info(self, url: str, latest: bool = True) -> list: "zipball_url": release["zipball_url"], }, ) + except httpx.HTTPStatusError as e: + response_body = "" + if e.response is not None: + response_body = self._truncate_response_body(e.response.text) + logger.error( + f"请求 {url} 失败,状态码: {e.response.status_code}, 内容: {response_body}", + ) + raise Exception("解析版本信息失败") from e except Exception as e: logger.error(f"解析版本信息时发生异常: {e}") - raise Exception("解析版本信息失败") + raise Exception("解析版本信息失败") from e return ret def github_api_release_parser(self, releases: list) -> list: @@ -186,7 +211,7 @@ async def download_from_repo_url( f"检查到设置了镜像站,将使用镜像站下载 {author}/{repo} 仓库源码: {release_url}", ) - await download_file(release_url, target_path + ".zip") + await self._download_file(release_url, target_path + ".zip") def parse_github_url(self, url: str): """使用正则表达式解析 GitHub 仓库 URL,支持 `.git` 后缀和 `tree/branch` 结构 diff --git a/tests/test_updator_socks.py b/tests/test_updator_socks.py new file mode 100644 index 0000000000..6dec6fe0f4 --- /dev/null +++ b/tests/test_updator_socks.py @@ -0,0 +1,359 @@ +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace + +import certifi +import httpx +import pytest + +from astrbot.core.zip_updator import RepoZipUpdator + + +class _FakeJSONResponse: + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self): + return self._payload + + +class _FakeStreamResponse: + def __init__(self, payload: bytes): + self._payload = payload + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + def raise_for_status(self) -> None: + return None + + async def aiter_bytes(self, chunk_size: int = 8192): + for start in range(0, len(self._payload), chunk_size): + yield self._payload[start : start + chunk_size] + + +class _FakeFailingStreamResponse: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + def raise_for_status(self) -> None: + return None + + async def aiter_bytes(self, chunk_size: int = 8192): # noqa: ARG002 + yield b"partial" + raise RuntimeError("stream interrupted") + + +class _FakeStatusErrorResponse: + def __init__(self, status_code: int, body: str, url: str): + self._status_code = status_code + self._body = body + self._url = url + + def raise_for_status(self) -> None: + request = httpx.Request("GET", self._url) + response = httpx.Response( + self._status_code, + text=self._body, + request=request, + ) + raise httpx.HTTPStatusError( + "status error", + request=request, + response=response, + ) + + +@dataclass +class _FakeAsyncClientState: + json_payload: list[dict] = field(default_factory=list) + stream_payload: bytes = b"" + init_kwargs: dict | None = None + requested_urls: list[str] = field(default_factory=list) + stream_urls: list[str] = field(default_factory=list) + + +class _FakeStatusErrorAsyncClient: + def __init__(self, response: _FakeStatusErrorResponse): + self._response = response + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + async def get(self, url: str): + return self._response + + +class _FakeFailingStreamAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + def stream(self, method: str, url: str): # noqa: ARG002 + return _FakeFailingStreamResponse() + + +def _build_fake_httpx_module(state: _FakeAsyncClientState) -> SimpleNamespace: + class _FakeAsyncClient: + def __init__(self, **kwargs): + state.init_kwargs = kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + async def get(self, url: str): + state.requested_urls.append(url) + return _FakeJSONResponse(state.json_payload) + + def stream(self, method: str, url: str): + assert method == "GET" + state.stream_urls.append(url) + return _FakeStreamResponse(state.stream_payload) + + return SimpleNamespace( + AsyncClient=_FakeAsyncClient, + HTTPStatusError=httpx.HTTPStatusError, + ) + + +@pytest.fixture +def fake_async_client_state() -> _FakeAsyncClientState: + return _FakeAsyncClientState() + + +@pytest.mark.asyncio +async def test_fetch_release_info_uses_httpx_client_with_env_proxy_support( + monkeypatch: pytest.MonkeyPatch, + fake_async_client_state: _FakeAsyncClientState, +) -> None: + import astrbot.core.zip_updator as zip_updator_module + + fake_async_client_state.json_payload = [ + { + "name": "AstrBot v4.23.2", + "published_at": "2026-04-16T00:00:00Z", + "body": "fix updater socks proxy support", + "tag_name": "v4.23.2", + "zipball_url": "https://example.com/astrbot.zip", + } + ] + + monkeypatch.setattr( + zip_updator_module, + "aiohttp", + SimpleNamespace( + ClientSession=lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError( + "fetch_release_info should not use aiohttp.ClientSession" + ) + ) + ), + raising=False, + ) + monkeypatch.setattr( + zip_updator_module, + "httpx", + _build_fake_httpx_module(fake_async_client_state), + raising=False, + ) + + release_info = await RepoZipUpdator().fetch_release_info( + "https://api.soulter.top/releases" + ) + + assert release_info == [ + { + "version": "AstrBot v4.23.2", + "published_at": "2026-04-16T00:00:00Z", + "body": "fix updater socks proxy support", + "tag_name": "v4.23.2", + "zipball_url": "https://example.com/astrbot.zip", + } + ] + assert fake_async_client_state.requested_urls == ["https://api.soulter.top/releases"] + assert fake_async_client_state.init_kwargs is not None + assert fake_async_client_state.init_kwargs["follow_redirects"] is True + assert fake_async_client_state.init_kwargs["timeout"] == 30.0 + assert fake_async_client_state.init_kwargs["trust_env"] is True + assert fake_async_client_state.init_kwargs["verify"] == certifi.where() + + +@pytest.mark.asyncio +async def test_download_from_repo_url_uses_httpx_stream_for_zip_download( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + fake_async_client_state: _FakeAsyncClientState, +) -> None: + import astrbot.core.zip_updator as zip_updator_module + + fake_async_client_state.stream_payload = b"zip-data" + + async def fake_fetch_release_info(self, url: str, latest: bool = True): # noqa: ARG001 + return [ + { + "version": "AstrBot v4.23.2", + "published_at": "2026-04-16T00:00:00Z", + "body": "fix updater socks proxy support", + "tag_name": "v4.23.2", + "zipball_url": "https://example.com/archive.zip", + } + ] + + monkeypatch.setattr(RepoZipUpdator, "fetch_release_info", fake_fetch_release_info) + monkeypatch.setattr( + zip_updator_module, + "download_file", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("download_from_repo_url should not use aiohttp download_file") + ), + raising=False, + ) + monkeypatch.setattr( + zip_updator_module, + "httpx", + _build_fake_httpx_module(fake_async_client_state), + raising=False, + ) + + target_path = tmp_path / "AstrBot" + await RepoZipUpdator().download_from_repo_url( + str(target_path), + "https://github.com/AstrBotDevs/AstrBot", + ) + + assert (tmp_path / "AstrBot.zip").read_bytes() == b"zip-data" + assert fake_async_client_state.stream_urls == ["https://example.com/archive.zip"] + assert fake_async_client_state.init_kwargs is not None + assert fake_async_client_state.init_kwargs["follow_redirects"] is True + assert fake_async_client_state.init_kwargs["timeout"] == 1800.0 + assert fake_async_client_state.init_kwargs["trust_env"] is True + assert fake_async_client_state.init_kwargs["verify"] == certifi.where() + + +def test_create_httpx_client_uses_custom_verify_setting( + monkeypatch: pytest.MonkeyPatch, + fake_async_client_state: _FakeAsyncClientState, +) -> None: + import astrbot.core.zip_updator as zip_updator_module + + custom_verify = "/tmp/custom-ca.pem" + + monkeypatch.setattr( + zip_updator_module, + "httpx", + _build_fake_httpx_module(fake_async_client_state), + raising=False, + ) + + RepoZipUpdator(verify=custom_verify)._create_httpx_client(timeout=45.0) + + assert fake_async_client_state.init_kwargs is not None + assert fake_async_client_state.init_kwargs["follow_redirects"] is True + assert fake_async_client_state.init_kwargs["timeout"] == 45.0 + assert fake_async_client_state.init_kwargs["trust_env"] is True + assert fake_async_client_state.init_kwargs["verify"] == custom_verify + + +@pytest.mark.asyncio +async def test_fetch_release_info_logs_status_code_and_truncated_body_on_http_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import astrbot.core.zip_updator as zip_updator_module + + url = "https://api.soulter.top/releases" + body = "x" * 1005 + log_messages: list[str] = [] + + monkeypatch.setattr( + RepoZipUpdator, + "_create_httpx_client", + staticmethod( + lambda timeout=30.0: _FakeStatusErrorAsyncClient( # noqa: ARG005 + _FakeStatusErrorResponse(502, body, url) + ) + ), + ) + monkeypatch.setattr( + zip_updator_module.logger, + "error", + lambda message: log_messages.append(message), + ) + + with pytest.raises(Exception, match="解析版本信息失败"): + await RepoZipUpdator().fetch_release_info(url) + + assert any("状态码: 502" in message for message in log_messages) + assert any("内容: " in message for message in log_messages) + assert any("...[truncated]" in message for message in log_messages) + + +@pytest.mark.asyncio +async def test_download_file_removes_partial_file_when_stream_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr( + RepoZipUpdator, + "_create_httpx_client", + staticmethod( + lambda timeout=30.0: _FakeFailingStreamAsyncClient() # noqa: ARG005 + ), + ) + + target_path = tmp_path / "partial.zip" + + with pytest.raises(RuntimeError, match="stream interrupted"): + await RepoZipUpdator()._download_file( + "https://example.com/archive.zip", + str(target_path), + ) + + assert not target_path.exists() + + +@pytest.mark.asyncio +async def test_download_file_logs_url_and_target_path_on_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import astrbot.core.zip_updator as zip_updator_module + + url = "https://example.com/archive.zip" + target_path = tmp_path / "logged-partial.zip" + log_messages: list[str] = [] + + monkeypatch.setattr( + RepoZipUpdator, + "_create_httpx_client", + staticmethod( + lambda timeout=30.0: _FakeFailingStreamAsyncClient() # noqa: ARG005 + ), + ) + monkeypatch.setattr( + zip_updator_module.logger, + "error", + lambda message: log_messages.append(message), + ) + + with pytest.raises(RuntimeError, match="stream interrupted"): + await RepoZipUpdator()._download_file(url, str(target_path)) + + assert any(url in message for message in log_messages) + assert any(str(target_path) in message for message in log_messages)