@neoline/hostbridge 2.0.1 → 2.0.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@neoline/hostbridge",
-  "version": "2.0.1",
+  "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.1"
+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.1"
+__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.1" 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:
@@ -121,9 +197,9 @@ class HostBridgeSSHProcess:
             process.exit(status)
 
     async def _read_exec_stdin(self, process: asyncssh.SSHServerProcess[bytes], command: str) -> bytes | None:
-        if process.stdin.at_eof():
-            return b""
         waits_for_eof = _command_reads_stdin(command)
+        if process.stdin.at_eof():
+            return b"" if waits_for_eof else None
         try:
             first_chunk = await asyncio.wait_for(
                 process.stdin.read(4096),
@@ -134,7 +210,7 @@ class HostBridgeSSHProcess:
                 raise HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin") from None
             return None
         if not first_chunk:
-            return b""
+            return b"" if waits_for_eof else None
 
         first_chunk_bytes = _to_bytes(first_chunk)
         total = len(first_chunk_bytes)
@@ -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

package/src/server_control_mcp/transports/pty.py

@@ -163,8 +163,7 @@ class PtyTransport:
                 raise TransportError("PTY transport is not alive")
             if stdin is not None:
                 stdin_path = self._stage_stdin_locked(stdin, timeout=timeout)
-                command = f"{command} < {shlex.quote(stdin_path)}"
-            return self._exec_locked(command, timeout=timeout, output_limit=output_limit)
+            return self._exec_locked(command, timeout=timeout, output_limit=output_limit, stdin_path=stdin_path)
         finally:
             if stdin_path is not None and self.child.isalive():
                 with suppress(Exception):
@@ -205,15 +204,23 @@ class PtyTransport:
             raise TransportError("failed to stage PTY stdin")
         return path
 
-    def _exec_locked(self, command: str, *, timeout: float, output_limit: int) -> ExecResult:
+    def _exec_locked(
+        self,
+        command: str,
+        *,
+        timeout: float,
+        output_limit: int,
+        stdin_path: str | None = None,
+    ) -> ExecResult:
         token = self._token_factory()
         start_marker = f"__{self._marker_prefix}_START_{token}__"
         exit_pattern = re.compile(rf"__{self._marker_prefix}_EXIT_{re.escape(token)}__:(\d+)")
         encoded_command = base64.b64encode(command.encode("utf-8")).decode("ascii")
+        stdin_redirect = f" < {shlex.quote(stdin_path)}" if stdin_path is not None else ""
         remote_line = (
             f"__scm_cmd=$(printf %s {shlex.quote(encoded_command)} | {_REMOTE_BASE64_DECODE}); "
             f"printf '\n{start_marker}\n'; "
-            f"{self.remote_shell} -lc \"$__scm_cmd\"; "
+            f"{self.remote_shell} -lc \"$__scm_cmd\"{stdin_redirect}; "
             "__scm_status=$?; "
             f"printf '\n__{self._marker_prefix}_EXIT_{token}__:%s\n' \"$__scm_status\""
         )