mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 10:57:03 +08:00
refactor(proxy): 将 aether-proxy 从 HMAC 正向代理迁移到 WebSocket 隧道模式
移除 HMAC 认证、TLS 自签名证书、HTTP CONNECT 代理和代发(delegate)模式, 改为 aether-proxy 主动通过 WebSocket 连接 Aether 服务端建立隧道。 Aether 服务端新增: - WebSocket 隧道端点 (proxy_tunnel.py) - TunnelManager 管理隧道连接和请求分发 - TunnelTransport 作为 httpx 自定义 transport 层 - 基于二进制帧的隧道协议 (tunnel_protocol.py) aether-proxy (Rust) 重构: - 新增 tunnel 模块 (client/dispatcher/stream_handler/protocol) - 支持多 Aether 服务端连接 ([[servers]] 配置) - 移除 proxy/auth/delegate 模块和 hyper 依赖 - 改用 tokio-tungstenite 实现 WebSocket 客户端 同时: - 添加浏览器指纹 Headers 绕过 Cloudflare 防护 - 删除节点时自动清理 Provider/Endpoint 的代理引用 - 数据库迁移: 新增 tunnel_mode/tunnel_connected/tunnel_connected_at 字段
This commit is contained in:
@@ -32,7 +32,7 @@ pipeline = ApiRequestPipeline()
|
||||
class ProxyNodeRegisterRequest(BaseModel):
|
||||
name: str = Field(..., min_length=1, max_length=100, description="节点名")
|
||||
ip: str = Field(..., description="公网 IP(IPv4/IPv6)")
|
||||
port: int = Field(..., ge=1, le=65535, description="代理端口")
|
||||
port: int = Field(0, ge=0, le=65535, description="代理端口(tunnel 模式下为 0)")
|
||||
region: str | None = Field(None, max_length=100, description="区域标签")
|
||||
heartbeat_interval: int = Field(30, ge=5, le=600, description="心跳间隔(秒)")
|
||||
|
||||
@@ -41,16 +41,13 @@ class ProxyNodeRegisterRequest(BaseModel):
|
||||
total_requests: int | None = Field(None, ge=0, description="累计请求数")
|
||||
avg_latency_ms: float | None = Field(None, ge=0, description="平均延迟(毫秒)")
|
||||
|
||||
# TLS
|
||||
tls_enabled: bool = Field(False, description="是否启用 TLS 加密")
|
||||
tls_cert_fingerprint: str | None = Field(
|
||||
None, max_length=128, description="TLS 证书 SHA-256 指纹"
|
||||
)
|
||||
|
||||
# 硬件信息
|
||||
hardware_info: dict | None = Field(None, description="硬件信息 JSON")
|
||||
estimated_max_concurrency: int | None = Field(None, ge=0, description="估算最大并发连接数")
|
||||
|
||||
# Tunnel 模式
|
||||
tunnel_mode: bool = Field(False, description="是否使用 tunnel 模式连接")
|
||||
|
||||
@field_validator("ip")
|
||||
@classmethod
|
||||
def validate_ip(cls, v: str) -> str:
|
||||
@@ -82,9 +79,6 @@ class ProxyNodeRemoteConfigRequest(BaseModel):
|
||||
allowed_ports: list[int] | None = Field(None, description="允许代理的目标端口")
|
||||
log_level: str | None = Field(None, description="日志级别 (trace/debug/info/warn/error)")
|
||||
heartbeat_interval: int | None = Field(None, ge=5, le=600, description="心跳间隔(秒)")
|
||||
timestamp_tolerance: int | None = Field(
|
||||
None, ge=10, le=3600, description="HMAC 时间戳容差(秒)"
|
||||
)
|
||||
|
||||
@field_validator("allowed_ports")
|
||||
@classmethod
|
||||
@@ -216,12 +210,6 @@ async def test_proxy_node(node_id: str, request: Request, db: Session = Depends(
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.get("/hmac-key")
|
||||
async def get_proxy_hmac_key(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||
adapter = AdminGetProxyHmacKeyAdapter()
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.post("/test-url")
|
||||
async def test_proxy_url(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||
adapter = AdminTestProxyUrlAdapter()
|
||||
@@ -273,14 +261,13 @@ class AdminRegisterProxyNodeAdapter(AdminApiAdapter):
|
||||
port=req.port,
|
||||
region=req.region,
|
||||
heartbeat_interval=req.heartbeat_interval,
|
||||
tls_enabled=req.tls_enabled,
|
||||
tls_cert_fingerprint=req.tls_cert_fingerprint,
|
||||
hardware_info=req.hardware_info,
|
||||
estimated_max_concurrency=req.estimated_max_concurrency,
|
||||
active_connections=req.active_connections,
|
||||
total_requests=req.total_requests,
|
||||
avg_latency_ms=req.avg_latency_ms,
|
||||
registered_by=context.user.id if context.user else None,
|
||||
tunnel_mode=req.tunnel_mode,
|
||||
)
|
||||
|
||||
context.add_audit_metadata(
|
||||
@@ -376,11 +363,24 @@ class AdminDeleteProxyNodeAdapter(AdminApiAdapter):
|
||||
)
|
||||
|
||||
was_system_proxy = result["cleared_system_proxy"]
|
||||
msg = "deleted, system default proxy cleared" if was_system_proxy else "deleted"
|
||||
cleared_providers = result.get("cleared_providers", 0)
|
||||
cleared_endpoints = result.get("cleared_endpoints", 0)
|
||||
|
||||
parts = ["deleted"]
|
||||
if was_system_proxy:
|
||||
parts.append("system default proxy cleared")
|
||||
if cleared_providers or cleared_endpoints:
|
||||
parts.append(
|
||||
f"cleared proxy from {cleared_providers} provider(s) "
|
||||
f"and {cleared_endpoints} endpoint(s)"
|
||||
)
|
||||
|
||||
return {
|
||||
"message": msg,
|
||||
"message": ", ".join(parts),
|
||||
"node_id": self.node_id,
|
||||
"cleared_system_proxy": was_system_proxy,
|
||||
"cleared_providers": cleared_providers,
|
||||
"cleared_endpoints": cleared_endpoints,
|
||||
}
|
||||
|
||||
|
||||
@@ -478,8 +478,6 @@ class AdminUpdateProxyNodeConfigAdapter(AdminApiAdapter):
|
||||
config_updates["log_level"] = req.log_level
|
||||
if req.heartbeat_interval is not None:
|
||||
config_updates["heartbeat_interval"] = req.heartbeat_interval
|
||||
if req.timestamp_tolerance is not None:
|
||||
config_updates["timestamp_tolerance"] = req.timestamp_tolerance
|
||||
|
||||
node = ProxyNodeService.update_node_config(
|
||||
context.db, node_id=self.node_id, config_updates=config_updates
|
||||
@@ -499,23 +497,6 @@ class AdminUpdateProxyNodeConfigAdapter(AdminApiAdapter):
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdminGetProxyHmacKeyAdapter(AdminApiAdapter):
|
||||
"""获取 proxy_hmac_key 供管理员复制到 aether-proxy 部署"""
|
||||
|
||||
name: str = "admin_get_proxy_hmac_key"
|
||||
|
||||
async def handle(self, context: ApiRequestContext) -> Any:
|
||||
from src.config.settings import config
|
||||
|
||||
key = config.proxy_hmac_key
|
||||
if not key:
|
||||
raise InvalidRequestException(
|
||||
"PROXY_HMAC_KEY 未配置(也未设置 ENCRYPTION_KEY 用于自动派生)"
|
||||
)
|
||||
return {"proxy_hmac_key": key}
|
||||
|
||||
|
||||
class TestProxyUrlRequest(BaseModel):
|
||||
proxy_url: str = Field(..., min_length=1, max_length=500)
|
||||
username: str | None = Field(None, max_length=255)
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
WebSocket 隧道端点
|
||||
|
||||
aether-proxy 通过此端点建立 tunnel 连接。
|
||||
路径: /api/internal/proxy-tunnel
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.proxy_node.tunnel_manager import (
|
||||
TunnelConnection,
|
||||
get_tunnel_manager,
|
||||
)
|
||||
from src.services.proxy_node.tunnel_protocol import Frame
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# 单帧最大 64 MB -- AI API 请求体可能包含多张 base64 图片,需要足够余量
|
||||
_MAX_FRAME_SIZE = 64 * 1024 * 1024
|
||||
|
||||
# WebSocket 空闲超时(秒)-- proxy 端 ping 间隔默认 15s,3 倍余量
|
||||
_IDLE_TIMEOUT = 90.0
|
||||
|
||||
|
||||
async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
||||
"""验证 WebSocket 连接的认证信息,返回 (node_id, node_name) 或 None
|
||||
|
||||
认证方式:Bearer <management_token>,通过 Management Token 系统验证。
|
||||
authenticate_management_token 是 async 方法(内部有 Redis 速率限制),
|
||||
因此直接 await 调用。节点存在性检查复用同一 session。
|
||||
"""
|
||||
auth = ws.headers.get("authorization", "")
|
||||
if not auth.startswith("Bearer "):
|
||||
return None
|
||||
|
||||
token = auth[7:]
|
||||
if not token or not token.startswith("ae_"):
|
||||
return None
|
||||
|
||||
client_ip = getattr(ws.client, "host", "unknown") if ws.client else "unknown"
|
||||
node_id_header = ws.headers.get("x-node-id", "").strip()
|
||||
node_name_header = ws.headers.get("x-node-name", "").strip()
|
||||
|
||||
if not node_id_header:
|
||||
return None
|
||||
|
||||
from src.database import create_session
|
||||
from src.models.database import ProxyNode
|
||||
from src.services.auth.service import AuthService
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
result = await AuthService.authenticate_management_token(db, token, client_ip)
|
||||
if not result:
|
||||
return None
|
||||
|
||||
# 节点存在性检查(复用同一 session,避免额外连接开销)
|
||||
exists = db.query(
|
||||
db.query(ProxyNode).filter(ProxyNode.id == node_id_header).exists()
|
||||
).scalar()
|
||||
if not exists:
|
||||
logger.warning("tunnel auth: node_id={} not found in DB", node_id_header)
|
||||
return None
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
return node_id_header, node_name_header or node_id_header
|
||||
|
||||
|
||||
@router.websocket("/api/internal/proxy-tunnel")
|
||||
async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
||||
"""aether-proxy tunnel WebSocket 端点"""
|
||||
try:
|
||||
auth = await _authenticate(ws)
|
||||
except Exception as e:
|
||||
logger.warning("tunnel auth error: {}", e)
|
||||
await ws.accept()
|
||||
await ws.close(code=4002, reason="authentication error")
|
||||
return
|
||||
|
||||
if not auth:
|
||||
await ws.accept()
|
||||
await ws.close(code=4001, reason="unauthorized")
|
||||
return
|
||||
|
||||
node_id: str = auth[0]
|
||||
node_name: str = auth[1]
|
||||
await ws.accept()
|
||||
|
||||
manager = get_tunnel_manager()
|
||||
conn = TunnelConnection(node_id, node_name, ws)
|
||||
manager.register(conn)
|
||||
|
||||
# 更新 DB: tunnel_connected = True
|
||||
await _update_tunnel_status(node_id, connected=True)
|
||||
|
||||
try:
|
||||
oversized_count = 0
|
||||
while True:
|
||||
try:
|
||||
data = await asyncio.wait_for(ws.receive_bytes(), timeout=_IDLE_TIMEOUT)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("tunnel idle timeout for node_id={}", node_id)
|
||||
await ws.close(code=4004, reason="idle timeout")
|
||||
break
|
||||
if len(data) > _MAX_FRAME_SIZE:
|
||||
oversized_count += 1
|
||||
logger.warning("tunnel frame too large from {}: {} bytes", node_id, len(data))
|
||||
if oversized_count >= 5:
|
||||
logger.warning("too many oversized frames from {}, closing", node_id)
|
||||
await ws.close(code=4003, reason="too many oversized frames")
|
||||
break
|
||||
continue
|
||||
oversized_count = 0 # 正常帧重置计数
|
||||
try:
|
||||
frame = Frame.decode(data)
|
||||
except ValueError as e:
|
||||
logger.warning("tunnel frame decode error from {}: {}", node_id, e)
|
||||
continue
|
||||
|
||||
await manager.handle_incoming_frame(node_id, frame)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info("tunnel WebSocket disconnected: node_id={}", node_id)
|
||||
except Exception as e:
|
||||
logger.error("tunnel WebSocket error for node_id={}: {}", node_id, e)
|
||||
finally:
|
||||
manager.unregister(node_id)
|
||||
await _update_tunnel_status(node_id, connected=False)
|
||||
|
||||
|
||||
async def _update_tunnel_status(node_id: str, *, connected: bool) -> None:
|
||||
"""更新 ProxyNode 的 tunnel 连接状态(在线程池中执行,避免阻塞 event loop)"""
|
||||
|
||||
def _sync_update() -> None:
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from src.database import create_session
|
||||
from src.models.database import ProxyNode, ProxyNodeStatus
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
|
||||
if node:
|
||||
node.tunnel_connected = connected
|
||||
now = datetime.now(timezone.utc)
|
||||
if connected:
|
||||
node.tunnel_connected_at = now
|
||||
node.status = ProxyNodeStatus.ONLINE
|
||||
else:
|
||||
# 记录断开时刻,供 health_scheduler 计算 UNHEALTHY 缓冲期
|
||||
node.tunnel_connected_at = now
|
||||
node.status = ProxyNodeStatus.UNHEALTHY
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(_sync_update)
|
||||
except Exception as e:
|
||||
logger.warning("failed to update tunnel status for {}: {}", node_id, e)
|
||||
|
||||
# 清除节点信息缓存,确保后续请求能立即感知连接状态变化
|
||||
from src.services.proxy_node.resolver import invalidate_proxy_node_cache
|
||||
|
||||
invalidate_proxy_node_cache(node_id)
|
||||
@@ -16,6 +16,7 @@ from src.api.handlers.base.chat_adapter_base import ChatAdapterBase, register_ad
|
||||
from src.api.handlers.base.chat_handler_base import ChatHandlerBase
|
||||
from src.core.api_format import ApiFamily, get_auth_handler
|
||||
from src.core.api_format.enums import AuthMethod
|
||||
from src.core.api_format.headers import BROWSER_FINGERPRINT_HEADERS
|
||||
from src.core.logger import logger
|
||||
from src.models.gemini import GeminiRequest
|
||||
from src.services.provider.transport import redact_url_for_log
|
||||
@@ -197,7 +198,7 @@ class GeminiChatAdapter(ChatAdapterBase):
|
||||
else:
|
||||
models_url = f"{base_url_clean}/v1beta/models?key={api_key}"
|
||||
|
||||
headers: dict[str, str] = {}
|
||||
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
|
||||
+55
-98
@@ -51,8 +51,8 @@ class HTTPClientPool:
|
||||
_proxy_clients: dict[str, tuple[httpx.AsyncClient, float]] = {}
|
||||
# 代理客户端缓存上限(避免内存泄漏)
|
||||
_max_proxy_clients: int = 50
|
||||
# 代发客户端缓存:{tls: client, plain: client}
|
||||
_delegate_clients: dict[str, httpx.AsyncClient] = {}
|
||||
# Tunnel 客户端缓存:{node_id: client}
|
||||
_tunnel_clients: dict[str, httpx.AsyncClient] = {}
|
||||
|
||||
def __new__(cls) -> "HTTPClientPool":
|
||||
if cls._instance is None:
|
||||
@@ -289,15 +289,15 @@ class HTTPClientPool:
|
||||
|
||||
cls._proxy_clients.clear()
|
||||
|
||||
# 关闭代发客户端缓存
|
||||
for cache_key, client in cls._delegate_clients.items():
|
||||
# 关闭 tunnel 客户端缓存
|
||||
for nid, client in cls._tunnel_clients.items():
|
||||
try:
|
||||
await client.aclose()
|
||||
logger.debug("代发客户端已关闭: {}", cache_key)
|
||||
logger.debug("tunnel 客户端已关闭: {}", nid)
|
||||
except Exception as e:
|
||||
logger.warning("关闭代发客户端失败: {}", e)
|
||||
logger.warning("关闭 tunnel 客户端失败: {}", e)
|
||||
|
||||
cls._delegate_clients.clear()
|
||||
cls._tunnel_clients.clear()
|
||||
logger.info("所有HTTP客户端已关闭")
|
||||
|
||||
@classmethod
|
||||
@@ -380,88 +380,6 @@ class HTTPClientPool:
|
||||
client_config.update(kwargs)
|
||||
return httpx.AsyncClient(**client_config) # type: ignore[arg-type]
|
||||
|
||||
@classmethod
|
||||
def create_delegate_stream_client(
|
||||
cls,
|
||||
delegate_config: dict[str, Any],
|
||||
timeout: httpx.Timeout | None = None,
|
||||
) -> httpx.AsyncClient:
|
||||
"""
|
||||
创建用于代发流式请求的 httpx 客户端
|
||||
|
||||
代发模式下不配置 proxy,直接 POST 到 proxy 的 /_aether/delegate 端点。
|
||||
调用者需要负责关闭返回的客户端。
|
||||
"""
|
||||
client_config: dict[str, Any] = {
|
||||
"http2": False,
|
||||
"follow_redirects": False,
|
||||
}
|
||||
|
||||
if timeout:
|
||||
client_config["timeout"] = timeout
|
||||
else:
|
||||
client_config["timeout"] = httpx.Timeout(
|
||||
connect=config.http_connect_timeout,
|
||||
read=config.http_read_timeout,
|
||||
write=config.http_write_timeout,
|
||||
pool=config.http_pool_timeout,
|
||||
)
|
||||
|
||||
if delegate_config.get("tls_enabled"):
|
||||
from src.utils.ssl_utils import get_proxy_ssl_context
|
||||
|
||||
client_config["verify"] = get_proxy_ssl_context()
|
||||
else:
|
||||
client_config["verify"] = get_ssl_context()
|
||||
|
||||
return httpx.AsyncClient(**client_config)
|
||||
|
||||
@classmethod
|
||||
async def get_delegate_client(
|
||||
cls,
|
||||
delegate_config: dict[str, Any],
|
||||
) -> httpx.AsyncClient:
|
||||
"""
|
||||
获取可复用的代发客户端(非流式请求用)
|
||||
|
||||
根据 TLS 状态缓存两个客户端(tls / plain),避免每次请求创建新客户端。
|
||||
当 tls_enabled=True 时使用 get_proxy_ssl_context()(信任自签名证书)。
|
||||
"""
|
||||
cache_key = "tls" if delegate_config.get("tls_enabled") else "plain"
|
||||
|
||||
lock = cls._get_proxy_clients_lock()
|
||||
async with lock:
|
||||
existing = cls._delegate_clients.get(cache_key)
|
||||
if existing and not existing.is_closed:
|
||||
return existing
|
||||
|
||||
if cache_key == "tls":
|
||||
from src.utils.ssl_utils import get_proxy_ssl_context
|
||||
|
||||
verify: Any = get_proxy_ssl_context()
|
||||
else:
|
||||
verify = get_ssl_context()
|
||||
|
||||
client = httpx.AsyncClient(
|
||||
http2=False,
|
||||
verify=verify,
|
||||
follow_redirects=False,
|
||||
timeout=httpx.Timeout(
|
||||
connect=config.http_connect_timeout,
|
||||
read=config.http_read_timeout,
|
||||
write=config.http_write_timeout,
|
||||
pool=config.http_pool_timeout,
|
||||
),
|
||||
limits=httpx.Limits(
|
||||
max_connections=config.http_max_connections,
|
||||
max_keepalive_connections=config.http_keepalive_connections,
|
||||
keepalive_expiry=config.http_keepalive_expiry,
|
||||
),
|
||||
)
|
||||
cls._delegate_clients[cache_key] = client
|
||||
logger.debug("创建代发客户端(缓存): {}", cache_key)
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
async def get_upstream_client(
|
||||
cls,
|
||||
@@ -469,31 +387,70 @@ class HTTPClientPool:
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
) -> httpx.AsyncClient:
|
||||
"""
|
||||
获取可复用的上游请求客户端(自动选择代发或代理模式)
|
||||
获取可复用的上游请求客户端(自动选择 tunnel/代理模式)
|
||||
|
||||
代发模式(delegate_cfg非空):返回代发客户端
|
||||
tunnel 模式(delegate_cfg.tunnel=True):返回 TunnelTransport 客户端
|
||||
直连/代理模式:返回代理客户端(含系统默认代理回退)
|
||||
"""
|
||||
if delegate_cfg:
|
||||
return await cls.get_delegate_client(delegate_cfg)
|
||||
if delegate_cfg and delegate_cfg.get("tunnel"):
|
||||
return await cls._get_tunnel_client(delegate_cfg["node_id"])
|
||||
return await cls.get_proxy_client(proxy_config=proxy_config)
|
||||
|
||||
@classmethod
|
||||
def create_upstream_stream_client(
|
||||
async def create_upstream_stream_client(
|
||||
cls,
|
||||
delegate_cfg: dict[str, Any] | None,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
timeout: httpx.Timeout | None = None,
|
||||
) -> httpx.AsyncClient:
|
||||
"""
|
||||
创建上游流式请求客户端(自动选择代发或代理模式)
|
||||
创建上游流式请求客户端(自动选择 tunnel/代理模式)
|
||||
|
||||
调用者需负责关闭返回的客户端。
|
||||
"""
|
||||
if delegate_cfg:
|
||||
return cls.create_delegate_stream_client(delegate_cfg, timeout=timeout)
|
||||
if delegate_cfg and delegate_cfg.get("tunnel"):
|
||||
return await cls._get_tunnel_client(delegate_cfg["node_id"], timeout=timeout)
|
||||
return cls.create_client_with_proxy(proxy_config=proxy_config, timeout=timeout)
|
||||
|
||||
@classmethod
|
||||
async def _get_tunnel_client(
|
||||
cls,
|
||||
node_id: str,
|
||||
timeout: httpx.Timeout | None = None,
|
||||
) -> httpx.AsyncClient:
|
||||
"""获取使用 TunnelTransport 的 httpx 客户端
|
||||
|
||||
当 timeout 为 None 时(非流式请求),返回按 node_id 缓存的 client,
|
||||
调用方不应关闭此 client,其生命周期由 HTTPClientPool 管理。
|
||||
当 timeout 非 None 时(流式请求),每次创建新 client,由调用方负责关闭。
|
||||
"""
|
||||
from src.services.proxy_node.tunnel_transport import TunnelTransport
|
||||
|
||||
t = timeout or httpx.Timeout(
|
||||
connect=config.http_connect_timeout,
|
||||
read=config.http_read_timeout,
|
||||
write=config.http_write_timeout,
|
||||
pool=config.http_pool_timeout,
|
||||
)
|
||||
timeout_secs = t.read if isinstance(t, httpx.Timeout) else 60.0
|
||||
|
||||
# 流式请求:每次创建新 client(调用方负责关闭)
|
||||
if timeout is not None:
|
||||
transport = TunnelTransport(node_id, timeout=timeout_secs or 60.0)
|
||||
return httpx.AsyncClient(transport=transport, timeout=t)
|
||||
|
||||
# 非流式请求:复用缓存的 client(加锁与 proxy_clients 保持一致)
|
||||
lock = cls._get_proxy_clients_lock()
|
||||
async with lock:
|
||||
existing = cls._tunnel_clients.get(node_id)
|
||||
if existing and not existing.is_closed:
|
||||
return existing
|
||||
|
||||
transport = TunnelTransport(node_id, timeout=timeout_secs or 60.0)
|
||||
client = httpx.AsyncClient(transport=transport, timeout=t)
|
||||
cls._tunnel_clients[node_id] = client
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def get_pool_stats(cls) -> dict[str, Any]:
|
||||
"""获取连接池统计信息"""
|
||||
@@ -502,7 +459,7 @@ class HTTPClientPool:
|
||||
"named_clients_count": len(cls._clients),
|
||||
"proxy_clients_count": len(cls._proxy_clients),
|
||||
"max_proxy_clients": cls._max_proxy_clients,
|
||||
"delegate_clients_count": len(cls._delegate_clients),
|
||||
"tunnel_clients_count": len(cls._tunnel_clients),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
从环境变量或 .env 文件加载配置
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -48,13 +46,6 @@ class Config:
|
||||
# 加密密钥配置(独立于JWT密钥,用于敏感数据加密)
|
||||
self.encryption_key = os.getenv("ENCRYPTION_KEY", None)
|
||||
|
||||
# 代理节点 HMAC 密钥(用于 aether-proxy 认证)
|
||||
proxy_hmac_key_env = os.getenv("PROXY_HMAC_KEY")
|
||||
if proxy_hmac_key_env and proxy_hmac_key_env.strip():
|
||||
self.proxy_hmac_key = proxy_hmac_key_env.strip()
|
||||
else:
|
||||
self.proxy_hmac_key = self._derive_proxy_hmac_key()
|
||||
|
||||
# 环境配置 - 智能检测
|
||||
# Docker 部署默认为生产环境,本地开发默认为开发环境
|
||||
is_docker = (
|
||||
@@ -327,20 +318,6 @@ class Config:
|
||||
# 验证连接池配置
|
||||
self._validate_pool_config()
|
||||
|
||||
def _derive_proxy_hmac_key(self) -> str:
|
||||
"""
|
||||
从 ENCRYPTION_KEY 派生 PROXY_HMAC_KEY
|
||||
|
||||
目的:避免把 ENCRYPTION_KEY 直接下发到 VPS(aether-proxy)。
|
||||
"""
|
||||
if not self.encryption_key:
|
||||
return ""
|
||||
return hmac.new(
|
||||
self.encryption_key.encode("utf-8"),
|
||||
b"aether-proxy-hmac-key-v1",
|
||||
hashlib.sha256,
|
||||
).hexdigest()
|
||||
|
||||
def _auto_pool_size(self) -> int:
|
||||
"""
|
||||
智能计算连接池大小 - 根据 Worker 数量和 PostgreSQL 限制计算
|
||||
|
||||
@@ -15,10 +15,12 @@ from __future__ import annotations
|
||||
from collections.abc import Set as AbstractSet
|
||||
from typing import Any
|
||||
|
||||
from src.core.api_format.enums import ApiFamily
|
||||
from src.core.api_format.metadata import (
|
||||
get_auth_config_for_endpoint,
|
||||
get_extra_headers_for_endpoint,
|
||||
get_protected_keys_for_endpoint,
|
||||
resolve_endpoint_definition,
|
||||
)
|
||||
from src.core.api_format.signature import EndpointSignature, parse_signature_key
|
||||
from src.core.logger import logger
|
||||
@@ -27,6 +29,37 @@ from src.core.logger import logger
|
||||
# 头部常量定义
|
||||
# =============================================================================
|
||||
|
||||
# 通用浏览器指纹 Headers,用于绕过 Cloudflare 等反爬防护
|
||||
# 基于 Electron 桌面客户端的真实请求头构建,作为所有 adapter 请求的底层默认值
|
||||
BROWSER_FINGERPRINT_HEADERS: dict[str, str] = {
|
||||
"User-Agent": (
|
||||
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
"Chrome/140.0.7339.249 Electron/38.7.0 Safari/537.36"
|
||||
),
|
||||
"Accept": "application/json",
|
||||
"Accept-Language": "zh-CN",
|
||||
"sec-ch-ua": '"Not=A?Brand";v="24", "Chromium";v="140"',
|
||||
"sec-ch-ua-mobile": "?0",
|
||||
"sec-ch-ua-platform": '"macOS"',
|
||||
"Sec-Fetch-Site": "cross-site",
|
||||
"Sec-Fetch-Mode": "cors",
|
||||
"Sec-Fetch-Dest": "empty",
|
||||
}
|
||||
|
||||
# Anthropic/Claude 专属 Headers(仅 Claude API family 使用)
|
||||
# 包含 Stainless SDK 指纹和 direct-browser-access 标记
|
||||
_ANTHROPIC_EXTRA_HEADERS: dict[str, str] = {
|
||||
"anthropic-dangerous-direct-browser-access": "true",
|
||||
"x-stainless-os": "Unknown",
|
||||
"x-stainless-runtime": "browser:chrome",
|
||||
"x-stainless-arch": "unknown",
|
||||
"x-stainless-lang": "js",
|
||||
"x-stainless-package-version": "0.41.0",
|
||||
"x-stainless-runtime-version": "140.0.7339",
|
||||
"x-stainless-retry-count": "0",
|
||||
}
|
||||
|
||||
# 转发给上游时需要剔除的头部(系统管理 + 认证替换 + 客户端/代理元数据)
|
||||
UPSTREAM_DROP_HEADERS: frozenset[str] = frozenset(
|
||||
{
|
||||
@@ -503,14 +536,23 @@ def build_adapter_base_headers_for_endpoint(
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
新模式:根据 endpoint signature 构建基础请求头。
|
||||
|
||||
浏览器指纹 headers 作为底层默认值注入,Claude API family 额外注入 Anthropic 专属 header。
|
||||
认证头和 extra_headers 会覆盖它们。
|
||||
"""
|
||||
auth_header, auth_type = get_auth_config_for_endpoint(endpoint)
|
||||
auth_value = f"Bearer {api_key}" if auth_type == "bearer" else api_key
|
||||
|
||||
headers: dict[str, str] = {
|
||||
auth_header: auth_value,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
# 以浏览器指纹为底层默认值,绕过 Cloudflare 等反爬防护
|
||||
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
|
||||
|
||||
# Claude API family 额外注入 Anthropic 专属 header
|
||||
definition = resolve_endpoint_definition(endpoint)
|
||||
if definition and definition.api_family == ApiFamily.CLAUDE:
|
||||
headers.update(_ANTHROPIC_EXTRA_HEADERS)
|
||||
|
||||
headers[auth_header] = auth_value
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
if include_extra:
|
||||
extra = get_extra_headers_for_endpoint(endpoint)
|
||||
|
||||
@@ -497,6 +497,11 @@ app.include_router(dashboard_router) # 仪表盘端点
|
||||
app.include_router(public_router) # 公开API端点(用户可查看提供商和模型)
|
||||
app.include_router(monitoring_router) # 监控端点
|
||||
|
||||
# WebSocket 隧道端点(aether-proxy tunnel 模式)
|
||||
from src.api.admin.proxy_tunnel import router as proxy_tunnel_router
|
||||
|
||||
app.include_router(proxy_tunnel_router)
|
||||
|
||||
|
||||
def main() -> Any:
|
||||
# 初始化新日志系统
|
||||
|
||||
@@ -920,12 +920,6 @@ class ProxyNode(Base):
|
||||
total_requests = Column(BigInteger, default=0, nullable=False)
|
||||
avg_latency_ms = Column(Float, nullable=True)
|
||||
|
||||
# TLS 加密
|
||||
tls_enabled = Column(Boolean, default=False, nullable=False, comment="是否启用 TLS 加密")
|
||||
tls_cert_fingerprint = Column(
|
||||
String(128), nullable=True, comment="TLS 证书 SHA-256 指纹(hex)"
|
||||
)
|
||||
|
||||
# 硬件信息(注册时上报,JSON 可扩展)
|
||||
hardware_info = Column(
|
||||
JSON,
|
||||
@@ -936,6 +930,15 @@ class ProxyNode(Base):
|
||||
Integer, nullable=True, comment="基于硬件估算的最大并发连接数"
|
||||
)
|
||||
|
||||
# 隧道模式(proxy 主动连接 Aether 的 WebSocket 隧道)
|
||||
tunnel_mode = Column(
|
||||
Boolean, default=False, nullable=False, comment="是否使用 WebSocket 隧道模式"
|
||||
)
|
||||
tunnel_connected = Column(Boolean, default=False, nullable=False, comment="隧道是否已连接")
|
||||
tunnel_connected_at = Column(
|
||||
DateTime(timezone=True), nullable=True, comment="隧道最近一次建立时间"
|
||||
)
|
||||
|
||||
# 管理端远程配置(通过心跳下发给 aether-proxy)
|
||||
remote_config = Column(
|
||||
JSON,
|
||||
|
||||
+10
-7
@@ -270,12 +270,6 @@ class ProxyNode(Base):
|
||||
total_requests = Column(BigInteger, default=0, nullable=False)
|
||||
avg_latency_ms = Column(Float, nullable=True)
|
||||
|
||||
# TLS 加密
|
||||
tls_enabled = Column(Boolean, default=False, nullable=False, comment="是否启用 TLS 加密")
|
||||
tls_cert_fingerprint = Column(
|
||||
String(128), nullable=True, comment="TLS 证书 SHA-256 指纹(hex)"
|
||||
)
|
||||
|
||||
# 硬件信息(注册时上报,JSON 可扩展)
|
||||
hardware_info = Column(
|
||||
JSON,
|
||||
@@ -286,11 +280,20 @@ class ProxyNode(Base):
|
||||
Integer, nullable=True, comment="基于硬件估算的最大并发连接数"
|
||||
)
|
||||
|
||||
# 隧道模式(proxy 主动连接 Aether 的 WebSocket 隧道)
|
||||
tunnel_mode = Column(
|
||||
Boolean, default=False, nullable=False, comment="是否使用 WebSocket 隧道模式"
|
||||
)
|
||||
tunnel_connected = Column(Boolean, default=False, nullable=False, comment="隧道是否已连接")
|
||||
tunnel_connected_at = Column(
|
||||
DateTime(timezone=True), nullable=True, comment="隧道最近一次建立时间"
|
||||
)
|
||||
|
||||
# 管理端远程配置(通过心跳下发给 aether-proxy)
|
||||
remote_config = Column(
|
||||
JSON,
|
||||
nullable=True,
|
||||
comment="管理端下发的远程配置 (allowed_ports, log_level, heartbeat_interval, timestamp_tolerance)",
|
||||
comment="管理端下发的远程配置 (allowed_ports, log_level, heartbeat_interval)",
|
||||
)
|
||||
config_version = Column(
|
||||
Integer, default=0, nullable=False, comment="远程配置版本号,每次更新 +1"
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
代理节点模块
|
||||
|
||||
提供海外 VPS 代理节点的注册、心跳、管理功能。
|
||||
aether-proxy 部署在海外 VPS 上自动注册节点,Aether 通过 HMAC 签名认证转发请求。
|
||||
aether-proxy 部署在海外 VPS 上,通过 WebSocket 隧道连接 Aether 转发 API 请求。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -77,15 +77,7 @@ async def _health_check() -> ModuleHealth:
|
||||
|
||||
|
||||
def _validate_config(db: Session) -> tuple[bool, str]:
|
||||
"""
|
||||
验证配置
|
||||
|
||||
代理节点模块需要 PROXY_HMAC_KEY 配置
|
||||
"""
|
||||
from src.config.settings import config
|
||||
|
||||
if not config.proxy_hmac_key:
|
||||
return False, "PROXY_HMAC_KEY 未配置(也未设置 ENCRYPTION_KEY 用于自动派生)"
|
||||
"""验证配置(tunnel 模式无需额外密钥配置)"""
|
||||
return True, ""
|
||||
|
||||
|
||||
@@ -93,7 +85,7 @@ proxy_nodes_module = ModuleDefinition(
|
||||
metadata=ModuleMetadata(
|
||||
name="proxy_nodes",
|
||||
display_name="代理节点",
|
||||
description="海外 VPS 代理节点管理,通过 HMAC 签名认证转发 API 请求",
|
||||
description="海外 VPS 代理节点管理,通过 WebSocket 隧道转发 API 请求",
|
||||
category=ModuleCategory.INTEGRATION,
|
||||
env_key="PROXY_NODES_AVAILABLE",
|
||||
default_available=True,
|
||||
|
||||
@@ -51,9 +51,10 @@ class ProviderConnector(ABC):
|
||||
self._last_error: str | None = None
|
||||
|
||||
# 代理配置(支持 proxy_node_id 和旧的 proxy URL)
|
||||
from src.services.proxy_node.resolver import resolve_ops_proxy
|
||||
from src.services.proxy_node.resolver import resolve_ops_proxy, resolve_ops_tunnel_node_id
|
||||
|
||||
self._proxy: str | httpx.Proxy | None = resolve_ops_proxy(self.config)
|
||||
self._tunnel_node_id: str | None = resolve_ops_tunnel_node_id(self.config)
|
||||
|
||||
# HTTP 客户端配置
|
||||
self._timeout = self.config.get("timeout", 30)
|
||||
@@ -117,13 +118,18 @@ class ProviderConnector(ABC):
|
||||
"""
|
||||
获取已认证的 HTTP 客户端
|
||||
|
||||
使用 context manager 确保资源正确释放
|
||||
使用 context manager 确保资源正确释放。
|
||||
tunnel 模式下使用 TunnelTransport 替代 proxy transport。
|
||||
|
||||
Yields:
|
||||
已配置认证信息的 AsyncClient
|
||||
"""
|
||||
transport = None
|
||||
if self._proxy:
|
||||
if self._tunnel_node_id:
|
||||
from src.services.proxy_node.tunnel_transport import TunnelTransport
|
||||
|
||||
transport = TunnelTransport(self._tunnel_node_id, timeout=self._timeout)
|
||||
elif self._proxy:
|
||||
transport = httpx.AsyncHTTPTransport(proxy=self._proxy)
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
|
||||
@@ -8,6 +8,7 @@ from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.core.api_format.headers import BROWSER_FINGERPRINT_HEADERS
|
||||
from src.services.provider_ops.actions import (
|
||||
NewApiBalanceAction,
|
||||
ProviderAction,
|
||||
@@ -75,6 +76,10 @@ class NewApiConnector(ProviderConnector):
|
||||
|
||||
def _apply_auth(self, request: httpx.Request) -> httpx.Request:
|
||||
"""为请求应用认证信息"""
|
||||
# 添加浏览器指纹 Headers 以绕过 Cloudflare 等防护
|
||||
for key, value in BROWSER_FINGERPRINT_HEADERS.items():
|
||||
request.headers.setdefault(key, value)
|
||||
|
||||
if self._api_key:
|
||||
request.headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
if self._user_id:
|
||||
@@ -206,7 +211,8 @@ class NewApiArchitecture(ProviderArchitecture):
|
||||
|
||||
New API 特有:需要 New-Api-User Header 传递用户 ID
|
||||
"""
|
||||
headers: dict[str, str] = {}
|
||||
# 以浏览器指纹 Headers 为基础,绕过 Cloudflare 等防护
|
||||
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
|
||||
|
||||
# Bearer Token 认证
|
||||
api_key = credentials.get("api_key", "")
|
||||
|
||||
@@ -1033,10 +1033,11 @@ class ProviderOpsService:
|
||||
list(headers.keys()),
|
||||
)
|
||||
|
||||
# 获取代理配置(支持 proxy_node_id 和旧的 proxy URL)
|
||||
from src.services.proxy_node.resolver import resolve_ops_proxy
|
||||
# 获取代理配置(支持 proxy_node_id、tunnel 模式和旧的 proxy URL)
|
||||
from src.services.proxy_node.resolver import resolve_ops_proxy, resolve_ops_tunnel_node_id
|
||||
|
||||
proxy = resolve_ops_proxy(config)
|
||||
tunnel_node_id = resolve_ops_tunnel_node_id(config)
|
||||
|
||||
try:
|
||||
# 构建 httpx client 参数
|
||||
@@ -1044,7 +1045,12 @@ class ProviderOpsService:
|
||||
"timeout": 30.0,
|
||||
"verify": get_ssl_context(),
|
||||
}
|
||||
if proxy:
|
||||
if tunnel_node_id:
|
||||
from src.services.proxy_node.tunnel_transport import TunnelTransport
|
||||
|
||||
client_kwargs["transport"] = TunnelTransport(tunnel_node_id, timeout=30.0)
|
||||
logger.debug("使用 tunnel 代理: node_id={}", tunnel_node_id)
|
||||
elif proxy:
|
||||
client_kwargs["proxy"] = proxy
|
||||
logger.debug("使用代理: {}", proxy)
|
||||
|
||||
|
||||
@@ -2,9 +2,6 @@
|
||||
|
||||
from .health_scheduler import ProxyNodeHealthScheduler, get_proxy_node_health_scheduler
|
||||
from .resolver import (
|
||||
build_delegate_post_kwargs,
|
||||
build_delegate_stream_kwargs,
|
||||
build_hmac_proxy_url,
|
||||
build_post_kwargs,
|
||||
build_proxy_url,
|
||||
build_stream_kwargs,
|
||||
@@ -12,10 +9,12 @@ from .resolver import (
|
||||
get_proxy_label,
|
||||
get_system_proxy_config,
|
||||
inject_auth_into_proxy_url,
|
||||
invalidate_proxy_node_cache,
|
||||
invalidate_system_proxy_cache,
|
||||
make_proxy_param,
|
||||
resolve_delegate_config,
|
||||
resolve_ops_proxy,
|
||||
resolve_ops_tunnel_node_id,
|
||||
resolve_proxy_info,
|
||||
)
|
||||
from .service import ProxyNodeService, node_to_dict
|
||||
@@ -25,9 +24,6 @@ __all__ = [
|
||||
"get_proxy_node_health_scheduler",
|
||||
"ProxyNodeService",
|
||||
"node_to_dict",
|
||||
"build_delegate_post_kwargs",
|
||||
"build_delegate_stream_kwargs",
|
||||
"build_hmac_proxy_url",
|
||||
"build_post_kwargs",
|
||||
"build_proxy_url",
|
||||
"build_stream_kwargs",
|
||||
@@ -36,8 +32,10 @@ __all__ = [
|
||||
"make_proxy_param",
|
||||
"get_proxy_label",
|
||||
"get_system_proxy_config",
|
||||
"invalidate_proxy_node_cache",
|
||||
"invalidate_system_proxy_cache",
|
||||
"resolve_delegate_config",
|
||||
"resolve_ops_proxy",
|
||||
"resolve_ops_tunnel_node_id",
|
||||
"resolve_proxy_info",
|
||||
]
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
"""
|
||||
ProxyNode 心跳检测调度器
|
||||
|
||||
定期检查 proxy_nodes 的 last_heartbeat_at,更新节点状态:
|
||||
- elapsed > interval * 3 -> unhealthy
|
||||
- elapsed > interval * 10 -> offline
|
||||
定期检查 proxy_nodes 的 tunnel 连接状态,更新节点状态:
|
||||
- tunnel_connected=True -> ONLINE
|
||||
- tunnel 刚断开 (<60s) -> UNHEALTHY(缓冲期,避免正在进行的请求被立即切走)
|
||||
- tunnel 断开超过 60s -> OFFLINE
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -56,6 +57,7 @@ class ProxyNodeHealthScheduler:
|
||||
try:
|
||||
now = datetime.now(timezone.utc)
|
||||
# 仅检查非手动节点(手动节点无心跳,始终保持 ONLINE)
|
||||
# 非手动节点均为 tunnel 模式,由 tunnel 连接状态决定
|
||||
nodes = (
|
||||
db.query(ProxyNode)
|
||||
.filter(
|
||||
@@ -69,19 +71,16 @@ class ProxyNodeHealthScheduler:
|
||||
|
||||
changed = 0
|
||||
for node in nodes:
|
||||
interval = int(node.heartbeat_interval or 30)
|
||||
last = node.last_heartbeat_at
|
||||
|
||||
if last is None:
|
||||
new_status = ProxyNodeStatus.OFFLINE
|
||||
if node.tunnel_connected:
|
||||
new_status = ProxyNodeStatus.ONLINE
|
||||
elif node.tunnel_connected_at:
|
||||
# tunnel 刚断开:给 60s 缓冲期标记为 UNHEALTHY
|
||||
elapsed = (now - node.tunnel_connected_at).total_seconds()
|
||||
new_status = (
|
||||
ProxyNodeStatus.UNHEALTHY if elapsed < 60 else ProxyNodeStatus.OFFLINE
|
||||
)
|
||||
else:
|
||||
elapsed = (now - last).total_seconds()
|
||||
if elapsed > interval * 10:
|
||||
new_status = ProxyNodeStatus.OFFLINE
|
||||
elif elapsed > interval * 3:
|
||||
new_status = ProxyNodeStatus.UNHEALTHY
|
||||
else:
|
||||
new_status = ProxyNodeStatus.ONLINE
|
||||
new_status = ProxyNodeStatus.OFFLINE
|
||||
|
||||
if node.status != new_status:
|
||||
node.status = new_status
|
||||
|
||||
@@ -7,18 +7,13 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import gzip as _gzip
|
||||
import hashlib
|
||||
import hmac as _hmac
|
||||
import json as _json
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config import config
|
||||
from src.core.exceptions import ProxyNodeUnavailableError
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -81,8 +76,8 @@ def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
|
||||
"name": node.name,
|
||||
"ip": node.ip,
|
||||
"port": node.port,
|
||||
"tls_enabled": bool(node.tls_enabled),
|
||||
"tls_cert_fingerprint": node.tls_cert_fingerprint,
|
||||
"tunnel_mode": bool(node.tunnel_mode),
|
||||
"tunnel_connected": bool(node.tunnel_connected),
|
||||
}
|
||||
|
||||
_proxy_node_cache[node_id] = (value, now + _PROXY_NODE_CACHE_TTL_SECONDS)
|
||||
@@ -91,40 +86,6 @@ def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
|
||||
db.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HMAC 签名
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_hmac_proxy_url(ip: str, port: int, *, tls_enabled: bool = False) -> str:
|
||||
"""
|
||||
构建带 HMAC BasicAuth 的 httpx proxy URL
|
||||
|
||||
格式: http(s)://hmac:{timestamp}.{signature}@{ip}:{port}
|
||||
signature = HMAC-SHA256(PROXY_HMAC_KEY, "{timestamp}") 的 hex
|
||||
|
||||
签名不再包含 node_id,避免 proxy 重新注册后 Aether 端缓存的旧 node_id
|
||||
与 proxy 端新 node_id 不一致导致的认证失败窗口。
|
||||
|
||||
当 tls_enabled=True 时使用 https:// scheme。
|
||||
"""
|
||||
if not config.proxy_hmac_key:
|
||||
logger.error("PROXY_HMAC_KEY 未配置,无法使用 ProxyNode 代理")
|
||||
raise ProxyNodeUnavailableError("PROXY_HMAC_KEY 未配置,无法使用 ProxyNode 代理")
|
||||
|
||||
timestamp = str(int(time.time()))
|
||||
payload = timestamp.encode("utf-8")
|
||||
signature = _hmac.new(
|
||||
config.proxy_hmac_key.encode("utf-8"),
|
||||
payload,
|
||||
hashlib.sha256,
|
||||
).hexdigest()
|
||||
|
||||
host = f"[{ip}]" if ":" in ip else ip
|
||||
scheme = "https" if tls_enabled else "http"
|
||||
return f"{scheme}://hmac:{timestamp}.{signature}@{host}:{int(port)}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 系统默认代理
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -132,6 +93,11 @@ _system_proxy_cache: tuple[dict[str, Any] | None, float] | None = None
|
||||
_SYSTEM_PROXY_CACHE_TTL = 60.0
|
||||
|
||||
|
||||
def invalidate_proxy_node_cache(node_id: str) -> None:
|
||||
"""主动清除指定节点的信息缓存(tunnel 断开时调用,避免使用过期的连接状态)"""
|
||||
_proxy_node_cache.pop(node_id, None)
|
||||
|
||||
|
||||
def invalidate_system_proxy_cache() -> None:
|
||||
"""手动失效系统代理缓存(在删除节点等操作后调用)"""
|
||||
global _system_proxy_cache
|
||||
@@ -224,6 +190,32 @@ def make_proxy_param(proxy_url: str | None) -> str | httpx.Proxy | None:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _resolve_effective_node(
|
||||
connector_config: dict[str, Any] | None,
|
||||
) -> tuple[str | None, dict[str, Any] | None]:
|
||||
"""
|
||||
从 connector_config 或系统默认代理中解析有效的 proxy_node_id 及其信息。
|
||||
|
||||
Returns:
|
||||
(node_id, node_info) 或 (None, None)
|
||||
"""
|
||||
if connector_config:
|
||||
node_id = connector_config.get("proxy_node_id")
|
||||
if isinstance(node_id, str) and node_id.strip():
|
||||
nid = node_id.strip()
|
||||
return nid, _get_proxy_node_info(nid)
|
||||
|
||||
# 回退:系统默认代理
|
||||
system_proxy = get_system_proxy_config()
|
||||
if system_proxy:
|
||||
node_id_sys = system_proxy.get("node_id")
|
||||
if isinstance(node_id_sys, str) and node_id_sys.strip():
|
||||
nid = node_id_sys.strip()
|
||||
return nid, _get_proxy_node_info(nid)
|
||||
|
||||
return None, None
|
||||
|
||||
|
||||
def resolve_ops_proxy(
|
||||
connector_config: dict[str, Any] | None,
|
||||
) -> str | httpx.Proxy | None:
|
||||
@@ -235,37 +227,50 @@ def resolve_ops_proxy(
|
||||
2. connector_config.proxy(旧格式 URL 字符串)
|
||||
3. 系统默认代理节点
|
||||
|
||||
tunnel 模式节点不返回代理 URL(由 resolve_ops_tunnel_node_id 处理)。
|
||||
|
||||
Args:
|
||||
connector_config: connector 的 config 字典
|
||||
|
||||
Returns:
|
||||
httpx 可接受的代理参数(str 或 httpx.Proxy),或 None
|
||||
"""
|
||||
if connector_config:
|
||||
# 新格式:proxy_node_id -> 通过 build_proxy_url 解析
|
||||
node_id = connector_config.get("proxy_node_id")
|
||||
if isinstance(node_id, str) and node_id.strip():
|
||||
try:
|
||||
url = build_proxy_url({"node_id": node_id.strip(), "enabled": True})
|
||||
return make_proxy_param(url)
|
||||
except Exception as exc:
|
||||
logger.warning("解析 proxy_node_id={} 失败,回退到直连: {}", node_id, exc)
|
||||
return None
|
||||
from .tunnel_transport import is_tunnel_node
|
||||
|
||||
# 旧格式:直接返回 proxy URL 字符串
|
||||
node_id, node_info = _resolve_effective_node(connector_config)
|
||||
if node_id and node_info:
|
||||
if is_tunnel_node(node_info):
|
||||
return None # tunnel 模式不使用 proxy URL
|
||||
try:
|
||||
url = build_proxy_url({"node_id": node_id, "enabled": True})
|
||||
return make_proxy_param(url)
|
||||
except Exception as exc:
|
||||
logger.warning("解析 proxy_node_id={} 失败,回退到直连: {}", node_id, exc)
|
||||
return None
|
||||
|
||||
# 旧格式:直接返回 proxy URL 字符串
|
||||
if connector_config:
|
||||
proxy = connector_config.get("proxy")
|
||||
if isinstance(proxy, str) and proxy.strip():
|
||||
return proxy
|
||||
|
||||
# 回退:系统默认代理
|
||||
system_proxy = get_system_proxy_config()
|
||||
if system_proxy:
|
||||
try:
|
||||
url = build_proxy_url(system_proxy)
|
||||
return make_proxy_param(url)
|
||||
except Exception as exc:
|
||||
logger.warning("构建系统默认代理 URL 失败: {}", exc)
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def resolve_ops_tunnel_node_id(
|
||||
connector_config: dict[str, Any] | None,
|
||||
) -> str | None:
|
||||
"""
|
||||
解析 ops connector 的 tunnel 节点 ID
|
||||
|
||||
如果配置的代理节点是 tunnel 模式且已连接,返回 node_id。
|
||||
否则返回 None(含系统默认代理回退)。
|
||||
"""
|
||||
from .tunnel_transport import is_tunnel_node
|
||||
|
||||
node_id, node_info = _resolve_effective_node(connector_config)
|
||||
if node_id and node_info and is_tunnel_node(node_info):
|
||||
return node_id
|
||||
|
||||
return None
|
||||
|
||||
@@ -401,12 +406,15 @@ def build_proxy_url(proxy_config: dict[str, Any]) -> str | None:
|
||||
return inject_auth_into_proxy_url(manual_url, username, password)
|
||||
return manual_url
|
||||
|
||||
# aether-proxy 节点:使用 HMAC 认证
|
||||
return build_hmac_proxy_url(
|
||||
node_info["ip"],
|
||||
node_info["port"],
|
||||
tls_enabled=node_info.get("tls_enabled", False),
|
||||
)
|
||||
# tunnel 模式节点:不构建 proxy URL(通过 TunnelTransport 处理)
|
||||
from .tunnel_transport import is_tunnel_node
|
||||
|
||||
if is_tunnel_node(node_info):
|
||||
return None
|
||||
|
||||
# aether-proxy 节点均为 tunnel 模式,不应走到这里
|
||||
logger.warning("非 tunnel 模式的 aether-proxy 节点不再支持: node_id={}", node_id)
|
||||
return None
|
||||
|
||||
proxy_url: str | None = proxy_config.get("url")
|
||||
if not proxy_url:
|
||||
@@ -504,11 +512,10 @@ def compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
|
||||
if not proxy_config.get("enabled", True):
|
||||
return "__no_proxy__"
|
||||
|
||||
# ProxyNode 模式:基于 node_id + 时间桶缓存,避免签名随时间变化导致 cache key 爆炸
|
||||
# ProxyNode 模式:基于 node_id 缓存
|
||||
node_id = proxy_config.get("node_id")
|
||||
if isinstance(node_id, str) and node_id.strip():
|
||||
time_bucket = int(time.time() / 240) # 240s bucket, within 300s HMAC tolerance
|
||||
return f"proxy_node:{node_id.strip()}:{time_bucket}"
|
||||
return f"proxy_node:{node_id.strip()}"
|
||||
|
||||
# 构建代理 URL 作为缓存键的基础
|
||||
proxy_url = build_proxy_url(proxy_config)
|
||||
@@ -520,46 +527,20 @@ def compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 代发模式 (Delegate API)
|
||||
# Tunnel 代理配置解析
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_hmac_auth_header() -> str:
|
||||
"""
|
||||
构建代发请求的 Authorization 头
|
||||
|
||||
格式: Basic base64(hmac:{timestamp}.{signature})
|
||||
签名算法与 build_hmac_proxy_url 相同(仅使用 timestamp,不含 node_id)。
|
||||
"""
|
||||
if not config.proxy_hmac_key:
|
||||
raise ProxyNodeUnavailableError("PROXY_HMAC_KEY 未配置,无法使用代发模式")
|
||||
|
||||
timestamp = str(int(time.time()))
|
||||
payload = timestamp.encode("utf-8")
|
||||
signature = _hmac.new(
|
||||
config.proxy_hmac_key.encode("utf-8"),
|
||||
payload,
|
||||
hashlib.sha256,
|
||||
).hexdigest()
|
||||
|
||||
cred = f"hmac:{timestamp}.{signature}"
|
||||
encoded = base64.b64encode(cred.encode()).decode()
|
||||
return f"Basic {encoded}"
|
||||
|
||||
|
||||
def resolve_delegate_config(proxy_config: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
"""
|
||||
解析代发配置(仅 aether-proxy 节点支持,手动节点/旧格式 URL 不支持)
|
||||
解析 tunnel 代理配置(仅 aether-proxy tunnel 节点支持)
|
||||
|
||||
无特定代理时自动回退到系统默认代理。
|
||||
auth_header 延迟生成:通过 ``fresh_auth_header()`` 闭包在每次请求 / 重试时
|
||||
获取新鲜的 HMAC 签名,避免长生命周期内时间戳过期。
|
||||
tunnel 模式节点返回 {"tunnel": True, "node_id": str},
|
||||
调用方应使用 TunnelTransport。
|
||||
|
||||
Returns:
|
||||
{"delegate_url": str, "node_id": str, "tls_enabled": bool,
|
||||
"auth_header": str, # 首次生成的签名(兼容旧调用)
|
||||
"fresh_auth_header": Callable} # 延迟生成签名的闭包
|
||||
或 None
|
||||
{"tunnel": True, "node_id": str} 或 None
|
||||
"""
|
||||
effective_config = proxy_config
|
||||
|
||||
@@ -571,148 +552,28 @@ def resolve_delegate_config(proxy_config: dict[str, Any] | None) -> dict[str, An
|
||||
|
||||
node_id = effective_config.get("node_id")
|
||||
if not isinstance(node_id, str) or not node_id.strip():
|
||||
return None # 旧格式 URL 模式不支持代发
|
||||
return None
|
||||
|
||||
node_id = node_id.strip()
|
||||
node_info = _get_proxy_node_info(node_id)
|
||||
if not node_info or node_info.get("is_manual"):
|
||||
return None # 手动节点不支持代发
|
||||
return None
|
||||
|
||||
tls_enabled = node_info.get("tls_enabled", False)
|
||||
host = f"[{node_info['ip']}]" if ":" in node_info["ip"] else node_info["ip"]
|
||||
scheme = "https" if tls_enabled else "http"
|
||||
delegate_url = f"{scheme}://{host}:{int(node_info['port'])}/_aether/delegate"
|
||||
from .tunnel_transport import is_tunnel_node
|
||||
|
||||
# 每次调用生成新鲜签名(避免长连接内时间戳过期)
|
||||
def _fresh() -> str:
|
||||
return _build_hmac_auth_header()
|
||||
if is_tunnel_node(node_info):
|
||||
return {"tunnel": True, "node_id": node_id}
|
||||
|
||||
return {
|
||||
"delegate_url": delegate_url,
|
||||
"auth_header": _fresh(), # 立即生成一份,兼容旧调用方
|
||||
"fresh_auth_header": _fresh,
|
||||
"node_id": node_id,
|
||||
"tls_enabled": tls_enabled,
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 代发请求参数构建(消除 handler 层重复代码)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_JSON_CT = "application/json"
|
||||
|
||||
|
||||
def _build_delegate_kwargs_core(
|
||||
delegate_cfg: dict[str, Any],
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
payload: Any,
|
||||
timeout: float,
|
||||
refresh_auth: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
构建代发请求的核心参数(post/stream 共用)
|
||||
|
||||
元数据通过 HTTP headers 传递(X-Delegate-Method/Url/Headers),
|
||||
上游请求体 gzip 压缩后直接作为 HTTP body 发送,大幅减少跨国传输耗时。
|
||||
|
||||
Args:
|
||||
delegate_cfg: resolve_delegate_config 返回的配置
|
||||
url: 上游实际 URL
|
||||
headers: 上游请求头
|
||||
payload: 上游 JSON body(可以为 None)
|
||||
timeout: 上游超时秒数
|
||||
refresh_auth: 为 True 时重新生成 HMAC 签名(用于 retry)
|
||||
"""
|
||||
auth = (
|
||||
delegate_cfg["fresh_auth_header"]()
|
||||
if refresh_auth
|
||||
else delegate_cfg.get("auth_header") or delegate_cfg["fresh_auth_header"]()
|
||||
)
|
||||
|
||||
# 上游 headers base64 编码
|
||||
headers_b64 = base64.b64encode(_json.dumps(headers, ensure_ascii=False).encode("utf-8")).decode(
|
||||
"ascii"
|
||||
)
|
||||
|
||||
# 构建代发请求 headers(元数据)
|
||||
delegate_headers: dict[str, str] = {
|
||||
"Authorization": auth,
|
||||
"X-Delegate-Method": "POST",
|
||||
"X-Delegate-Url": url,
|
||||
"X-Delegate-Headers": headers_b64,
|
||||
"X-Delegate-Timeout": str(int(timeout)),
|
||||
}
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"url": delegate_cfg["delegate_url"],
|
||||
"headers": delegate_headers,
|
||||
"timeout": httpx.Timeout(timeout + 10),
|
||||
}
|
||||
|
||||
# body gzip 压缩后直接作为 HTTP content
|
||||
if payload is not None:
|
||||
body_bytes = _json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
compressed = _gzip.compress(body_bytes)
|
||||
kwargs["content"] = compressed
|
||||
kwargs["headers"]["Content-Encoding"] = "gzip"
|
||||
kwargs["headers"]["Content-Type"] = _JSON_CT
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
def build_delegate_post_kwargs(
|
||||
delegate_cfg: dict[str, Any],
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
payload: Any,
|
||||
timeout: float,
|
||||
refresh_auth: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""构建代发 POST 请求的 httpx kwargs(非流式,传给 client.post)"""
|
||||
return _build_delegate_kwargs_core(
|
||||
delegate_cfg,
|
||||
url=url,
|
||||
headers=headers,
|
||||
payload=payload,
|
||||
timeout=timeout,
|
||||
refresh_auth=refresh_auth,
|
||||
)
|
||||
|
||||
|
||||
def build_delegate_stream_kwargs(
|
||||
delegate_cfg: dict[str, Any],
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
payload: Any,
|
||||
timeout: float,
|
||||
refresh_auth: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""构建代发 stream 请求的 httpx kwargs(传给 client.stream)"""
|
||||
kwargs = _build_delegate_kwargs_core(
|
||||
delegate_cfg,
|
||||
url=url,
|
||||
headers=headers,
|
||||
payload=payload,
|
||||
timeout=timeout,
|
||||
refresh_auth=refresh_auth,
|
||||
)
|
||||
# stream() 需要显式 method 参数
|
||||
kwargs["method"] = "POST"
|
||||
return kwargs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 统一上游请求参数构建(消除 handler 层 delegate/直连 分支重复)
|
||||
# 统一上游请求参数构建
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_post_kwargs(
|
||||
delegate_cfg: dict[str, Any] | None,
|
||||
_delegate_cfg: dict[str, Any] | None = None,
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
@@ -721,19 +582,13 @@ def build_post_kwargs(
|
||||
refresh_auth: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
构建上游 POST 请求的 httpx kwargs(自动选择代发或直连模式)
|
||||
构建上游 POST 请求的 httpx kwargs
|
||||
|
||||
返回的 dict 可直接传给 ``http_client.post(**kwargs)``。
|
||||
|
||||
``_delegate_cfg`` 和 ``refresh_auth`` 已废弃(tunnel 模式下认证由 transport 层处理),
|
||||
保留仅为兼容现有调用方签名。
|
||||
"""
|
||||
if delegate_cfg:
|
||||
return build_delegate_post_kwargs(
|
||||
delegate_cfg,
|
||||
url=url,
|
||||
headers=headers,
|
||||
payload=payload,
|
||||
timeout=timeout,
|
||||
refresh_auth=refresh_auth,
|
||||
)
|
||||
return {
|
||||
"url": url,
|
||||
"json": payload,
|
||||
@@ -743,7 +598,7 @@ def build_post_kwargs(
|
||||
|
||||
|
||||
def build_stream_kwargs(
|
||||
delegate_cfg: dict[str, Any] | None,
|
||||
_delegate_cfg: dict[str, Any] | None = None,
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
@@ -751,21 +606,13 @@ def build_stream_kwargs(
|
||||
timeout: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
构建上游 stream 请求的 httpx kwargs(自动选择代发或直连模式)
|
||||
构建上游 stream 请求的 httpx kwargs
|
||||
|
||||
返回的 dict 可直接传给 ``http_client.stream(**kwargs)``。
|
||||
当 ``timeout`` 为 None 时由外层 asyncio.wait_for 控制超时。
|
||||
|
||||
当 ``timeout`` 为 None(直连模式下由外层 asyncio.wait_for 控制超时),
|
||||
直连分支不设置 timeout;代发分支始终携带 timeout(proxy 协议需要)。
|
||||
``_delegate_cfg`` 已废弃,保留仅为兼容现有调用方签名。
|
||||
"""
|
||||
if delegate_cfg:
|
||||
return build_delegate_stream_kwargs(
|
||||
delegate_cfg,
|
||||
url=url,
|
||||
headers=headers,
|
||||
payload=payload,
|
||||
timeout=timeout or 60,
|
||||
)
|
||||
kwargs: dict[str, Any] = {
|
||||
"method": "POST",
|
||||
"url": url,
|
||||
|
||||
@@ -17,10 +17,9 @@ import httpx
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.models.database import ProxyNode, ProxyNodeStatus, SystemConfig
|
||||
from src.models.database import Provider, ProviderEndpoint, ProxyNode, ProxyNodeStatus, SystemConfig
|
||||
|
||||
from .resolver import (
|
||||
build_hmac_proxy_url,
|
||||
inject_auth_into_proxy_url,
|
||||
invalidate_system_proxy_cache,
|
||||
make_proxy_param,
|
||||
@@ -50,14 +49,15 @@ def node_to_dict(node: ProxyNode) -> dict[str, Any]:
|
||||
"region": node.region,
|
||||
"status": node.status.value if node.status else None,
|
||||
"is_manual": bool(node.is_manual),
|
||||
"tunnel_mode": bool(node.tunnel_mode),
|
||||
"tunnel_connected": bool(node.tunnel_connected),
|
||||
"tunnel_connected_at": node.tunnel_connected_at,
|
||||
"registered_by": node.registered_by,
|
||||
"last_heartbeat_at": node.last_heartbeat_at,
|
||||
"heartbeat_interval": node.heartbeat_interval,
|
||||
"active_connections": node.active_connections,
|
||||
"total_requests": node.total_requests,
|
||||
"avg_latency_ms": node.avg_latency_ms,
|
||||
"tls_enabled": bool(node.tls_enabled),
|
||||
"tls_cert_fingerprint": node.tls_cert_fingerprint,
|
||||
"hardware_info": node.hardware_info,
|
||||
"estimated_max_concurrency": node.estimated_max_concurrency,
|
||||
"remote_config": node.remote_config,
|
||||
@@ -166,8 +166,8 @@ def _build_test_proxy_url(node: ProxyNode) -> str:
|
||||
)
|
||||
return proxy_url
|
||||
else:
|
||||
# aether-proxy: 使用 HMAC 认证构建代理 URL
|
||||
return build_hmac_proxy_url(node.ip, node.port, tls_enabled=bool(node.tls_enabled))
|
||||
# aether-proxy 节点均为 tunnel 模式,不支持通过代理 URL 测试
|
||||
raise InvalidRequestException("aether-proxy tunnel 节点不支持代理 URL 连通性测试")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -187,27 +187,36 @@ class ProxyNodeService:
|
||||
port: int,
|
||||
region: str | None = None,
|
||||
heartbeat_interval: int = 30,
|
||||
tls_enabled: bool = False,
|
||||
tls_cert_fingerprint: str | None = None,
|
||||
hardware_info: dict[str, Any] | None = None,
|
||||
estimated_max_concurrency: int | None = None,
|
||||
active_connections: int | None = None,
|
||||
total_requests: int | None = None,
|
||||
avg_latency_ms: float | None = None,
|
||||
registered_by: str | None = None,
|
||||
tunnel_mode: bool = False,
|
||||
) -> ProxyNode:
|
||||
"""注册或更新 aether-proxy 节点(按 ip+port upsert)"""
|
||||
"""注册或更新 aether-proxy 节点
|
||||
|
||||
tunnel 模式按 name upsert(port 固定为 0,同 IP 可能有多个实例);
|
||||
旧模式按 ip+port upsert(向后兼容)。
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
node = db.query(ProxyNode).filter(ProxyNode.ip == ip, ProxyNode.port == port).first()
|
||||
if tunnel_mode:
|
||||
node = (
|
||||
db.query(ProxyNode)
|
||||
.filter(ProxyNode.name == name, ProxyNode.is_manual == False) # noqa: E712
|
||||
.first()
|
||||
)
|
||||
else:
|
||||
node = db.query(ProxyNode).filter(ProxyNode.ip == ip, ProxyNode.port == port).first()
|
||||
if node:
|
||||
node.name = name
|
||||
node.region = region
|
||||
node.status = ProxyNodeStatus.ONLINE
|
||||
node.last_heartbeat_at = now
|
||||
node.heartbeat_interval = heartbeat_interval
|
||||
node.tls_enabled = tls_enabled
|
||||
node.tls_cert_fingerprint = tls_cert_fingerprint
|
||||
node.tunnel_mode = tunnel_mode
|
||||
if hardware_info is not None:
|
||||
node.hardware_info = hardware_info
|
||||
if estimated_max_concurrency is not None:
|
||||
@@ -232,10 +241,9 @@ class ProxyNodeService:
|
||||
active_connections=active_connections or 0,
|
||||
total_requests=total_requests or 0,
|
||||
avg_latency_ms=avg_latency_ms,
|
||||
tls_enabled=tls_enabled,
|
||||
tls_cert_fingerprint=tls_cert_fingerprint,
|
||||
hardware_info=hardware_info,
|
||||
estimated_max_concurrency=estimated_max_concurrency,
|
||||
tunnel_mode=tunnel_mode,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
@@ -424,6 +432,21 @@ class ProxyNodeService:
|
||||
sys_cfg.value = None
|
||||
was_system_proxy = True
|
||||
|
||||
# 清理引用该节点的 Provider / ProviderEndpoint 的 proxy 字段(批量 SQL 更新)
|
||||
cleared_providers = (
|
||||
db.query(Provider)
|
||||
.filter(Provider.proxy.isnot(None), Provider.proxy["node_id"].as_string() == node_id)
|
||||
.update({"proxy": None}, synchronize_session="fetch")
|
||||
)
|
||||
cleared_endpoints = (
|
||||
db.query(ProviderEndpoint)
|
||||
.filter(
|
||||
ProviderEndpoint.proxy.isnot(None),
|
||||
ProviderEndpoint.proxy["node_id"].as_string() == node_id,
|
||||
)
|
||||
.update({"proxy": None}, synchronize_session="fetch")
|
||||
)
|
||||
|
||||
node_info = {"proxy_node_ip": node.ip, "proxy_node_port": node.port}
|
||||
db.delete(node)
|
||||
db.commit()
|
||||
@@ -435,6 +458,8 @@ class ProxyNodeService:
|
||||
"node_id": node_id,
|
||||
"node_info": node_info,
|
||||
"cleared_system_proxy": was_system_proxy,
|
||||
"cleared_providers": cleared_providers,
|
||||
"cleared_endpoints": cleared_endpoints,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -0,0 +1,334 @@
|
||||
"""
|
||||
WebSocket 隧道管理器
|
||||
|
||||
管理所有活跃的 aether-proxy tunnel 连接,提供通过隧道发送 HTTP 请求的能力。
|
||||
每个 proxy node 最多一条 tunnel 连接。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
from starlette.websockets import WebSocket, WebSocketState
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
from .tunnel_protocol import Frame, FrameFlags, MsgType
|
||||
|
||||
|
||||
class TunnelConnection:
|
||||
"""单条 tunnel 连接"""
|
||||
|
||||
__slots__ = (
|
||||
"node_id",
|
||||
"node_name",
|
||||
"ws",
|
||||
"connected_at",
|
||||
"_pending_streams",
|
||||
"_write_lock",
|
||||
"_next_stream_id",
|
||||
)
|
||||
|
||||
def __init__(self, node_id: str, node_name: str, ws: WebSocket) -> None:
|
||||
self.node_id = node_id
|
||||
self.node_name = node_name
|
||||
self.ws = ws
|
||||
self.connected_at = time.time()
|
||||
self._pending_streams: dict[int, _StreamState] = {}
|
||||
self._write_lock = asyncio.Lock()
|
||||
# Per-connection stream ID 分配器(Aether 端使用偶数,从 2 开始)
|
||||
self._next_stream_id: int = 2
|
||||
|
||||
@property
|
||||
def is_alive(self) -> bool:
|
||||
return self.ws.client_state == WebSocketState.CONNECTED
|
||||
|
||||
async def send_frame(self, frame: Frame) -> None:
|
||||
async with self._write_lock:
|
||||
await self.ws.send_bytes(frame.encode())
|
||||
|
||||
def create_stream(self, stream_id: int) -> _StreamState:
|
||||
state = _StreamState(stream_id)
|
||||
self._pending_streams[stream_id] = state
|
||||
return state
|
||||
|
||||
def get_stream(self, stream_id: int) -> _StreamState | None:
|
||||
return self._pending_streams.get(stream_id)
|
||||
|
||||
def remove_stream(self, stream_id: int) -> None:
|
||||
self._pending_streams.pop(stream_id, None)
|
||||
|
||||
@property
|
||||
def stream_count(self) -> int:
|
||||
return len(self._pending_streams)
|
||||
|
||||
def has_stream(self, stream_id: int) -> bool:
|
||||
return stream_id in self._pending_streams
|
||||
|
||||
def alloc_stream_id(self, max_streams: int) -> int:
|
||||
"""分配一个未被占用的偶数 stream_id,回绕时跳过飞行中的 ID"""
|
||||
# 最多尝试 max_streams + 16 次(飞行中的 stream 数量不超过 max_streams)
|
||||
for _ in range(max_streams + 16):
|
||||
sid = self._next_stream_id
|
||||
self._next_stream_id += 2
|
||||
if self._next_stream_id > 0xFFFF_FFFE:
|
||||
self._next_stream_id = 2
|
||||
if sid not in self._pending_streams:
|
||||
return sid
|
||||
raise TunnelStreamError("stream ID space exhausted")
|
||||
|
||||
def cancel_all_streams(self) -> None:
|
||||
for state in self._pending_streams.values():
|
||||
state.set_error("tunnel disconnected")
|
||||
self._pending_streams.clear()
|
||||
|
||||
|
||||
class _StreamState:
|
||||
"""跟踪单个 stream 的响应状态"""
|
||||
|
||||
__slots__ = (
|
||||
"stream_id",
|
||||
"status",
|
||||
"headers",
|
||||
"_header_event",
|
||||
"_body_chunks",
|
||||
"_done_event",
|
||||
"_error",
|
||||
)
|
||||
|
||||
def __init__(self, stream_id: int) -> None:
|
||||
self.stream_id = stream_id
|
||||
self.status: int = 0
|
||||
self.headers: list[list[str]] = []
|
||||
self._header_event = asyncio.Event()
|
||||
self._body_chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
|
||||
self._done_event = asyncio.Event()
|
||||
self._error: str | None = None
|
||||
|
||||
def set_response_headers(self, status: int, headers: list[list[str]] | dict[str, str]) -> None:
|
||||
self.status = status
|
||||
# headers 可能是 [[k, v], ...] (多值) 或 {k: v} (旧格式兼容)
|
||||
if isinstance(headers, list):
|
||||
self.headers = headers # type: ignore[assignment]
|
||||
else:
|
||||
self.headers = list(headers.items()) # type: ignore[assignment]
|
||||
self._header_event.set()
|
||||
|
||||
def push_body_chunk(self, data: bytes) -> None:
|
||||
self._body_chunks.put_nowait(data)
|
||||
|
||||
def set_done(self) -> None:
|
||||
self._body_chunks.put_nowait(None) # sentinel
|
||||
self._done_event.set()
|
||||
|
||||
def set_error(self, msg: str) -> None:
|
||||
self._error = msg
|
||||
self._header_event.set()
|
||||
self._body_chunks.put_nowait(None)
|
||||
self._done_event.set()
|
||||
|
||||
async def wait_headers(self, timeout: float = 60.0) -> None:
|
||||
await asyncio.wait_for(self._header_event.wait(), timeout=timeout)
|
||||
if self._error:
|
||||
raise TunnelStreamError(self._error)
|
||||
|
||||
async def iter_body(self, chunk_timeout: float = 60.0) -> AsyncGenerator[bytes, None]:
|
||||
while True:
|
||||
try:
|
||||
chunk = await asyncio.wait_for(self._body_chunks.get(), timeout=chunk_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
self._error = "body chunk timeout"
|
||||
self._done_event.set()
|
||||
raise TunnelStreamError("body chunk timeout")
|
||||
if chunk is None:
|
||||
if self._error:
|
||||
raise TunnelStreamError(self._error)
|
||||
return
|
||||
yield chunk
|
||||
|
||||
|
||||
class TunnelStreamError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 全局 TunnelManager 单例
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TunnelManager:
|
||||
"""管理所有活跃的 tunnel 连接"""
|
||||
|
||||
# 单条 tunnel 上允许的最大并发 stream 数(超出时拒绝新请求)
|
||||
MAX_STREAMS_PER_CONN = 2048
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._connections: dict[str, TunnelConnection] = {} # node_id -> conn
|
||||
|
||||
@property
|
||||
def active_count(self) -> int:
|
||||
return len(self._connections)
|
||||
|
||||
def get_connection(self, node_id: str) -> TunnelConnection | None:
|
||||
conn = self._connections.get(node_id)
|
||||
if conn and not conn.is_alive:
|
||||
self._connections.pop(node_id, None)
|
||||
conn.cancel_all_streams()
|
||||
return None
|
||||
return conn
|
||||
|
||||
def register(self, conn: TunnelConnection) -> None:
|
||||
old = self._connections.get(conn.node_id)
|
||||
if old:
|
||||
old.cancel_all_streams()
|
||||
self._connections[conn.node_id] = conn
|
||||
logger.info("tunnel connected: node_id={}, name={}", conn.node_id, conn.node_name)
|
||||
|
||||
def unregister(self, node_id: str) -> None:
|
||||
conn = self._connections.pop(node_id, None)
|
||||
if conn:
|
||||
conn.cancel_all_streams()
|
||||
logger.info("tunnel disconnected: node_id={}, name={}", node_id, conn.node_name)
|
||||
|
||||
def has_tunnel(self, node_id: str) -> bool:
|
||||
conn = self.get_connection(node_id)
|
||||
return conn is not None
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
node_id: str,
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: bytes | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> _StreamState:
|
||||
"""
|
||||
通过 tunnel 发送 HTTP 请求,返回 StreamState 用于读取响应。
|
||||
"""
|
||||
conn = self.get_connection(node_id)
|
||||
if not conn:
|
||||
raise TunnelStreamError(f"tunnel not connected for node {node_id}")
|
||||
|
||||
if conn.stream_count >= self.MAX_STREAMS_PER_CONN:
|
||||
raise TunnelStreamError(
|
||||
f"tunnel stream limit reached ({self.MAX_STREAMS_PER_CONN}) for node {node_id}"
|
||||
)
|
||||
|
||||
stream_id = conn.alloc_stream_id(self.MAX_STREAMS_PER_CONN)
|
||||
stream_state = conn.create_stream(stream_id)
|
||||
|
||||
try:
|
||||
# 发送 REQUEST_HEADERS
|
||||
meta = json.dumps(
|
||||
{
|
||||
"method": method,
|
||||
"url": url,
|
||||
"headers": headers,
|
||||
"timeout": int(timeout),
|
||||
}
|
||||
).encode()
|
||||
await conn.send_frame(Frame(stream_id, MsgType.REQUEST_HEADERS, 0, meta))
|
||||
|
||||
# 发送 REQUEST_BODY + END_STREAM
|
||||
body_data = body or b""
|
||||
await conn.send_frame(
|
||||
Frame(stream_id, MsgType.REQUEST_BODY, FrameFlags.END_STREAM, body_data)
|
||||
)
|
||||
except Exception:
|
||||
conn.remove_stream(stream_id)
|
||||
raise
|
||||
|
||||
return stream_state
|
||||
|
||||
async def handle_incoming_frame(self, node_id: str, frame: Frame) -> None:
|
||||
"""处理从 proxy 收到的响应帧"""
|
||||
conn = self.get_connection(node_id)
|
||||
if not conn:
|
||||
return
|
||||
|
||||
stream = conn.get_stream(frame.stream_id)
|
||||
|
||||
if frame.msg_type == MsgType.RESPONSE_HEADERS:
|
||||
if not stream:
|
||||
return
|
||||
try:
|
||||
meta = json.loads(frame.payload)
|
||||
stream.set_response_headers(meta["status"], meta.get("headers", []))
|
||||
except Exception as e:
|
||||
stream.set_error(f"invalid response headers: {e}")
|
||||
|
||||
elif frame.msg_type == MsgType.RESPONSE_BODY:
|
||||
if stream:
|
||||
stream.push_body_chunk(frame.payload)
|
||||
|
||||
elif frame.msg_type == MsgType.STREAM_END:
|
||||
if stream:
|
||||
stream.set_done()
|
||||
conn.remove_stream(frame.stream_id)
|
||||
|
||||
elif frame.msg_type == MsgType.STREAM_ERROR:
|
||||
if stream:
|
||||
msg = frame.payload.decode(errors="replace") if frame.payload else "stream error"
|
||||
stream.set_error(msg)
|
||||
conn.remove_stream(frame.stream_id)
|
||||
|
||||
elif frame.msg_type == MsgType.HEARTBEAT_DATA:
|
||||
await self._handle_heartbeat(conn, frame)
|
||||
|
||||
elif frame.msg_type == MsgType.PING:
|
||||
await conn.send_frame(Frame(0, MsgType.PONG, 0, frame.payload))
|
||||
|
||||
async def _handle_heartbeat(self, conn: TunnelConnection, frame: Frame) -> None:
|
||||
"""处理 proxy 上报的心跳数据,更新 DB,返回 ACK"""
|
||||
try:
|
||||
data = json.loads(frame.payload) if frame.payload else {}
|
||||
except Exception:
|
||||
data = {}
|
||||
|
||||
def _sync_heartbeat() -> dict[str, Any]:
|
||||
from src.database import create_session
|
||||
from src.services.proxy_node.service import ProxyNodeService
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
node = ProxyNodeService.heartbeat(
|
||||
db,
|
||||
node_id=conn.node_id,
|
||||
active_connections=data.get("active_connections"),
|
||||
total_requests=data.get("total_requests"),
|
||||
avg_latency_ms=data.get("avg_latency_ms"),
|
||||
)
|
||||
result: dict[str, Any] = {}
|
||||
if node.remote_config:
|
||||
result["remote_config"] = node.remote_config
|
||||
result["config_version"] = node.config_version or 0
|
||||
return result
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
try:
|
||||
ack = await asyncio.to_thread(_sync_heartbeat)
|
||||
except Exception as e:
|
||||
logger.warning("tunnel heartbeat DB update failed: {}", e)
|
||||
ack = {}
|
||||
|
||||
await conn.send_frame(Frame(0, MsgType.HEARTBEAT_ACK, 0, json.dumps(ack).encode()))
|
||||
|
||||
|
||||
# 全局单例
|
||||
_tunnel_manager: TunnelManager | None = None
|
||||
|
||||
|
||||
def get_tunnel_manager() -> TunnelManager:
|
||||
global _tunnel_manager
|
||||
if _tunnel_manager is None:
|
||||
_tunnel_manager = TunnelManager()
|
||||
return _tunnel_manager
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
WebSocket \u96a7\u9053\u4e8c\u8fdb\u5236\u5e27\u534f\u8bae
|
||||
|
||||
\u5e27\u683c\u5f0f:
|
||||
| stream_id (4B) | msg_type (1B) | flags (1B) | payload_len (4B) | payload (NB) |
|
||||
|
||||
\u7528\u4e8e Aether \u4e0e aether-proxy \u4e4b\u95f4\u7684 WebSocket \u96a7\u9053\u591a\u8def\u590d\u7528\u901a\u4fe1\u3002
|
||||
"""
|
||||
|
||||
import struct
|
||||
from enum import IntEnum
|
||||
from typing import Self
|
||||
|
||||
HEADER_SIZE = 10 # 4 + 1 + 1 + 4 bytes
|
||||
|
||||
|
||||
class MsgType(IntEnum):
|
||||
"""\u6d88\u606f\u7c7b\u578b"""
|
||||
|
||||
REQUEST_HEADERS = 0x01 # Aether -> Proxy: \u8bf7\u6c42\u5143\u6570\u636e (JSON)
|
||||
REQUEST_BODY = 0x02 # Aether -> Proxy: \u8bf7\u6c42\u4f53
|
||||
RESPONSE_HEADERS = 0x03 # Proxy -> Aether: \u54cd\u5e94\u72b6\u6001\u7801 + headers (JSON)
|
||||
RESPONSE_BODY = 0x04 # Proxy -> Aether: \u54cd\u5e94\u4f53\uff08\u6d41\u5f0f\u5206\u5757\uff09
|
||||
STREAM_END = 0x05 # \u53cc\u5411: \u6d41\u7ed3\u675f
|
||||
STREAM_ERROR = 0x06 # \u53cc\u5411: \u6d41\u9519\u8bef
|
||||
|
||||
PING = 0x10 # \u53cc\u5411: \u5fc3\u8df3 (stream_id=0)
|
||||
PONG = 0x11 # \u53cc\u5411: \u5fc3\u8df3\u54cd\u5e94 (stream_id=0)
|
||||
GOAWAY = 0x12 # \u53cc\u5411: \u4f18\u96c5\u5173\u95ed (stream_id=0)
|
||||
HEARTBEAT_DATA = 0x13 # Proxy -> Aether: \u6307\u6807\u4e0a\u62a5
|
||||
HEARTBEAT_ACK = 0x14 # Aether -> Proxy: \u5fc3\u8df3\u786e\u8ba4 + \u8fdc\u7a0b\u914d\u7f6e
|
||||
|
||||
|
||||
class FrameFlags:
|
||||
"""\u5e27\u6807\u5fd7\u4f4d"""
|
||||
|
||||
END_STREAM = 0x01
|
||||
GZIP_COMPRESSED = 0x02
|
||||
|
||||
|
||||
class Frame:
|
||||
"""WebSocket \u96a7\u9053\u5e27"""
|
||||
|
||||
__slots__ = ("stream_id", "msg_type", "flags", "payload")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream_id: int,
|
||||
msg_type: MsgType,
|
||||
flags: int = 0,
|
||||
payload: bytes = b"",
|
||||
) -> None:
|
||||
self.stream_id = stream_id
|
||||
self.msg_type = msg_type
|
||||
self.flags = flags
|
||||
self.payload = payload
|
||||
|
||||
def encode(self) -> bytes:
|
||||
header = struct.pack(
|
||||
"!IBBI",
|
||||
self.stream_id,
|
||||
self.msg_type,
|
||||
self.flags,
|
||||
len(self.payload),
|
||||
)
|
||||
return header + self.payload
|
||||
|
||||
@classmethod
|
||||
def decode(cls, data: bytes) -> Self:
|
||||
if len(data) < HEADER_SIZE:
|
||||
raise ValueError(
|
||||
f"\u5e27\u6570\u636e\u592a\u77ed: \u9700\u8981 {HEADER_SIZE} \u5b57\u8282, \u5b9e\u9645 {len(data)}"
|
||||
)
|
||||
stream_id, msg_type_raw, flags, payload_len = struct.unpack("!IBBI", data[:HEADER_SIZE])
|
||||
expected_total = HEADER_SIZE + payload_len
|
||||
if len(data) < expected_total:
|
||||
raise ValueError(
|
||||
f"\u5e27\u6570\u636e\u4e0d\u5b8c\u6574: \u9700\u8981 {expected_total} \u5b57\u8282, \u5b9e\u9645 {len(data)}"
|
||||
)
|
||||
try:
|
||||
msg_type = MsgType(msg_type_raw)
|
||||
except ValueError:
|
||||
raise ValueError(f"\u672a\u77e5\u6d88\u606f\u7c7b\u578b: 0x{msg_type_raw:02x}")
|
||||
payload = data[HEADER_SIZE:expected_total]
|
||||
return cls(stream_id, msg_type, flags, payload)
|
||||
|
||||
@property
|
||||
def is_end_stream(self) -> bool:
|
||||
return bool(self.flags & FrameFlags.END_STREAM)
|
||||
|
||||
@property
|
||||
def is_gzip(self) -> bool:
|
||||
return bool(self.flags & FrameFlags.GZIP_COMPRESSED)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"Frame(stream={self.stream_id}, type={self.msg_type.name}, "
|
||||
f"flags=0x{self.flags:02x}, payload_len={len(self.payload)})"
|
||||
)
|
||||
@@ -0,0 +1,133 @@
|
||||
"""
|
||||
Tunnel httpx Transport
|
||||
|
||||
自定义 httpx AsyncBaseTransport,将 HTTP 请求通过 WebSocket tunnel 发送到 aether-proxy。
|
||||
对 handler 层完全透明 -- 只需在创建 httpx.AsyncClient 时使用此 transport。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
import httpx
|
||||
|
||||
from .tunnel_manager import TunnelManager, TunnelStreamError, _StreamState, get_tunnel_manager
|
||||
|
||||
_HOP_BY_HOP_HEADERS = frozenset(
|
||||
{
|
||||
"host",
|
||||
"transfer-encoding",
|
||||
"content-length",
|
||||
"connection",
|
||||
"upgrade",
|
||||
"keep-alive",
|
||||
"proxy-authorization",
|
||||
"proxy-connection",
|
||||
"te",
|
||||
"trailer",
|
||||
}
|
||||
)
|
||||
|
||||
# bytes 版本,用于直接比较 httpx raw headers(key 已经是小写 bytes)
|
||||
_HOP_BY_HOP_HEADERS_BYTES = frozenset(h.encode("ascii") for h in _HOP_BY_HOP_HEADERS)
|
||||
|
||||
|
||||
class TunnelTransport(httpx.AsyncBaseTransport):
|
||||
"""通过 WebSocket tunnel 发送请求的 httpx transport"""
|
||||
|
||||
def __init__(self, node_id: str, timeout: float = 60.0) -> None:
|
||||
self._node_id = node_id
|
||||
self._timeout = timeout
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
manager = get_tunnel_manager()
|
||||
|
||||
# 构建 headers dict(跳过 hop-by-hop 和 httpx 内部 headers)
|
||||
# request.headers.raw 返回 (bytes, bytes) 元组,key 已经是小写
|
||||
headers: dict[str, str] = {}
|
||||
for key, value in request.headers.raw:
|
||||
if key not in _HOP_BY_HOP_HEADERS_BYTES:
|
||||
headers[key.decode("latin-1")] = value.decode("latin-1")
|
||||
|
||||
# 读取 body -- request.content 在 json= 传参时已由 httpx 序列化好;
|
||||
# 对 stream 类型的 request 需要先 read() 才能拿到完整 content。
|
||||
body = request.content or await request.aread() or None
|
||||
|
||||
stream_state: _StreamState | None = None
|
||||
try:
|
||||
stream_state = await manager.send_request(
|
||||
self._node_id,
|
||||
method=request.method,
|
||||
url=str(request.url),
|
||||
headers=headers,
|
||||
body=body,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
|
||||
# 等待响应头
|
||||
await stream_state.wait_headers(timeout=self._timeout)
|
||||
|
||||
# 构建 httpx.Response(流式 body)
|
||||
resp_headers = httpx.Headers(stream_state.headers)
|
||||
|
||||
return httpx.Response(
|
||||
status_code=stream_state.status,
|
||||
headers=resp_headers,
|
||||
stream=TunnelResponseStream(
|
||||
manager, self._node_id, stream_state, timeout=self._timeout
|
||||
),
|
||||
)
|
||||
|
||||
except TunnelStreamError as e:
|
||||
self._cleanup_stream(manager, stream_state)
|
||||
# 区分连接阶段和响应阶段的错误
|
||||
if stream_state and stream_state.status > 0:
|
||||
raise httpx.ReadError(str(e)) from e
|
||||
raise httpx.ConnectError(str(e)) from e
|
||||
except asyncio.TimeoutError:
|
||||
self._cleanup_stream(manager, stream_state)
|
||||
raise httpx.ReadTimeout("tunnel request timeout") from None
|
||||
|
||||
def _cleanup_stream(self, manager: TunnelManager, stream_state: _StreamState | None) -> None:
|
||||
if stream_state is None:
|
||||
return
|
||||
conn = manager.get_connection(self._node_id)
|
||||
if conn:
|
||||
conn.remove_stream(stream_state.stream_id)
|
||||
|
||||
|
||||
class TunnelResponseStream(httpx.AsyncByteStream):
|
||||
"""将 tunnel stream 的 body chunks 包装为 httpx AsyncByteStream"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manager: TunnelManager,
|
||||
node_id: str,
|
||||
stream_state: _StreamState,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
self._manager = manager
|
||||
self._node_id = node_id
|
||||
self._stream_state = stream_state
|
||||
self._timeout = timeout
|
||||
|
||||
async def __aiter__(self) -> AsyncGenerator[bytes, None]:
|
||||
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
|
||||
yield chunk
|
||||
|
||||
async def aclose(self) -> None:
|
||||
# 确保 stream 从 connection 的 pending 列表中移除,防止内存泄漏
|
||||
conn = self._manager.get_connection(self._node_id)
|
||||
if conn:
|
||||
conn.remove_stream(self._stream_state.stream_id)
|
||||
|
||||
|
||||
def is_tunnel_node(node_info: dict[str, Any] | None) -> bool:
|
||||
"""检查节点是否为 tunnel 模式且已连接"""
|
||||
if not node_info:
|
||||
return False
|
||||
return bool(node_info.get("tunnel_mode")) and bool(node_info.get("tunnel_connected"))
|
||||
Reference in New Issue
Block a user