mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-09-01 15:32:49 +08:00
fix: support SOCKS proxies in updater requests (#7615)
* fix: support SOCKS proxies in updater requests * fix: log updater HTTP status details * fix: clean partial updater downloads on failure * test: lock updater httpx client options * refactor: harden updater httpx configuration
This commit is contained in:
エイカク
@@ -11,8 +11,8 @@ from ..updator import RepoZipUpdator
|
||||
|
||||
|
||||
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:
|
||||
|
||||
@@ -7,7 +7,6 @@ import psutil
|
||||
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 @@ class AstrBotUpdator(RepoZipUpdator):
|
||||
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:
|
||||
|
||||
+52
-27
@@ -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 @@ class ReleaseInfo:
|
||||
|
||||
|
||||
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 @@ class RepoZipUpdator:
|
||||
"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 @@ class RepoZipUpdator:
|
||||
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` 结构
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user