feat(channel/gateway): 新增完整网关通道模块
新增设备身份管理、认证限流、并发通道、Webhook路由、RBAC权限控制、SSE/轮询降级等全套网关通道功能,包含: 1. 设备身份生成与签名验证 2. 设备令牌认证与速率限制 3. 内存+数据库双重设备注册表 4. 并发通道限流管理 5. Webhook安全处理与路由 6. RBAC权限校验系统 7. OpenAI API兼容适配层 8. Tailscale认证支持 9. HTTP轮询降级机制
This commit is contained in:
parent
c9d0a26f52
commit
ecd3c90e80
177
backend/package/yuxi/channel/gateway/__init__.py
Normal file
177
backend/package/yuxi/channel/gateway/__init__.py
Normal file
@ -0,0 +1,177 @@
|
||||
from yuxi.channel.gateway.auth import (
|
||||
GatewayAuthMode,
|
||||
GatewayAuthResult,
|
||||
TrustedProxyConfig,
|
||||
authenticate_gateway_connect,
|
||||
authenticate_tailscale,
|
||||
authenticate_trusted_proxy,
|
||||
clear_bootstrap_token,
|
||||
generate_bootstrap_token,
|
||||
get_bootstrap_token,
|
||||
set_bootstrap_token,
|
||||
)
|
||||
from yuxi.channel.gateway.auth_rate_limiter import AuthRateLimiter, AuthRateScope, auth_rate_limiter
|
||||
from yuxi.channel.gateway.device_auth import build_device_token, parse_device_token
|
||||
from yuxi.channel.gateway.device_identity import (
|
||||
derive_device_id,
|
||||
generate_device_identity,
|
||||
sign_challenge,
|
||||
verify_signature,
|
||||
)
|
||||
from yuxi.channel.gateway.device_registry import (
|
||||
list_devices,
|
||||
lookup_public_key,
|
||||
register_device,
|
||||
unregister_device,
|
||||
)
|
||||
from yuxi.channel.gateway.lanes import ChannelLane, LaneManager, lane_manager
|
||||
from yuxi.channel.gateway.net_utils import (
|
||||
GatewayBindMode,
|
||||
build_ws_security_error,
|
||||
can_bind_to_host,
|
||||
can_bind_to_host_cached,
|
||||
has_forwarded_request_headers,
|
||||
has_tailscale_proxy_headers,
|
||||
is_local_direct_request,
|
||||
is_localish_host,
|
||||
is_loopback_address,
|
||||
is_loopback_host,
|
||||
is_private_host,
|
||||
is_private_or_loopback_host,
|
||||
is_secure_ws_url,
|
||||
is_tailscale_proxy_request,
|
||||
is_trusted_proxy_address,
|
||||
resolve_client_ip,
|
||||
resolve_forwarded_client_ip,
|
||||
resolve_gateway_bind_host,
|
||||
resolve_tailscale_client_ip,
|
||||
)
|
||||
from yuxi.channel.gateway.protocol import (
|
||||
DeliveryMode,
|
||||
FrameType,
|
||||
GatewayErrorCode,
|
||||
GatewayRpcMethod,
|
||||
HelloOk,
|
||||
RpcEvent,
|
||||
RpcFrame,
|
||||
RpcRequest,
|
||||
RpcResponse,
|
||||
marshal_frame,
|
||||
unmarshal_frame,
|
||||
)
|
||||
from yuxi.channel.gateway.routes import (
|
||||
WebhookRegistry,
|
||||
build_health_response,
|
||||
build_webhook_path,
|
||||
webhook_registry,
|
||||
)
|
||||
from yuxi.channel.gateway.openai_adapter import (
|
||||
OPENAI_CHAT_COMPLETIONS_PATH,
|
||||
OPENAI_MODELS_PATH,
|
||||
OpenAIMessageConverter,
|
||||
format_sse_stream,
|
||||
)
|
||||
from yuxi.channel.gateway.rpc_dispatcher import RpcDispatcher, rpc_dispatcher
|
||||
from yuxi.channel.gateway.server import GatewayWsServer, gateway_ws_server
|
||||
from yuxi.channel.gateway.sse import GatewaySseEndpoint, gateway_sse_endpoint
|
||||
from yuxi.channel.gateway.polling import PollingFallback, polling_fallback, start_polling_cleanup
|
||||
from yuxi.channel.gateway.tailscale_auth import (
|
||||
TailscaleWhoisIdentity,
|
||||
get_tailscale_user_from_headers,
|
||||
read_tailscale_whois_identity,
|
||||
)
|
||||
from yuxi.channel.gateway.webhook_security import (
|
||||
WebhookAnomalyTracker,
|
||||
WebhookConcurrencyGuard,
|
||||
WebhookGuard,
|
||||
WebhookGuardConfig,
|
||||
WebhookGuardResult,
|
||||
WebhookSigner,
|
||||
webhook_guard,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AuthRateLimiter",
|
||||
"AuthRateScope",
|
||||
"ChannelLane",
|
||||
"DeliveryMode",
|
||||
"FrameType",
|
||||
"GatewayAuthMode",
|
||||
"GatewayAuthResult",
|
||||
"GatewayBindMode",
|
||||
"GatewayErrorCode",
|
||||
"GatewayRpcMethod",
|
||||
"HelloOk",
|
||||
"GatewayWsServer",
|
||||
"GatewaySseEndpoint",
|
||||
"LaneManager",
|
||||
"OPENAI_CHAT_COMPLETIONS_PATH",
|
||||
"OPENAI_MODELS_PATH",
|
||||
"OpenAIMessageConverter",
|
||||
"PollingFallback",
|
||||
"RpcDispatcher",
|
||||
"RpcEvent",
|
||||
"RpcFrame",
|
||||
"RpcRequest",
|
||||
"RpcResponse",
|
||||
"TailscaleWhoisIdentity",
|
||||
"TrustedProxyConfig",
|
||||
"WebhookAnomalyTracker",
|
||||
"WebhookConcurrencyGuard",
|
||||
"WebhookGuard",
|
||||
"WebhookGuardConfig",
|
||||
"WebhookGuardResult",
|
||||
"WebhookRegistry",
|
||||
"WebhookSigner",
|
||||
"authenticate_gateway_connect",
|
||||
"authenticate_tailscale",
|
||||
"authenticate_trusted_proxy",
|
||||
"auth_rate_limiter",
|
||||
"build_device_token",
|
||||
"build_health_response",
|
||||
"build_webhook_path",
|
||||
"build_ws_security_error",
|
||||
"can_bind_to_host",
|
||||
"can_bind_to_host_cached",
|
||||
"clear_bootstrap_token",
|
||||
"derive_device_id",
|
||||
"gateway_ws_server",
|
||||
"gateway_sse_endpoint",
|
||||
"format_sse_stream",
|
||||
"polling_fallback",
|
||||
"generate_bootstrap_token",
|
||||
"generate_device_identity",
|
||||
"get_bootstrap_token",
|
||||
"get_tailscale_user_from_headers",
|
||||
"has_forwarded_request_headers",
|
||||
"has_tailscale_proxy_headers",
|
||||
"is_local_direct_request",
|
||||
"is_localish_host",
|
||||
"is_loopback_address",
|
||||
"is_loopback_host",
|
||||
"is_private_host",
|
||||
"is_private_or_loopback_host",
|
||||
"is_secure_ws_url",
|
||||
"is_tailscale_proxy_request",
|
||||
"is_trusted_proxy_address",
|
||||
"lane_manager",
|
||||
"list_devices",
|
||||
"lookup_public_key",
|
||||
"marshal_frame",
|
||||
"parse_device_token",
|
||||
"read_tailscale_whois_identity",
|
||||
"register_device",
|
||||
"resolve_client_ip",
|
||||
"resolve_forwarded_client_ip",
|
||||
"resolve_gateway_bind_host",
|
||||
"resolve_tailscale_client_ip",
|
||||
"rpc_dispatcher",
|
||||
"set_bootstrap_token",
|
||||
"sign_challenge",
|
||||
"start_polling_cleanup",
|
||||
"unmarshal_frame",
|
||||
"unregister_device",
|
||||
"verify_signature",
|
||||
"webhook_guard",
|
||||
"webhook_registry",
|
||||
]
|
||||
387
backend/package/yuxi/channel/gateway/auth.py
Normal file
387
backend/package/yuxi/channel/gateway/auth.py
Normal file
@ -0,0 +1,387 @@
|
||||
import hmac
|
||||
import logging
|
||||
import secrets
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
|
||||
from server.utils.auth_utils import AuthUtils
|
||||
from yuxi.channel.gateway.auth_rate_limiter import (
|
||||
AUTH_LOCKOUT_SECONDS,
|
||||
AuthRateScope,
|
||||
auth_rate_limiter,
|
||||
)
|
||||
from yuxi.channel.gateway.rbac import GatewayRole, map_user_role
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AUTH_RATE_LOCKOUT_MINUTES = int(AUTH_LOCKOUT_SECONDS / 60)
|
||||
|
||||
LookupPublicKey = Callable[[str], Awaitable[str | None]]
|
||||
|
||||
|
||||
class GatewayAuthMode(StrEnum):
|
||||
TOKEN = "token"
|
||||
PASSWORD = "password"
|
||||
DEVICE_TOKEN = "device_token"
|
||||
BOOTSTRAP = "bootstrap"
|
||||
TAILSCALE = "tailscale"
|
||||
TRUSTED_PROXY = "trusted_proxy"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrustedProxyConfig:
|
||||
user_header: str
|
||||
required_headers: list[str] = field(default_factory=list)
|
||||
allow_users: list[str] = field(default_factory=list)
|
||||
allow_loopback: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class GatewayAuthResult:
|
||||
authenticated: bool
|
||||
user_id: str | None = None
|
||||
mode: GatewayAuthMode | None = None
|
||||
error: str | None = None
|
||||
metadata: dict = field(default_factory=dict)
|
||||
roles: list[GatewayRole] = field(default_factory=list)
|
||||
|
||||
|
||||
import threading
|
||||
|
||||
_bootstrap_token: str | None = None
|
||||
_bootstrap_token_lock = threading.Lock()
|
||||
|
||||
|
||||
def set_bootstrap_token(token: str | None) -> None:
|
||||
global _bootstrap_token
|
||||
with _bootstrap_token_lock:
|
||||
_bootstrap_token = token
|
||||
|
||||
|
||||
def get_bootstrap_token() -> str | None:
|
||||
with _bootstrap_token_lock:
|
||||
return _bootstrap_token
|
||||
|
||||
|
||||
def clear_bootstrap_token() -> None:
|
||||
global _bootstrap_token
|
||||
with _bootstrap_token_lock:
|
||||
_bootstrap_token = None
|
||||
|
||||
|
||||
def generate_bootstrap_token() -> str:
|
||||
return secrets.token_hex(24)
|
||||
|
||||
|
||||
def _timing_safe_compare(a: str, b: str) -> bool:
|
||||
if len(a) != len(b):
|
||||
return False
|
||||
return hmac.compare_digest(a.encode(), b.encode())
|
||||
|
||||
|
||||
def _authenticate_bootstrap_token(token: str, client_ip: str | None) -> GatewayAuthResult:
|
||||
scope_key = client_ip or "unknown"
|
||||
stored = get_bootstrap_token()
|
||||
|
||||
is_locked = auth_rate_limiter.is_locked(scope_key, AuthRateScope.BOOTSTRAP)
|
||||
if is_locked:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error=f"Bootstrap 认证尝试次数过多,请 {AUTH_RATE_LOCKOUT_MINUTES} 分钟后重试。",
|
||||
metadata={"rate_limited": True, "retry_after_seconds": AUTH_LOCKOUT_SECONDS},
|
||||
)
|
||||
|
||||
if not stored:
|
||||
auth_rate_limiter.record_failure(scope_key, AuthRateScope.BOOTSTRAP)
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="Bootstrap token 未配置或已过期。",
|
||||
)
|
||||
|
||||
if not _timing_safe_compare(token, stored):
|
||||
auth_rate_limiter.record_failure(scope_key, AuthRateScope.BOOTSTRAP)
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="Bootstrap token 验证失败。",
|
||||
)
|
||||
|
||||
auth_rate_limiter.reset(scope_key, AuthRateScope.BOOTSTRAP)
|
||||
return GatewayAuthResult(
|
||||
authenticated=True,
|
||||
user_id="bootstrap",
|
||||
mode=GatewayAuthMode.BOOTSTRAP,
|
||||
roles=[GatewayRole.ADMIN],
|
||||
metadata={"bootstrap": True},
|
||||
)
|
||||
|
||||
|
||||
async def authenticate_tailscale(
|
||||
headers: dict | None,
|
||||
client_ip: str | None = None,
|
||||
remote_addr: str | None = None,
|
||||
tailscale_whois: Callable[[str], Awaitable] | None = None,
|
||||
) -> GatewayAuthResult:
|
||||
from yuxi.channel.gateway.net_utils import (
|
||||
is_tailscale_proxy_request,
|
||||
resolve_tailscale_client_ip,
|
||||
)
|
||||
from yuxi.channel.gateway.tailscale_auth import (
|
||||
get_tailscale_user_from_headers,
|
||||
read_tailscale_whois_identity,
|
||||
)
|
||||
|
||||
tailscale_user = get_tailscale_user_from_headers(headers)
|
||||
if not tailscale_user:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="缺少 Tailscale 用户身份头信息。",
|
||||
)
|
||||
|
||||
login, name = tailscale_user
|
||||
|
||||
if not is_tailscale_proxy_request(remote_addr, headers):
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="请求未经过 Tailscale Serve 代理。",
|
||||
)
|
||||
|
||||
real_client_ip = resolve_tailscale_client_ip(remote_addr, headers)
|
||||
if not real_client_ip:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="无法解析 Tailscale 客户端 IP。",
|
||||
)
|
||||
|
||||
whois_func = tailscale_whois or read_tailscale_whois_identity
|
||||
whois = await whois_func(real_client_ip)
|
||||
if not whois or not whois.login:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="Tailscale whois 查询失败。",
|
||||
)
|
||||
|
||||
if whois.login != login:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="Tailscale 用户身份不匹配。",
|
||||
)
|
||||
|
||||
return GatewayAuthResult(
|
||||
authenticated=True,
|
||||
user_id=f"tailscale:{whois.login}",
|
||||
mode=GatewayAuthMode.TAILSCALE,
|
||||
metadata={"tailscale_login": whois.login, "tailscale_name": whois.name},
|
||||
roles=[GatewayRole.OPERATOR],
|
||||
)
|
||||
|
||||
|
||||
def authenticate_trusted_proxy(
|
||||
headers: dict | None,
|
||||
client_ip: str | None,
|
||||
remote_addr: str | None,
|
||||
config: TrustedProxyConfig,
|
||||
trusted_proxies: list[str],
|
||||
) -> GatewayAuthResult:
|
||||
from yuxi.channel.gateway.net_utils import (
|
||||
is_loopback_address,
|
||||
is_trusted_proxy_address,
|
||||
)
|
||||
|
||||
if not headers:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="缺少请求头信息,无法进行受信代理认证。",
|
||||
)
|
||||
|
||||
if not remote_addr or not is_trusted_proxy_address(remote_addr, trusted_proxies):
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="请求来源 IP 不在受信代理列表中。",
|
||||
)
|
||||
|
||||
remote_is_loopback = is_loopback_address(remote_addr)
|
||||
if remote_is_loopback and not config.allow_loopback:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="受信代理不允许回环地址来源,请启用 allow_loopback 配置。",
|
||||
)
|
||||
|
||||
for header_name in config.required_headers:
|
||||
value = headers.get(header_name.lower())
|
||||
if not value or not value.strip():
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error=f"受信代理缺少必要的请求头: {header_name}。",
|
||||
)
|
||||
|
||||
user_header_name = config.user_header.lower()
|
||||
user_header_value = headers.get(user_header_name)
|
||||
if not user_header_value or not user_header_value.strip():
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error=f"受信代理未提供用户身份请求头: {config.user_header}。",
|
||||
)
|
||||
|
||||
user = user_header_value.strip()
|
||||
|
||||
if config.allow_users and user not in config.allow_users:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error=f"用户 {user} 不在受信代理的允许列表中。",
|
||||
)
|
||||
|
||||
return GatewayAuthResult(
|
||||
authenticated=True,
|
||||
user_id=f"proxy:{user}",
|
||||
mode=GatewayAuthMode.TRUSTED_PROXY,
|
||||
metadata={"proxy_user": user, "proxy_remote": remote_addr},
|
||||
roles=[GatewayRole.ADMIN],
|
||||
)
|
||||
|
||||
|
||||
async def authenticate_gateway_connect(
|
||||
auth_header: str | None,
|
||||
query_token: str | None = None,
|
||||
client_ip: str | None = None,
|
||||
shared_password: str | None = None,
|
||||
lookup_public_key: LookupPublicKey | None = None,
|
||||
bootstrap_token_param: str | None = None,
|
||||
headers: dict | None = None,
|
||||
remote_addr: str | None = None,
|
||||
trusted_proxy_config: TrustedProxyConfig | None = None,
|
||||
trusted_proxies: list[str] | None = None,
|
||||
allow_tailscale: bool = False,
|
||||
) -> GatewayAuthResult:
|
||||
scope_key = client_ip or "unknown"
|
||||
|
||||
if trusted_proxy_config and trusted_proxies:
|
||||
result = authenticate_trusted_proxy(
|
||||
headers=headers,
|
||||
client_ip=client_ip,
|
||||
remote_addr=remote_addr,
|
||||
config=trusted_proxy_config,
|
||||
trusted_proxies=trusted_proxies,
|
||||
)
|
||||
if result.authenticated:
|
||||
return result
|
||||
|
||||
if bootstrap_token_param:
|
||||
return _authenticate_bootstrap_token(bootstrap_token_param, client_ip)
|
||||
|
||||
token = None
|
||||
if auth_header and auth_header.startswith("Bearer "):
|
||||
token = auth_header[7:]
|
||||
elif query_token:
|
||||
token = query_token
|
||||
|
||||
if allow_tailscale:
|
||||
from yuxi.channel.gateway.net_utils import is_local_direct_request
|
||||
from yuxi.channel.gateway.tailscale_auth import get_tailscale_user_from_headers
|
||||
|
||||
if not is_local_direct_request(remote_addr, headers):
|
||||
tailscale_user = get_tailscale_user_from_headers(headers)
|
||||
if tailscale_user and not (shared_password and token):
|
||||
result = await authenticate_tailscale(
|
||||
headers=headers,
|
||||
client_ip=client_ip,
|
||||
remote_addr=remote_addr,
|
||||
)
|
||||
if result.authenticated:
|
||||
return result
|
||||
|
||||
if token and token.startswith("dv."):
|
||||
return await _authenticate_device_token(token, client_ip, lookup_public_key)
|
||||
|
||||
if shared_password and token:
|
||||
if not _timing_safe_compare(token, shared_password):
|
||||
auth_rate_limiter.record_failure(scope_key, AuthRateScope.PASSWORD)
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="密码认证失败。",
|
||||
)
|
||||
auth_rate_limiter.reset(scope_key, AuthRateScope.PASSWORD)
|
||||
return GatewayAuthResult(
|
||||
authenticated=True,
|
||||
user_id="shared",
|
||||
mode=GatewayAuthMode.PASSWORD,
|
||||
roles=[GatewayRole.ADMIN],
|
||||
)
|
||||
|
||||
if shared_password and not token:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="缺少认证 token。",
|
||||
)
|
||||
|
||||
if not token:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="缺少认证 token。WebSocket 头部须携带 Bearer token,或连接 URL 附加 ?token=xxx 参数。",
|
||||
)
|
||||
|
||||
is_locked = auth_rate_limiter.is_locked(scope_key, AuthRateScope.TOKEN)
|
||||
if is_locked:
|
||||
remaining = auth_rate_limiter.remaining_attempts(scope_key, AuthRateScope.TOKEN)
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error=f"认证尝试次数过多,请 {AUTH_RATE_LOCKOUT_MINUTES} 分钟后重试。",
|
||||
metadata={"rate_limited": True, "retry_after_seconds": AUTH_LOCKOUT_SECONDS},
|
||||
)
|
||||
|
||||
try:
|
||||
payload = AuthUtils.verify_access_token(token)
|
||||
user_id = payload.get("sub")
|
||||
if not user_id:
|
||||
raise ValueError("Token 缺少 subject 声明")
|
||||
user_role = payload.get("role")
|
||||
roles = [map_user_role(user_role)]
|
||||
auth_rate_limiter.reset(scope_key, AuthRateScope.TOKEN)
|
||||
return GatewayAuthResult(
|
||||
authenticated=True,
|
||||
user_id=user_id,
|
||||
mode=GatewayAuthMode.TOKEN,
|
||||
metadata={"token_payload": payload},
|
||||
roles=roles,
|
||||
)
|
||||
except Exception as e:
|
||||
auth_rate_limiter.record_failure(scope_key, AuthRateScope.TOKEN)
|
||||
logger.debug("Token verification failed for scope=%s: %s", scope_key, e)
|
||||
remaining = auth_rate_limiter.remaining_attempts(scope_key, AuthRateScope.TOKEN)
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error="Token 验证失败,请检查凭据是否有效。",
|
||||
metadata={"remaining_attempts": remaining},
|
||||
)
|
||||
|
||||
|
||||
async def _authenticate_device_token(
|
||||
token: str,
|
||||
client_ip: str | None,
|
||||
lookup_public_key: LookupPublicKey | None,
|
||||
) -> GatewayAuthResult:
|
||||
from yuxi.channel.gateway.device_auth import authenticate_device
|
||||
from yuxi.channel.gateway.device_registry import lookup_public_key as default_lookup
|
||||
|
||||
scope_key = client_ip or "unknown"
|
||||
|
||||
is_locked = auth_rate_limiter.is_locked(scope_key, AuthRateScope.DEVICE_TOKEN)
|
||||
if is_locked:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
error=f"设备认证尝试次数过多,请 {AUTH_RATE_LOCKOUT_MINUTES} 分钟后重试。",
|
||||
metadata={"rate_limited": True, "retry_after_seconds": AUTH_LOCKOUT_SECONDS},
|
||||
)
|
||||
|
||||
result = await authenticate_device(
|
||||
token=token,
|
||||
lookup_public_key=lookup_public_key or default_lookup,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
|
||||
if not result.authenticated:
|
||||
auth_rate_limiter.record_failure(scope_key, AuthRateScope.DEVICE_TOKEN)
|
||||
else:
|
||||
auth_rate_limiter.reset(scope_key, AuthRateScope.DEVICE_TOKEN)
|
||||
|
||||
return result
|
||||
111
backend/package/yuxi/channel/gateway/auth_rate_limiter.py
Normal file
111
backend/package/yuxi/channel/gateway/auth_rate_limiter.py
Normal file
@ -0,0 +1,111 @@
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from enum import StrEnum
|
||||
from typing import NamedTuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AUTH_FAIL_LIMIT = 10
|
||||
AUTH_LOCKOUT_SECONDS = 300
|
||||
AUTH_RATE_SCOPE_BUCKETS = 10_000
|
||||
|
||||
|
||||
class AuthRateScope(StrEnum):
|
||||
TOKEN = "token"
|
||||
PASSWORD = "password"
|
||||
DEVICE_TOKEN = "device_token"
|
||||
BOOTSTRAP = "bootstrap"
|
||||
|
||||
|
||||
class _ScopeState(NamedTuple):
|
||||
failures: int
|
||||
locked_until: float
|
||||
|
||||
|
||||
class AuthRateLimiter:
|
||||
def __init__(
|
||||
self,
|
||||
fail_limit: int = AUTH_FAIL_LIMIT,
|
||||
lockout_seconds: float = AUTH_LOCKOUT_SECONDS,
|
||||
max_buckets: int = AUTH_RATE_SCOPE_BUCKETS,
|
||||
):
|
||||
self._fail_limit = fail_limit
|
||||
self._lockout_seconds = lockout_seconds
|
||||
self._max_buckets = max_buckets
|
||||
self._buckets: dict[tuple[str, AuthRateScope], _ScopeState] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _cleanup_expired(self) -> None:
|
||||
if len(self._buckets) <= self._max_buckets:
|
||||
return
|
||||
now = time.monotonic()
|
||||
expired = [
|
||||
key
|
||||
for key, state in self._buckets.items()
|
||||
if now >= state.locked_until and state.failures < self._fail_limit
|
||||
]
|
||||
for key in expired:
|
||||
del self._buckets[key]
|
||||
|
||||
def record_failure(self, identifier: str, scope: AuthRateScope) -> bool:
|
||||
with self._lock:
|
||||
self._cleanup_expired()
|
||||
key = (identifier, scope)
|
||||
now = time.monotonic()
|
||||
state = self._buckets.get(key)
|
||||
|
||||
if state is not None and now < state.locked_until:
|
||||
return False
|
||||
|
||||
new_failures = (state.failures + 1) if state else 1
|
||||
locked_until = now + self._lockout_seconds if new_failures >= self._fail_limit else 0.0
|
||||
self._buckets[key] = _ScopeState(new_failures, locked_until)
|
||||
|
||||
if new_failures >= self._fail_limit:
|
||||
logger.warning(
|
||||
"Auth rate limit locked: identifier=%s scope=%s",
|
||||
identifier,
|
||||
scope.value,
|
||||
)
|
||||
|
||||
if len(self._buckets) > self._max_buckets * 2:
|
||||
self._force_evict(now)
|
||||
|
||||
return True
|
||||
|
||||
def reset(self, identifier: str, scope: AuthRateScope) -> None:
|
||||
key = (identifier, scope)
|
||||
with self._lock:
|
||||
self._buckets.pop(key, None)
|
||||
|
||||
def is_locked(self, identifier: str, scope: AuthRateScope) -> bool:
|
||||
key = (identifier, scope)
|
||||
with self._lock:
|
||||
state = self._buckets.get(key)
|
||||
if state is None:
|
||||
return False
|
||||
return time.monotonic() < state.locked_until
|
||||
|
||||
def remaining_attempts(self, identifier: str, scope: AuthRateScope) -> int:
|
||||
key = (identifier, scope)
|
||||
with self._lock:
|
||||
state = self._buckets.get(key)
|
||||
if state is None:
|
||||
return self._fail_limit
|
||||
now = time.monotonic()
|
||||
if state.locked_until > 0.0 and now >= state.locked_until:
|
||||
return self._fail_limit
|
||||
return max(0, self._fail_limit - state.failures)
|
||||
|
||||
def _force_evict(self, now: float) -> None:
|
||||
by_expiry = sorted(
|
||||
self._buckets.items(),
|
||||
key=lambda item: item[1].locked_until,
|
||||
)
|
||||
to_remove = len(self._buckets) - self._max_buckets
|
||||
for i in range(to_remove):
|
||||
del self._buckets[by_expiry[i][0]]
|
||||
|
||||
|
||||
auth_rate_limiter = AuthRateLimiter()
|
||||
359
backend/package/yuxi/channel/gateway/broadcaster.py
Normal file
359
backend/package/yuxi/channel/gateway/broadcaster.py
Normal file
@ -0,0 +1,359 @@
|
||||
import logging
|
||||
import time as _time
|
||||
from collections.abc import Collection
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.gateway.auth import GatewayAuthResult
|
||||
from yuxi.channel.gateway.protocol import RpcEvent, marshal_frame
|
||||
from yuxi.channel.gateway.rbac import ROLE_HIERARCHY, GatewayRole
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_BROADCAST_RATE = 100
|
||||
DEFAULT_BROADCAST_BURST = 200
|
||||
DEFAULT_MAX_BUFFERED_FRAMES = 256
|
||||
DEFAULT_SLOW_CONSUMER_THRESHOLD = 50
|
||||
|
||||
PUBLIC_EVENTS: set[str] = {
|
||||
"heartbeat",
|
||||
"ping",
|
||||
"server.shutdown",
|
||||
"server.tick",
|
||||
}
|
||||
|
||||
READ_EVENTS: set[str] = PUBLIC_EVENTS | {
|
||||
"channel.created",
|
||||
"channel.updated",
|
||||
"channel.deleted",
|
||||
"channel.state_change",
|
||||
"channel.message",
|
||||
"session.created",
|
||||
"session.updated",
|
||||
"session.deleted",
|
||||
"session.message",
|
||||
"session.tool",
|
||||
"agent.status",
|
||||
"agent.event",
|
||||
"cron.status",
|
||||
}
|
||||
|
||||
WRITE_EVENTS: set[str] = READ_EVENTS | {
|
||||
"config.changed",
|
||||
}
|
||||
|
||||
ADMIN_EVENTS: set[str] = WRITE_EVENTS | {
|
||||
"plugin.installed",
|
||||
"plugin.removed",
|
||||
"device.pair.requested",
|
||||
"device.pair.resolved",
|
||||
}
|
||||
|
||||
EVENT_MIN_ROLE: dict[str, GatewayRole] = {}
|
||||
|
||||
for _event in PUBLIC_EVENTS:
|
||||
EVENT_MIN_ROLE[_event] = GatewayRole.VIEWER
|
||||
for _event in READ_EVENTS - PUBLIC_EVENTS:
|
||||
EVENT_MIN_ROLE[_event] = GatewayRole.VIEWER
|
||||
for _event in WRITE_EVENTS - READ_EVENTS:
|
||||
EVENT_MIN_ROLE[_event] = GatewayRole.OPERATOR
|
||||
for _event in ADMIN_EVENTS - WRITE_EVENTS:
|
||||
EVENT_MIN_ROLE[_event] = GatewayRole.ADMIN
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ConnInfo:
|
||||
conn_id: str
|
||||
user_id: str | None = None
|
||||
roles: list[GatewayRole] = field(default_factory=list)
|
||||
subscribed_events: set[str] = field(default_factory=set)
|
||||
pending_frames: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class BroadcastFilter:
|
||||
conn_ids: Collection[str] | None = None
|
||||
user_ids: Collection[str] | None = None
|
||||
roles: Collection[GatewayRole] | None = None
|
||||
drop_if_slow: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class BroadcastResult:
|
||||
sent: int = 0
|
||||
dropped: int = 0
|
||||
filtered: int = 0
|
||||
|
||||
|
||||
class TokenBucket:
|
||||
def __init__(self, rate: float, burst: float):
|
||||
self.rate = rate
|
||||
self.burst = burst
|
||||
self.tokens = burst
|
||||
self.last_refill = _time.monotonic()
|
||||
|
||||
def consume(self, n: float = 1.0) -> bool:
|
||||
self._refill()
|
||||
if self.tokens >= n:
|
||||
self.tokens -= n
|
||||
return True
|
||||
return False
|
||||
|
||||
def _refill(self):
|
||||
now = _time.monotonic()
|
||||
elapsed = now - self.last_refill
|
||||
self.tokens = min(self.burst, self.tokens + elapsed * self.rate)
|
||||
self.last_refill = now
|
||||
|
||||
|
||||
class GatewayBroadcaster:
|
||||
def __init__(
|
||||
self,
|
||||
broadcast_rate: float = DEFAULT_BROADCAST_RATE,
|
||||
broadcast_burst: float = DEFAULT_BROADCAST_BURST,
|
||||
max_buffered_frames: int = DEFAULT_MAX_BUFFERED_FRAMES,
|
||||
slow_consumer_threshold: int = DEFAULT_SLOW_CONSUMER_THRESHOLD,
|
||||
):
|
||||
self._connections: dict[str, _ConnInfo] = {}
|
||||
self._rate_limiter = TokenBucket(broadcast_rate, broadcast_burst)
|
||||
self._max_buffered_frames = max_buffered_frames
|
||||
self._slow_consumer_threshold = slow_consumer_threshold
|
||||
self._send_func: dict[str, Any] = {}
|
||||
self._global_seq = 0
|
||||
|
||||
@property
|
||||
def global_seq(self) -> int:
|
||||
return self._global_seq
|
||||
|
||||
@property
|
||||
def connection_count(self) -> int:
|
||||
return len(self._connections)
|
||||
|
||||
def register_connection(
|
||||
self,
|
||||
conn_id: str,
|
||||
auth: GatewayAuthResult,
|
||||
send_fn: Any,
|
||||
) -> None:
|
||||
info = _ConnInfo(
|
||||
conn_id=conn_id,
|
||||
user_id=auth.user_id,
|
||||
roles=list(auth.roles) if auth.roles else [],
|
||||
)
|
||||
self._connections[conn_id] = info
|
||||
self._send_func[conn_id] = send_fn
|
||||
|
||||
def unregister_connection(self, conn_id: str) -> None:
|
||||
self._connections.pop(conn_id, None)
|
||||
self._send_func.pop(conn_id, None)
|
||||
|
||||
def subscribe(self, conn_id: str, events: Collection[str]) -> None:
|
||||
info = self._connections.get(conn_id)
|
||||
if info is None:
|
||||
return
|
||||
info.subscribed_events.update(events)
|
||||
|
||||
def unsubscribe(self, conn_id: str, events: Collection[str]) -> None:
|
||||
info = self._connections.get(conn_id)
|
||||
if info is None:
|
||||
return
|
||||
info.subscribed_events.difference_update(events)
|
||||
if not events or "*" in events:
|
||||
info.subscribed_events.clear()
|
||||
|
||||
def clear_subscriptions(self, conn_id: str) -> None:
|
||||
info = self._connections.get(conn_id)
|
||||
if info is None:
|
||||
return
|
||||
info.subscribed_events.clear()
|
||||
|
||||
def get_subscriptions(self, conn_id: str) -> set[str]:
|
||||
info = self._connections.get(conn_id)
|
||||
if info is None:
|
||||
return set()
|
||||
return set(info.subscribed_events)
|
||||
|
||||
async def broadcast(
|
||||
self,
|
||||
event: str,
|
||||
data: dict | None = None,
|
||||
filter_: BroadcastFilter | None = None,
|
||||
) -> BroadcastResult:
|
||||
return await self._broadcast_internal(event, data, filter_)
|
||||
|
||||
async def broadcast_to_conn_ids(
|
||||
self,
|
||||
event: str,
|
||||
data: dict | None,
|
||||
conn_ids: Collection[str],
|
||||
) -> BroadcastResult:
|
||||
return await self._broadcast_internal(event, data, BroadcastFilter(conn_ids=conn_ids))
|
||||
|
||||
async def broadcast_to_user(
|
||||
self,
|
||||
event: str,
|
||||
data: dict | None,
|
||||
user_id: str,
|
||||
) -> BroadcastResult:
|
||||
return await self._broadcast_internal(event, data, BroadcastFilter(user_ids={user_id}))
|
||||
|
||||
async def broadcast_to_role(
|
||||
self,
|
||||
event: str,
|
||||
data: dict | None,
|
||||
role: GatewayRole,
|
||||
) -> BroadcastResult:
|
||||
return await self._broadcast_internal(event, data, BroadcastFilter(roles={role}))
|
||||
|
||||
async def send_event(
|
||||
self,
|
||||
conn_id: str,
|
||||
event: str,
|
||||
data: dict | None = None,
|
||||
) -> bool:
|
||||
send_fn = self._send_func.get(conn_id)
|
||||
if send_fn is None:
|
||||
return False
|
||||
self._global_seq += 1
|
||||
rpc_event = RpcEvent(
|
||||
event=event,
|
||||
data=data,
|
||||
seq=self._global_seq,
|
||||
state_version=self._resolve_state_version(event),
|
||||
)
|
||||
try:
|
||||
await send_fn(marshal_frame(rpc_event))
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def _broadcast_internal(
|
||||
self,
|
||||
event: str,
|
||||
data: dict | None,
|
||||
filter_: BroadcastFilter | None,
|
||||
) -> BroadcastResult:
|
||||
|
||||
result = BroadcastResult()
|
||||
|
||||
if not self._connections:
|
||||
return result
|
||||
|
||||
candidates = [
|
||||
conn_id
|
||||
for conn_id, info in self._connections.items()
|
||||
if self._passes_filter(info, event, filter_) and self._has_event_scope(info, event)
|
||||
]
|
||||
candidate_count = len(candidates)
|
||||
result.filtered = len(self._connections) - candidate_count
|
||||
|
||||
if not candidates:
|
||||
return result
|
||||
|
||||
if not self._rate_limiter.consume(candidate_count):
|
||||
logger.warning("Broadcast rate limited: event=%s candidates=%d", event, candidate_count)
|
||||
return result
|
||||
|
||||
self._global_seq += 1
|
||||
rpc_event = RpcEvent(
|
||||
event=event,
|
||||
data=data,
|
||||
seq=self._global_seq,
|
||||
state_version=self._resolve_state_version(event),
|
||||
)
|
||||
frame = marshal_frame(rpc_event)
|
||||
|
||||
for conn_id in candidates:
|
||||
info = self._connections.get(conn_id)
|
||||
if info is None:
|
||||
continue
|
||||
|
||||
info.pending_frames += 1
|
||||
is_slow = info.pending_frames > self._slow_consumer_threshold
|
||||
drop = filter_ is not None and filter_.drop_if_slow
|
||||
|
||||
if is_slow and drop:
|
||||
info.pending_frames = max(0, info.pending_frames - 1)
|
||||
result.dropped += 1
|
||||
continue
|
||||
|
||||
if info.pending_frames > self._max_buffered_frames:
|
||||
logger.warning(
|
||||
"Slow consumer exceeded max buffer: %s (event=%s pending=%d)",
|
||||
conn_id,
|
||||
event,
|
||||
info.pending_frames,
|
||||
)
|
||||
info.pending_frames = max(0, info.pending_frames - 1)
|
||||
result.dropped += 1
|
||||
continue
|
||||
|
||||
send_fn = self._send_func.get(conn_id)
|
||||
if send_fn is None:
|
||||
continue
|
||||
|
||||
try:
|
||||
await send_fn(frame)
|
||||
info.pending_frames = max(0, info.pending_frames - 1)
|
||||
result.sent += 1
|
||||
except Exception:
|
||||
logger.debug("Failed to send event %s to %s", event, conn_id)
|
||||
|
||||
return result
|
||||
|
||||
def _passes_filter(
|
||||
self,
|
||||
info: _ConnInfo,
|
||||
event: str,
|
||||
filter_: BroadcastFilter | None,
|
||||
) -> bool:
|
||||
if info.subscribed_events:
|
||||
if event not in info.subscribed_events and "*" not in info.subscribed_events:
|
||||
return False
|
||||
|
||||
if filter_ is None:
|
||||
return True
|
||||
|
||||
if filter_.conn_ids is not None and info.conn_id not in filter_.conn_ids:
|
||||
return False
|
||||
if filter_.user_ids is not None and info.user_id not in filter_.user_ids:
|
||||
return False
|
||||
if filter_.roles is not None:
|
||||
user_has_role = any(r in filter_.roles for r in info.roles)
|
||||
if not user_has_role:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _has_event_scope(info: _ConnInfo, event: str) -> bool:
|
||||
min_role = EVENT_MIN_ROLE.get(event)
|
||||
if min_role is None:
|
||||
return True
|
||||
|
||||
if not info.roles:
|
||||
return min_role == GatewayRole.VIEWER
|
||||
|
||||
max_level = max(ROLE_HIERARCHY.get(r, 0) for r in info.roles)
|
||||
required_level = ROLE_HIERARCHY.get(min_role, 0)
|
||||
return max_level >= required_level
|
||||
|
||||
def pending_frame_count(self, conn_id: str) -> int:
|
||||
info = self._connections.get(conn_id)
|
||||
if info is None:
|
||||
return 0
|
||||
return info.pending_frames
|
||||
|
||||
def _resolve_state_version(self, event: str) -> int | None:
|
||||
versioned_events = {
|
||||
"channel.state_change",
|
||||
"config.changed",
|
||||
"plugin.installed",
|
||||
"plugin.removed",
|
||||
}
|
||||
if event in versioned_events:
|
||||
return self._global_seq
|
||||
return None
|
||||
|
||||
|
||||
gateway_broadcaster = GatewayBroadcaster()
|
||||
192
backend/package/yuxi/channel/gateway/channel_plugin.py
Normal file
192
backend/package/yuxi/channel/gateway/channel_plugin.py
Normal file
@ -0,0 +1,192 @@
|
||||
import logging
|
||||
|
||||
from yuxi.channel.runtime.backoff import BackoffConfig, ErrorBackoff
|
||||
from yuxi.channel.extensions.base import BaseChannelPlugin
|
||||
from yuxi.channel.gateway.client import GatewayClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_RECONNECT_LEVELS = [1.0, 2.0, 4.0, 8.0, 15.0, 30.0]
|
||||
|
||||
|
||||
class GatewayChannelPlugin(BaseChannelPlugin):
|
||||
"""使用 Gateway Client SDK 的渠道插件基类。
|
||||
|
||||
封装常用连接模式:
|
||||
- 自动重连(基于 ErrorBackoff 的状态管理)
|
||||
- 请求重试(基于 ErrorBackoff.execute)
|
||||
- 设备身份管理(自动生成/注册/刷新)
|
||||
- 生命周期回调(on_connect / on_disconnect)
|
||||
|
||||
用法::
|
||||
|
||||
class MyChannelPlugin(GatewayChannelPlugin):
|
||||
id = "my_channel"
|
||||
name = "MyChannel"
|
||||
|
||||
async def on_connect(self, client: GatewayClient) -> None:
|
||||
await client.request("channel.start")
|
||||
|
||||
async def on_event(self, client: GatewayClient, event) -> None:
|
||||
...
|
||||
"""
|
||||
|
||||
id: str = ""
|
||||
name: str = ""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
gateway_url: str = "ws://127.0.0.1:9001",
|
||||
token: str | None = None,
|
||||
device_token: str | None = None,
|
||||
request_timeout: float = 30.0,
|
||||
reconnect_base_ms: float = 1000.0,
|
||||
reconnect_max_ms: float = 30000.0,
|
||||
reconnect_factor: float = 2.0,
|
||||
reconnect_jitter: float = 0.1,
|
||||
):
|
||||
self._gateway_url = gateway_url
|
||||
self._token = token
|
||||
self._device_token = device_token
|
||||
self._request_timeout = request_timeout
|
||||
self._reconnect_base_ms = reconnect_base_ms
|
||||
self._reconnect_max_ms = reconnect_max_ms
|
||||
self._reconnect_factor = reconnect_factor
|
||||
self._reconnect_jitter = reconnect_jitter
|
||||
|
||||
self._client: GatewayClient | None = None
|
||||
self._backoff = ErrorBackoff(
|
||||
config=BackoffConfig(
|
||||
base_delay=reconnect_base_ms / 1000.0,
|
||||
max_delay=reconnect_max_ms / 1000.0,
|
||||
exponent=reconnect_factor,
|
||||
jitter=True,
|
||||
jitter_factor=reconnect_jitter,
|
||||
max_retries=0,
|
||||
),
|
||||
)
|
||||
|
||||
@property
|
||||
def client(self) -> GatewayClient | None:
|
||||
return self._client
|
||||
|
||||
@property
|
||||
def connected(self) -> bool:
|
||||
return self._client is not None and self._client.connected
|
||||
|
||||
async def start(self, ctx) -> object:
|
||||
self._client = GatewayClient(
|
||||
url=self._gateway_url,
|
||||
token=self._token,
|
||||
device_token=self._device_token,
|
||||
request_timeout=self._request_timeout,
|
||||
reconnect_base_ms=self._reconnect_base_ms,
|
||||
reconnect_max_ms=self._reconnect_max_ms,
|
||||
reconnect_factor=self._reconnect_factor,
|
||||
reconnect_jitter=self._reconnect_jitter,
|
||||
on_connect=self._on_gw_connect,
|
||||
on_disconnect=self._on_gw_disconnect,
|
||||
on_event=self._on_gw_event,
|
||||
on_error=self._on_gw_error,
|
||||
on_close=self._on_gw_close,
|
||||
)
|
||||
|
||||
try:
|
||||
await self._client.start()
|
||||
except ValueError as e:
|
||||
logger.error("[%s] Gateway client start failed: %s", self.id, e)
|
||||
raise
|
||||
|
||||
logger.info("[%s] Gateway channel plugin started", self.id)
|
||||
return self._client
|
||||
|
||||
async def stop(self, ctx) -> None:
|
||||
if self._client:
|
||||
await self._client.stop()
|
||||
self._client = None
|
||||
self._backoff.reset(f"ch-{self.id}")
|
||||
logger.info("[%s] Gateway channel plugin stopped", self.id)
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
method: str,
|
||||
params: dict | None = None,
|
||||
*,
|
||||
timeout: float | None = None,
|
||||
max_attempts: int = 3,
|
||||
) -> dict:
|
||||
"""发送 RPC 请求,带自动重试。"""
|
||||
if self._client is None:
|
||||
raise RuntimeError(f"[{self.id}] gateway client not started")
|
||||
|
||||
return await self._client.request_with_retry(
|
||||
method,
|
||||
params=params,
|
||||
timeout=timeout,
|
||||
max_attempts=max_attempts,
|
||||
)
|
||||
|
||||
async def init_device_identity(self) -> str:
|
||||
"""初始化设备身份,返回 device_token。"""
|
||||
if self._client is None:
|
||||
self._client = GatewayClient(url=self._gateway_url, token=self._token)
|
||||
return await self._client.init_device_identity()
|
||||
|
||||
async def refresh_device_token(self) -> str | None:
|
||||
"""刷新设备令牌。"""
|
||||
if self._client is None:
|
||||
return None
|
||||
return await self._client.refresh_device_token()
|
||||
|
||||
def create_backoff(self, name: str, levels: list[float] | None = None) -> ErrorBackoff:
|
||||
"""创建 ErrorBackoff 实例,供子类管理独立退避状态。"""
|
||||
return ErrorBackoff(
|
||||
config=BackoffConfig(
|
||||
base_delay=levels[0] if levels else 1.0,
|
||||
max_delay=levels[-1] if levels else 30.0,
|
||||
exponent=2.0,
|
||||
jitter=True,
|
||||
jitter_factor=0.1,
|
||||
max_retries=0,
|
||||
),
|
||||
)
|
||||
|
||||
async def on_connect(self, client: GatewayClient) -> None:
|
||||
"""子类可重写:Gateway 连接成功时回调。"""
|
||||
|
||||
async def on_disconnect(self, client: GatewayClient, code: int, reason: str) -> None:
|
||||
"""子类可重写:Gateway 断开连接时回调。"""
|
||||
|
||||
async def on_event(self, client: GatewayClient, event) -> None:
|
||||
"""子类可重写:收到 Gateway 事件时回调。"""
|
||||
|
||||
async def on_error(self, client: GatewayClient, error: Exception) -> None:
|
||||
"""子类可重写:Gateway 错误时回调。"""
|
||||
|
||||
async def _on_gw_connect(self, client: GatewayClient) -> None:
|
||||
try:
|
||||
await self.on_connect(client)
|
||||
except Exception:
|
||||
logger.exception("[%s] on_connect callback failed", self.id)
|
||||
|
||||
async def _on_gw_disconnect(self, client: GatewayClient, code: int, reason: str) -> None:
|
||||
try:
|
||||
await self.on_disconnect(client, code, reason)
|
||||
except Exception:
|
||||
logger.exception("[%s] on_disconnect callback failed", self.id)
|
||||
|
||||
async def _on_gw_event(self, client: GatewayClient, event) -> None:
|
||||
try:
|
||||
await self.on_event(client, event)
|
||||
except Exception:
|
||||
logger.exception("[%s] on_event callback failed", self.id)
|
||||
|
||||
async def _on_gw_error(self, client: GatewayClient, error: Exception) -> None:
|
||||
try:
|
||||
await self.on_error(client, error)
|
||||
except Exception:
|
||||
logger.exception("[%s] on_error callback failed", self.id)
|
||||
|
||||
async def _on_gw_close(self, client: GatewayClient, code: int, reason: str) -> None:
|
||||
logger.info("[%s] Gateway connection closed: code=%d reason=%s", self.id, code, reason)
|
||||
557
backend/package/yuxi/channel/gateway/client.py
Normal file
557
backend/package/yuxi/channel/gateway/client.py
Normal file
@ -0,0 +1,557 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
import aiohttp
|
||||
|
||||
from yuxi.channel.config.defaults import TIMEOUT
|
||||
from yuxi.channel.runtime.backoff import (
|
||||
BackoffConfig,
|
||||
ErrorBackoff,
|
||||
)
|
||||
from yuxi.channel.gateway.net_utils import is_secure_ws_url
|
||||
from yuxi.channel.gateway.protocol import (
|
||||
GatewayErrorCode,
|
||||
RpcEvent,
|
||||
RpcRequest,
|
||||
RpcResponse,
|
||||
marshal_frame,
|
||||
unmarshal_frame,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_GATEWAY_URL = "ws://127.0.0.1:9001"
|
||||
DEFAULT_REQUEST_TIMEOUT = TIMEOUT.http.normal
|
||||
DEFAULT_RECONNECT_BASE_MS = 1000.0
|
||||
DEFAULT_RECONNECT_MAX_MS = 30000.0
|
||||
DEFAULT_RECONNECT_FACTOR = 2.0
|
||||
DEFAULT_RECONNECT_JITTER = 0.1
|
||||
|
||||
OnConnectCallback = Callable[["GatewayClient"], Awaitable[None]]
|
||||
OnDisconnectCallback = Callable[["GatewayClient", int, str], Awaitable[None]]
|
||||
OnEventCallback = Callable[["GatewayClient", RpcEvent], Awaitable[None]]
|
||||
OnErrorCallback = Callable[["GatewayClient", Exception], Awaitable[None]]
|
||||
OnCloseCallback = Callable[["GatewayClient", int, str], Awaitable[None]]
|
||||
|
||||
|
||||
class GatewayClientError(Exception):
|
||||
def __init__(self, code: GatewayErrorCode | str, message: str, details: dict | None = None):
|
||||
self.code = code
|
||||
self.details = details or {}
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class _PendingRequest:
|
||||
__slots__ = ("future", "timeout_handle")
|
||||
future: asyncio.Future
|
||||
timeout_handle: asyncio.TimerHandle | None
|
||||
|
||||
def __init__(self):
|
||||
self.future = asyncio.get_running_loop().create_future()
|
||||
self.timeout_handle = None
|
||||
|
||||
|
||||
class GatewayClient:
|
||||
"""统一 Gateway WebSocket 客户端 SDK。
|
||||
|
||||
连接 → 认证 → 收发 RPC 请求/事件 → 自动重连。
|
||||
对齐 OpenClaw GatewayClient 语义,适配项目现有 Gateway 协议。
|
||||
|
||||
用法::
|
||||
|
||||
client = GatewayClient(
|
||||
url="ws://127.0.0.1:9001",
|
||||
token="my-token",
|
||||
on_event=my_event_handler,
|
||||
)
|
||||
await client.start()
|
||||
|
||||
resp = await client.request("system.health")
|
||||
...
|
||||
|
||||
await client.stop()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
url: str = DEFAULT_GATEWAY_URL,
|
||||
token: str | None = None,
|
||||
password: str | None = None,
|
||||
device_token: str | None = None,
|
||||
request_timeout: float = DEFAULT_REQUEST_TIMEOUT,
|
||||
reconnect_base_ms: float = DEFAULT_RECONNECT_BASE_MS,
|
||||
reconnect_max_ms: float = DEFAULT_RECONNECT_MAX_MS,
|
||||
reconnect_factor: float = DEFAULT_RECONNECT_FACTOR,
|
||||
reconnect_jitter: float = DEFAULT_RECONNECT_JITTER,
|
||||
on_connect: OnConnectCallback | None = None,
|
||||
on_disconnect: OnDisconnectCallback | None = None,
|
||||
on_event: OnEventCallback | None = None,
|
||||
on_error: OnErrorCallback | None = None,
|
||||
on_close: OnCloseCallback | None = None,
|
||||
):
|
||||
self._url = url
|
||||
self._token = token
|
||||
self._password = password
|
||||
self._device_token = device_token
|
||||
self._request_timeout = max(1.0, request_timeout)
|
||||
self._reconnect_base_ms = max(100.0, reconnect_base_ms)
|
||||
self._reconnect_max_ms = max(reconnect_base_ms, reconnect_max_ms)
|
||||
self._reconnect_factor = max(1.0, reconnect_factor)
|
||||
self._reconnect_jitter = float(reconnect_jitter)
|
||||
|
||||
self._on_connect = on_connect
|
||||
self._on_disconnect = on_disconnect
|
||||
self._on_event = on_event
|
||||
self._on_error = on_error
|
||||
self._on_close = on_close
|
||||
|
||||
self._ws: aiohttp.ClientWebSocketResponse | None = None
|
||||
self._session: aiohttp.ClientSession | None = None
|
||||
self._pending: dict[str, _PendingRequest] = {}
|
||||
self._running = False
|
||||
self._close_event = asyncio.Event()
|
||||
self._reconnect_task: asyncio.Task | None = None
|
||||
|
||||
self._device_id: str | None = None
|
||||
self._device_private_key_pem: str | None = None
|
||||
|
||||
self._last_tick: float | None = None
|
||||
self._tick_interval_ms = 30_000
|
||||
self._tick_watch_task: asyncio.Task | None = None
|
||||
|
||||
self._reconnect_backoff = ErrorBackoff(
|
||||
config=BackoffConfig(
|
||||
base_delay=reconnect_base_ms / 1000.0,
|
||||
max_delay=reconnect_max_ms / 1000.0,
|
||||
exponent=reconnect_factor,
|
||||
jitter=True,
|
||||
jitter_factor=reconnect_jitter,
|
||||
max_retries=0,
|
||||
),
|
||||
)
|
||||
|
||||
@property
|
||||
def connected(self) -> bool:
|
||||
return self._ws is not None and not self._ws.closed
|
||||
|
||||
@property
|
||||
def url(self) -> str:
|
||||
return self._url
|
||||
|
||||
async def start(self) -> None:
|
||||
"""启动客户端,开始连接和自动重连循环。"""
|
||||
if self._running:
|
||||
return
|
||||
self._running = True
|
||||
self._close_event.clear()
|
||||
self._validate_url()
|
||||
self._reconnect_task = asyncio.create_task(self._reconnect_loop())
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""停止客户端,断开连接并等待清理完成。"""
|
||||
self._running = False
|
||||
self._close_event.set()
|
||||
self._stop_tick_watch()
|
||||
self._cancel_all_pending(RuntimeError("gateway client stopped"))
|
||||
await self._disconnect()
|
||||
if self._session and not self._session.closed:
|
||||
await self._session.close()
|
||||
self._session = None
|
||||
if self._reconnect_task:
|
||||
self._reconnect_task.cancel()
|
||||
try:
|
||||
await self._reconnect_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._reconnect_task = None
|
||||
|
||||
async def request(self, method: str, params: dict | None = None, timeout: float | None = None) -> dict:
|
||||
"""发送 RPC 请求,返回响应结果。
|
||||
|
||||
Raises:
|
||||
GatewayClientError: 服务端返回错误
|
||||
RuntimeError: 未连接
|
||||
TimeoutError: 请求超时
|
||||
"""
|
||||
if not self.connected:
|
||||
raise RuntimeError("gateway not connected")
|
||||
|
||||
req = RpcRequest(method=method, params=params)
|
||||
frame = marshal_frame(req)
|
||||
request_id = req.id
|
||||
effective_timeout = timeout if timeout is not None else self._request_timeout
|
||||
|
||||
pending = _PendingRequest()
|
||||
self._pending[request_id] = pending
|
||||
|
||||
timeout_coro: asyncio.Task | None = None
|
||||
try:
|
||||
timeout_coro = asyncio.create_task(asyncio.sleep(effective_timeout))
|
||||
send_task = asyncio.create_task(self._ws.send_str(frame))
|
||||
|
||||
done, _pending_set = await asyncio.wait(
|
||||
[pending.future, send_task, timeout_coro],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in _pending_set:
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
if timeout_coro in done and not pending.future.done():
|
||||
raise TimeoutError(f"gateway request timeout for {method}")
|
||||
|
||||
if send_task in done:
|
||||
send_exc = send_task.exception()
|
||||
if send_exc:
|
||||
raise send_exc
|
||||
|
||||
if pending.future in done:
|
||||
exc = pending.future.exception()
|
||||
if exc:
|
||||
raise exc
|
||||
return pending.future.result()
|
||||
|
||||
raise RuntimeError(f"gateway request failed for {method}: unexpected state")
|
||||
|
||||
finally:
|
||||
self._pending.pop(request_id, None)
|
||||
if timeout_coro and not timeout_coro.done():
|
||||
timeout_coro.cancel()
|
||||
try:
|
||||
await timeout_coro
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def request_with_retry(
|
||||
self,
|
||||
method: str,
|
||||
params: dict | None = None,
|
||||
*,
|
||||
timeout: float | None = None,
|
||||
max_attempts: int = 3,
|
||||
) -> dict:
|
||||
"""发送 RPC 请求,失败时自动重试。
|
||||
|
||||
Raises:
|
||||
GatewayClientError: 服务端返回错误
|
||||
RuntimeError: 未连接
|
||||
TimeoutError: 请求超时
|
||||
"""
|
||||
backoff = ErrorBackoff(
|
||||
config=BackoffConfig(
|
||||
base_delay=0.4,
|
||||
max_delay=5.0,
|
||||
exponent=2.0,
|
||||
jitter=True,
|
||||
jitter_factor=0.1,
|
||||
max_retries=max_attempts - 1,
|
||||
),
|
||||
)
|
||||
|
||||
async def _do_request() -> dict:
|
||||
return await self.request(method, params=params, timeout=timeout)
|
||||
|
||||
return await backoff.execute(f"gw-rpc:{method}", _do_request)
|
||||
|
||||
async def init_device_identity(self) -> str:
|
||||
"""生成并注册设备身份,返回 device_token。
|
||||
|
||||
首次调用生成 Ed25519 密钥对,注册到设备注册表,
|
||||
并构建可用于认证的 device_token。
|
||||
"""
|
||||
from yuxi.channel.gateway.device_auth import build_device_token as _build_token
|
||||
from yuxi.channel.gateway.device_identity import generate_device_identity
|
||||
from yuxi.channel.gateway.device_registry import register_device
|
||||
|
||||
device_id, public_key_pem, private_key_pem = generate_device_identity()
|
||||
self._device_id = device_id
|
||||
self._device_private_key_pem = private_key_pem
|
||||
|
||||
register_device(device_id, public_key_pem)
|
||||
|
||||
token = self._build_device_token_internal()
|
||||
self._device_token = token
|
||||
|
||||
logger.info("GatewayClient device identity initialized: device_id=%s", device_id)
|
||||
return token
|
||||
|
||||
async def refresh_device_token(self) -> str | None:
|
||||
"""刷新设备令牌(重新签名时间戳),返回新的 device_token。"""
|
||||
if not self._device_id or not self._device_private_key_pem:
|
||||
return None
|
||||
token = self._build_device_token_internal()
|
||||
self._device_token = token
|
||||
return token
|
||||
|
||||
@property
|
||||
def device_id(self) -> str | None:
|
||||
return self._device_id
|
||||
|
||||
# ---- internal ----
|
||||
|
||||
def _validate_url(self) -> None:
|
||||
if not is_secure_ws_url(self._url):
|
||||
raise ValueError(
|
||||
f"不安全连接: 不允许通过明文 ws:// 连接到非回环地址 "
|
||||
f"(CWE-319)。请使用 wss:// 或将 gateway.bind 设为 loopback。"
|
||||
f" 当前 URL: {self._url}"
|
||||
)
|
||||
|
||||
async def _ensure_session(self) -> None:
|
||||
if self._session is None or self._session.closed:
|
||||
self._session = aiohttp.ClientSession()
|
||||
|
||||
async def _connect(self) -> bool:
|
||||
await self._ensure_session()
|
||||
headers = {}
|
||||
auth_value = self._resolve_auth_header()
|
||||
if auth_value:
|
||||
headers["Authorization"] = auth_value
|
||||
|
||||
try:
|
||||
self._ws = await self._session.ws_connect(self._url, headers=headers, heartbeat=30.0)
|
||||
except aiohttp.ClientError as e:
|
||||
logger.warning("Gateway client connect failed: %s", e)
|
||||
await self._fire_error(e)
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.warning("Gateway client connect error: %s", e)
|
||||
await self._fire_error(e)
|
||||
return False
|
||||
|
||||
try:
|
||||
request_id = f"auth-check-{id(self):x}"
|
||||
await asyncio.wait_for(
|
||||
self._ws.send_str(marshal_frame(RpcRequest(id=request_id, method="system.health", params={}))),
|
||||
timeout=TIMEOUT.gateway.connect,
|
||||
)
|
||||
auth_ok = False
|
||||
async for msg in self._ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
frame = unmarshal_frame(msg.data)
|
||||
if isinstance(frame, RpcResponse) and frame.id == request_id:
|
||||
auth_ok = frame.ok
|
||||
break
|
||||
elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
|
||||
break
|
||||
if not auth_ok:
|
||||
logger.warning("Gateway client auth check failed")
|
||||
await self._ws.close(code=4001, message=b"auth failed")
|
||||
self._ws = None
|
||||
return False
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("Gateway client auth check error: %s", e)
|
||||
ws = self._ws
|
||||
self._ws = None
|
||||
if ws and not ws.closed:
|
||||
try:
|
||||
await ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
def _resolve_auth_header(self) -> str | None:
|
||||
if self._device_token:
|
||||
return f"Bearer {self._device_token}"
|
||||
if self._password:
|
||||
return f"Bearer {self._password}"
|
||||
if self._token:
|
||||
return f"Bearer {self._token}"
|
||||
return None
|
||||
|
||||
async def _disconnect(self) -> None:
|
||||
self._stop_tick_watch()
|
||||
self._cancel_all_pending(RuntimeError("gateway disconnected"))
|
||||
ws = self._ws
|
||||
self._ws = None
|
||||
if ws and not ws.closed:
|
||||
try:
|
||||
await ws.close(code=1000, message=b"client shutdown")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _start_tick_watch(self) -> None:
|
||||
self._last_tick = None
|
||||
|
||||
async def _watch():
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(self._tick_interval_ms / 1000.0 * 2)
|
||||
if self._last_tick is None:
|
||||
continue
|
||||
gap_ms = (asyncio.get_event_loop().time() - self._last_tick) * 1000
|
||||
if gap_ms > self._tick_interval_ms * 3:
|
||||
logger.warning("Gateway tick timeout: gap=%.0fms", gap_ms)
|
||||
await self._disconnect()
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
self._tick_watch_task = asyncio.create_task(_watch())
|
||||
|
||||
def _stop_tick_watch(self) -> None:
|
||||
if self._tick_watch_task and not self._tick_watch_task.done():
|
||||
self._tick_watch_task.cancel()
|
||||
self._tick_watch_task = None
|
||||
self._last_tick = None
|
||||
|
||||
async def _reconnect_loop(self) -> None:
|
||||
backoff_ms = self._reconnect_base_ms
|
||||
|
||||
while self._running:
|
||||
connected = await self._connect()
|
||||
if not connected:
|
||||
delay = self._reconnect_backoff.config.compute_delay(0)
|
||||
try:
|
||||
await asyncio.wait_for(self._close_event.wait(), timeout=delay)
|
||||
return
|
||||
except TimeoutError:
|
||||
backoff_ms = min(backoff_ms * self._reconnect_factor, self._reconnect_max_ms)
|
||||
continue
|
||||
|
||||
backoff_ms = self._reconnect_base_ms
|
||||
self._start_tick_watch()
|
||||
await self._fire_connect()
|
||||
|
||||
try:
|
||||
await self._recv_loop()
|
||||
except aiohttp.ClientError as e:
|
||||
logger.warning("Gateway client recv error: %s", e)
|
||||
await self._fire_error(e)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
except Exception as e:
|
||||
logger.exception("Gateway client recv unexpected error")
|
||||
await self._fire_error(e)
|
||||
|
||||
close_code = 1006
|
||||
close_reason = "connection lost"
|
||||
if self._ws is not None:
|
||||
close_code = self._ws.close_code or close_code
|
||||
close_reason = self._ws._close_reason or close_reason
|
||||
|
||||
await self._disconnect()
|
||||
await self._fire_close(close_code, close_reason)
|
||||
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
delay = self._reconnect_backoff.config.compute_delay(0)
|
||||
try:
|
||||
await asyncio.wait_for(self._close_event.wait(), timeout=delay)
|
||||
return
|
||||
except TimeoutError:
|
||||
backoff_ms = min(backoff_ms * self._reconnect_factor, self._reconnect_max_ms)
|
||||
|
||||
async def _recv_loop(self) -> None:
|
||||
ws = self._ws
|
||||
if ws is None:
|
||||
return
|
||||
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
frame = unmarshal_frame(msg.data)
|
||||
except Exception:
|
||||
logger.warning("Gateway client: invalid frame received")
|
||||
continue
|
||||
|
||||
if isinstance(frame, RpcResponse):
|
||||
self._handle_response(frame)
|
||||
elif isinstance(frame, RpcEvent):
|
||||
self._handle_event(frame)
|
||||
|
||||
elif msg.type == aiohttp.WSMsgType.CLOSED:
|
||||
break
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
break
|
||||
|
||||
def _handle_response(self, resp: RpcResponse) -> None:
|
||||
pending = self._pending.get(resp.id)
|
||||
if pending is None:
|
||||
return
|
||||
self._pending.pop(resp.id, None)
|
||||
|
||||
if resp.ok:
|
||||
pending.future.set_result(resp.result or {})
|
||||
else:
|
||||
code = resp.error_code or GatewayErrorCode.INTERNAL_ERROR
|
||||
msg = resp.error_message or "unknown error"
|
||||
err = GatewayClientError(code=code, message=msg)
|
||||
pending.future.set_exception(err)
|
||||
|
||||
def _handle_event(self, evt: RpcEvent) -> None:
|
||||
if evt.event == "tick":
|
||||
self._last_tick = asyncio.get_event_loop().time()
|
||||
if self._on_event:
|
||||
asyncio.create_task(self._fire_event(evt))
|
||||
|
||||
def _cancel_all_pending(self, exc: Exception) -> None:
|
||||
for pending in list(self._pending.values()):
|
||||
if not pending.future.done():
|
||||
pending.future.set_exception(exc)
|
||||
self._pending.clear()
|
||||
|
||||
def _compute_reconnect_levels(self) -> list[float]:
|
||||
levels = []
|
||||
current = self._reconnect_base_ms / 1000.0
|
||||
for _ in range(10):
|
||||
if current >= self._reconnect_max_ms / 1000.0:
|
||||
break
|
||||
levels.append(current)
|
||||
current *= self._reconnect_factor
|
||||
if not levels or levels[-1] < self._reconnect_max_ms / 1000.0:
|
||||
levels.append(self._reconnect_max_ms / 1000.0)
|
||||
return levels
|
||||
|
||||
def _build_device_token_internal(self) -> str:
|
||||
from yuxi.channel.gateway.device_auth import build_device_token as _build_token, build_challenge
|
||||
from yuxi.channel.gateway.device_identity import sign_challenge
|
||||
|
||||
timestamp_ms = int(time.time() * 1000)
|
||||
challenge = build_challenge(self._device_id, timestamp_ms)
|
||||
signature = sign_challenge(self._device_private_key_pem, challenge)
|
||||
return _build_token(self._device_id, timestamp_ms, signature)
|
||||
|
||||
async def _fire_connect(self) -> None:
|
||||
if self._on_connect:
|
||||
try:
|
||||
await self._on_connect(self)
|
||||
except Exception:
|
||||
logger.exception("on_connect callback failed")
|
||||
|
||||
async def _fire_close(self, code: int, reason: str) -> None:
|
||||
await self._fire_disconnect(code, reason)
|
||||
if self._on_close:
|
||||
try:
|
||||
await self._on_close(self, code, reason)
|
||||
except Exception:
|
||||
logger.exception("on_close callback failed")
|
||||
|
||||
async def _fire_disconnect(self, code: int, reason: str) -> None:
|
||||
if self._on_disconnect:
|
||||
try:
|
||||
await self._on_disconnect(self, code, reason)
|
||||
except Exception:
|
||||
logger.exception("on_disconnect callback failed")
|
||||
|
||||
async def _fire_event(self, evt: RpcEvent) -> None:
|
||||
if self._on_event:
|
||||
try:
|
||||
await self._on_event(self, evt)
|
||||
except Exception:
|
||||
logger.exception("on_event callback failed")
|
||||
|
||||
async def _fire_error(self, err: Exception) -> None:
|
||||
if self._on_error:
|
||||
try:
|
||||
await self._on_error(self, err)
|
||||
except Exception:
|
||||
logger.exception("on_error callback failed")
|
||||
108
backend/package/yuxi/channel/gateway/device_auth.py
Normal file
108
backend/package/yuxi/channel/gateway/device_auth.py
Normal file
@ -0,0 +1,108 @@
|
||||
import base64
|
||||
import logging
|
||||
import time
|
||||
|
||||
from yuxi.channel.gateway.auth import GatewayAuthMode, GatewayAuthResult
|
||||
from yuxi.channel.gateway.device_identity import verify_signature
|
||||
from yuxi.channel.gateway.rbac import GatewayRole
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEVICE_TOKEN_PREFIX = "dv."
|
||||
DEVICE_TOKEN_TIMESTAMP_TOLERANCE_SECONDS = 300
|
||||
|
||||
|
||||
def parse_device_token(token: str) -> tuple[str, int, bytes] | None:
|
||||
"""解析设备令牌,格式: dv.{base64(device_id)}.{base64(timestamp_ms)}.{base64(signature)}"""
|
||||
if not token.startswith(DEVICE_TOKEN_PREFIX):
|
||||
return None
|
||||
parts = token[len(DEVICE_TOKEN_PREFIX) :].split(".")
|
||||
if len(parts) != 3:
|
||||
return None
|
||||
try:
|
||||
def _add_padding(s: str) -> str:
|
||||
missing = (4 - len(s) % 4) % 4
|
||||
return s + "=" * missing if missing else s
|
||||
|
||||
device_id = base64.urlsafe_b64decode(_add_padding(parts[0])).decode("utf-8")
|
||||
timestamp_ms = int(base64.urlsafe_b64decode(_add_padding(parts[1])).decode("utf-8"))
|
||||
signature = base64.urlsafe_b64decode(_add_padding(parts[2]))
|
||||
except (ValueError, UnicodeDecodeError, base64.binascii.Error):
|
||||
return None
|
||||
return device_id, timestamp_ms, signature
|
||||
|
||||
|
||||
def build_device_token(device_id: str, timestamp_ms: int, signature: bytes) -> str:
|
||||
"""构建设备令牌"""
|
||||
device_id_b64 = base64.urlsafe_b64encode(device_id.encode("utf-8")).rstrip(b"=").decode("ascii")
|
||||
ts_b64 = base64.urlsafe_b64encode(str(timestamp_ms).encode("utf-8")).rstrip(b"=").decode("ascii")
|
||||
sig_b64 = base64.urlsafe_b64encode(signature).rstrip(b"=").decode("ascii")
|
||||
return f"{DEVICE_TOKEN_PREFIX}{device_id_b64}.{ts_b64}.{sig_b64}"
|
||||
|
||||
|
||||
def build_challenge(device_id: str, timestamp_ms: int) -> bytes:
|
||||
return f"{device_id}:{timestamp_ms}".encode("utf-8")
|
||||
|
||||
|
||||
def validate_timestamp(timestamp_ms: int) -> bool:
|
||||
now_ms = int(time.time() * 1000)
|
||||
if timestamp_ms > now_ms:
|
||||
return False
|
||||
diff_ms = now_ms - timestamp_ms
|
||||
return diff_ms <= DEVICE_TOKEN_TIMESTAMP_TOLERANCE_SECONDS * 1000
|
||||
|
||||
|
||||
async def authenticate_device(
|
||||
token: str,
|
||||
lookup_public_key,
|
||||
client_ip: str | None = None,
|
||||
) -> GatewayAuthResult:
|
||||
parsed = parse_device_token(token)
|
||||
if parsed is None:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
mode=GatewayAuthMode.DEVICE_TOKEN,
|
||||
error="设备令牌格式无效。预期格式: dv.{deviceId}.{ts}.{sig}",
|
||||
)
|
||||
|
||||
device_id, timestamp_ms, signature = parsed
|
||||
|
||||
if not validate_timestamp(timestamp_ms):
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
mode=GatewayAuthMode.DEVICE_TOKEN,
|
||||
error="设备令牌时间戳超出容忍范围(±5 分钟)。",
|
||||
metadata={"device_id": device_id},
|
||||
)
|
||||
|
||||
public_key_pem = await lookup_public_key(device_id)
|
||||
if public_key_pem is None:
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
mode=GatewayAuthMode.DEVICE_TOKEN,
|
||||
error=f"未知设备: {device_id}",
|
||||
metadata={"device_id": device_id},
|
||||
)
|
||||
|
||||
challenge = build_challenge(device_id, timestamp_ms)
|
||||
if not verify_signature(public_key_pem, challenge, signature):
|
||||
return GatewayAuthResult(
|
||||
authenticated=False,
|
||||
mode=GatewayAuthMode.DEVICE_TOKEN,
|
||||
error="设备签名验证失败。",
|
||||
metadata={"device_id": device_id},
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Device authenticated: device_id=%s ip=%s",
|
||||
device_id,
|
||||
client_ip or "unknown",
|
||||
)
|
||||
|
||||
return GatewayAuthResult(
|
||||
authenticated=True,
|
||||
user_id=f"device:{device_id}",
|
||||
mode=GatewayAuthMode.DEVICE_TOKEN,
|
||||
metadata={"device_id": device_id},
|
||||
roles=[GatewayRole.OPERATOR],
|
||||
)
|
||||
68
backend/package/yuxi/channel/gateway/device_identity.py
Normal file
68
backend/package/yuxi/channel/gateway/device_identity.py
Normal file
@ -0,0 +1,68 @@
|
||||
import hashlib
|
||||
import logging
|
||||
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEVICE_ID_HEX_LENGTH = 16
|
||||
|
||||
|
||||
def generate_device_identity() -> tuple[str, str, str]:
|
||||
"""生成 Ed25519 设备身份。
|
||||
|
||||
Returns:
|
||||
(device_id, public_key_pem, private_key_pem)
|
||||
"""
|
||||
private_key = Ed25519PrivateKey.generate()
|
||||
public_key = private_key.public_key()
|
||||
|
||||
public_key_pem = public_key.public_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
).decode("utf-8")
|
||||
|
||||
private_key_pem = private_key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.BestAvailableEncryption(b"yuxi-device-key"),
|
||||
).decode("utf-8")
|
||||
|
||||
device_id = _derive_device_id_from_key(public_key)
|
||||
|
||||
logger.info("Generated device identity: device_id=%s", device_id)
|
||||
return device_id, public_key_pem, private_key_pem
|
||||
|
||||
|
||||
def derive_device_id(public_key_pem: str) -> str:
|
||||
public_key = serialization.load_pem_public_key(public_key_pem.encode("utf-8"))
|
||||
return _derive_device_id_from_key(public_key)
|
||||
|
||||
|
||||
_PRIVATE_KEY_PASSWORD = b"yuxi-device-key"
|
||||
|
||||
|
||||
def sign_challenge(private_key_pem: str, challenge: bytes) -> bytes:
|
||||
private_key = serialization.load_pem_private_key(
|
||||
private_key_pem.encode("utf-8"),
|
||||
password=_PRIVATE_KEY_PASSWORD,
|
||||
)
|
||||
return private_key.sign(challenge)
|
||||
|
||||
|
||||
def verify_signature(public_key_pem: str, challenge: bytes, signature: bytes) -> bool:
|
||||
try:
|
||||
public_key = serialization.load_pem_public_key(public_key_pem.encode("utf-8"))
|
||||
public_key.verify(signature, challenge)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _derive_device_id_from_key(public_key) -> str:
|
||||
public_key_raw = public_key.public_bytes(
|
||||
encoding=serialization.Encoding.Raw,
|
||||
format=serialization.PublicFormat.Raw,
|
||||
)
|
||||
return hashlib.sha256(public_key_raw).hexdigest()[:DEVICE_ID_HEX_LENGTH]
|
||||
61
backend/package/yuxi/channel/gateway/device_registry.py
Normal file
61
backend/package/yuxi/channel/gateway/device_registry.py
Normal file
@ -0,0 +1,61 @@
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_registry: dict[str, str] = {}
|
||||
_registry_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_device(device_id: str, public_key_pem: str) -> None:
|
||||
with _registry_lock:
|
||||
_registry[device_id] = public_key_pem
|
||||
logger.info("Device registered: device_id=%s", device_id)
|
||||
try:
|
||||
from yuxi.channel.gateway.device_registry_db import db_save_device
|
||||
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.create_task(db_save_device(device_id, public_key_pem))
|
||||
except RuntimeError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.debug("DB device save skipped: device_id=%s", device_id)
|
||||
|
||||
|
||||
def unregister_device(device_id: str) -> bool:
|
||||
with _registry_lock:
|
||||
if device_id in _registry:
|
||||
del _registry[device_id]
|
||||
logger.info("Device unregistered: device_id=%s", device_id)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def lookup_public_key(device_id: str) -> str | None:
|
||||
with _registry_lock:
|
||||
key = _registry.get(device_id)
|
||||
if key is not None:
|
||||
return key
|
||||
key = await _db_lookup(device_id)
|
||||
if key is not None:
|
||||
with _registry_lock:
|
||||
_registry[device_id] = key
|
||||
return key
|
||||
|
||||
|
||||
async def _db_lookup(device_id: str) -> str | None:
|
||||
try:
|
||||
from yuxi.channel.gateway.device_registry_db import db_lookup_public_key
|
||||
|
||||
return await db_lookup_public_key(device_id)
|
||||
except Exception:
|
||||
logger.debug("DB device lookup skipped: device_id=%s", device_id)
|
||||
return None
|
||||
|
||||
|
||||
def list_devices() -> list[str]:
|
||||
with _registry_lock:
|
||||
return list(_registry.keys())
|
||||
31
backend/package/yuxi/channel/gateway/device_registry_db.py
Normal file
31
backend/package/yuxi/channel/gateway/device_registry_db.py
Normal file
@ -0,0 +1,31 @@
|
||||
import logging
|
||||
|
||||
from yuxi.repositories.channel_device_identity_repo import DeviceIdentityRepository
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_repo = DeviceIdentityRepository()
|
||||
|
||||
|
||||
async def db_lookup_public_key(device_id: str) -> str | None:
|
||||
record = await _repo.get_by_device_id(device_id)
|
||||
if record is None or record.status != "active":
|
||||
return None
|
||||
|
||||
await _repo.mark_used(device_id)
|
||||
return record.public_key_pem
|
||||
|
||||
|
||||
async def db_save_device(device_id: str, public_key_pem: str) -> None:
|
||||
existing = await _repo.get_by_device_id(device_id)
|
||||
if existing is not None:
|
||||
await _repo.update(device_id, {"public_key_pem": public_key_pem, "status": "active"})
|
||||
else:
|
||||
await _repo.create(
|
||||
{
|
||||
"device_id": device_id,
|
||||
"public_key_pem": public_key_pem,
|
||||
"status": "active",
|
||||
}
|
||||
)
|
||||
logger.debug("Device saved to DB: device_id=%s", device_id)
|
||||
114
backend/package/yuxi/channel/gateway/lanes.py
Normal file
114
backend/package/yuxi/channel/gateway/lanes.py
Normal file
@ -0,0 +1,114 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MAX_LANES = 5
|
||||
DEFAULT_GLOBAL_MAX_LANES = 50
|
||||
|
||||
|
||||
class ChannelLane:
|
||||
"""单个渠道账户的并发通道。
|
||||
|
||||
每个账户一个独立 lane,控制该账户下同时运行的 Agent 数量。
|
||||
"""
|
||||
|
||||
def __init__(self, max_lanes: int = DEFAULT_MAX_LANES):
|
||||
self._semaphore = asyncio.Semaphore(max_lanes)
|
||||
self._active_runs: int = 0
|
||||
self._busy: bool = False
|
||||
self._max_lanes = max_lanes
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[None]:
|
||||
async with self._semaphore:
|
||||
self._active_runs += 1
|
||||
if self._active_runs >= self._max_lanes:
|
||||
self._busy = True
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._active_runs = max(0, self._active_runs - 1)
|
||||
if self._active_runs < self._max_lanes:
|
||||
self._busy = False
|
||||
|
||||
@property
|
||||
def active_runs(self) -> int:
|
||||
return self._active_runs
|
||||
|
||||
@property
|
||||
def is_busy(self) -> bool:
|
||||
return self._busy
|
||||
|
||||
@property
|
||||
def available_slots(self) -> int:
|
||||
return max(0, self._max_lanes - self._active_runs)
|
||||
|
||||
|
||||
class LaneManager:
|
||||
"""并发通道管理器。
|
||||
|
||||
按账户维度管理并发通道,同时提供全局上限保护。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
per_account_max: int = DEFAULT_MAX_LANES,
|
||||
global_max: int = DEFAULT_GLOBAL_MAX_LANES,
|
||||
):
|
||||
self._per_account_max = per_account_max
|
||||
self._lanes: dict[str, ChannelLane] = {}
|
||||
self._global_semaphore = asyncio.Semaphore(global_max)
|
||||
|
||||
def _key(self, channel_type: str, account_id: str) -> str:
|
||||
return f"{channel_type}:{account_id}"
|
||||
|
||||
def get_lane(self, channel_type: str, account_id: str) -> ChannelLane:
|
||||
key = self._key(channel_type, account_id)
|
||||
if key not in self._lanes:
|
||||
self._lanes[key] = ChannelLane(max_lanes=self._per_account_max)
|
||||
return self._lanes[key]
|
||||
|
||||
@asynccontextmanager
|
||||
async def run_with_lane(self, channel_type: str, account_id: str) -> AsyncIterator[ChannelLane]:
|
||||
lane = self.get_lane(channel_type, account_id)
|
||||
if lane.is_busy:
|
||||
logger.warning(
|
||||
"Channel lane busy: %s/%s, active=%d",
|
||||
channel_type,
|
||||
account_id,
|
||||
lane.active_runs,
|
||||
)
|
||||
async with lane.acquire():
|
||||
async with self._global_semaphore:
|
||||
yield lane
|
||||
|
||||
def remove_lane(self, channel_type: str, account_id: str) -> None:
|
||||
key = self._key(channel_type, account_id)
|
||||
self._lanes.pop(key, None)
|
||||
|
||||
def get_stats(self) -> dict[str, dict]:
|
||||
return {
|
||||
key: {
|
||||
"active_runs": lane.active_runs,
|
||||
"busy": lane.is_busy,
|
||||
"available_slots": lane.available_slots,
|
||||
}
|
||||
for key, lane in self._lanes.items()
|
||||
}
|
||||
|
||||
def get_total_active(self) -> int:
|
||||
return sum(lane.active_runs for lane in self._lanes.values())
|
||||
|
||||
def cleanup_idle(self) -> int:
|
||||
idle_keys = [key for key, lane in self._lanes.items() if lane.active_runs == 0]
|
||||
for key in idle_keys:
|
||||
del self._lanes[key]
|
||||
if idle_keys:
|
||||
logger.debug("LaneManager cleaned up %d idle lanes", len(idle_keys))
|
||||
return len(idle_keys)
|
||||
|
||||
|
||||
lane_manager = LaneManager()
|
||||
373
backend/package/yuxi/channel/gateway/net_utils.py
Normal file
373
backend/package/yuxi/channel/gateway/net_utils.py
Normal file
@ -0,0 +1,373 @@
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
from enum import StrEnum
|
||||
from urllib.parse import urlparse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_LOOPBACK_HOSTS = frozenset(["localhost", "127.0.0.1", "::1"])
|
||||
|
||||
_PRIVATE_IP_NETS = [
|
||||
ipaddress.ip_network("10.0.0.0/8"),
|
||||
ipaddress.ip_network("172.16.0.0/12"),
|
||||
ipaddress.ip_network("192.168.0.0/16"),
|
||||
ipaddress.ip_network("127.0.0.0/8"),
|
||||
ipaddress.ip_network("::1/128"),
|
||||
ipaddress.ip_network("fc00::/7"),
|
||||
ipaddress.ip_network("fe80::/10"),
|
||||
]
|
||||
|
||||
_TAILNET_IPV4_NET = ipaddress.ip_network("100.64.0.0/10")
|
||||
|
||||
_ENV_ALLOW_INSECURE_PRIVATE_WS = "YUXI_ALLOW_INSECURE_PRIVATE_WS"
|
||||
|
||||
|
||||
class GatewayBindMode(StrEnum):
|
||||
LOOPBACK = "loopback"
|
||||
LAN = "lan"
|
||||
TAILNET = "tailnet"
|
||||
AUTO = "auto"
|
||||
CUSTOM = "custom"
|
||||
|
||||
|
||||
def is_loopback_address(ip: str | None) -> bool:
|
||||
if not ip:
|
||||
return False
|
||||
try:
|
||||
return ipaddress.ip_address(ip.strip()).is_loopback
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def is_trusted_proxy_address(ip: str | None, trusted_proxies: list[str] | None) -> bool:
|
||||
if not ip or not trusted_proxies:
|
||||
return False
|
||||
try:
|
||||
addr = ipaddress.ip_address(ip.strip())
|
||||
except ValueError:
|
||||
return False
|
||||
for proxy in trusted_proxies:
|
||||
candidate = proxy.strip()
|
||||
if not candidate:
|
||||
continue
|
||||
try:
|
||||
net = ipaddress.ip_network(candidate, strict=False)
|
||||
except ValueError:
|
||||
continue
|
||||
if addr in net:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
TAILSCALE_TRUSTED_PROXIES = ["127.0.0.1", "::1"]
|
||||
|
||||
|
||||
def _parse_ip_literal(raw: str | None) -> str | None:
|
||||
if not raw:
|
||||
return None
|
||||
trimmed = raw.strip()
|
||||
if not trimmed:
|
||||
return None
|
||||
if trimmed.startswith("[") and "]" in trimmed:
|
||||
trimmed = trimmed[1 : trimmed.index("]")]
|
||||
if ":" in trimmed and "." in trimmed:
|
||||
last_colon = trimmed.rfind(":")
|
||||
candidate = trimmed[:last_colon]
|
||||
try:
|
||||
ipaddress.IPv4Address(candidate)
|
||||
trimmed = candidate
|
||||
except ValueError:
|
||||
pass
|
||||
try:
|
||||
ipaddress.ip_address(trimmed)
|
||||
return trimmed
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def resolve_forwarded_client_ip(
|
||||
forwarded_for: str | None,
|
||||
trusted_proxies: list[str] | None,
|
||||
) -> str | None:
|
||||
if not trusted_proxies:
|
||||
return None
|
||||
chain: list[str] = []
|
||||
for entry in (forwarded_for or "").split(","):
|
||||
normalized = _parse_ip_literal(entry)
|
||||
if normalized:
|
||||
chain.append(normalized)
|
||||
if not chain:
|
||||
return None
|
||||
for hop in reversed(chain):
|
||||
if is_loopback_address(hop):
|
||||
continue
|
||||
if not is_trusted_proxy_address(hop, trusted_proxies):
|
||||
return hop
|
||||
return None
|
||||
|
||||
|
||||
def resolve_client_ip(
|
||||
remote_addr: str | None,
|
||||
forwarded_for: str | None = None,
|
||||
real_ip: str | None = None,
|
||||
trusted_proxies: list[str] | None = None,
|
||||
allow_real_ip_fallback: bool = False,
|
||||
) -> str | None:
|
||||
remote = _parse_ip_literal(remote_addr)
|
||||
if not remote:
|
||||
return None
|
||||
if not is_trusted_proxy_address(remote, trusted_proxies):
|
||||
return remote
|
||||
forwarded = resolve_forwarded_client_ip(forwarded_for, trusted_proxies)
|
||||
if forwarded:
|
||||
return forwarded
|
||||
if allow_real_ip_fallback:
|
||||
return _parse_ip_literal(real_ip)
|
||||
return None
|
||||
|
||||
|
||||
def has_forwarded_request_headers(headers: dict) -> bool:
|
||||
return bool(
|
||||
headers.get("forwarded")
|
||||
or headers.get("x-forwarded-for")
|
||||
or headers.get("x-forwarded-proto")
|
||||
or headers.get("x-real-ip")
|
||||
or headers.get("x-forwarded-host")
|
||||
)
|
||||
|
||||
|
||||
def is_local_direct_request(
|
||||
remote_addr: str | None,
|
||||
headers: dict | None = None,
|
||||
) -> bool:
|
||||
if not remote_addr:
|
||||
return False
|
||||
if headers and has_forwarded_request_headers(headers):
|
||||
return False
|
||||
return is_loopback_address(remote_addr)
|
||||
|
||||
|
||||
def has_tailscale_proxy_headers(headers: dict | None) -> bool:
|
||||
if not headers:
|
||||
return False
|
||||
return bool(headers.get("x-forwarded-for") and headers.get("x-forwarded-proto") and headers.get("x-forwarded-host"))
|
||||
|
||||
|
||||
def is_tailscale_proxy_request(
|
||||
remote_addr: str | None,
|
||||
headers: dict | None = None,
|
||||
) -> bool:
|
||||
if not remote_addr:
|
||||
return False
|
||||
return is_loopback_address(remote_addr) and has_tailscale_proxy_headers(headers)
|
||||
|
||||
|
||||
def resolve_tailscale_client_ip(
|
||||
remote_addr: str | None,
|
||||
headers: dict | None = None,
|
||||
) -> str | None:
|
||||
return resolve_client_ip(
|
||||
remote_addr=remote_addr,
|
||||
forwarded_for=headers.get("x-forwarded-for") if headers else None,
|
||||
trusted_proxies=list(TAILSCALE_TRUSTED_PROXIES),
|
||||
)
|
||||
|
||||
|
||||
def is_loopback_host(host: str) -> bool:
|
||||
host = host.strip().lower().rstrip(".")
|
||||
if not host:
|
||||
return False
|
||||
if host in _LOOPBACK_HOSTS:
|
||||
return True
|
||||
try:
|
||||
addr = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
return False
|
||||
return addr.is_loopback
|
||||
|
||||
|
||||
def is_private_host(host: str) -> bool:
|
||||
host = host.strip().lower().rstrip(".")
|
||||
if not host:
|
||||
return False
|
||||
try:
|
||||
addr = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
return False
|
||||
return addr.is_private
|
||||
|
||||
|
||||
def is_private_or_loopback_host(host: str) -> bool:
|
||||
host = host.strip().lower().rstrip(".")
|
||||
if not host:
|
||||
return False
|
||||
try:
|
||||
addr = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
return False
|
||||
return addr.is_private or addr.is_loopback or addr.is_link_local
|
||||
|
||||
|
||||
def is_localish_host(host: str | None) -> bool:
|
||||
if not host:
|
||||
return False
|
||||
host = host.strip().lower().rstrip(".")
|
||||
return is_loopback_host(host) or host.endswith(".ts.net")
|
||||
|
||||
|
||||
def is_secure_ws_url(url: str, allow_private_ws: bool = False) -> bool:
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
protocol = parsed.scheme.lower()
|
||||
if protocol == "wss":
|
||||
return True
|
||||
if protocol not in ("ws", "http"):
|
||||
return False
|
||||
|
||||
hostname = _extract_ws_hostname(parsed)
|
||||
if not hostname:
|
||||
return False
|
||||
|
||||
if is_loopback_host(hostname):
|
||||
return True
|
||||
|
||||
if allow_private_ws and is_private_or_loopback_host(hostname):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _extract_ws_hostname(parsed) -> str:
|
||||
netloc = parsed.netloc or parsed.hostname or ""
|
||||
if "@" in netloc:
|
||||
netloc = netloc.rsplit("@", 1)[-1]
|
||||
if "[" in netloc and "]" in netloc:
|
||||
start = netloc.index("[") + 1
|
||||
end = netloc.index("]")
|
||||
return netloc[start:end]
|
||||
if ":" in netloc:
|
||||
return netloc.rsplit(":", 1)[0]
|
||||
return netloc
|
||||
|
||||
|
||||
async def can_bind_to_host(host: str) -> bool:
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(None, _sync_can_bind_to_host, host)
|
||||
|
||||
|
||||
def _sync_can_bind_to_host(host: str) -> bool:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
try:
|
||||
sock.bind((host, 0))
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
_can_bind_cache: dict[str, bool] = {}
|
||||
|
||||
|
||||
async def can_bind_to_host_cached(host: str) -> bool:
|
||||
cached = _can_bind_cache.get(host)
|
||||
if cached is not None:
|
||||
return cached
|
||||
result = await can_bind_to_host(host)
|
||||
_can_bind_cache[host] = result
|
||||
return result
|
||||
|
||||
|
||||
def _is_container_environment() -> bool:
|
||||
if os.path.exists("/.dockerenv"):
|
||||
return True
|
||||
try:
|
||||
with open("/proc/1/cgroup") as f:
|
||||
content = f.read()
|
||||
if "docker" in content or "kubepods" in content:
|
||||
return True
|
||||
except OSError:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _pick_primary_tailnet_ipv4() -> str | None:
|
||||
try:
|
||||
for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
|
||||
ip = info[4][0]
|
||||
try:
|
||||
addr = ipaddress.ip_address(ip)
|
||||
except ValueError:
|
||||
continue
|
||||
if addr in _TAILNET_IPV4_NET:
|
||||
return ip
|
||||
except OSError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
async def resolve_gateway_bind_host(
|
||||
mode: GatewayBindMode | None = None,
|
||||
custom_host: str | None = None,
|
||||
) -> str:
|
||||
mode = mode or GatewayBindMode.LOOPBACK
|
||||
|
||||
if mode == GatewayBindMode.LOOPBACK:
|
||||
if await can_bind_to_host("127.0.0.1"):
|
||||
return "127.0.0.1"
|
||||
return "0.0.0.0"
|
||||
|
||||
if mode == GatewayBindMode.TAILNET:
|
||||
tailnet_ip = _pick_primary_tailnet_ipv4()
|
||||
if tailnet_ip and await can_bind_to_host(tailnet_ip):
|
||||
return tailnet_ip
|
||||
if await can_bind_to_host("127.0.0.1"):
|
||||
return "127.0.0.1"
|
||||
return "0.0.0.0"
|
||||
|
||||
if mode == GatewayBindMode.LAN:
|
||||
return "0.0.0.0"
|
||||
|
||||
if mode == GatewayBindMode.CUSTOM:
|
||||
host = (custom_host or "").strip()
|
||||
if not host:
|
||||
return "0.0.0.0"
|
||||
try:
|
||||
ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
logger.warning("gateway bind=custom: invalid IP '%s', falling back to 0.0.0.0", host)
|
||||
return "0.0.0.0"
|
||||
if await can_bind_to_host(host):
|
||||
return host
|
||||
logger.warning("gateway bind=custom: cannot bind '%s', falling back to 0.0.0.0", host)
|
||||
return "0.0.0.0"
|
||||
|
||||
if mode == GatewayBindMode.AUTO:
|
||||
if _is_container_environment():
|
||||
return "0.0.0.0"
|
||||
if await can_bind_to_host("127.0.0.1"):
|
||||
return "127.0.0.1"
|
||||
return "0.0.0.0"
|
||||
|
||||
return "0.0.0.0"
|
||||
|
||||
|
||||
def build_ws_security_error(display_host: str) -> str:
|
||||
allow_private = os.environ.get(_ENV_ALLOW_INSECURE_PRIVATE_WS) == "1"
|
||||
msg = (
|
||||
f'SECURITY ERROR: Cannot connect to "{display_host}" over plaintext ws://. '
|
||||
"Both credentials and chat data would be exposed to network interception. "
|
||||
"Use wss:// for remote URLs. Safe defaults: keep gateway.bind=loopback and "
|
||||
"connect via SSH tunnel "
|
||||
"(ssh -N -L 18789:127.0.0.1:18789 user@gateway-host), or use Tailscale Serve/Funnel."
|
||||
)
|
||||
if not allow_private:
|
||||
msg += f" Break-glass (trusted private networks only): set {_ENV_ALLOW_INSECURE_PRIVATE_WS}=1."
|
||||
return msg
|
||||
161
backend/package/yuxi/channel/gateway/openai_adapter.py
Normal file
161
backend/package/yuxi/channel/gateway/openai_adapter.py
Normal file
@ -0,0 +1,161 @@
|
||||
"""OpenAI API 兼容层 — 请求格式转换与 SS 流式响应"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
OPENAI_CHAT_COMPLETIONS_PATH = "/v1/chat/completions"
|
||||
OPENAI_MODELS_PATH = "/v1/models"
|
||||
|
||||
|
||||
class OpenAIMessageConverter:
|
||||
@staticmethod
|
||||
def extract_system_prompt(messages: list[dict]) -> str | None:
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
return msg.get("content", "").strip() or None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def extract_user_query(messages: list[dict]) -> tuple[str, str | None]:
|
||||
user_texts: list[str] = []
|
||||
image_base64: str | None = None
|
||||
|
||||
for msg in messages:
|
||||
if msg.get("role") != "user":
|
||||
continue
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
user_texts.append(content.strip())
|
||||
elif isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict):
|
||||
if part.get("type") == "text":
|
||||
user_texts.append(part.get("text", "").strip())
|
||||
elif part.get("type") == "image_url":
|
||||
url = part.get("image_url", {}).get("url", "")
|
||||
if url.startswith("data:"):
|
||||
image_base64 = url.split(",", 1)[1] if "," in url else url
|
||||
|
||||
return "\n".join(t for t in user_texts if t), image_base64
|
||||
|
||||
@staticmethod
|
||||
def extract_conversation_history(messages: list[dict]) -> list[dict]:
|
||||
history: list[dict] = []
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
if role in ("user", "assistant") and content:
|
||||
if isinstance(content, str) and content.strip():
|
||||
history.append({"role": role, "content": content.strip()})
|
||||
return history
|
||||
|
||||
@staticmethod
|
||||
def extract_model_name(request_data: dict) -> str:
|
||||
return request_data.get("model", "default")
|
||||
|
||||
@staticmethod
|
||||
def extract_stream_flag(request_data: dict) -> bool:
|
||||
return bool(request_data.get("stream", False))
|
||||
|
||||
@staticmethod
|
||||
def extract_max_tokens(request_data: dict) -> int | None:
|
||||
return request_data.get("max_tokens")
|
||||
|
||||
@staticmethod
|
||||
def extract_temperature(request_data: dict) -> float | None:
|
||||
return request_data.get("temperature")
|
||||
|
||||
|
||||
def _format_sse_chunk(
|
||||
content: str | None = None,
|
||||
status: str = "streaming",
|
||||
*,
|
||||
finish_reason: str | None = None,
|
||||
model: str = "default",
|
||||
index: int = 0,
|
||||
) -> bytes:
|
||||
delta: dict = {}
|
||||
if content is not None:
|
||||
delta["content"] = content
|
||||
if status is not None:
|
||||
payload: dict = {
|
||||
"id": f"chatcmpl-{int(time.time() * 1000)}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": index,
|
||||
"delta": delta,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
}
|
||||
return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n".encode()
|
||||
|
||||
return b"data: [DONE]\n\n"
|
||||
|
||||
|
||||
def _format_openai_non_stream_response(
|
||||
content: str,
|
||||
*,
|
||||
model: str = "default",
|
||||
finish_reason: str = "stop",
|
||||
) -> dict:
|
||||
return {
|
||||
"id": f"chatcmpl-{int(time.time() * 1000)}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
},
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
async def format_sse_stream(
|
||||
raw_stream: AsyncIterator[bytes],
|
||||
*,
|
||||
model: str = "default",
|
||||
) -> AsyncIterator[bytes]:
|
||||
accumulated: list[str] = []
|
||||
async for raw in raw_stream:
|
||||
try:
|
||||
data = json.loads(raw.decode("utf-8"))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
continue
|
||||
|
||||
status = data.get("status", "")
|
||||
content = data.get("response", "")
|
||||
|
||||
if status == "streaming" and content:
|
||||
accumulated.append(content)
|
||||
yield _format_sse_chunk(content=content, status="streaming", model=model)
|
||||
elif status == "finished":
|
||||
pass
|
||||
elif status == "error":
|
||||
error_msg = data.get("error_message", content or "未知错误")
|
||||
yield _format_sse_chunk(
|
||||
content=error_msg,
|
||||
status="error",
|
||||
finish_reason="error",
|
||||
model=model,
|
||||
)
|
||||
return
|
||||
|
||||
yield _format_sse_chunk(content="", status=None, finish_reason="stop", model=model)
|
||||
102
backend/package/yuxi/channel/gateway/polling.py
Normal file
102
backend/package/yuxi/channel/gateway/polling.py
Normal file
@ -0,0 +1,102 @@
|
||||
"""HTTP Polling 降级机制 — SSE/WS 不可用时的备选路径。
|
||||
|
||||
按 session_id 维护响应队列,前端通过 HTTP GET 定时轮询获取回复。
|
||||
与 SSE 端点协同使用,形成三级降级链:SSE → Polling → 单次 fetch。
|
||||
|
||||
Usage:
|
||||
polling = PollingFallback()
|
||||
await polling.push("session_abc", {"type": "delta", "data": {"content": "hello"}})
|
||||
events = await polling.poll("session_abc")
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time as _time
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_POLL_TTL = 600
|
||||
|
||||
|
||||
class PollingFallback:
|
||||
def __init__(self, max_queue_size: int = 100, ttl_seconds: int = DEFAULT_POLL_TTL):
|
||||
self._queues: dict[str, asyncio.Queue[dict]] = {}
|
||||
self._ttl = ttl_seconds
|
||||
self._last_active: dict[str, float] = {}
|
||||
self._max_queue_size = max_queue_size
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def ensure(self, session_id: str) -> None:
|
||||
async with self._lock:
|
||||
if session_id not in self._queues:
|
||||
self._queues[session_id] = asyncio.Queue(maxsize=self._max_queue_size)
|
||||
self._last_active[session_id] = _time.monotonic()
|
||||
|
||||
async def push(self, session_id: str, event: dict) -> None:
|
||||
async with self._lock:
|
||||
self._last_active[session_id] = _time.monotonic()
|
||||
q = self._queues.get(session_id)
|
||||
if q is None:
|
||||
q = asyncio.Queue(maxsize=self._max_queue_size)
|
||||
self._queues[session_id] = q
|
||||
try:
|
||||
q.put_nowait(event)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning("Polling queue full for session %s, dropping event", session_id)
|
||||
|
||||
async def poll(self, session_id: str) -> list[dict]:
|
||||
async with self._lock:
|
||||
self._last_active[session_id] = _time.monotonic()
|
||||
q = self._queues.get(session_id)
|
||||
if q is None:
|
||||
return []
|
||||
|
||||
events: list[dict] = []
|
||||
while not q.empty():
|
||||
try:
|
||||
events.append(q.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
return events
|
||||
|
||||
async def cleanup_stale(self) -> int:
|
||||
async with self._lock:
|
||||
now = _time.monotonic()
|
||||
stale = [sid for sid, ts in self._last_active.items() if now - ts > self._ttl]
|
||||
for sid in stale:
|
||||
self._queues.pop(sid, None)
|
||||
self._last_active.pop(sid, None)
|
||||
if stale:
|
||||
logger.info("PollingFallback cleaned up %d stale sessions", len(stale))
|
||||
return len(stale)
|
||||
|
||||
def active_sessions(self) -> int:
|
||||
return len(self._queues)
|
||||
|
||||
|
||||
polling_fallback = PollingFallback()
|
||||
|
||||
|
||||
async def _polling_cleanup_loop(interval: int = 300) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
await polling_fallback.cleanup_stale()
|
||||
|
||||
|
||||
_cleanup_task: asyncio.Task | None = None
|
||||
|
||||
|
||||
def start_polling_cleanup(interval: int = 300) -> None:
|
||||
global _cleanup_task
|
||||
if _cleanup_task is None or _cleanup_task.done():
|
||||
_cleanup_task = asyncio.ensure_future(_polling_cleanup_loop(interval))
|
||||
logger.info("PollingFallback cleanup loop started (interval=%ds)", interval)
|
||||
|
||||
|
||||
def stop_polling_cleanup() -> None:
|
||||
global _cleanup_task
|
||||
if _cleanup_task and not _cleanup_task.done():
|
||||
_cleanup_task.cancel()
|
||||
_cleanup_task = None
|
||||
289
backend/package/yuxi/channel/gateway/probe.py
Normal file
289
backend/package/yuxi/channel/gateway/probe.py
Normal file
@ -0,0 +1,289 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
|
||||
import aiohttp
|
||||
|
||||
from yuxi.channel.gateway.protocol import (
|
||||
RpcRequest,
|
||||
RpcResponse,
|
||||
marshal_frame,
|
||||
unmarshal_frame,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_PROBE_TIMEOUT_MS = 10_000
|
||||
MIN_PROBE_TIMEOUT_MS = 250
|
||||
MAX_PROBE_TIMEOUT_MS = 30_000
|
||||
|
||||
|
||||
class ProbeCapability(StrEnum):
|
||||
UNKNOWN = "unknown"
|
||||
PAIRING_PENDING = "pairing_pending"
|
||||
CONNECTED_NO_OPERATOR_SCOPE = "connected_no_operator_scope"
|
||||
READ_ONLY = "read_only"
|
||||
WRITE_CAPABLE = "write_capable"
|
||||
ADMIN_CAPABLE = "admin_capable"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProbeAuthSummary:
|
||||
role: str | None = None
|
||||
scopes: list[str] = field(default_factory=list)
|
||||
capability: ProbeCapability = ProbeCapability.UNKNOWN
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProbeServerSummary:
|
||||
version: str | None = None
|
||||
conn_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProbeClose:
|
||||
code: int
|
||||
reason: str
|
||||
hint: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class GatewayProbeResult:
|
||||
ok: bool
|
||||
url: str
|
||||
connect_latency_ms: float | None = None
|
||||
error: str | None = None
|
||||
close: ProbeClose | None = None
|
||||
auth: ProbeAuthSummary = field(default_factory=ProbeAuthSummary)
|
||||
server: ProbeServerSummary = field(default_factory=ProbeServerSummary)
|
||||
health: dict | None = None
|
||||
status: dict | None = None
|
||||
presence: list | None = None
|
||||
config_snapshot: dict | None = None
|
||||
|
||||
|
||||
def _clamp_timeout_ms(timeout_ms: float | int) -> float:
|
||||
return max(MIN_PROBE_TIMEOUT_MS, min(float(timeout_ms), MAX_PROBE_TIMEOUT_MS))
|
||||
|
||||
|
||||
def _resolve_capability(
|
||||
scopes: list[str],
|
||||
connect_latency_ms: float | None,
|
||||
auth_metadata_present: bool,
|
||||
) -> ProbeCapability:
|
||||
if "operator.admin" in scopes:
|
||||
return ProbeCapability.ADMIN_CAPABLE
|
||||
if "operator.write" in scopes:
|
||||
return ProbeCapability.WRITE_CAPABLE
|
||||
if "operator.read" in scopes or "admin" in scopes:
|
||||
return ProbeCapability.READ_ONLY
|
||||
if connect_latency_ms is not None and auth_metadata_present:
|
||||
return ProbeCapability.CONNECTED_NO_OPERATOR_SCOPE
|
||||
return ProbeCapability.UNKNOWN
|
||||
|
||||
|
||||
async def probe_gateway(
|
||||
url: str,
|
||||
token: str | None = None,
|
||||
password: str | None = None,
|
||||
timeout_ms: float = DEFAULT_PROBE_TIMEOUT_MS,
|
||||
detail_level: str = "full",
|
||||
) -> GatewayProbeResult:
|
||||
connect_latency_ms: float | None = None
|
||||
connect_error: str | None = None
|
||||
close_info: ProbeClose | None = None
|
||||
auth_summary = ProbeAuthSummary()
|
||||
server_summary = ProbeServerSummary()
|
||||
auth_metadata_present = False
|
||||
|
||||
effective_timeout = _clamp_timeout_ms(timeout_ms)
|
||||
timeout_sec = effective_timeout / 1000.0
|
||||
|
||||
headers: dict[str, str] = {}
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
elif password:
|
||||
headers["Authorization"] = f"Bearer {password}"
|
||||
|
||||
conn_start = time.monotonic()
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(total=timeout_sec),
|
||||
) as session:
|
||||
async with session.ws_connect(
|
||||
url,
|
||||
headers=headers,
|
||||
heartbeat=15.0,
|
||||
) as ws:
|
||||
connect_latency_ms = (time.monotonic() - conn_start) * 1000
|
||||
|
||||
auth_metadata_present = True
|
||||
|
||||
if detail_level == "none":
|
||||
ws_close = ws.close_code
|
||||
if ws_close is not None:
|
||||
close_info = ProbeClose(
|
||||
code=ws_close,
|
||||
reason="connection closed by server",
|
||||
)
|
||||
return GatewayProbeResult(
|
||||
ok=ws_close is None,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
error=None,
|
||||
close=close_info,
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
)
|
||||
|
||||
rpc_request_id = "probe"
|
||||
rpc_request = RpcRequest(
|
||||
id=rpc_request_id,
|
||||
method="system.health",
|
||||
params={},
|
||||
)
|
||||
await ws.send_str(marshal_frame(rpc_request))
|
||||
|
||||
try:
|
||||
raw = await asyncio.wait_for(
|
||||
ws.receive_str(),
|
||||
timeout=effective_timeout / 1000.0,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return GatewayProbeResult(
|
||||
ok=False,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
error="timeout waiting for health response",
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
)
|
||||
|
||||
try:
|
||||
frame = unmarshal_frame(raw)
|
||||
except Exception:
|
||||
return GatewayProbeResult(
|
||||
ok=False,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
error="invalid response frame from gateway",
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
)
|
||||
|
||||
if not isinstance(frame, RpcResponse) or not frame.ok:
|
||||
error_msg = frame.error_message if isinstance(frame, RpcResponse) else "unexpected frame type"
|
||||
return GatewayProbeResult(
|
||||
ok=False,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
error=error_msg or "gateway returned error",
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
)
|
||||
|
||||
health_data = frame.result
|
||||
|
||||
auth_summary = ProbeAuthSummary(
|
||||
role="viewer",
|
||||
scopes=[],
|
||||
capability=ProbeCapability.READ_ONLY,
|
||||
)
|
||||
|
||||
if detail_level == "presence":
|
||||
return GatewayProbeResult(
|
||||
ok=True,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
health=health_data,
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
presence=health_data.get("channels") if health_data else None,
|
||||
)
|
||||
|
||||
status_data: dict | None = None
|
||||
try:
|
||||
status_request = RpcRequest(
|
||||
id="probe-status",
|
||||
method="channels.status",
|
||||
params={},
|
||||
)
|
||||
await ws.send_str(marshal_frame(status_request))
|
||||
raw_status = await asyncio.wait_for(
|
||||
ws.receive_str(),
|
||||
timeout=effective_timeout / 1000.0,
|
||||
)
|
||||
status_frame = unmarshal_frame(raw_status)
|
||||
if isinstance(status_frame, RpcResponse) and status_frame.ok:
|
||||
status_data = status_frame.result
|
||||
except Exception:
|
||||
logger.debug("Failed to get channel status during probe", exc_info=True)
|
||||
|
||||
return GatewayProbeResult(
|
||||
ok=True,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
health=health_data,
|
||||
status=status_data,
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
)
|
||||
|
||||
except aiohttp.ClientConnectorError as e:
|
||||
connect_error = f"连接失败: {e}"
|
||||
except aiohttp.WSServerHandshakeError as e:
|
||||
connect_error = f"WebSocket 握手失败: {e.status} {e.message}"
|
||||
except aiohttp.ClientError as e:
|
||||
connect_error = f"客户端错误: {e}"
|
||||
except asyncio.TimeoutError:
|
||||
connect_error = "连接超时"
|
||||
except Exception as e:
|
||||
connect_error = f"未知错误: {e}"
|
||||
|
||||
return GatewayProbeResult(
|
||||
ok=False,
|
||||
url=url,
|
||||
connect_latency_ms=connect_latency_ms,
|
||||
error=connect_error,
|
||||
close=close_info,
|
||||
auth=auth_summary,
|
||||
server=server_summary,
|
||||
)
|
||||
|
||||
|
||||
def format_probe_result(result: GatewayProbeResult, verbose: bool = False) -> str:
|
||||
lines: list[str] = []
|
||||
|
||||
if result.ok:
|
||||
lines.append(f"✓ 网关探测成功: {result.url}")
|
||||
if result.connect_latency_ms is not None:
|
||||
lines.append(f" 连接延迟: {result.connect_latency_ms:.1f}ms")
|
||||
lines.append(f" 能力级别: {result.auth.capability.value}")
|
||||
else:
|
||||
lines.append(f"✗ 网关探测失败: {result.url}")
|
||||
if result.connect_latency_ms is not None:
|
||||
lines.append(f" 连接延迟: {result.connect_latency_ms:.1f}ms")
|
||||
if result.error:
|
||||
lines.append(f" 错误: {result.error}")
|
||||
if result.close:
|
||||
lines.append(f" 关闭信息: code={result.close.code} reason={result.close.reason}")
|
||||
|
||||
if verbose and result.health:
|
||||
lines.append("")
|
||||
lines.append("健康报告:")
|
||||
health = result.health
|
||||
if isinstance(health, dict):
|
||||
lines.append(f" 状态: {health.get('status', 'unknown')}")
|
||||
summary = health.get("summary", {})
|
||||
if summary:
|
||||
lines.append(
|
||||
f" 频道: 总计={summary.get('total', 0)}, "
|
||||
f"运行={summary.get('running', 0)}, "
|
||||
f"停止={summary.get('stopped', 0)}, "
|
||||
f"异常={summary.get('unhealthy', 0)}"
|
||||
)
|
||||
|
||||
return "\n".join(lines)
|
||||
207
backend/package/yuxi/channel/gateway/protocol.py
Normal file
207
backend/package/yuxi/channel/gateway/protocol.py
Normal file
@ -0,0 +1,207 @@
|
||||
import json as _json
|
||||
import time as _time
|
||||
import uuid as _uuid
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
|
||||
|
||||
class FrameType(StrEnum):
|
||||
REQUEST = "request"
|
||||
RESPONSE = "response"
|
||||
EVENT = "event"
|
||||
|
||||
|
||||
class GatewayErrorCode(StrEnum):
|
||||
INVALID_REQUEST = "INVALID_REQUEST"
|
||||
METHOD_NOT_FOUND = "METHOD_NOT_FOUND"
|
||||
INVALID_PARAMS = "INVALID_PARAMS"
|
||||
INTERNAL_ERROR = "INTERNAL_ERROR"
|
||||
AUTH_ERROR = "AUTH_ERROR"
|
||||
PERMISSION_DENIED = "PERMISSION_DENIED"
|
||||
RATE_LIMITED = "RATE_LIMITED"
|
||||
TIMEOUT = "TIMEOUT"
|
||||
NOT_CONNECTED = "NOT_CONNECTED"
|
||||
UNAVAILABLE = "UNAVAILABLE"
|
||||
DEVICE_IDENTITY_REQUIRED = "DEVICE_IDENTITY_REQUIRED"
|
||||
|
||||
|
||||
class GatewayRpcMethod(StrEnum):
|
||||
CONNECT = "connect"
|
||||
CHAT_SEND = "chat.send"
|
||||
CHAT_STREAM = "chat.stream"
|
||||
CHANNELS_LIST = "channels.list"
|
||||
CHANNELS_STATUS = "channels.status"
|
||||
CHANNELS_CONFIGURE = "channels.configure"
|
||||
SESSIONS_LIST = "sessions.list"
|
||||
SESSIONS_HISTORY = "sessions.history"
|
||||
PLUGINS_LIST = "plugins.list"
|
||||
PLUGINS_INSTALL = "plugins.install"
|
||||
SYSTEM_HEALTH = "system.health"
|
||||
SYSTEM_CONFIG = "system.config"
|
||||
SYSTEM_VERSION = "system.version"
|
||||
CHANNELS_START = "channels.start"
|
||||
CHANNELS_STOP = "channels.stop"
|
||||
CHANNELS_RESTART = "channels.restart"
|
||||
PAIRING_LIST = "pairing.list"
|
||||
PAIRING_APPROVE = "pairing.approve"
|
||||
PAIRING_REJECT = "pairing.reject"
|
||||
ALLOWLIST_GET = "allowlist.get"
|
||||
ALLOWLIST_ADD = "allowlist.add"
|
||||
ALLOWLIST_REMOVE = "allowlist.remove"
|
||||
CHAT_HISTORY = "chat.history"
|
||||
CHAT_CANCEL = "chat.cancel"
|
||||
CRON_LIST = "cron.list"
|
||||
CRON_CREATE = "cron.create"
|
||||
CRON_UPDATE = "cron.update"
|
||||
CRON_DELETE = "cron.delete"
|
||||
CRON_FORCE_RUN = "cron.force_run"
|
||||
CRON_PAUSE = "cron.pause"
|
||||
CRON_RESUME = "cron.resume"
|
||||
CRON_DIAGNOSTICS = "cron.diagnostics"
|
||||
CRON_RUN_LOG = "cron.run_log"
|
||||
AGENT_TOOLS_LIST = "agentTools.list"
|
||||
AGENT_TOOLS_EXECUTE = "agentTools.execute"
|
||||
SYSTEM_LOG_TAIL = "system.log.tail"
|
||||
SESSIONS_CREATE = "sessions.create"
|
||||
SESSIONS_DELETE = "sessions.delete"
|
||||
CONFIG_GET = "config.get"
|
||||
CONFIG_SET = "config.set"
|
||||
MESSAGE_SEND = "message.send"
|
||||
MESSAGE_ACTION = "message.action"
|
||||
CHANNELS_PROBE = "channels.probe"
|
||||
IDENTITY_LINKS_LIST = "identity_links.list"
|
||||
IDENTITY_LINKS_ADD = "identity_links.add"
|
||||
IDENTITY_LINKS_REMOVE = "identity_links.remove"
|
||||
|
||||
|
||||
class DeliveryMode(StrEnum):
|
||||
DIRECT = "direct"
|
||||
GATEWAY = "gateway"
|
||||
HYBRID = "hybrid"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RpcFrame:
|
||||
frame_type: FrameType
|
||||
|
||||
|
||||
@dataclass
|
||||
class RpcRequest(RpcFrame):
|
||||
frame_type: FrameType = FrameType.REQUEST
|
||||
id: str = field(default_factory=lambda: _uuid.uuid4().hex[:12])
|
||||
method: str = ""
|
||||
params: dict | None = None
|
||||
session_id: str | None = None
|
||||
caller_user_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RpcResponse(RpcFrame):
|
||||
frame_type: FrameType = FrameType.RESPONSE
|
||||
id: str = ""
|
||||
ok: bool = True
|
||||
result: dict | None = None
|
||||
error_code: GatewayErrorCode | None = None
|
||||
error_message: str | None = None
|
||||
error_data: dict | None = None
|
||||
retryable: bool = False
|
||||
retry_after_ms: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class HelloOk(RpcResponse):
|
||||
frame_type: FrameType = FrameType.RESPONSE
|
||||
id: str = ""
|
||||
ok: bool = True
|
||||
result: dict | None = field(default_factory=lambda: {
|
||||
"protocolVersion": 1,
|
||||
"server": {"version": "1.0.0", "connId": ""},
|
||||
"features": {"methods": [], "events": []},
|
||||
"auth": {"role": "", "scopes": []},
|
||||
"policy": {
|
||||
"maxPayload": 25 * 1024 * 1024,
|
||||
"maxBufferedBytes": 256 * 1024,
|
||||
"tickIntervalMs": 30_000,
|
||||
},
|
||||
})
|
||||
|
||||
|
||||
@dataclass
|
||||
class RpcEvent(RpcFrame):
|
||||
frame_type: FrameType = FrameType.EVENT
|
||||
event: str = ""
|
||||
data: dict | None = None
|
||||
timestamp: float = field(default_factory=_time.time)
|
||||
seq: int | None = None
|
||||
state_version: int | None = None
|
||||
|
||||
|
||||
def marshal_frame(frame: RpcFrame) -> str:
|
||||
if isinstance(frame, RpcRequest):
|
||||
payload = {"type": frame.frame_type.value, "id": frame.id, "method": frame.method}
|
||||
if frame.params is not None:
|
||||
payload["params"] = frame.params
|
||||
if frame.session_id is not None:
|
||||
payload["sessionId"] = frame.session_id
|
||||
elif isinstance(frame, RpcResponse):
|
||||
payload = {"type": frame.frame_type.value, "id": frame.id, "ok": frame.ok}
|
||||
if frame.result is not None:
|
||||
payload["result"] = frame.result
|
||||
if frame.error_code is not None:
|
||||
error_payload = {
|
||||
"code": frame.error_code.value,
|
||||
"message": frame.error_message or "",
|
||||
}
|
||||
if frame.error_data is not None:
|
||||
error_payload["data"] = frame.error_data
|
||||
if frame.retryable:
|
||||
error_payload["retryable"] = True
|
||||
if frame.retry_after_ms is not None:
|
||||
error_payload["retryAfterMs"] = frame.retry_after_ms
|
||||
payload["error"] = error_payload
|
||||
elif isinstance(frame, RpcEvent):
|
||||
payload = {"type": frame.frame_type.value, "event": frame.event, "timestamp": frame.timestamp}
|
||||
if frame.data is not None:
|
||||
payload["data"] = frame.data
|
||||
if frame.seq is not None:
|
||||
payload["seq"] = frame.seq
|
||||
if frame.state_version is not None:
|
||||
payload["stateVersion"] = frame.state_version
|
||||
else:
|
||||
raise TypeError(f"Unknown frame type: {type(frame)}")
|
||||
return _json.dumps(payload, ensure_ascii=False, default=str)
|
||||
|
||||
|
||||
def unmarshal_frame(data: str | bytes) -> RpcFrame:
|
||||
if isinstance(data, bytes):
|
||||
data = data.decode("utf-8")
|
||||
raw = _json.loads(data)
|
||||
frame_type = raw.get("type", FrameType.REQUEST.value)
|
||||
if frame_type == FrameType.REQUEST.value:
|
||||
return RpcRequest(
|
||||
id=raw.get("id", ""),
|
||||
method=raw.get("method", ""),
|
||||
params=raw.get("params"),
|
||||
session_id=raw.get("sessionId"),
|
||||
)
|
||||
elif frame_type == FrameType.RESPONSE.value:
|
||||
error = raw.get("error")
|
||||
return RpcResponse(
|
||||
id=raw.get("id", ""),
|
||||
ok=raw.get("ok", True),
|
||||
result=raw.get("result"),
|
||||
error_code=GatewayErrorCode(error["code"]) if error else None,
|
||||
error_message=error.get("message") if error else None,
|
||||
error_data=error.get("data") if error else None,
|
||||
retryable=error.get("retryable", False) if error else False,
|
||||
retry_after_ms=error.get("retryAfterMs") if error else None,
|
||||
)
|
||||
elif frame_type == FrameType.EVENT.value:
|
||||
return RpcEvent(
|
||||
event=raw.get("event", ""),
|
||||
data=raw.get("data"),
|
||||
timestamp=raw.get("timestamp", _time.time()),
|
||||
seq=raw.get("seq"),
|
||||
state_version=raw.get("stateVersion"),
|
||||
)
|
||||
raise ValueError(f"Unknown frame type: {frame_type}")
|
||||
149
backend/package/yuxi/channel/gateway/rbac.py
Normal file
149
backend/package/yuxi/channel/gateway/rbac.py
Normal file
@ -0,0 +1,149 @@
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from enum import StrEnum
|
||||
from functools import wraps
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GatewayRole(StrEnum):
|
||||
SUPERADMIN = "superadmin"
|
||||
ADMIN = "admin"
|
||||
OPERATOR = "operator"
|
||||
VIEWER = "viewer"
|
||||
|
||||
|
||||
ROLE_HIERARCHY: dict[GatewayRole, int] = {
|
||||
GatewayRole.SUPERADMIN: 4,
|
||||
GatewayRole.ADMIN: 3,
|
||||
GatewayRole.OPERATOR: 2,
|
||||
GatewayRole.VIEWER: 1,
|
||||
}
|
||||
|
||||
|
||||
_USER_ROLE_MAP: dict[str, GatewayRole] = {
|
||||
"superadmin": GatewayRole.SUPERADMIN,
|
||||
"admin": GatewayRole.ADMIN,
|
||||
"user": GatewayRole.OPERATOR,
|
||||
}
|
||||
|
||||
|
||||
def map_user_role(user_role: str | None) -> GatewayRole:
|
||||
if not user_role:
|
||||
return GatewayRole.VIEWER
|
||||
return _USER_ROLE_MAP.get(user_role.lower(), GatewayRole.VIEWER)
|
||||
|
||||
|
||||
READ_ONLY_METHODS: set[str] = {
|
||||
"system.health",
|
||||
"system.version",
|
||||
"channels.list",
|
||||
"channels.status",
|
||||
"channels.probe",
|
||||
"chat.history",
|
||||
"sessions.list",
|
||||
"sessions.history",
|
||||
"plugins.list",
|
||||
"agentTools.list",
|
||||
"config.get",
|
||||
}
|
||||
|
||||
OPERATOR_METHODS: set[str] = READ_ONLY_METHODS | {
|
||||
"chat.send",
|
||||
"chat.stream",
|
||||
"chat.cancel",
|
||||
"sessions.create",
|
||||
"sessions.delete",
|
||||
"allowlist.get",
|
||||
"cron.list",
|
||||
"agentTools.execute",
|
||||
"message.send",
|
||||
"message.action",
|
||||
}
|
||||
|
||||
ADMIN_METHODS: set[str] = OPERATOR_METHODS | {
|
||||
"system.log.tail",
|
||||
"system.config",
|
||||
"channels.start",
|
||||
"channels.stop",
|
||||
"channels.restart",
|
||||
"channels.configure",
|
||||
"pairing.list",
|
||||
"pairing.approve",
|
||||
"pairing.reject",
|
||||
"allowlist.add",
|
||||
"allowlist.remove",
|
||||
"identity_links.list",
|
||||
"identity_links.add",
|
||||
"identity_links.remove",
|
||||
"plugins.install",
|
||||
"config.set",
|
||||
}
|
||||
|
||||
SUPERADMIN_METHODS: set[str] = ADMIN_METHODS | {
|
||||
"cron.create",
|
||||
"cron.delete",
|
||||
"cron.force_run",
|
||||
"cron.pause",
|
||||
"cron.resume",
|
||||
}
|
||||
|
||||
METHOD_MIN_ROLE: dict[str, GatewayRole] = {}
|
||||
|
||||
for _method in READ_ONLY_METHODS:
|
||||
METHOD_MIN_ROLE[_method] = GatewayRole.VIEWER
|
||||
|
||||
for _method in OPERATOR_METHODS - READ_ONLY_METHODS:
|
||||
METHOD_MIN_ROLE[_method] = GatewayRole.OPERATOR
|
||||
|
||||
for _method in ADMIN_METHODS - OPERATOR_METHODS:
|
||||
METHOD_MIN_ROLE[_method] = GatewayRole.ADMIN
|
||||
|
||||
for _method in {"cron.create", "cron.delete", "cron.force_run", "cron.pause", "cron.resume"}:
|
||||
METHOD_MIN_ROLE[_method] = GatewayRole.SUPERADMIN
|
||||
|
||||
|
||||
def check_permission(role: GatewayRole | None, method: str) -> bool:
|
||||
effective_role = role or GatewayRole.VIEWER
|
||||
min_role = METHOD_MIN_ROLE.get(method)
|
||||
if min_role is None:
|
||||
return True
|
||||
return ROLE_HIERARCHY[effective_role] >= ROLE_HIERARCHY[min_role]
|
||||
|
||||
|
||||
def require_role(min_role: GatewayRole):
|
||||
def decorator(
|
||||
func: Callable[..., Any],
|
||||
) -> Callable[..., Any]:
|
||||
@wraps(func)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
caller_role: GatewayRole | None = kwargs.pop("_caller_role", None)
|
||||
if caller_role is None:
|
||||
logger.warning(
|
||||
"RBAC: _caller_role not found in handler kwargs for %s, allowing by default",
|
||||
func.__name__,
|
||||
)
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
if ROLE_HIERARCHY[caller_role] < ROLE_HIERARCHY[min_role]:
|
||||
from yuxi.channel.gateway.protocol import GatewayErrorCode, RpcResponse
|
||||
|
||||
logger.warning(
|
||||
"RBAC: %s requires %s, caller has %s — denied",
|
||||
func.__name__,
|
||||
min_role.value,
|
||||
caller_role.value,
|
||||
)
|
||||
return RpcResponse(
|
||||
id=kwargs.get("request_id", ""),
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.PERMISSION_DENIED,
|
||||
error_message=f"需要 {min_role.value} 权限,当前角色为 {caller_role.value}",
|
||||
)
|
||||
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
161
backend/package/yuxi/channel/gateway/routes.py
Normal file
161
backend/package/yuxi/channel/gateway/routes.py
Normal file
@ -0,0 +1,161 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.gateway.protocol import DeliveryMode
|
||||
from yuxi.channel.gateway.webhook_security import (
|
||||
WebhookGuard,
|
||||
WebhookGuardConfig,
|
||||
WebhookGuardResult,
|
||||
)
|
||||
from yuxi.channel.gateway.webhook_security import (
|
||||
webhook_guard as _default_webhook_guard,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WebhookHandler = Callable[[dict[str, Any]], asyncio.Future[dict[str, Any]]]
|
||||
|
||||
|
||||
class WebhookRegistry:
|
||||
"""渠道 Webhook 处理器注册表。
|
||||
|
||||
每个渠道插件注册自己的 webhook 处理函数,
|
||||
HTTP 层统一通过 FastAPI router 将请求分发到此处。
|
||||
|
||||
集成了多层 Webhook 安全 Guard:
|
||||
- Method 检查(仅允许注册的 HTTP 方法)
|
||||
- Content-Type 验证
|
||||
- Body 大小限制
|
||||
- HMAC-SHA256 签名校验
|
||||
- 并发限制
|
||||
- 异常追踪告警
|
||||
"""
|
||||
|
||||
def __init__(self, webhook_guard: WebhookGuard | None = None):
|
||||
self._handlers: dict[str, WebhookHandler] = {}
|
||||
self._delivery_modes: dict[str, DeliveryMode] = {}
|
||||
self._guard = webhook_guard or _default_webhook_guard
|
||||
|
||||
def register(
|
||||
self,
|
||||
channel_type: str,
|
||||
handler: WebhookHandler,
|
||||
delivery_mode: DeliveryMode = DeliveryMode.DIRECT,
|
||||
guard_config: WebhookGuardConfig | None = None,
|
||||
) -> None:
|
||||
if channel_type in self._handlers:
|
||||
logger.warning("Overwriting webhook handler for channel: %s", channel_type)
|
||||
self._handlers[channel_type] = handler
|
||||
self._delivery_modes[channel_type] = delivery_mode
|
||||
self._guard.register_channel(channel_type, guard_config)
|
||||
|
||||
def get_handler(self, channel_type: str) -> WebhookHandler | None:
|
||||
return self._handlers.get(channel_type)
|
||||
|
||||
def get_delivery_mode(self, channel_type: str) -> DeliveryMode:
|
||||
return self._delivery_modes.get(channel_type, DeliveryMode.DIRECT)
|
||||
|
||||
def list_channels(self) -> list[dict[str, str]]:
|
||||
return [
|
||||
{
|
||||
"channel_type": ct,
|
||||
"delivery_mode": self._delivery_modes.get(ct, DeliveryMode.DIRECT).value,
|
||||
}
|
||||
for ct in self._handlers
|
||||
]
|
||||
|
||||
def remove(self, channel_type: str) -> None:
|
||||
self._handlers.pop(channel_type, None)
|
||||
self._delivery_modes.pop(channel_type, None)
|
||||
self._guard.unregister_channel(channel_type)
|
||||
|
||||
def clear(self) -> None:
|
||||
for ct in list(self._handlers):
|
||||
self._guard.unregister_channel(ct)
|
||||
self._handlers.clear()
|
||||
self._delivery_modes.clear()
|
||||
|
||||
async def dispatch_guarded(
|
||||
self,
|
||||
channel_type: str,
|
||||
method: str,
|
||||
content_type: str | None,
|
||||
body: bytes,
|
||||
signature: str | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> tuple[WebhookGuardResult | None, dict[str, Any] | None]:
|
||||
guard_result = self._guard.run_pipeline(
|
||||
channel_type,
|
||||
method,
|
||||
content_type,
|
||||
body,
|
||||
signature,
|
||||
)
|
||||
if not guard_result.allowed:
|
||||
return guard_result, None
|
||||
|
||||
handler = self._handlers.get(channel_type)
|
||||
if handler is None:
|
||||
self._guard.release_concurrency(channel_type)
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=None,
|
||||
error_message=f"No handler registered for channel: {channel_type}",
|
||||
http_status=404,
|
||||
), None
|
||||
|
||||
try:
|
||||
import json
|
||||
|
||||
payload = json.loads(body)
|
||||
except Exception:
|
||||
self._guard.release_concurrency(channel_type)
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=None,
|
||||
error_message="Invalid JSON payload",
|
||||
http_status=400,
|
||||
), None
|
||||
|
||||
if extra_headers:
|
||||
payload["_headers"] = extra_headers
|
||||
payload["_raw_body"] = body
|
||||
|
||||
try:
|
||||
result = await handler(payload)
|
||||
except Exception:
|
||||
logger.exception("Webhook handler failed for channel: %s", channel_type)
|
||||
result = None
|
||||
finally:
|
||||
self._guard.release_concurrency(channel_type)
|
||||
|
||||
return None, result
|
||||
|
||||
@property
|
||||
def guard(self) -> WebhookGuard:
|
||||
return self._guard
|
||||
|
||||
|
||||
webhook_registry = WebhookRegistry()
|
||||
|
||||
|
||||
import re
|
||||
|
||||
_WEBHOOK_PATH_RE = re.compile(r"^[a-zA-Z0-9_-]+$")
|
||||
|
||||
|
||||
def build_webhook_path(channel_type: str) -> str:
|
||||
if not channel_type or not _WEBHOOK_PATH_RE.match(channel_type):
|
||||
raise ValueError(f"Invalid channel_type: {channel_type!r}")
|
||||
return f"/webhook/{channel_type}"
|
||||
|
||||
|
||||
def build_health_response(status: str, channels: dict, summary: dict, ts: str | None = None) -> dict:
|
||||
return {
|
||||
"status": status,
|
||||
"timestamp": ts,
|
||||
"channels": channels,
|
||||
"summary": summary,
|
||||
}
|
||||
193
backend/package/yuxi/channel/gateway/rpc_dispatcher.py
Normal file
193
backend/package/yuxi/channel/gateway/rpc_dispatcher.py
Normal file
@ -0,0 +1,193 @@
|
||||
import logging
|
||||
from collections.abc import AsyncGenerator, Callable, Coroutine
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.gateway.protocol import (
|
||||
GatewayErrorCode,
|
||||
GatewayRpcMethod,
|
||||
RpcEvent,
|
||||
RpcRequest,
|
||||
RpcResponse,
|
||||
)
|
||||
from yuxi.channel.gateway.rbac import GatewayRole, check_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RpcHandler = Callable[[RpcRequest], Coroutine[Any, Any, RpcResponse]]
|
||||
RpcStreamHandler = Callable[[RpcRequest], AsyncGenerator[RpcEvent | RpcResponse, None]]
|
||||
|
||||
|
||||
class RpcDispatcher:
|
||||
def __init__(self):
|
||||
self._handlers: dict[str, RpcHandler] = {}
|
||||
self._stream_handlers: dict[str, RpcStreamHandler] = {}
|
||||
self._lazy_handlers: dict[str, Callable[[], RpcHandler]] = {}
|
||||
self._lazy_stream_handlers: dict[str, Callable[[], RpcStreamHandler]] = {}
|
||||
self._aliases: dict[str, str] = {}
|
||||
self._enforce_rbac: bool = True
|
||||
|
||||
@property
|
||||
def enforce_rbac(self) -> bool:
|
||||
return self._enforce_rbac
|
||||
|
||||
@enforce_rbac.setter
|
||||
def enforce_rbac(self, value: bool) -> None:
|
||||
self._enforce_rbac = value
|
||||
|
||||
def register(self, method: GatewayRpcMethod | str, handler: RpcHandler) -> None:
|
||||
key = method.value if isinstance(method, GatewayRpcMethod) else method
|
||||
if key in self._handlers or key in self._stream_handlers:
|
||||
logger.warning("RpcDispatcher: 覆盖已注册方法 %s", key)
|
||||
self._handlers[key] = handler
|
||||
|
||||
def register_stream(self, method: GatewayRpcMethod | str, handler: RpcStreamHandler) -> None:
|
||||
key = method.value if isinstance(method, GatewayRpcMethod) else method
|
||||
if key in self._handlers or key in self._stream_handlers:
|
||||
logger.warning("RpcDispatcher: 覆盖已注册方法 %s", key)
|
||||
self._stream_handlers[key] = handler
|
||||
|
||||
def register_lazy(self, method: GatewayRpcMethod | str, factory: Callable[[], RpcHandler]) -> None:
|
||||
key = method.value if isinstance(method, GatewayRpcMethod) else method
|
||||
if key in self._handlers or key in self._lazy_handlers:
|
||||
logger.warning("RpcDispatcher: 覆盖已有懒加载方法 %s", key)
|
||||
self._lazy_handlers[key] = factory
|
||||
|
||||
def register_stream_lazy(
|
||||
self, method: GatewayRpcMethod | str, factory: Callable[[], RpcStreamHandler]
|
||||
) -> None:
|
||||
key = method.value if isinstance(method, GatewayRpcMethod) else method
|
||||
if key in self._stream_handlers or key in self._lazy_stream_handlers:
|
||||
logger.warning("RpcDispatcher: 覆盖已有懒加载流方法 %s", key)
|
||||
self._lazy_stream_handlers[key] = factory
|
||||
|
||||
def add_alias(self, alias: str, method: str) -> None:
|
||||
self._aliases[alias] = method
|
||||
|
||||
def _resolve_method(self, method: str) -> str:
|
||||
return self._aliases.get(method, method)
|
||||
|
||||
def _resolve_handler(self, method: str) -> RpcHandler | None:
|
||||
handler = self._handlers.get(method)
|
||||
if handler is not None:
|
||||
return handler
|
||||
|
||||
factory = self._lazy_handlers.get(method)
|
||||
if factory is not None:
|
||||
handler = factory()
|
||||
self._handlers[method] = handler
|
||||
del self._lazy_handlers[method]
|
||||
return handler
|
||||
|
||||
return None
|
||||
|
||||
def _resolve_stream_handler(self, method: str) -> RpcStreamHandler | None:
|
||||
handler = self._stream_handlers.get(method)
|
||||
if handler is not None:
|
||||
return handler
|
||||
|
||||
factory = self._lazy_stream_handlers.get(method)
|
||||
if factory is not None:
|
||||
handler = factory()
|
||||
self._stream_handlers[method] = handler
|
||||
del self._lazy_stream_handlers[method]
|
||||
return handler
|
||||
|
||||
return None
|
||||
|
||||
def is_stream(self, method: str) -> bool:
|
||||
resolved = self._resolve_method(method)
|
||||
return resolved in self._stream_handlers or resolved in self._lazy_stream_handlers
|
||||
|
||||
async def dispatch(self, request: RpcRequest, caller_roles: list[GatewayRole] | None = None) -> RpcResponse:
|
||||
resolved_method = self._resolve_method(request.method)
|
||||
|
||||
if self._enforce_rbac and caller_roles:
|
||||
primary_role = caller_roles[0] if caller_roles else None
|
||||
if not check_permission(primary_role, resolved_method):
|
||||
logger.warning(
|
||||
"RpcDispatcher: 权限拒绝 method=%s role=%s",
|
||||
resolved_method,
|
||||
primary_role,
|
||||
)
|
||||
return RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.PERMISSION_DENIED,
|
||||
error_message=f"方法 '{resolved_method}' 需要更高权限,当前角色为 {primary_role}",
|
||||
)
|
||||
|
||||
handler = self._resolve_handler(resolved_method)
|
||||
if handler is None:
|
||||
logger.warning("RpcDispatcher: 未注册方法 %s", resolved_method)
|
||||
return RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.METHOD_NOT_FOUND,
|
||||
error_message=f"方法 '{resolved_method}' 未注册",
|
||||
)
|
||||
try:
|
||||
return await handler(request)
|
||||
except Exception:
|
||||
logger.exception("RpcDispatcher: 方法 %s 执行异常", resolved_method)
|
||||
return RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.INTERNAL_ERROR,
|
||||
error_message=f"方法 '{resolved_method}' 执行异常",
|
||||
)
|
||||
|
||||
async def dispatch_stream(
|
||||
self, request: RpcRequest, caller_roles: list[GatewayRole] | None = None
|
||||
) -> AsyncGenerator[RpcEvent | RpcResponse, None]:
|
||||
resolved_method = self._resolve_method(request.method)
|
||||
|
||||
if self._enforce_rbac and caller_roles:
|
||||
primary_role = caller_roles[0] if caller_roles else None
|
||||
if not check_permission(primary_role, resolved_method):
|
||||
logger.warning(
|
||||
"RpcDispatcher: 流式权限拒绝 method=%s role=%s",
|
||||
resolved_method,
|
||||
primary_role,
|
||||
)
|
||||
yield RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.PERMISSION_DENIED,
|
||||
error_message=f"方法 '{resolved_method}' 需要更高权限,当前角色为 {primary_role}",
|
||||
)
|
||||
return
|
||||
|
||||
handler = self._resolve_stream_handler(resolved_method)
|
||||
if handler is None:
|
||||
logger.warning("RpcDispatcher: 未注册流式方法 %s", resolved_method)
|
||||
yield RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.METHOD_NOT_FOUND,
|
||||
error_message=f"方法 '{resolved_method}' 未注册",
|
||||
)
|
||||
return
|
||||
try:
|
||||
async for item in handler(request):
|
||||
yield item
|
||||
except Exception:
|
||||
logger.exception("RpcDispatcher: 流式方法 %s 执行异常", resolved_method)
|
||||
yield RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.INTERNAL_ERROR,
|
||||
error_message=f"方法 '{resolved_method}' 执行异常",
|
||||
)
|
||||
|
||||
def list_methods(self) -> list[str]:
|
||||
methods = set(self._handlers.keys())
|
||||
methods.update(self._lazy_handlers.keys())
|
||||
methods.update(self._stream_handlers.keys())
|
||||
methods.update(self._lazy_stream_handlers.keys())
|
||||
return list(methods)
|
||||
|
||||
def list_aliases(self) -> dict[str, str]:
|
||||
return dict(self._aliases)
|
||||
|
||||
|
||||
rpc_dispatcher = RpcDispatcher()
|
||||
1444
backend/package/yuxi/channel/gateway/rpc_handlers.py
Normal file
1444
backend/package/yuxi/channel/gateway/rpc_handlers.py
Normal file
File diff suppressed because it is too large
Load Diff
382
backend/package/yuxi/channel/gateway/server.py
Normal file
382
backend/package/yuxi/channel/gateway/server.py
Normal file
@ -0,0 +1,382 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time as _time
|
||||
from collections.abc import Callable, Collection
|
||||
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
|
||||
from yuxi.channel.gateway.auth import GatewayAuthResult, authenticate_gateway_connect
|
||||
from yuxi.channel.gateway.broadcaster import BroadcastFilter, BroadcastResult, gateway_broadcaster
|
||||
from yuxi.channel.gateway.net_utils import (
|
||||
_ENV_ALLOW_INSECURE_PRIVATE_WS,
|
||||
is_loopback_address,
|
||||
is_secure_ws_url,
|
||||
)
|
||||
from yuxi.channel.gateway.protocol import (
|
||||
GatewayErrorCode,
|
||||
GatewayRpcMethod,
|
||||
HelloOk,
|
||||
RpcEvent,
|
||||
RpcRequest,
|
||||
RpcResponse,
|
||||
marshal_frame,
|
||||
unmarshal_frame,
|
||||
)
|
||||
from yuxi.channel.config.defaults import TIMEOUT
|
||||
from yuxi.channel.gateway.rbac import GatewayRole
|
||||
from yuxi.channel.gateway.rpc_dispatcher import rpc_dispatcher
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OnConnectCallback = Callable[[str, GatewayAuthResult], None]
|
||||
OnDisconnectCallback = Callable[[str, GatewayAuthResult], None]
|
||||
|
||||
HEARTBEAT_INTERVAL = TIMEOUT.gateway.heartbeat_interval
|
||||
HEARTBEAT_TIMEOUT = TIMEOUT.gateway.heartbeat
|
||||
SEND_TIMEOUT = TIMEOUT.gateway.send
|
||||
_SEND_SLOW_WARN_THRESHOLD = 5.0
|
||||
CONNECT_NEGOTIATION_TIMEOUT = TIMEOUT.gateway.negotiation
|
||||
TICK_INTERVAL_MS = 30_000
|
||||
|
||||
|
||||
async def _send_with_timeout(ws: WebSocket, text: str, timeout: float = SEND_TIMEOUT) -> None:
|
||||
try:
|
||||
await asyncio.wait_for(ws.send_text(text), timeout=timeout)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"Gateway WS send timed out after %.0fs, closing connection",
|
||||
timeout,
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class GatewayWsServer:
|
||||
def __init__(self, allow_private_ws: bool | None = None):
|
||||
self._active_connections: dict[str, WebSocket] = {}
|
||||
self._connection_auth: dict[str, GatewayAuthResult] = {}
|
||||
self._heartbeat_tasks: dict[str, asyncio.Task] = {}
|
||||
self._tick_tasks: dict[str, asyncio.Task] = {}
|
||||
self._on_connect: list[OnConnectCallback] = []
|
||||
self._on_disconnect: list[OnDisconnectCallback] = []
|
||||
self._shutting_down = False
|
||||
self._allow_private_ws = (
|
||||
allow_private_ws
|
||||
if allow_private_ws is not None
|
||||
else os.environ.get(_ENV_ALLOW_INSECURE_PRIVATE_WS) == "1"
|
||||
)
|
||||
|
||||
@property
|
||||
def active_count(self) -> int:
|
||||
return len(self._active_connections)
|
||||
|
||||
@property
|
||||
def shutting_down(self) -> bool:
|
||||
return self._shutting_down
|
||||
|
||||
@property
|
||||
def allow_private_ws(self) -> bool:
|
||||
return self._allow_private_ws
|
||||
|
||||
def set_allow_private_ws(self, value: bool) -> None:
|
||||
self._allow_private_ws = value
|
||||
|
||||
async def handle_connection(self, ws: WebSocket, token: str | None = None):
|
||||
client_ip = _resolve_ws_client_ip(ws)
|
||||
is_loopback_client = is_loopback_address(client_ip)
|
||||
|
||||
if not is_loopback_client and not is_secure_ws_url(str(ws.url), allow_private_ws=self._allow_private_ws):
|
||||
display_host = _format_ws_display_host(str(ws.url))
|
||||
logger.warning(
|
||||
"Gateway WS: rejected insecure ws:// connection from %s to %s",
|
||||
client_ip or "unknown",
|
||||
display_host,
|
||||
)
|
||||
await ws.close(code=4400, reason="insecure_ws_connection")
|
||||
return
|
||||
|
||||
await ws.accept()
|
||||
|
||||
auth_header = ws.headers.get("authorization")
|
||||
auth_result = await authenticate_gateway_connect(
|
||||
auth_header, token, client_ip=client_ip, headers=dict(ws.headers), remote_addr=client_ip
|
||||
)
|
||||
if not auth_result.authenticated:
|
||||
await _send_with_timeout(
|
||||
ws,
|
||||
marshal_frame(
|
||||
RpcResponse(
|
||||
id="auth",
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.AUTH_ERROR,
|
||||
error_message=auth_result.error or "认证失败",
|
||||
)
|
||||
),
|
||||
)
|
||||
await ws.close(code=4001, reason="unauthorized")
|
||||
return
|
||||
|
||||
conn_id = f"{auth_result.user_id}-{id(ws):x}"
|
||||
self._active_connections[conn_id] = ws
|
||||
self._connection_auth[conn_id] = auth_result
|
||||
|
||||
async def _send_to_ws(text: str) -> None:
|
||||
await ws.send_text(text)
|
||||
|
||||
gateway_broadcaster.register_connection(conn_id, auth_result, _send_to_ws)
|
||||
|
||||
heartbeat_task = asyncio.create_task(
|
||||
self._heartbeat_loop(conn_id, ws),
|
||||
name=f"gateway_heartbeat:{conn_id}",
|
||||
)
|
||||
self._heartbeat_tasks[conn_id] = heartbeat_task
|
||||
|
||||
tick_task = asyncio.create_task(
|
||||
self._tick_loop(conn_id, ws),
|
||||
name=f"gateway_tick:{conn_id}",
|
||||
)
|
||||
self._tick_tasks[conn_id] = tick_task
|
||||
|
||||
try:
|
||||
first_frame = await self._negotiate_connect(conn_id, ws, auth_result)
|
||||
except Exception:
|
||||
logger.warning("Gateway connect negotiation failed: %s", conn_id)
|
||||
first_frame = None
|
||||
|
||||
try:
|
||||
for cb in self._on_connect:
|
||||
try:
|
||||
cb(conn_id, auth_result)
|
||||
except Exception:
|
||||
logger.exception("on_connect callback failed: %s", conn_id)
|
||||
|
||||
logger.info(
|
||||
"Gateway WS connected: %s (user=%s, roles=%s)",
|
||||
conn_id,
|
||||
auth_result.user_id,
|
||||
[r.value for r in auth_result.roles],
|
||||
)
|
||||
|
||||
if first_frame is not None:
|
||||
await self._dispatch_frame(conn_id, ws, first_frame)
|
||||
|
||||
while True:
|
||||
raw = await ws.receive_text()
|
||||
try:
|
||||
frame = unmarshal_frame(raw)
|
||||
except Exception:
|
||||
logger.warning("Gateway WS: 无效帧 from %s", conn_id)
|
||||
await _send_with_timeout(
|
||||
ws,
|
||||
marshal_frame(
|
||||
RpcResponse(
|
||||
id="",
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.INVALID_REQUEST,
|
||||
error_message="无效的 RPC 帧格式",
|
||||
)
|
||||
),
|
||||
)
|
||||
continue
|
||||
|
||||
await self._dispatch_frame(conn_id, ws, frame)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info("Gateway WS disconnected: %s", conn_id)
|
||||
except Exception:
|
||||
logger.exception("Gateway WS error: %s", conn_id)
|
||||
finally:
|
||||
gateway_broadcaster.unregister_connection(conn_id)
|
||||
for task_key, task_dict in [("heartbeat", self._heartbeat_tasks), ("tick", self._tick_tasks)]:
|
||||
task = task_dict.get(conn_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
task_dict.pop(conn_id, None)
|
||||
self._active_connections.pop(conn_id, None)
|
||||
auth = self._connection_auth.pop(conn_id, None)
|
||||
for cb in self._on_disconnect:
|
||||
try:
|
||||
cb(conn_id, auth or GatewayAuthResult(authenticated=False))
|
||||
except Exception:
|
||||
logger.exception("on_disconnect callback failed: %s", conn_id)
|
||||
|
||||
async def _heartbeat_loop(self, conn_id: str, ws: WebSocket) -> None:
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
||||
try:
|
||||
await ws.send_text(marshal_frame(RpcEvent(event="ping")))
|
||||
except Exception:
|
||||
logger.warning("Gateway WS heartbeat failed: %s", conn_id)
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def _negotiate_connect(self, conn_id: str, ws: WebSocket, auth_result: GatewayAuthResult):
|
||||
try:
|
||||
raw = await asyncio.wait_for(ws.receive_text(), timeout=CONNECT_NEGOTIATION_TIMEOUT)
|
||||
frame = unmarshal_frame(raw)
|
||||
if isinstance(frame, RpcRequest) and frame.method == GatewayRpcMethod.CONNECT:
|
||||
hello_ok = self._build_hello_ok(conn_id, auth_result)
|
||||
await _send_with_timeout(ws, marshal_frame(hello_ok))
|
||||
return None
|
||||
return frame
|
||||
except TimeoutError:
|
||||
return None
|
||||
|
||||
def _build_hello_ok(self, conn_id: str, auth_result: GatewayAuthResult) -> HelloOk:
|
||||
hello = HelloOk()
|
||||
hello.result["server"]["connId"] = conn_id
|
||||
hello.result["auth"] = {
|
||||
"role": auth_result.roles[0].value if auth_result.roles else "viewer",
|
||||
"scopes": [r.value for r in auth_result.roles],
|
||||
}
|
||||
return hello
|
||||
|
||||
async def _dispatch_frame(self, conn_id: str, ws: WebSocket, frame):
|
||||
auth_result = self._connection_auth.get(conn_id)
|
||||
if auth_result is None:
|
||||
return
|
||||
if isinstance(frame, RpcRequest):
|
||||
frame.caller_user_id = auth_result.user_id
|
||||
if frame.method == "event.subscribe":
|
||||
resp = await self._handle_event_subscribe(conn_id, frame)
|
||||
await _send_with_timeout(ws, marshal_frame(resp))
|
||||
elif frame.method == "event.unsubscribe":
|
||||
resp = await self._handle_event_unsubscribe(conn_id, frame)
|
||||
await _send_with_timeout(ws, marshal_frame(resp))
|
||||
elif rpc_dispatcher.is_stream(frame.method):
|
||||
async for item in rpc_dispatcher.dispatch_stream(frame, auth_result.roles):
|
||||
await _send_with_timeout(ws, marshal_frame(item), timeout=TIMEOUT.gateway.stream_rpc)
|
||||
else:
|
||||
resp = await rpc_dispatcher.dispatch(frame, auth_result.roles)
|
||||
await _send_with_timeout(ws, marshal_frame(resp))
|
||||
|
||||
async def _tick_loop(self, conn_id: str, ws: WebSocket) -> None:
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(TICK_INTERVAL_MS / 1000.0)
|
||||
tick_event = RpcEvent(
|
||||
event="tick",
|
||||
data={"ts": int(_time.time() * 1000)},
|
||||
)
|
||||
await _send_with_timeout(ws, marshal_frame(tick_event))
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.warning("Gateway tick failed for %s", conn_id)
|
||||
|
||||
async def _handle_event_subscribe(self, conn_id: str, req: RpcRequest) -> RpcResponse:
|
||||
events = req.params.get("events", []) if req.params else []
|
||||
if not isinstance(events, list) or not all(isinstance(e, str) for e in events):
|
||||
return RpcResponse(
|
||||
id=req.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.INVALID_PARAMS,
|
||||
error_message="events 必须是字符串数组",
|
||||
)
|
||||
gateway_broadcaster.subscribe(conn_id, events)
|
||||
subscribed = gateway_broadcaster.get_subscriptions(conn_id)
|
||||
return RpcResponse(id=req.id, ok=True, result={"subscribed": list(subscribed)})
|
||||
|
||||
async def _handle_event_unsubscribe(self, conn_id: str, req: RpcRequest) -> RpcResponse:
|
||||
events = req.params.get("events", []) if req.params else []
|
||||
if not isinstance(events, list) or not all(isinstance(e, str) for e in events):
|
||||
return RpcResponse(
|
||||
id=req.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.INVALID_PARAMS,
|
||||
error_message="events 必须是字符串数组",
|
||||
)
|
||||
gateway_broadcaster.unsubscribe(conn_id, events)
|
||||
subscribed = gateway_broadcaster.get_subscriptions(conn_id)
|
||||
return RpcResponse(id=req.id, ok=True, result={"subscribed": list(subscribed)})
|
||||
|
||||
async def send_event(self, conn_id: str, event: RpcEvent) -> bool:
|
||||
return await gateway_broadcaster.send_event(conn_id, event.event, event.data)
|
||||
|
||||
async def broadcast_event(self, event: RpcEvent) -> int:
|
||||
result = await gateway_broadcaster.broadcast(event.event, event.data)
|
||||
return result.sent
|
||||
|
||||
async def broadcast_event_filtered(
|
||||
self,
|
||||
event: str,
|
||||
data: dict | None = None,
|
||||
*,
|
||||
conn_ids: Collection[str] | None = None,
|
||||
user_ids: Collection[str] | None = None,
|
||||
roles: Collection[GatewayRole] | None = None,
|
||||
drop_if_slow: bool = True,
|
||||
) -> BroadcastResult:
|
||||
filter_ = BroadcastFilter(
|
||||
conn_ids=conn_ids,
|
||||
user_ids=user_ids,
|
||||
roles=roles,
|
||||
drop_if_slow=drop_if_slow,
|
||||
)
|
||||
return await gateway_broadcaster.broadcast(event, data, filter_)
|
||||
|
||||
def register_on_connect(self, cb: OnConnectCallback) -> None:
|
||||
self._on_connect.append(cb)
|
||||
|
||||
def register_on_disconnect(self, cb: OnDisconnectCallback) -> None:
|
||||
self._on_disconnect.append(cb)
|
||||
|
||||
async def shutdown(self, timeout: float = 10.0) -> None:
|
||||
self._shutting_down = True
|
||||
logger.info("Gateway WS server shutting down (%d connections)", len(self._active_connections))
|
||||
|
||||
await gateway_broadcaster.broadcast("server.shutdown", {"reason": "server_shutdown"})
|
||||
|
||||
for task_dict in [self._heartbeat_tasks, self._tick_tasks]:
|
||||
for conn_id, task in list(task_dict.items()):
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
gateway_broadcaster.unregister_connection(conn_id)
|
||||
|
||||
async def _close_conn(conn_id: str, ws: WebSocket) -> None:
|
||||
try:
|
||||
await asyncio.wait_for(ws.close(code=1001, reason="server_shutdown"), timeout=timeout)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tasks = [_close_conn(cid, ws) for cid, ws in self._active_connections.items()]
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
self._active_connections.clear()
|
||||
self._connection_auth.clear()
|
||||
self._heartbeat_tasks.clear()
|
||||
self._tick_tasks.clear()
|
||||
logger.info("Gateway WS server shut down complete")
|
||||
|
||||
|
||||
gateway_ws_server = GatewayWsServer()
|
||||
|
||||
|
||||
def _resolve_ws_client_ip(ws: WebSocket) -> str | None:
|
||||
client = getattr(ws, "client", None)
|
||||
if client and hasattr(client, "host"):
|
||||
return client.host
|
||||
return None
|
||||
|
||||
|
||||
def _format_ws_display_host(url_str: str) -> str:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
try:
|
||||
parsed = urlparse(url_str)
|
||||
host = parsed.hostname or parsed.netloc or "unknown"
|
||||
port = parsed.port
|
||||
if port and port not in (80, 443):
|
||||
return f"{host}:{port}"
|
||||
return host
|
||||
except Exception:
|
||||
return url_str
|
||||
277
backend/package/yuxi/channel/gateway/sse.py
Normal file
277
backend/package/yuxi/channel/gateway/sse.py
Normal file
@ -0,0 +1,277 @@
|
||||
"""Gateway HTTP SSE 端点 — 为不支持 WebSocket 的客户端提供备选流式通道。
|
||||
|
||||
复用 stream_agent_chat 的流式输出能力,通过 HTTP SSE (text/event-stream) 推送给浏览器客户端。
|
||||
|
||||
SSE 事件类型对齐 CowAgent 标准:
|
||||
reasoning / delta / tool_start / tool_end / message_end /
|
||||
phase / image / file / video / done / error
|
||||
|
||||
路由:
|
||||
POST /api/sse/chat — 发送聊天消息,返回 request_id
|
||||
GET /api/sse/stream — 订阅 SSE 事件流
|
||||
GET /api/poll — Polling 降级轮询
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from fastapi import APIRouter, Header, Query, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from yuxi.channel.protocols import SseEventType
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SSE_KEEP_ALIVE_SEC = 15.0
|
||||
SSE_QUEUE_TTL_SEC = 600
|
||||
MAX_QUEUE_SIZE = 256
|
||||
|
||||
|
||||
def _sse_frame(event_type: str, data: dict | str) -> str:
|
||||
content = json.dumps(data, ensure_ascii=False) if isinstance(data, dict) else data
|
||||
return f"event: {event_type}\ndata: {content}\n\n"
|
||||
|
||||
|
||||
def _sse_comment(comment: str) -> str:
|
||||
return f": {comment}\n\n"
|
||||
|
||||
|
||||
class GatewaySseEndpoint:
|
||||
def __init__(self):
|
||||
self._queues: dict[str, asyncio.Queue[dict]] = {}
|
||||
self._tasks: dict[str, asyncio.Task] = {}
|
||||
self._last_active: dict[str, float] = {}
|
||||
|
||||
def create(self, rid: str, q: asyncio.Queue[dict], task: asyncio.Task) -> None:
|
||||
self._queues[rid] = q
|
||||
self._tasks[rid] = task
|
||||
self._last_active[rid] = time.monotonic()
|
||||
|
||||
def get(self, rid: str) -> asyncio.Queue[dict] | None:
|
||||
return self._queues.get(rid)
|
||||
|
||||
def touch(self, rid: str) -> None:
|
||||
self._last_active[rid] = time.monotonic()
|
||||
|
||||
def remove(self, rid: str) -> None:
|
||||
self._queues.pop(rid, None)
|
||||
task = self._tasks.pop(rid, None)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
self._last_active.pop(rid, None)
|
||||
|
||||
def cleanup_stale(self) -> int:
|
||||
now = time.monotonic()
|
||||
stale = [rid for rid, ts in self._last_active.items() if now - ts > SSE_QUEUE_TTL_SEC]
|
||||
for rid in stale:
|
||||
self.remove(rid)
|
||||
return len(stale)
|
||||
|
||||
|
||||
gateway_sse_endpoint = GatewaySseEndpoint()
|
||||
|
||||
|
||||
async def _run_chat_feed(
|
||||
query: str,
|
||||
agent_config_id: int,
|
||||
thread_id: str | None,
|
||||
image_content: str | None,
|
||||
current_user,
|
||||
db,
|
||||
q: asyncio.Queue[dict],
|
||||
rid: str,
|
||||
) -> None:
|
||||
from yuxi.services.chat_service import stream_agent_chat
|
||||
|
||||
meta = {
|
||||
"source": "gateway_sse",
|
||||
"channel_type": "sse",
|
||||
"account_id": "default",
|
||||
"request_id": rid,
|
||||
}
|
||||
|
||||
try:
|
||||
async for chunk in stream_agent_chat(
|
||||
query=query,
|
||||
agent_config_id=agent_config_id,
|
||||
thread_id=thread_id,
|
||||
meta=meta,
|
||||
image_content=image_content,
|
||||
current_user=current_user,
|
||||
db=db,
|
||||
):
|
||||
try:
|
||||
data = json.loads(chunk.decode("utf-8"))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
continue
|
||||
|
||||
status = data.get("status", "")
|
||||
content = data.get("response", "")
|
||||
|
||||
if status == "error":
|
||||
await q.put({"type": SseEventType.ERROR, "data": {"message": data.get("error_message", content)}})
|
||||
await q.put({"type": SseEventType.DONE, "data": {}})
|
||||
return
|
||||
|
||||
if status == "init":
|
||||
continue
|
||||
|
||||
if status in ("streaming", "reasoning"):
|
||||
event_type = SseEventType.REASONING if data.get("reasoning") else SseEventType.DELTA
|
||||
if content:
|
||||
await q.put({"type": event_type, "data": {"content": content}})
|
||||
|
||||
if status == "thinking":
|
||||
await q.put({"type": SseEventType.PHASE, "data": {"content": content or "thinking"}})
|
||||
|
||||
if status == "finished":
|
||||
final_content = data.get("final_response", content)
|
||||
await q.put(
|
||||
{
|
||||
"type": SseEventType.MESSAGE_END,
|
||||
"data": {"content": final_content, "thread_id": data.get("thread_id", thread_id)},
|
||||
}
|
||||
)
|
||||
|
||||
await q.put({"type": SseEventType.DONE, "data": {}})
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception("SSE chat stream failed for %s", rid)
|
||||
await q.put({"type": SseEventType.ERROR, "data": {"message": "SSE stream error"}})
|
||||
await q.put({"type": SseEventType.DONE, "data": {}})
|
||||
|
||||
|
||||
async def _resolve_sse_user(db, authorization: str | None):
|
||||
from sqlalchemy import select
|
||||
|
||||
from server.utils.auth_utils import AuthUtils
|
||||
from yuxi.storage.postgres.models_business import User
|
||||
|
||||
if authorization and authorization.startswith("Bearer "):
|
||||
token = authorization[7:]
|
||||
try:
|
||||
payload = AuthUtils.verify_access_token(token)
|
||||
user_id = payload.get("sub")
|
||||
if user_id:
|
||||
result = await db.execute(select(User).where(User.id == int(user_id)))
|
||||
user = result.scalar_one_or_none()
|
||||
if user:
|
||||
return user
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
result = await db.execute(
|
||||
select(User).where(User.is_deleted == 0).order_by(User.id).limit(1)
|
||||
)
|
||||
user = result.scalar_one_or_none()
|
||||
if user:
|
||||
logger.warning("SSE endpoint using fallback user id=%s (no valid auth provided)", user.id)
|
||||
return user
|
||||
|
||||
|
||||
router = APIRouter(prefix="/api/sse", tags=["sse"])
|
||||
|
||||
|
||||
@router.post("/chat")
|
||||
async def sse_chat(
|
||||
query: str = Query(...),
|
||||
agent_config_id: int = Query(...),
|
||||
thread_id: str | None = Query(None),
|
||||
image_content: str | None = Query(None),
|
||||
authorization: str | None = Header(None),
|
||||
):
|
||||
rid = str(uuid.uuid4())
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
current_user = await _resolve_sse_user(db, authorization)
|
||||
|
||||
q: asyncio.Queue[dict] = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_run_chat_feed(
|
||||
query=query,
|
||||
agent_config_id=agent_config_id,
|
||||
thread_id=thread_id,
|
||||
image_content=image_content,
|
||||
current_user=current_user,
|
||||
db=db,
|
||||
q=q,
|
||||
rid=rid,
|
||||
),
|
||||
name=f"sse_chat:{rid}",
|
||||
)
|
||||
gateway_sse_endpoint.create(rid, q, task)
|
||||
|
||||
return {"request_id": rid, "stream_url": f"/api/sse/stream?request_id={rid}"}
|
||||
|
||||
|
||||
@router.get("/stream")
|
||||
async def sse_stream(request: Request, request_id: str = Query(...)):
|
||||
q = gateway_sse_endpoint.get(request_id)
|
||||
if q is None:
|
||||
return StreamingResponse(
|
||||
iter([_sse_frame(SseEventType.ERROR, {"message": "request_id invalid or expired"})]),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
gateway_sse_endpoint.touch(request_id)
|
||||
|
||||
async def generate():
|
||||
try:
|
||||
while True:
|
||||
if await request.is_disconnected():
|
||||
logger.info("SSE client disconnected: %s", request_id)
|
||||
break
|
||||
try:
|
||||
event = await asyncio.wait_for(q.get(), timeout=SSE_KEEP_ALIVE_SEC)
|
||||
yield _sse_frame(str(event["type"]), event["data"])
|
||||
if event["type"] == SseEventType.DONE:
|
||||
break
|
||||
except TimeoutError:
|
||||
yield _sse_comment("keepalive")
|
||||
finally:
|
||||
gateway_sse_endpoint.remove(request_id)
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ── OPT-7: 实时日志 SSE ──────────────────────────────────
|
||||
|
||||
|
||||
@router.get("/logs")
|
||||
async def sse_logs(
|
||||
request: Request,
|
||||
levels: str | None = Query(None, description="comma-separated: DEBUG,INFO,WARNING,ERROR"),
|
||||
):
|
||||
logger.info("SSE log stream requested (not yet implemented), levels=%s", levels)
|
||||
|
||||
async def generate():
|
||||
yield _sse_frame(SseEventType.ERROR, {"message": "SSE log streaming is not yet implemented"})
|
||||
yield _sse_frame(SseEventType.DONE, {})
|
||||
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
173
backend/package/yuxi/channel/gateway/tailscale_auth.py
Normal file
173
backend/package/yuxi/channel/gateway/tailscale_auth.py
Normal file
@ -0,0 +1,173 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WHOIS_CACHE_TTL_MS = 60_000
|
||||
WHOIS_ERROR_TTL_MS = 5_000
|
||||
WHOIS_TIMEOUT_S = 5.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class TailscaleWhoisIdentity:
|
||||
login: str
|
||||
name: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CacheEntry:
|
||||
value: TailscaleWhoisIdentity | None
|
||||
expires_at: float
|
||||
|
||||
|
||||
_whois_cache: dict[str, _CacheEntry] = {}
|
||||
|
||||
|
||||
def _read_cached_whois(ip: str, now_ms: float) -> TailscaleWhoisIdentity | None | bool:
|
||||
"""Returns None=cache miss, False=cached null, or the identity."""
|
||||
entry = _whois_cache.get(ip)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry.expires_at <= now_ms:
|
||||
del _whois_cache[ip]
|
||||
return None
|
||||
return entry.value if entry.value is not None else False # type: ignore[return-value]
|
||||
|
||||
|
||||
def _write_cached_whois(ip: str, value: TailscaleWhoisIdentity | None, ttl_ms: float) -> None:
|
||||
_whois_cache[ip] = _CacheEntry(value=value, expires_at=time.time() * 1000 + ttl_ms)
|
||||
|
||||
|
||||
def _parse_whois_identity(payload: dict) -> TailscaleWhoisIdentity | None:
|
||||
user_profile = payload.get("UserProfile") or payload.get("userProfile") or payload.get("User")
|
||||
if not isinstance(user_profile, dict):
|
||||
return None
|
||||
login = (
|
||||
user_profile.get("LoginName")
|
||||
or user_profile.get("Login")
|
||||
or user_profile.get("loginName")
|
||||
or user_profile.get("login")
|
||||
)
|
||||
if not login or not isinstance(login, str):
|
||||
return None
|
||||
name = (
|
||||
user_profile.get("DisplayName")
|
||||
or user_profile.get("displayName")
|
||||
or user_profile.get("Name")
|
||||
or user_profile.get("name")
|
||||
)
|
||||
return TailscaleWhoisIdentity(
|
||||
login=login.lower().strip(),
|
||||
name=name.strip() if isinstance(name, str) and name.strip() else None,
|
||||
)
|
||||
|
||||
|
||||
def _parse_possibly_noisy_json(stdout: str) -> dict:
|
||||
trimmed = stdout.strip()
|
||||
start = trimmed.find("{")
|
||||
end = trimmed.rfind("}")
|
||||
if start >= 0 and end > start:
|
||||
return json.loads(trimmed[start : end + 1])
|
||||
return json.loads(trimmed)
|
||||
|
||||
|
||||
async def _get_tailscale_binary() -> str | None:
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
"which",
|
||||
"tailscale",
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
)
|
||||
stdout, _ = await proc.communicate()
|
||||
if proc.returncode == 0 and stdout:
|
||||
return stdout.decode().strip()
|
||||
except FileNotFoundError:
|
||||
logger.debug("tailscale binary not found via 'which'")
|
||||
except Exception:
|
||||
logger.debug("tailscale binary lookup failed via 'which'", exc_info=True)
|
||||
|
||||
import os
|
||||
|
||||
mac_path = "/Applications/Tailscale.app/Contents/MacOS/Tailscale"
|
||||
if os.path.exists(mac_path) and os.access(mac_path, os.X_OK):
|
||||
return mac_path
|
||||
return None
|
||||
|
||||
|
||||
async def read_tailscale_whois_identity(
|
||||
ip: str,
|
||||
timeout_s: float = WHOIS_TIMEOUT_S,
|
||||
cache_ttl_ms: float = WHOIS_CACHE_TTL_MS,
|
||||
error_ttl_ms: float = WHOIS_ERROR_TTL_MS,
|
||||
) -> TailscaleWhoisIdentity | None:
|
||||
normalized = ip.strip()
|
||||
if not normalized:
|
||||
return None
|
||||
|
||||
now_ms = time.time() * 1000
|
||||
cached = _read_cached_whois(normalized, now_ms)
|
||||
if cached is not None:
|
||||
return cached if cached is not False else None
|
||||
|
||||
tailscale_bin = await _get_tailscale_binary()
|
||||
if not tailscale_bin:
|
||||
logger.debug("tailscale binary not found, whois unavailable")
|
||||
_write_cached_whois(normalized, None, error_ttl_ms)
|
||||
return None
|
||||
|
||||
proc: asyncio.subprocess.Process | None = None
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
tailscale_bin,
|
||||
"whois",
|
||||
"--json",
|
||||
normalized,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
proc.communicate(),
|
||||
timeout=timeout_s,
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
logger.debug("tailscale whois failed for %s: %s", normalized, stderr.decode().strip())
|
||||
_write_cached_whois(normalized, None, error_ttl_ms)
|
||||
return None
|
||||
payload = _parse_possibly_noisy_json(stdout.decode())
|
||||
identity = _parse_whois_identity(payload)
|
||||
_write_cached_whois(normalized, identity, cache_ttl_ms)
|
||||
return identity
|
||||
except TimeoutError:
|
||||
logger.debug("tailscale whois timeout for %s", normalized)
|
||||
if proc is not None:
|
||||
try:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
_write_cached_whois(normalized, None, error_ttl_ms)
|
||||
return None
|
||||
except Exception:
|
||||
logger.debug("tailscale whois error for %s", normalized, exc_info=True)
|
||||
if proc is not None:
|
||||
try:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except Exception:
|
||||
pass
|
||||
_write_cached_whois(normalized, None, error_ttl_ms)
|
||||
return None
|
||||
|
||||
|
||||
def get_tailscale_user_from_headers(headers: dict | None) -> tuple[str, str] | None:
|
||||
if not headers:
|
||||
return None
|
||||
login = headers.get("tailscale-user-login")
|
||||
if not login:
|
||||
return None
|
||||
name = headers.get("tailscale-user-name") or login
|
||||
return (login.lower().strip(), name.strip() if isinstance(name, str) else login)
|
||||
110
backend/package/yuxi/channel/gateway/validation.py
Normal file
110
backend/package/yuxi/channel/gateway/validation.py
Normal file
@ -0,0 +1,110 @@
|
||||
from functools import wraps
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
|
||||
from yuxi.channel.gateway.protocol import GatewayErrorCode, RpcRequest, RpcResponse
|
||||
|
||||
|
||||
class ErrorShape(BaseModel):
|
||||
code: str
|
||||
message: str = ""
|
||||
data: Any | None = None
|
||||
retryable: bool = False
|
||||
retryAfterMs: int | None = None
|
||||
|
||||
|
||||
class RequestFrameSchema(BaseModel):
|
||||
type: str = "request"
|
||||
id: str
|
||||
method: str
|
||||
params: dict | None = None
|
||||
sessionId: str | None = None
|
||||
|
||||
|
||||
class ResponseFrameSchema(BaseModel):
|
||||
type: str = "response"
|
||||
id: str
|
||||
ok: bool = True
|
||||
result: dict | None = None
|
||||
error: ErrorShape | None = None
|
||||
|
||||
|
||||
class EventFrameSchema(BaseModel):
|
||||
type: str = "event"
|
||||
event: str
|
||||
data: dict | None = None
|
||||
timestamp: float = 0.0
|
||||
seq: int | None = None
|
||||
stateVersion: int | None = None
|
||||
|
||||
|
||||
class StartAccountParams(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
channel_type: str = Field(..., min_length=1, max_length=50)
|
||||
account_id: str = Field(default="default", min_length=1, max_length=100)
|
||||
config: dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
class StopAccountParams(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
channel_type: str = Field(..., min_length=1, max_length=50)
|
||||
account_id: str = Field(default="default", min_length=1, max_length=100)
|
||||
force: bool = False
|
||||
|
||||
|
||||
class SendMessageParams(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
channel_type: str = Field(..., min_length=1, max_length=50)
|
||||
account_id: str = Field(default="default", min_length=1, max_length=100)
|
||||
target_id: str = Field(..., min_length=1, max_length=200)
|
||||
text: str = Field(..., min_length=1, max_length=4096)
|
||||
msg_type: str = Field(default="text")
|
||||
media_url: str | None = None
|
||||
extra: dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ProbeParams(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
channel_type: str = Field(..., min_length=1, max_length=50)
|
||||
account_id: str = Field(default="default", min_length=1, max_length=100)
|
||||
|
||||
|
||||
class DiagnoseParams(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
channel_type: str = Field(..., min_length=1, max_length=50)
|
||||
account_id: str = Field(default="default", min_length=1, max_length=100)
|
||||
|
||||
|
||||
class RepairParams(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
channel_type: str = Field(..., min_length=1, max_length=50)
|
||||
step_id: str = Field(..., min_length=1, max_length=100)
|
||||
account_id: str = Field(default="default", min_length=1, max_length=100)
|
||||
|
||||
|
||||
def validate_params(schema_cls: type[BaseModel]):
|
||||
def decorator(func):
|
||||
@wraps(func)
|
||||
async def wrapper(request: RpcRequest, *args, **kwargs):
|
||||
try:
|
||||
validated = schema_cls(**(request.params or {}))
|
||||
except ValidationError as e:
|
||||
return RpcResponse(
|
||||
id=request.id,
|
||||
ok=False,
|
||||
error_code=GatewayErrorCode.INVALID_PARAMS,
|
||||
error_message=f"参数校验失败: {e}",
|
||||
)
|
||||
request._validated_params = validated
|
||||
return await func(request, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
304
backend/package/yuxi/channel/gateway/webhook_security.py
Normal file
304
backend/package/yuxi/channel/gateway/webhook_security.py
Normal file
@ -0,0 +1,304 @@
|
||||
import hashlib
|
||||
import hmac
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PRE_AUTH_BODY_SIZE_LIMIT = 64 * 1024
|
||||
POST_AUTH_BODY_SIZE_LIMIT = 1 * 1024 * 1024
|
||||
PRE_AUTH_READ_TIMEOUT = 5.0
|
||||
POST_AUTH_READ_TIMEOUT = 30.0
|
||||
MAX_CONCURRENT_PER_KEY = 8
|
||||
MAX_TRACKING_KEYS = 4096
|
||||
|
||||
_VALID_WEBHOOK_METHODS = frozenset({"POST", "PUT", "PATCH"})
|
||||
|
||||
|
||||
class WebhookErrorCode(StrEnum):
|
||||
METHOD_NOT_ALLOWED = "METHOD_NOT_ALLOWED"
|
||||
CONTENT_TYPE_UNSUPPORTED = "CONTENT_TYPE_UNSUPPORTED"
|
||||
BODY_TOO_LARGE = "BODY_TOO_LARGE"
|
||||
TOO_MANY_REQUESTS = "TOO_MANY_REQUESTS"
|
||||
SIGNATURE_MISSING = "SIGNATURE_MISSING"
|
||||
SIGNATURE_INVALID = "SIGNATURE_INVALID"
|
||||
PAYLOAD_INVALID = "PAYLOAD_INVALID"
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebhookGuardConfig:
|
||||
channel_type: str
|
||||
allowed_methods: tuple[str, ...] = ("POST",)
|
||||
allowed_content_types: tuple[str, ...] = ("application/json",)
|
||||
pre_auth: bool = True
|
||||
verify_signature: bool = True
|
||||
body_size_limit: int = PRE_AUTH_BODY_SIZE_LIMIT
|
||||
|
||||
|
||||
import threading
|
||||
|
||||
class WebhookAnomalyTracker:
|
||||
def __init__(self, sample_window: float = 60.0, alert_threshold: int = 50):
|
||||
self._sample_window = sample_window
|
||||
self._alert_threshold = alert_threshold
|
||||
self._counters: dict[str, dict[str, int]] = {}
|
||||
self._last_reset: float = time.monotonic()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def record(self, channel_type: str, error_code: str) -> None:
|
||||
now = time.monotonic()
|
||||
with self._lock:
|
||||
if now - self._last_reset > self._sample_window:
|
||||
self._counters.clear()
|
||||
self._last_reset = now
|
||||
|
||||
channel_counters = self._counters.setdefault(channel_type, {})
|
||||
channel_counters[error_code] = channel_counters.get(error_code, 0) + 1
|
||||
|
||||
total = sum(channel_counters.values())
|
||||
if total >= self._alert_threshold and total % self._alert_threshold == 0:
|
||||
self._emit_alert(channel_type, channel_counters)
|
||||
|
||||
def _emit_alert(self, channel_type: str, counters: dict[str, int]) -> None:
|
||||
details = ", ".join(f"{k}={v}" for k, v in counters.items())
|
||||
logger.warning(
|
||||
"Webhook anomaly alert [%s]: %d anomalies in window — %s",
|
||||
channel_type,
|
||||
sum(counters.values()),
|
||||
details,
|
||||
)
|
||||
|
||||
def get_stats(self, channel_type: str | None = None) -> dict[str, dict[str, int]]:
|
||||
with self._lock:
|
||||
if channel_type:
|
||||
return {channel_type: self._counters.get(channel_type, {})}
|
||||
return dict(self._counters)
|
||||
|
||||
|
||||
class WebhookSigner:
|
||||
def __init__(self):
|
||||
self._signing_keys: dict[str, str] = {}
|
||||
|
||||
def set_channel_key(self, channel_type: str, signing_key: str) -> None:
|
||||
self._signing_keys[channel_type] = signing_key
|
||||
|
||||
def remove_channel_key(self, channel_type: str) -> None:
|
||||
self._signing_keys.pop(channel_type, None)
|
||||
|
||||
def compute_signature(self, channel_type: str, body: bytes) -> str | None:
|
||||
key = self._signing_keys.get(channel_type)
|
||||
if not key:
|
||||
return None
|
||||
return hmac.new(key.encode(), body, hashlib.sha256).hexdigest()
|
||||
|
||||
def verify(
|
||||
self,
|
||||
channel_type: str,
|
||||
body: bytes,
|
||||
signature: str,
|
||||
) -> bool:
|
||||
expected = self.compute_signature(channel_type, body)
|
||||
if expected is None:
|
||||
return False
|
||||
return hmac.compare_digest(expected, signature)
|
||||
|
||||
@staticmethod
|
||||
def compute_sha256_signature(secret: str, body: bytes) -> str:
|
||||
return hmac.new(secret.encode(), body, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
class WebhookConcurrencyGuard:
|
||||
def __init__(self, max_per_key: int = MAX_CONCURRENT_PER_KEY, max_keys: int = MAX_TRACKING_KEYS):
|
||||
self._max_per_key = max_per_key
|
||||
self._max_keys = max_keys
|
||||
self._inflight: dict[str, int] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def try_acquire(self, channel_type: str) -> bool:
|
||||
with self._lock:
|
||||
current = self._inflight.get(channel_type, 0)
|
||||
if current >= self._max_per_key:
|
||||
return False
|
||||
if len(self._inflight) >= self._max_keys and channel_type not in self._inflight:
|
||||
return False
|
||||
self._inflight[channel_type] = current + 1
|
||||
return True
|
||||
|
||||
def release(self, channel_type: str) -> None:
|
||||
with self._lock:
|
||||
current = self._inflight.get(channel_type, 0)
|
||||
if current <= 1:
|
||||
self._inflight.pop(channel_type, None)
|
||||
else:
|
||||
self._inflight[channel_type] = current - 1
|
||||
|
||||
@property
|
||||
def inflight_count(self) -> dict[str, int]:
|
||||
with self._lock:
|
||||
return dict(self._inflight)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebhookGuardResult:
|
||||
allowed: bool
|
||||
error_code: WebhookErrorCode | None = None
|
||||
error_message: str | None = None
|
||||
http_status: int = 200
|
||||
|
||||
|
||||
class WebhookGuard:
|
||||
def __init__(
|
||||
self,
|
||||
signer: WebhookSigner | None = None,
|
||||
concurrency_guard: WebhookConcurrencyGuard | None = None,
|
||||
anomaly_tracker: WebhookAnomalyTracker | None = None,
|
||||
):
|
||||
self._signer = signer or WebhookSigner()
|
||||
self._concurrency = concurrency_guard or WebhookConcurrencyGuard()
|
||||
self._anomaly = anomaly_tracker or WebhookAnomalyTracker()
|
||||
self._channel_configs: dict[str, WebhookGuardConfig] = {}
|
||||
|
||||
def register_channel(self, channel_type: str, config: WebhookGuardConfig | None = None) -> None:
|
||||
self._channel_configs[channel_type] = config or WebhookGuardConfig(channel_type=channel_type)
|
||||
|
||||
def unregister_channel(self, channel_type: str) -> None:
|
||||
self._channel_configs.pop(channel_type, None)
|
||||
|
||||
def check_method(self, channel_type: str, method: str) -> WebhookGuardResult:
|
||||
cfg = self._channel_configs.get(channel_type)
|
||||
allowed = list(cfg.allowed_methods) if cfg else ["POST"]
|
||||
if method not in _VALID_WEBHOOK_METHODS or method not in allowed:
|
||||
self._anomaly.record(channel_type, "405")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.METHOD_NOT_ALLOWED,
|
||||
error_message=f"Method {method} not allowed",
|
||||
http_status=405,
|
||||
)
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
def check_content_type(self, channel_type: str, content_type: str | None) -> WebhookGuardResult:
|
||||
cfg = self._channel_configs.get(channel_type)
|
||||
allowed = list(cfg.allowed_content_types) if cfg else ["application/json"]
|
||||
if not content_type:
|
||||
self._anomaly.record(channel_type, "415")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.CONTENT_TYPE_UNSUPPORTED,
|
||||
error_message="Content-Type header missing",
|
||||
http_status=415,
|
||||
)
|
||||
ctype_main = content_type.split(";")[0].strip().lower()
|
||||
if ctype_main not in allowed:
|
||||
self._anomaly.record(channel_type, "415")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.CONTENT_TYPE_UNSUPPORTED,
|
||||
error_message=f"Content-Type {content_type} not supported",
|
||||
http_status=415,
|
||||
)
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
def check_body_size(self, channel_type: str, body: bytes) -> WebhookGuardResult:
|
||||
cfg = self._channel_configs.get(channel_type)
|
||||
limit = cfg.body_size_limit if cfg else PRE_AUTH_BODY_SIZE_LIMIT
|
||||
if len(body) > limit:
|
||||
self._anomaly.record(channel_type, "413")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.BODY_TOO_LARGE,
|
||||
error_message=f"Body size {len(body)} exceeds limit {limit}",
|
||||
http_status=413,
|
||||
)
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
def check_concurrency(self, channel_type: str) -> WebhookGuardResult:
|
||||
if not self._concurrency.try_acquire(channel_type):
|
||||
self._anomaly.record(channel_type, "429")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.TOO_MANY_REQUESTS,
|
||||
error_message="Too many concurrent requests",
|
||||
http_status=429,
|
||||
)
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
def release_concurrency(self, channel_type: str) -> None:
|
||||
self._concurrency.release(channel_type)
|
||||
|
||||
def check_signature(
|
||||
self,
|
||||
channel_type: str,
|
||||
body: bytes,
|
||||
signature: str | None,
|
||||
) -> WebhookGuardResult:
|
||||
cfg = self._channel_configs.get(channel_type)
|
||||
if cfg is None or not cfg.verify_signature:
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
if not signature:
|
||||
self._anomaly.record(channel_type, "401")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.SIGNATURE_MISSING,
|
||||
error_message="Webhook signature missing",
|
||||
http_status=401,
|
||||
)
|
||||
|
||||
if not self._signer.verify(channel_type, body, signature):
|
||||
self._anomaly.record(channel_type, "401")
|
||||
return WebhookGuardResult(
|
||||
allowed=False,
|
||||
error_code=WebhookErrorCode.SIGNATURE_INVALID,
|
||||
error_message="Webhook signature invalid",
|
||||
http_status=401,
|
||||
)
|
||||
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
def run_pipeline(
|
||||
self,
|
||||
channel_type: str,
|
||||
method: str,
|
||||
content_type: str | None,
|
||||
body: bytes,
|
||||
signature: str | None = None,
|
||||
) -> WebhookGuardResult:
|
||||
result = self.check_method(channel_type, method)
|
||||
if not result.allowed:
|
||||
return result
|
||||
|
||||
result = self.check_content_type(channel_type, content_type)
|
||||
if not result.allowed:
|
||||
return result
|
||||
|
||||
result = self.check_body_size(channel_type, body)
|
||||
if not result.allowed:
|
||||
return result
|
||||
|
||||
result = self.check_signature(channel_type, body, signature)
|
||||
if not result.allowed:
|
||||
return result
|
||||
|
||||
result = self.check_concurrency(channel_type)
|
||||
if not result.allowed:
|
||||
return result
|
||||
|
||||
return WebhookGuardResult(allowed=True)
|
||||
|
||||
@property
|
||||
def signer(self) -> WebhookSigner:
|
||||
return self._signer
|
||||
|
||||
@property
|
||||
def anomaly_tracker(self) -> WebhookAnomalyTracker:
|
||||
return self._anomaly
|
||||
|
||||
@property
|
||||
def concurrency_guard(self) -> WebhookConcurrencyGuard:
|
||||
return self._concurrency
|
||||
|
||||
|
||||
webhook_guard = WebhookGuard()
|
||||
Loading…
Reference in New Issue
Block a user