@neoline/hostbridge 2.0.2 → 2.0.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@neoline/hostbridge",
-  "version": "2.0.2",
+  "version": "2.0.3",
   "description": "HostBridge MCP server for persistent allowlisted SSH and shell sessions.",
   "type": "module",
   "license": "MIT",

package/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hostbridge"
-version = "2.0.2"
+version = "2.0.3"
 description = "HostBridge MCP server for persistent allowlisted SSH and shell sessions."
 readme = "README.md"
 requires-python = ">=3.11"

package/src/server_control_mcp/__init__.py

@@ -8,7 +8,7 @@ import sys
 from typing import NoReturn
 
 __all__ = ["__version__", "ensure_runtime"]
-__version__ = "2.0.2"
+__version__ = "2.0.3"
 
 # (import_name, pip_name) pairs for runtime dependency probing.
 _REQUIRED_DEPENDENCIES = (

package/src/server_control_mcp/doctor.py

@@ -60,7 +60,7 @@ def _version_diagnostic() -> DiagnosticResult:
         "npm": npm.get("version"),
         "python": pyproject.get("project", {}).get("version"),
     }
-    status = "pass" if len(set(versions.values())) == 1 and __version__ == "2.0.2" else "fail"
+    status = "pass" if len(set(versions.values())) == 1 and __version__ == "2.0.3" else "fail"
     return DiagnosticResult("version", status, "version sources agree" if status == "pass" else "version drift detected", versions)
 
 

package/src/server_control_mcp/mock_ssh.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import asyncio
+import base64
+import binascii
 import json
 import posixpath
 import re
@@ -9,6 +11,9 @@ import sys
 import tempfile
 import threading
 import time
+import uuid
+from collections import deque
+from collections.abc import Awaitable
 from contextlib import suppress
 from dataclasses import dataclass, field
 from functools import partial
@@ -34,6 +39,61 @@ STDIN_READING_COMMANDS = [
     re.compile(r"(?:^|\s)scp(?:\s+\S+)*\s+-[A-Za-z]*(?:t|f)[A-Za-z]*(?:\s|$)"),
 ]
 MAX_IDLE_BACKEND_SESSIONS = 4
+TUNNEL_FRAME_BYTES = 2048
+TUNNEL_READ_BYTES = 32 * 1024
+TUNNEL_MAX_BUFFER_BYTES = 128 * 1024
+
+_REMOTE_TUNNEL_SCRIPT = r"""
+import base64
+import os
+import socket
+import sys
+import threading
+
+host, port, prefix = sys.argv[1], int(sys.argv[2]), sys.argv[3]
+write_lock = threading.Lock()
+
+def emit(kind, payload=b""):
+    suffix = ":" + base64.b64encode(payload).decode("ascii") if payload else ""
+    with write_lock:
+        print(prefix + ":" + kind + suffix, flush=True)
+
+try:
+    connection = socket.create_connection((host, port), timeout=15)
+    connection.settimeout(None)
+except Exception as exc:
+    emit("ERROR", str(exc).encode("utf-8", errors="replace"))
+    raise SystemExit(1)
+
+emit("READY")
+
+def receive():
+    try:
+        while True:
+            data = connection.recv(2048)
+            if not data:
+                emit("EOF")
+                os._exit(0)
+            emit("DATA", data)
+    except Exception as exc:
+        emit("ERROR", str(exc).encode("utf-8", errors="replace"))
+        os._exit(1)
+
+receiver = threading.Thread(target=receive, daemon=True)
+receiver.start()
+try:
+    for raw_line in sys.stdin.buffer:
+        line = raw_line.strip()
+        if line == (prefix + ":EOF").encode("ascii"):
+            connection.shutdown(socket.SHUT_WR)
+            receiver.join()
+            break
+        data_prefix = (prefix + ":DATA:").encode("ascii")
+        if line.startswith(data_prefix):
+            connection.sendall(base64.b64decode(line[len(data_prefix):], validate=True))
+finally:
+    connection.close()
+"""
 
 
 def _to_bytes(chunk: str | bytes) -> bytes:
@@ -47,8 +107,9 @@ def _command_reads_stdin(command: str | None) -> bool:
 
 
 class HostBridgeSSHServer(asyncssh.SSHServer):
-    def __init__(self, authorized_public_key: str) -> None:
+    def __init__(self, authorized_public_key: str, tcp_forwarder: HostBridgeTCPForwarder | None = None) -> None:
         self._authorized_public_key = authorized_public_key.strip()
