@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/tunnel_native_ssh.py

@@ -0,0 +1,454 @@
+"""Daemon-loop-owned native AsyncSSH TCP forwarding provider."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncGenerator, Callable
+from contextlib import suppress
+
+import asyncssh
+
+from .async_lifecycle import finish_task_before_cancellation
+from .hosts import ForwardingPolicy, HostConfig
+from .transports.base import ByteStream, ExecResult, ShellAttachment, ShellReadResult, TransferResult, TransportError
+from .transports.ssh import AsyncSshConnection
+from .tunnel_manager import TunnelBackendError
+
+DEFAULT_RECEIVE_WINDOW = 2 * 1024 * 1024
+
+
+class _ConnectionObserver(asyncssh.SSHClient):
+    def __init__(self, lost: Callable[[BaseException | None], None]) -> None:
+        self._lost = lost
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        self._lost(exc)
+
+
+class NativeSshTunnelEndpoint(asyncssh.SSHTCPSession[bytes]):
+    def __init__(
+        self,
+        provider_stream_id: int,
+        *,
+        receive_window: int,
+        closed_handler: Callable[[int], None],
+        write_stall_handler: Callable[[], None],
+    ) -> None:
+        if receive_window <= 0:
+            raise ValueError("receive_window must be positive")
+        self.provider_stream_id = provider_stream_id
+        self._receive_window = receive_window
+        self._high_water = max(1, receive_window * 3 // 4)
+        self._low_water = receive_window // 2
+        self._closed_handler = closed_handler
+        self._write_stall_handler = write_stall_handler
+        self._channel: asyncssh.SSHTCPChannel[bytes] | None = None
+        self._buffer = bytearray()
+        self._read_ready = asyncio.Event()
+        self._write_ready = asyncio.Event()
+        self._write_ready.set()
+        self._remote_eof = False
+        self._local_eof = False
+        self._reading_paused = False
+        self._error: TunnelBackendError | None = None
+        self._closed = False
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    @property
+    def queued_bytes(self) -> int:
+        return len(self._buffer)
+
+    def connection_made(self, chan: asyncssh.SSHTCPChannel[bytes]) -> None:
+        self._channel = chan
+
+    def session_started(self) -> None:
+        return None
+
+    def data_received(self, data: bytes, datatype: asyncssh.DataType) -> None:  # noqa: ARG002
+        if self._closed or not data:
+            return
+        payload = bytes(data)
+        if len(self._buffer) + len(payload) > self._receive_window:
+            self._buffer.clear()
+            error = TunnelBackendError("native SSH stream exceeded its receive window")
+            self._error = error
+            if self._channel is not None:
+                self._channel.abort()
+            self._finish()
+            return
+        self._buffer.extend(payload)
+        if len(self._buffer) >= self._high_water and not self._reading_paused and self._channel is not None:
+            self._channel.pause_reading()
+            self._reading_paused = True
+        self._read_ready.set()
+
+    def eof_received(self) -> bool:
+        self._remote_eof = True
+        self._read_ready.set()
+        return True
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        if exc is not None and self._error is None:
+            self._error = TunnelBackendError(str(exc) or "native SSH stream connection lost")
+        self._remote_eof = True
+        self._finish()
+
+    def pause_writing(self) -> None:
+        if self._write_ready.is_set():
+            self._write_stall_handler()
+        self._write_ready.clear()
+
+    def resume_writing(self) -> None:
+        self._write_ready.set()
+
+    async def send(self, data: bytes) -> None:
+        if not isinstance(data, bytes):
+            raise TypeError("tunnel stream data must be bytes")
+        if not data:
+            return
+        if self._local_eof:
+            raise TunnelBackendError("native SSH stream input is closed")
+        self._raise_if_failed()
+        await self._write_ready.wait()
+        self._raise_if_failed()
+        channel = self._require_channel()
+        channel.write(data)
+        await asyncio.sleep(0)
+
+    async def send_eof(self) -> None:
+        if self._local_eof or self._closed:
+            return
+        self._raise_if_failed()
+        self._local_eof = True
+        self._require_channel().write_eof()
+
+    async def receive(self, max_bytes: int = 64 * 1024) -> bytes | None:
+        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes <= 0:
+            raise ValueError("max_bytes must be a positive integer")
+        while not (self._buffer or self._remote_eof or self._error is not None or self._closed):
+            self._read_ready.clear()
+            await self._read_ready.wait()
+        if self._error is not None:
+            raise self._error
+        if self._buffer:
+            size = min(max_bytes, len(self._buffer))
+            data = bytes(self._buffer[:size])
+            del self._buffer[:size]
+            self._resume_reading_if_ready()
+            return data
+        return None
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        if self._channel is not None:
+            self._channel.close()
+        self._remote_eof = True
+        self._finish()
+
+    async def reset(self, reason: str) -> None:
+        if self._closed:
+            return
+        self._error = TunnelBackendError(reason or "native SSH stream reset")
+        if self._channel is not None:
+            self._channel.abort()
+        self._remote_eof = True
+        self._finish()
+
+    def backend_lost(self, error: BaseException | None) -> None:
+        if self._closed:
+            return
+        if error is not None:
+            self._error = TunnelBackendError(str(error) or "native SSH connection lost")
+        if self._channel is not None:
+            self._channel.abort()
+        self._remote_eof = True
+        self._finish()
+
+    def _require_channel(self) -> asyncssh.SSHTCPChannel[bytes]:
+        if self._channel is None:
+            raise TunnelBackendError("native SSH stream is not open")
+        return self._channel
+
+    def _raise_if_failed(self) -> None:
+        if self._error is not None:
+            raise self._error
+        if self._closed:
+            raise TunnelBackendError("native SSH stream is closed")
+
+    def _resume_reading_if_ready(self) -> None:
+        if self._reading_paused and len(self._buffer) <= self._low_water and self._channel is not None:
+            self._channel.resume_reading()
+            self._reading_paused = False
+
+    def _finish(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._write_ready.set()
+        self._read_ready.set()
+        self._closed_handler(self.provider_stream_id)
+
+
+class NativeSshHostRuntime:
+    provider_name = "native_ssh"
+    capabilities = AsyncSshConnection.capabilities
+
+    @property
+    def agent_pid(self) -> int | None:
+        return None
+
+    def __init__(
+        self,
+        host: HostConfig,
+        *,
+        generation: int,
+        connect_timeout: float = 30,
+        receive_window: int = DEFAULT_RECEIVE_WINDOW,
+    ) -> None:
+        if generation <= 0:
+            raise ValueError("generation must be positive")
+        if connect_timeout <= 0:
+            raise ValueError("connect_timeout must be positive")
+        if receive_window <= 0:
+            raise ValueError("receive_window must be positive")
+        self.host = host
+        self.generation = generation
+        self.connect_timeout = float(connect_timeout)
+        self.receive_window = receive_window
+        self._connection: asyncssh.SSHClientConnection | None = None
+        self._operations: AsyncSshConnection | None = None
+        self._closed_future: asyncio.Future[BaseException | None] | None = None
+        self._endpoints: dict[int, NativeSshTunnelEndpoint] = {}
+        self._next_stream_id = 1
+        self._closing = False
+        self._close_lock = asyncio.Lock()
+        self._window_stalls = 0
+
+    @property
+    def queued_bytes(self) -> int:
+        return sum(endpoint.queued_bytes for endpoint in self._endpoints.values())
+
+    @property
+    def window_stalls(self) -> int:
+        return self._window_stalls
+
+    @property
+    def protocol_failures(self) -> int:
+        return 0
+
+    def _record_window_stall(self) -> None:
+        self._window_stalls += 1
+
+    async def start(self) -> None:
+        if self._connection is not None:
+            return
+        ssh = self.host.ssh
+        if ssh is None:
+            raise TunnelBackendError("native SSH tunnel requires an SSH host configuration")
+        self._closed_future = asyncio.get_running_loop().create_future()
+        observer = _ConnectionObserver(self._connection_lost)
+        options: dict[str, object] = {
+            "port": ssh.port,
+            "username": ssh.username,
+            "tunnel": ssh.proxy_jump,
+            "password": self.host.secrets.get("password"),
+            "connect_timeout": self.connect_timeout,
+            "encoding": None,
+            "client_factory": lambda: observer,
+        }
+        if ssh.identity_file is not None:
+            options["client_keys"] = [str(ssh.identity_file)]
+        try:
+            self._connection = await asyncssh.connect(ssh.host, **options)
+            self._operations = AsyncSshConnection(self._connection)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            raise TunnelBackendError(f"native SSH connection failed: {exc}") from exc
+
+    async def exec(
+        self,
+        command: str,
+        *,
+        stdin: ByteStream | None,
+        timeout: float,
+        output_limit: int,
+    ) -> ExecResult:
+        return await self._require_operations().exec(
+            command,
+            stdin=stdin,
+            timeout=timeout,
+            output_limit=output_limit,
+        )
+
+    async def upload(
+        self,
+        chunks: ByteStream,
+        remote_path: str,
+        *,
+        mode: int,
+        expected_size: int | None,
+        expected_sha256: str | None,
+    ) -> TransferResult:
+        try:
+            return await self._require_operations().upload(
+                chunks,
+                remote_path,
+                mode=mode,
+                expected_size=expected_size,
+                expected_sha256=expected_sha256,
+            )
+        except TransportError as exc:
+            raise TunnelBackendError(str(exc)) from exc
+
+    async def download(self, remote_path: str, *, chunk_size: int) -> AsyncGenerator[bytes, None]:
+        async for chunk in self._require_operations().download(remote_path, chunk_size=chunk_size):
+            yield chunk
+
+    async def shell_write(self, shell_id: str, data: bytes) -> None:
+        await self._require_operations().shell_write(shell_id, data)
+
+    async def shell_read(self, shell_id: str, *, timeout: float | None, max_bytes: int) -> ShellReadResult:
+        return await self._require_operations().shell_read(shell_id, timeout=timeout, max_bytes=max_bytes)
+
+    async def attach_shell(self, shell_id: str, stdin: ByteStream, *, max_bytes: int) -> ShellAttachment:
+        return await self._require_operations().attach_shell(shell_id, stdin, max_bytes=max_bytes)
+
+    async def shell_close(self, shell_id: str) -> None:
+        await self._require_operations().shell_close(shell_id)
+
+    async def open_stream(
+        self,
+        host: str,
+        port: int,
+        policy: ForwardingPolicy,
+    ) -> NativeSshTunnelEndpoint:
+        if self._closing or self._connection is None or self._backend_closed():
+            raise TunnelBackendError("native SSH runtime is not connected")
+        if policy != self.host.forwarding:
+            raise TunnelBackendError("native SSH stream policy does not match the host runtime")
+        stream_id = self._next_stream_id
+        self._next_stream_id += 1
+        endpoint = NativeSshTunnelEndpoint(
+            stream_id,
+            receive_window=self.receive_window,
+            closed_handler=self._endpoint_closed,
+            write_stall_handler=self._record_window_stall,
+        )
+        try:
+            _, session = await self._connection.create_connection(
+                lambda: endpoint,
+                host,
+                port,
+                orig_host="",
+                orig_port=0,
+                encoding=None,
+            )
+        except asyncio.CancelledError:
+            await endpoint.reset("native SSH stream open cancelled")
+            raise
+        except Exception as exc:
+            await endpoint.reset(str(exc) or "native SSH stream open failed")
+            raise TunnelBackendError(f"native SSH destination open failed: {exc}") from exc
+        if session is not endpoint:
+            await endpoint.reset("native SSH returned an unexpected stream session")
+            raise TunnelBackendError("native SSH returned an unexpected stream session")
+        if self._closing or self._backend_closed():
+            message = self._backend_failure_message()
+            await endpoint.reset(message)
+            raise TunnelBackendError(message)
+        self._endpoints[stream_id] = endpoint
+        return endpoint
+
+    async def wait_closed(self) -> BaseException | None:
+        if self._closed_future is None:
+            raise TunnelBackendError("native SSH runtime has not started")
+        return await asyncio.shield(self._closed_future)
+
+    async def close(self) -> None:
+        async with self._close_lock:
+            self._closing = True
+            cleanup = asyncio.create_task(self._finish_close())
+            await finish_task_before_cancellation(cleanup)
+
+    def abort(self) -> None:
+        self._closing = True
+        error = TunnelBackendError("native SSH runtime cleanup timed out")
+        for endpoint in list(self._endpoints.values()):
+            with suppress(Exception):
+                endpoint.backend_lost(error)
+        self._endpoints.clear()
+        connection = self._connection
+        self._connection = None
+        self._operations = None
+        abort_error: Exception | None = None
+        if connection is not None:
+            try:
+                connection.abort()
+            except Exception as exc:
+                try:
+                    connection.close()
+                except Exception as fallback_error:
+                    fallback_error.add_note(f"connection abort failed: {exc}")
+                    abort_error = fallback_error
+        if self._closed_future is not None and not self._closed_future.done():
+            with suppress(Exception):
+                self._closed_future.set_result(error)
+        if abort_error is not None:
+            raise abort_error
+
+    async def _finish_close(self) -> None:
+        endpoints = list(self._endpoints.values())
+        await asyncio.gather(*(endpoint.close() for endpoint in endpoints), return_exceptions=True)
+        for endpoint in endpoints:
+            self._endpoints.pop(endpoint.provider_stream_id, None)
+        connection = self._connection
+        self._operations = None
+        if connection is not None:
+            connection.close()
+            with suppress(Exception):
+                await connection.wait_closed()
+            if self._connection is connection:
+                self._connection = None
+        if self._closed_future is not None and not self._closed_future.done():
+            self._closed_future.set_result(None)
+
+    def _connection_lost(self, error: BaseException | None) -> None:
+        if error is None and not self._closing:
+            error = TunnelBackendError("native SSH connection closed")
+        for endpoint in list(self._endpoints.values()):
+            endpoint.backend_lost(error)
+        self._endpoints.clear()
+        if self._closed_future is not None and not self._closed_future.done():
+            self._closed_future.set_result(error)
+
+    def _endpoint_closed(self, stream_id: int) -> None:
+        self._endpoints.pop(stream_id, None)
+
+    def _require_operations(self) -> AsyncSshConnection:
+        if self._closing or self._backend_closed() or self._operations is None:
+            raise TunnelBackendError("native SSH runtime is not connected")
+        return self._operations
+
+    def _backend_closed(self) -> bool:
+        return self._closed_future is not None and self._closed_future.done()
+
+    def _backend_failure_message(self) -> str:
+        if self._closed_future is None or not self._closed_future.done():
+            return "native SSH connection closed"
+        error = self._closed_future.result()
+        return str(error) if error is not None else "native SSH connection closed"
+
+
+def native_ssh_provider_factory(
+    host: HostConfig,
+    provider_name: str,
+    generation: int,
+) -> NativeSshHostRuntime:
+    if provider_name != "native_ssh":
+        raise TunnelBackendError(f"native SSH factory cannot create provider {provider_name!r}")
+    return NativeSshHostRuntime(host, generation=generation)

package/src/server_control_mcp/tunnel_providers.py

@@ -0,0 +1,20 @@
+"""Strict tunnel provider registry owned by the HostBridge daemon."""
+
+from __future__ import annotations
+
+from .hosts import HostConfig
+from .tunnel_manager import TunnelBackendError, TunnelProviderRuntime
+from .tunnel_native_ssh import native_ssh_provider_factory
+from .tunnel_pty_agent import pty_agent_provider_factory
+
+
+def tunnel_provider_factory(
+    host: HostConfig,
+    provider_name: str,
+    generation: int,
+) -> TunnelProviderRuntime:
+    if provider_name == "native_ssh":
+        return native_ssh_provider_factory(host, provider_name, generation)
+    if provider_name == "pty_agent":
+        return pty_agent_provider_factory(host, provider_name, generation)
+    raise TunnelBackendError(f"unsupported tunnel provider {provider_name!r}")