"""Nexus — asyncssh Connection Pool (Reference Counting Mode) ADR-002: BFF mode with asyncssh embedded (single process, full Python). ADR-004: Reference counting connection pool — reuse connections, reduce handshakes. W1: Extracted from ssh_direct.py, provides asyncssh-based connection pooling. W2: Replaces paramiko pool.py with asyncssh throughout the codebase. """ import asyncio import logging import secrets import shlex import time from contextlib import asynccontextmanager from dataclasses import dataclass, field from typing import AsyncIterator, Dict, Optional import asyncssh from server.domain.models import Server from server.utils.files_elevation import FilesElevation from server.utils.posix_paths import posix_dirname from server.infrastructure.database.crypto import decrypt_value from server.utils.sync_error_message import translate_ssh_error_message logger = logging.getLogger("nexus.asyncssh_pool") def _localized_ssh_error(exc: Exception, *, max_len: int = 200) -> str: raw = str(exc)[:max_len] return translate_ssh_error_message(raw) or raw async def try_ssh_login( host: str, port: int, username: str, *, password: str | None = None, private_key: str | None = None, timeout: float = 8.0, ) -> tuple[bool, str]: """One-shot SSH login probe (not pooled). Returns (ok, error_message).""" if not password and not private_key: return False, "未提供密码或私钥" connect_kwargs: dict = { "host": host, "port": port or 22, "username": username or "root", "connect_timeout": timeout, # 凭据轮询探测未知主机:跳过主机密钥校验(运维平台标准;持久连接见 SSH_STRICT_HOST_CHECKING) "known_hosts": None, } if password: connect_kwargs["password"] = password elif private_key: connect_kwargs["client_keys"] = [asyncssh.import_private_key(private_key)] conn: asyncssh.SSHClientConnection | None = None try: conn = await asyncssh.connect(**connect_kwargs) return True, "" except asyncssh.Error as e: return False, _localized_ssh_error(e) except Exception as e: return False, _localized_ssh_error(e) finally: if conn is not None: try: conn.close() except Exception: # noqa: S110 — best-effort cleanup pass @dataclass class PooledConnection: """A tracked asyncssh connection with reference counting""" conn: asyncssh.SSHClientConnection ref_count: int = 0 created_at: float = field(default_factory=time.monotonic) last_used: float = field(default_factory=time.monotonic) server_id: int = 0 @property def is_idle(self) -> bool: return self.ref_count == 0 @property def idle_seconds(self) -> float: return time.monotonic() - self.last_used if self.is_idle else 0 class AsyncSSHPool: """asyncssh connection pool with reference counting (ADR-004) Features: - Reference counting: multiple operations can share one SSH connection - Auto-cleanup: idle connections closed after IDLE_TIMEOUT - Per-server pooling: one connection per server at a time - Thread-safe: uses asyncio.Lock for concurrent access """ # Pool configuration MAX_CONNECTIONS = 100 # Global max connections IDLE_TIMEOUT = 300 # 5 minutes idle → close CONNECT_TIMEOUT = 10 # 10 seconds to establish connection CLEANUP_INTERVAL = 60 # Check for idle connections every 60s def __init__(self): self._pool: Dict[int, PooledConnection] = {} # server_id → connection self._lock = asyncio.Lock() self._cleanup_task: Optional[asyncio.Task] = None async def start(self): """Start the connection pool and cleanup task""" self._cleanup_task = asyncio.create_task(self._cleanup_loop(), name="asyncssh_cleanup") logger.info(f"AsyncSSH pool started (max={self.MAX_CONNECTIONS}, idle_timeout={self.IDLE_TIMEOUT}s)") async def stop(self): """Stop the pool and close all connections""" if self._cleanup_task: self._cleanup_task.cancel() try: await self._cleanup_task except asyncio.CancelledError: pass async with self._lock: for _server_id, pooled in self._pool.items(): try: pooled.conn.close() except Exception: # noqa: S110 — best-effort cleanup pass self._pool.clear() logger.info("AsyncSSH pool stopped — all connections closed") async def acquire(self, server: Server) -> asyncssh.SSHClientConnection: """Acquire a connection to a server (reference counting). If an idle connection exists for this server, reuse it. Otherwise, create a new connection. SSH handshake is done outside the lock to avoid blocking other pool operations. """ async with self._lock: # Check for existing idle connection pooled = self._pool.get(server.id) if pooled and not pooled.conn.is_closed(): pooled.ref_count += 1 pooled.last_used = time.monotonic() logger.debug(f"Reused SSH connection: server={server.id}, refs={pooled.ref_count}") return pooled.conn # Need to create — evict if at capacity if len(self._pool) >= self.MAX_CONNECTIONS: evicted = await self._evict_one() if not evicted: # All connections are busy — cannot exceed MAX_CONNECTIONS raise ConnectionError( f"SSH pool at capacity ({self.MAX_CONNECTIONS}) and all connections busy. " f"Try again later." ) # Create connection outside lock (handshake can take seconds) conn = await self._create_connection(server) async with self._lock: # Re-check: another coroutine may have created one while we were unlocked pooled = self._pool.get(server.id) if pooled and not pooled.conn.is_closed(): # Close our new connection, reuse the existing one try: conn.close() except Exception: # noqa: S110 — best-effort cleanup pass pooled.ref_count += 1 pooled.last_used = time.monotonic() return pooled.conn pooled = PooledConnection( conn=conn, ref_count=1, server_id=server.id, ) self._pool[server.id] = pooled logger.debug(f"New SSH connection: server={server.id}, pool_size={len(self._pool)}") return conn async def release(self, server_id: int): """Release a connection (decrement reference count). If ref_count reaches 0, the connection stays in pool for reuse. It will be cleaned up by the idle timeout task. """ async with self._lock: pooled = self._pool.get(server_id) if pooled: pooled.ref_count = max(0, pooled.ref_count - 1) pooled.last_used = time.monotonic() logger.debug(f"Released SSH connection: server={server_id}, refs={pooled.ref_count}") async def close_connection(self, server_id: int): """Force-close a connection to a specific server""" async with self._lock: pooled = self._pool.pop(server_id, None) if pooled: try: pooled.conn.close() except Exception: # noqa: S110 — best-effort cleanup pass logger.debug(f"Force-closed SSH connection: server={server_id}") async def _create_connection(self, server: Server) -> asyncssh.SSHClientConnection: """Create a new asyncssh connection to a server""" from server.config import settings connect_kwargs = { "host": server.domain, "port": server.port or 22, "username": server.username or "root", "connect_timeout": self.CONNECT_TIMEOUT, } # Host key verification: respect SSH_STRICT_HOST_CHECKING setting # Default is "true" — reject unknown hosts for security if settings.SSH_STRICT_HOST_CHECKING.lower() in ("false", "0", "no"): connect_kwargs["known_hosts"] = None # Skip verification (insecure) # If strict mode, asyncssh will use default known_hosts behavior # (rejects unknown hosts — caller should configure ~/.ssh/known_hosts) if server.auth_method == "password" and server.password: password = decrypt_value(server.password) connect_kwargs["password"] = password elif server.ssh_key_configured and server.ssh_key_private: key_content = decrypt_value(server.ssh_key_private) connect_kwargs["client_keys"] = [asyncssh.import_private_key(key_content)] elif server.ssh_key_path: connect_kwargs["client_keys"] = [server.ssh_key_path] else: raise ValueError(f"服务器 {server.id}({server.name})没有可用的 SSH 凭据") try: conn = await asyncssh.connect(**connect_kwargs) return conn except asyncssh.Error as e: msg = _localized_ssh_error(e) raise ConnectionError( f"SSH 连接失败 {server.domain}:{server.port or 22}:{msg}" ) from e async def _evict_one(self) -> bool: """Evict the oldest idle connection from the pool. Returns True if an idle connection was evicted, False if all are busy. """ oldest_idle = None oldest_time = float("inf") for server_id, pooled in self._pool.items(): if pooled.is_idle and pooled.last_used < oldest_time: oldest_idle = server_id oldest_time = pooled.last_used if oldest_idle is not None: pooled = self._pool.pop(oldest_idle) try: pooled.conn.close() except Exception: # noqa: S110 — best-effort cleanup pass logger.debug(f"Evicted idle SSH connection: server={oldest_idle}") return True return False async def _cleanup_loop(self): """Periodically clean up idle connections""" while True: await asyncio.sleep(self.CLEANUP_INTERVAL) try: async with self._lock: to_remove = [] for server_id, pooled in self._pool.items(): if pooled.is_idle and pooled.idle_seconds > self.IDLE_TIMEOUT: to_remove.append(server_id) elif not pooled.is_idle and pooled.conn.is_closed(): # Connection died while in use — remove from pool to_remove.append(server_id) for server_id in to_remove: pooled = self._pool.pop(server_id) try: pooled.conn.close() except Exception: # noqa: S110 — best-effort cleanup pass if to_remove: logger.info(f"SSH pool cleanup: removed {len(to_remove)} idle/dead connections") except asyncio.CancelledError: break except Exception as e: logger.error(f"SSH pool cleanup error: {e}") @property def stats(self) -> dict: """Get pool statistics""" total = len(self._pool) active = sum(1 for p in self._pool.values() if not p.is_idle) idle = total - active return {"total": total, "active": active, "idle": idle} # Global pool instance ssh_pool = AsyncSSHPool() @asynccontextmanager async def sftp_session(conn: asyncssh.SSHClientConnection) -> AsyncIterator[asyncssh.SFTPClient]: """Open an SFTP client on a pooled SSH connection (asyncssh 2.x: start_sftp_client).""" async with await conn.start_sftp_client() as sftp: yield sftp async def sftp_write_bytes(sftp: asyncssh.SFTPClient, remote_path: str, data: bytes) -> None: """Write in-memory bytes to a remote file (asyncssh put() only accepts local paths).""" async with sftp.open(remote_path, "wb") as remote_file: await remote_file.write(data) def _sftp_permission_denied(exc: Exception) -> bool: if isinstance(exc, asyncssh.PermissionDenied): return True if isinstance(exc, asyncssh.SFTPError): # SSH_FX_PERMISSION_DENIED = 3 if getattr(exc, "code", None) == 3: return True return "permission denied" in str(exc).lower() async def _run_tee_stdin( conn: asyncssh.SSHClientConnection, command: str, data: bytes, ) -> asyncssh.SSHCompletedProcess: """Feed raw bytes to stdin; default SSH session encoding expects str, not bytes.""" return await conn.run(command, input=data, encoding=None, check=False) async def _ensure_remote_parent_dir( conn: asyncssh.SSHClientConnection, remote_path: str, *, use_sudo: bool, ) -> None: parent = posix_dirname(remote_path) if not parent or parent == "/": return prefix = "sudo -n " if use_sudo else "" await conn.run(f"{prefix}mkdir -p {shlex.quote(parent)}", check=False) async def _ssh_write_bytes_via_tee( conn: asyncssh.SSHClientConnection, remote_path: str, data: bytes, *, use_sudo: bool, ) -> None: """Shell write fallback when SFTP cannot open the file (e.g. file not writable but directory is).""" prefix = "sudo -n " if use_sudo else "" await _ensure_remote_parent_dir(conn, remote_path, use_sudo=use_sudo) direct = await _run_tee_stdin( conn, f"{prefix}tee {shlex.quote(remote_path)}", data, ) if direct.exit_status == 0: return tmp_path = f"/tmp/nexus-write-{secrets.token_hex(12)}" staged = await _run_tee_stdin( conn, f"{prefix}tee {shlex.quote(tmp_path)}", data, ) if staged.exit_status != 0: stderr = (staged.stderr or staged.stdout or direct.stderr or "")[:500] raise PermissionError(stderr or "tee 写入失败") moved = await conn.run( f"{prefix}mv -f {shlex.quote(tmp_path)} {shlex.quote(remote_path)}", check=False, ) if moved.exit_status != 0: await conn.run(f"{prefix}rm -f {shlex.quote(tmp_path)}", check=False) stderr = (moved.stderr or "")[:500] raise PermissionError(stderr or "无法覆盖目标文件(目录可能不可写)") def _sudo_attempt_order(elevation: FilesElevation | None) -> tuple[bool, ...]: mode = elevation if elevation is not None else FilesElevation.AUTO if mode == FilesElevation.OFF: return (False,) if mode == FilesElevation.ALWAYS: return (True,) return (False, True) async def remote_write_bytes( conn: asyncssh.SSHClientConnection, remote_path: str, data: bytes, *, elevation: FilesElevation | None = None, ) -> None: """Write bytes to a remote path: SFTP first, then SSH tee / tmp+mv (incl. optional sudo -n).""" eff = elevation if elevation is not None else FilesElevation.AUTO skip_sftp = eff == FilesElevation.ALWAYS if not skip_sftp: try: parent = posix_dirname(remote_path) if parent and parent != "/": await _ensure_remote_parent_dir(conn, remote_path, use_sudo=False) async with await conn.start_sftp_client() as sftp: await sftp_write_bytes(sftp, remote_path, data) return except (asyncssh.PermissionDenied, asyncssh.SFTPError) as exc: if not _sftp_permission_denied(exc): raise last_error: Exception | None = None for use_sudo in _sudo_attempt_order(eff): try: await _ssh_write_bytes_via_tee(conn, remote_path, data, use_sudo=use_sudo) return except PermissionError as exc: last_error = exc reason = str(last_error) if last_error else "Permission denied" raise asyncssh.SFTPError(3, reason) async def exec_ssh_command(server: Server, command: str, timeout: int = 30) -> dict: """Execute a command via SSH using the asyncssh connection pool. W2: Replaces paramiko-based exec_command from pool.py. Returns structured result dict with stdout/stderr/exit_code. """ conn = None try: conn = await ssh_pool.acquire(server) result = await asyncio.wait_for( conn.run(command, timeout=timeout), timeout=timeout + 5, ) await ssh_pool.release(server.id) return { "status": "success" if result.exit_status == 0 else "failed", "stdout": result.stdout[:10000], "stderr": result.stderr[:10000], "exit_code": result.exit_status, } except asyncio.TimeoutError: # Timed-out connection may be in an unknown state — evict so it is NOT # returned to the pool and potentially reused for a future command. if conn is not None: await ssh_pool.close_connection(server.id) return { "status": "timeout", "stdout": "", "stderr": f"命令执行超时({timeout} 秒)", "exit_code": -1, } except ConnectionError as e: if conn is not None: await ssh_pool.close_connection(server.id) return { "status": "connection_error", "stdout": "", "stderr": str(e), "exit_code": -1, } except Exception as e: if conn is not None: await ssh_pool.release(server.id) return { "status": "error", "stdout": "", "stderr": str(e), "exit_code": -1, }