+        self._tcp_forwarder = tcp_forwarder
 
     def begin_auth(self, username: str) -> bool:  # noqa: ARG002
         return True
@@ -62,6 +123,21 @@ class HostBridgeSSHServer(asyncssh.SSHServer):
     def validate_public_key(self, username: str, key: asyncssh.SSHKey) -> bool:  # noqa: ARG002
         return key.export_public_key().decode("utf-8").strip() == self._authorized_public_key
 
+    def connection_requested(
+        self,
+        dest_host: str,
+        dest_port: int,
+        orig_host: str,  # noqa: ARG002
+        orig_port: int,  # noqa: ARG002
+    ) -> Awaitable[asyncssh.SSHTCPSession[bytes]] | bool:
+        if self._tcp_forwarder is None:
+            return False
+        try:
+            self._tcp_forwarder.validate_destination(dest_host, dest_port)
+        except ValueError as exc:
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
+        return self._tcp_forwarder.connect(dest_host, dest_port)
+
 
 class HostBridgeSSHProcess:
     def __init__(self, client: HostBridgeClient, session_factory: HostBridgeSessionFactory, timeout: int, output_limit: int) -> None:
@@ -305,6 +381,287 @@ class HostBridgeSessionFactory:
         )
 
 
+class _TunnelFrameReader:
+    def __init__(self, client: HostBridgeClient, session_id: str, prefix: str) -> None:
+        self._client = client
+        self._session_id = session_id
+        self._prefix = prefix.encode("ascii") + b":"
+        self._buffer = bytearray()
+        self._pending: deque[tuple[str, bytes]] = deque()
+
+    async def next(self) -> tuple[str, bytes]:
+        while not self._pending:
+            result = await asyncio.to_thread(
+                self._client.shell_read,
+                self._session_id,
+                timeout=0.2,
+                max_bytes=4096,
+                owner="mock-ssh",
+            )
+            if result.data:
+                self._buffer.extend(result.data)
+                self._parse_lines()
+            if not result.alive and not self._pending:
+                raise HostBridgeError("connection_lost", "remote TCP tunnel closed unexpectedly", retryable=True)
+        return self._pending.popleft()
+
+    def _parse_lines(self) -> None:
+        while b"\n" in self._buffer:
+            raw_line, _, remainder = self._buffer.partition(b"\n")
+            self._buffer = bytearray(remainder)
+            line = raw_line.rstrip(b"\r")
+            marker = line.find(self._prefix)
+            if marker < 0:
+                continue
+            frame = line[marker + len(self._prefix) :]
+            kind, separator, encoded = frame.partition(b":")
+            try:
+                kind_text = kind.decode("ascii")
+            except UnicodeDecodeError:
+                continue
+            if kind_text not in {"READY", "DATA", "EOF", "ERROR", "CLOSED"}:
+                continue
+            payload = b""
+            if separator:
+                try:
+                    payload = base64.b64decode(encoded, validate=True)
+                except (binascii.Error, ValueError) as exc:
+                    raise HostBridgeError("protocol_mismatch", "remote TCP tunnel returned invalid base64") from exc
+            self._pending.append((kind_text, payload))
+        if len(self._buffer) > TUNNEL_MAX_BUFFER_BYTES:
+            marker = self._buffer.rfind(self._prefix)
+            self._buffer = self._buffer[marker:] if marker >= 0 else bytearray()
+
+
+class HostBridgeTCPForwarder:
+    def __init__(
+        self,
+        client: HostBridgeClient,
+        session_factory: HostBridgeSessionFactory,
+        *,
+        connect_timeout: int,
+    ) -> None:
+        self._client = client
+        self._session_factory = session_factory
+        self._connect_timeout = connect_timeout
+
+    @staticmethod
+    def validate_destination(dest_host: str, dest_port: int) -> None:
+        if not isinstance(dest_host, str) or not dest_host.strip() or "\x00" in dest_host or len(dest_host) > 253:
+            raise ValueError("TCP forwarding destination host is invalid")
+        if not isinstance(dest_port, int) or isinstance(dest_port, bool) or not 1 <= dest_port <= 65535:
+            raise ValueError("TCP forwarding destination port must be between 1 and 65535")
+
+    async def connect(self, dest_host: str, dest_port: int) -> asyncssh.SSHTCPSession[bytes]:
+        self.validate_destination(dest_host, dest_port)
+        session_id = await asyncio.to_thread(self._session_factory.open)
+        prefix = f"__HB_TUNNEL_{uuid.uuid4().hex}__"
+        frames = _TunnelFrameReader(self._client, session_id, prefix)
+        try:
+            await self._start_remote(session_id, dest_host, dest_port, prefix)
+            kind, payload = await asyncio.wait_for(frames.next(), timeout=self._connect_timeout)
+            if kind == "ERROR":
+                raise HostBridgeError("connection_lost", payload.decode("utf-8", errors="replace"), retryable=True)
+            if kind != "READY":
+                raise HostBridgeError("protocol_mismatch", f"remote TCP tunnel returned {kind} before READY")
+        except asyncio.CancelledError:
+            await asyncio.to_thread(self._session_factory.close, session_id, reusable=False)
+            raise
+        except Exception as exc:
+            await asyncio.to_thread(self._session_factory.close, session_id, reusable=False)
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
+        return HostBridgeTCPSession(
+            self,
+            session_id=session_id,
+            prefix=prefix,
+            frames=frames,
+            dest_host=dest_host,
+            dest_port=dest_port,
+        )
+
+    async def _start_remote(self, session_id: str, dest_host: str, dest_port: int, prefix: str) -> None:
+        command = (
+            f"__hb_prefix={shlex.quote(prefix)}; stty -echo; "
+            f"python3 -u -c {shlex.quote(_REMOTE_TUNNEL_SCRIPT)} "
+            f"{shlex.quote(dest_host)} {dest_port} \"$__hb_prefix\"; "
+            "__hb_status=$?; stty echo; printf '\n%s:CLOSED\n' \"$__hb_prefix\"\n"
+        )
+        await asyncio.to_thread(
+            self._client.shell_write,
+            session_id,
+            command.encode("utf-8"),
+            owner="mock-ssh",
+        )
+
+    async def _send_data(self, session_id: str, prefix: str, payload: bytes) -> None:
+        lines = [
+            f"{prefix}:DATA:{base64.b64encode(payload[offset : offset + TUNNEL_FRAME_BYTES]).decode('ascii')}\n"
+            for offset in range(0, len(payload), TUNNEL_FRAME_BYTES)
+        ]
+        await self._send_lines(session_id, lines)
+
+    async def _send_lines(self, session_id: str, lines: list[str]) -> None:
+        await asyncio.to_thread(
+            self._client.shell_write,
+            session_id,
+            "".join(lines).encode("ascii"),
+            owner="mock-ssh",
+        )
+
+
+
+_TUNNEL_INPUT_EOF = object()
+
+
+class HostBridgeTCPSession(asyncssh.SSHTCPSession[bytes]):
+    def __init__(
+        self,
+        forwarder: HostBridgeTCPForwarder,
+        *,
+        session_id: str,
+        prefix: str,
+        frames: _TunnelFrameReader,
+        dest_host: str,
+        dest_port: int,
+    ) -> None:
+        self._forwarder = forwarder
+        self._session_id = session_id
+        self._prefix = prefix
+        self._frames = frames
+        self._dest_host = dest_host
+        self._dest_port = dest_port
+        self._channel: asyncssh.SSHTCPChannel[bytes] | None = None
+        self._input: asyncio.Queue[bytes | object] = asyncio.Queue()
+        self._write_ready = asyncio.Event()
+        self._write_ready.set()
+        self._runner: asyncio.Task[None] | None = None
+        self._orphan_cleanup: asyncio.Task[None] | None = None
+        self._released = False
+        self._finished = asyncio.Event()
+
+    def connection_made(self, chan: asyncssh.SSHTCPChannel[bytes]) -> None:
+        self._channel = chan
+
+    def session_started(self) -> None:
+        self._runner = asyncio.create_task(self._run())
+        self._runner.add_done_callback(self._runner_done)
+
+    def data_received(self, data: bytes, datatype: asyncssh.DataType) -> None:  # noqa: ARG002
+        if data:
+            self._input.put_nowait(bytes(data))
+
+    def eof_received(self) -> bool:
+        self._input.put_nowait(_TUNNEL_INPUT_EOF)
+        return True
+
+    def connection_lost(self, exc: Exception | None) -> None:  # noqa: ARG002
+        self.abort()
+
+    def pause_writing(self) -> None:
+        self._write_ready.clear()
+
+    def resume_writing(self) -> None:
+        self._write_ready.set()
+
+    def abort(self) -> None:
+        if self._runner is not None and self._runner is not asyncio.current_task() and not self._runner.done():
+            self._runner.cancel()
+        elif self._runner is None:
+            self._schedule_orphan_cleanup()
+
+    async def wait_finished(self) -> None:
+        await self._finished.wait()
+
+    def _runner_done(self, task: asyncio.Task[None]) -> None:  # noqa: ARG002
+        if not self._finished.is_set():
+            self._schedule_orphan_cleanup()
+
+    def _schedule_orphan_cleanup(self) -> None:
+        if self._orphan_cleanup is None:
+            self._orphan_cleanup = asyncio.create_task(self._finish_orphaned_session())
+
+    async def _finish_orphaned_session(self) -> None:
+        await self._release(reusable=False)
+        if self._channel is not None:
+            with suppress(Exception):
+                self._channel.close()
+        self._finished.set()
+
+    async def _run(self) -> None:
+        input_task = asyncio.create_task(self._pump_input())
+        output_task = asyncio.create_task(self._pump_output())
+        reusable = False
+        try:
+            done, _ = await asyncio.wait({input_task, output_task}, return_when=asyncio.FIRST_COMPLETED)
+            if output_task in done:
+                input_task.cancel()
+                await asyncio.gather(input_task, return_exceptions=True)
+                await output_task
+            else:
+                await input_task
+                await output_task
+            reusable = True
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            print(
+                f"hostbridge mock-ssh TCP forwarding to {self._dest_host}:{self._dest_port} closed: {exc}",
+                file=sys.stderr,
+                flush=True,
+            )
+        finally:
+            for task in (input_task, output_task):
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(input_task, output_task, return_exceptions=True)
+            await self._release(reusable=reusable)
+            if self._channel is not None:
+                with suppress(Exception):
+                    self._channel.close()
+            self._finished.set()
+
+    async def _release(self, *, reusable: bool) -> None:
+        if self._released:
+            return
+        self._released = True
+        await asyncio.to_thread(
+            self._forwarder._session_factory.close,
+            self._session_id,
+            reusable=reusable,
+        )
+
+    async def _pump_input(self) -> None:
+        while True:
+            item = await self._input.get()
+            if item is _TUNNEL_INPUT_EOF:
+                await self._forwarder._send_lines(self._session_id, [f"{self._prefix}:EOF\n"])
+                return
+            await self._forwarder._send_data(self._session_id, self._prefix, cast(bytes, item))
+
+    async def _pump_output(self) -> None:
+        if self._channel is None:
+            raise RuntimeError("TCP tunnel channel was not initialized")
+        saw_eof = False
+        while True:
+            kind, payload = await self._frames.next()
+            if kind == "DATA":
+                await self._write_ready.wait()
+                self._channel.write(payload)
+            elif kind == "EOF":
+                self._channel.write_eof()
+                saw_eof = True
+            elif kind == "CLOSED":
+                if not saw_eof:
+                    raise HostBridgeError("protocol_mismatch", "remote TCP tunnel closed without EOF")
+                return
+            elif kind == "ERROR":
+                raise HostBridgeError(
+                    "connection_lost",
+                    payload.decode("utf-8", errors="replace") or "remote TCP tunnel failed",
+                    retryable=True,
+                )
+
+
 class HostBridgeSFTPServer(asyncssh.SFTPServer):
     def __init__(
         self,
@@ -781,8 +1138,9 @@ async def serve_mock_ssh(
     client = HostBridgeClient()
     client.hello()
     session_factory = HostBridgeSessionFactory(client, host_id, connect_timeout)
+    tcp_forwarder = HostBridgeTCPForwarder(client, session_factory, connect_timeout=connect_timeout)
     server, selected_port = await create_server_on_available_port(
-        server_factory=lambda: HostBridgeSSHServer(public_key),
+        server_factory=lambda: HostBridgeSSHServer(public_key, tcp_forwarder),
         listen_host=listen_host,
         port=port,
         server_host_keys=[str(host_key_path)],
@@ -804,7 +1162,7 @@ async def serve_mock_ssh(
     if installed_config_path is not None:
         print(f"ssh config: {installed_config_path}", flush=True)
     print(f"identity file: {client_key_path}", flush=True)
-    print("supports SSH exec, shell, SFTP, and SCP; press Ctrl-C to stop", flush=True)
+    print("supports SSH exec, shell, SFTP, SCP, and TCP forwarding; press Ctrl-C to stop", flush=True)
     try:
         await server.wait_closed()
         return 0