@neoline/hostbridge 2.0.2 → 2.0.4

package/src/server_control_mcp/mock_ssh_tunnel.py

@@ -0,0 +1,592 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import binascii
+import sys
+import uuid
+from collections import deque
+from contextlib import suppress
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Protocol
+
+import asyncssh
+
+from .client import HostBridgeClient, HostBridgeError
+from .remote_tunnel_agent import build_remote_tunnel_command
+
+TUNNEL_PROTOCOL_VERSION = 1
+UPSTREAM_FRAME_BYTES = 2048
+DOWNSTREAM_FRAME_BYTES = 16 * 1024
+TUNNEL_READ_BYTES = 64 * 1024
+TUNNEL_WRITE_BATCH_BYTES = 32 * 1024
+TUNNEL_HIGH_WATER_BYTES = 128 * 1024
+TUNNEL_LOW_WATER_BYTES = 64 * 1024
+TUNNEL_MAX_LINE_BYTES = 128 * 1024
+
+
+class TunnelProtocolError(ValueError):
+    pass
+
+
+class Direction(StrEnum):
+    UPSTREAM = "U"
+    DOWNSTREAM = "D"
+
+
+class FrameKind(StrEnum):
+    READY = "READY"
+    DATA = "DATA"
+    EOF = "EOF"
+    ERROR = "ERROR"
+    CLOSED = "CLOSED"
+
+
+@dataclass(frozen=True, slots=True)
+class TunnelFrame:
+    direction: Direction
+    sequence: int
+    kind: FrameKind
+    payload: bytes = b""
+    status: int | None = None
+
+
+class TunnelCodec:
+    def __init__(self, prefix: str, *, inbound_direction: Direction) -> None:
+        if not prefix or any(character in prefix for character in "\r\n:"):
+            raise ValueError("tunnel prefix must be non-empty and must not contain CR, LF, or colon")
+        try:
+            prefix_bytes = prefix.encode("ascii")
+        except UnicodeEncodeError as exc:
+            raise ValueError("tunnel prefix must be ASCII") from exc
+        self._prefix = prefix_bytes
+        self._marker = prefix_bytes + b":V"
+        self._inbound_direction = inbound_direction
+        self._next_sequence = 0
+        self._buffer = bytearray()
+
+    def encode(self, frame: TunnelFrame) -> bytes:
+        self._validate_frame(frame)
+        status = b"-" if frame.status is None else str(frame.status).encode("ascii")
+        payload = base64.b64encode(frame.payload)
+        return (
+            b":".join(
+                (
+                    self._prefix,
+                    f"V{TUNNEL_PROTOCOL_VERSION}".encode("ascii"),
+                    frame.direction.value.encode("ascii"),
+                    str(frame.sequence).encode("ascii"),
+                    frame.kind.value.encode("ascii"),
+                    status,
+                    payload,
+                )
+            )
+            + b"\n"
+        )
+
+    def feed(self, data: bytes) -> list[TunnelFrame]:
+        self._buffer.extend(data)
+        frames: list[TunnelFrame] = []
+        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._marker)
+            if marker < 0:
+                continue
+            frames.append(self._decode_line(bytes(line[marker:])))
+        if len(self._buffer) > TUNNEL_MAX_LINE_BYTES:
+            raise TunnelProtocolError("tunnel frame exceeds maximum line size")
+        return frames
+
+    def _decode_line(self, line: bytes) -> TunnelFrame:
+        fields = line.split(b":", 6)
+        if len(fields) != 7:
+            raise TunnelProtocolError("malformed tunnel frame")
+        prefix, version, direction_raw, sequence_raw, kind_raw, status_raw, payload_raw = fields
+        if prefix != self._prefix:
+            raise TunnelProtocolError("tunnel frame prefix mismatch")
+        if version != f"V{TUNNEL_PROTOCOL_VERSION}".encode("ascii"):
+            raise TunnelProtocolError("unsupported tunnel protocol version")
+        try:
+            direction = Direction(direction_raw.decode("ascii"))
+        except (UnicodeDecodeError, ValueError) as exc:
+            raise TunnelProtocolError("invalid tunnel frame direction") from exc
+        if direction is not self._inbound_direction:
+            raise TunnelProtocolError("unexpected tunnel frame direction")
+        try:
+            sequence = int(sequence_raw)
+        except ValueError as exc:
+            raise TunnelProtocolError("invalid tunnel frame sequence") from exc
+        if sequence != self._next_sequence:
+            raise TunnelProtocolError(f"tunnel frame sequence expected {self._next_sequence}, received {sequence}")
+        try:
+            kind = FrameKind(kind_raw.decode("ascii"))
+        except (UnicodeDecodeError, ValueError) as exc:
+            raise TunnelProtocolError("invalid tunnel frame kind") from exc
+        if status_raw == b"-":
+            status = None
+        else:
+            try:
+                status = int(status_raw)
+            except ValueError as exc:
+                raise TunnelProtocolError("invalid tunnel frame status") from exc
+        try:
+            payload = base64.b64decode(payload_raw, validate=True)
+        except (binascii.Error, ValueError) as exc:
+            raise TunnelProtocolError("invalid tunnel frame base64 payload") from exc
+        frame = TunnelFrame(direction, sequence, kind, payload, status)
+        self._validate_frame(frame)
+        self._next_sequence += 1
+        return frame
+
+    @staticmethod
+    def _validate_frame(frame: TunnelFrame) -> None:
+        if not isinstance(frame.direction, Direction):
+            raise TunnelProtocolError("invalid tunnel frame direction")
+        if not isinstance(frame.sequence, int) or isinstance(frame.sequence, bool) or frame.sequence < 0:
+            raise TunnelProtocolError("invalid tunnel frame sequence")
+        if not isinstance(frame.kind, FrameKind):
+            raise TunnelProtocolError("invalid tunnel frame kind")
+        if not isinstance(frame.payload, bytes):
+            raise TunnelProtocolError("tunnel frame payload must be bytes")
+        if frame.kind is FrameKind.CLOSED:
+            if not isinstance(frame.status, int) or isinstance(frame.status, bool) or not 0 <= frame.status <= 255:
+                raise TunnelProtocolError("CLOSED tunnel frame requires an exit status")
+        elif frame.status is not None:
+            raise TunnelProtocolError("only CLOSED tunnel frames may contain a status")
+        if frame.kind not in {FrameKind.DATA, FrameKind.ERROR} and frame.payload:
+            raise TunnelProtocolError(f"{frame.kind.value} tunnel frame must not contain a payload")
+
+
+class TunnelState(StrEnum):
+    CONNECTING = "CONNECTING"
+    OPEN = "OPEN"
+    LOCAL_EOF = "LOCAL_EOF"
+    REMOTE_EOF = "REMOTE_EOF"
+    DRAINING = "DRAINING"
+    ABORTING = "ABORTING"
+    CLOSED = "CLOSED"
+
+
+class TunnelLifecycle:
+    def __init__(self) -> None:
+        self.state = TunnelState.CONNECTING
+        self.reusable = False
+
+    def mark_ready(self) -> None:
+        self._require(TunnelState.CONNECTING)
+        self.state = TunnelState.OPEN
+
+    def mark_local_eof(self) -> None:
+        if self.state is TunnelState.OPEN:
+            self.state = TunnelState.LOCAL_EOF
+        elif self.state is TunnelState.REMOTE_EOF:
+            self.state = TunnelState.DRAINING
+        else:
+            self._invalid("local EOF")
+
+    def mark_remote_eof(self) -> None:
+        if self.state is TunnelState.OPEN:
+            self.state = TunnelState.REMOTE_EOF
+        elif self.state is TunnelState.LOCAL_EOF:
+            self.state = TunnelState.DRAINING
+        else:
+            self._invalid("remote EOF")
+
+    def mark_closed(self, status: int) -> None:
+        self._require(TunnelState.DRAINING)
+        if not isinstance(status, int) or isinstance(status, bool) or not 0 <= status <= 255:
+            raise TunnelProtocolError("terminal tunnel status is invalid")
+        self.state = TunnelState.CLOSED
+        self.reusable = status == 0
+
+    def abort(self) -> None:
+        if self.state is not TunnelState.CLOSED:
+            self.state = TunnelState.ABORTING
+            self.reusable = False
+
+    def finish_abort(self) -> None:
+        self._require(TunnelState.ABORTING)
+        self.state = TunnelState.CLOSED
+
+    def _require(self, expected: TunnelState) -> None:
+        if self.state is not expected:
+            raise TunnelProtocolError(f"tunnel state must be {expected.value}, found {self.state.value}")
+
+    def _invalid(self, event: str) -> None:
+        raise TunnelProtocolError(f"cannot accept {event} while tunnel state is {self.state.value}")
+
+
+class SessionPool(Protocol):
+    def open(self) -> str: ...
+
+    def close(self, session_id: str, *, reusable: bool = True) -> None: ...
+
+
+class TunnelLease:
+    def __init__(self, pool: SessionPool, session_id: str) -> None:
+        self._pool = pool
+        self.session_id = session_id
+        self._close_task: asyncio.Task[None] | None = None
+
+    @classmethod
+    async def acquire(cls, pool: SessionPool) -> TunnelLease:
+        open_task = asyncio.create_task(asyncio.to_thread(pool.open))
+        try:
+            session_id = await asyncio.shield(open_task)
+        except asyncio.CancelledError:
+            await cls._wait_ignoring_cancellation(open_task)
+            if not open_task.cancelled() and open_task.exception() is None:
+                lease = cls(pool, open_task.result())
+                await lease.close(reusable=False)
+            raise
+        return cls(pool, session_id)
+
+    async def close(self, *, reusable: bool) -> None:
+        if self._close_task is None:
+            self._close_task = asyncio.create_task(
+                asyncio.to_thread(self._pool.close, self.session_id, reusable=reusable)
+            )
+        await self._wait_ignoring_cancellation(self._close_task)
+        await self._close_task
+
+    @staticmethod
+    async def _wait_ignoring_cancellation(task: asyncio.Task[object]) -> None:
+        while not task.done():
+            try:
+                await asyncio.shield(task)
+            except asyncio.CancelledError:
+                continue
+
+
+class _TunnelFrameStream:
+    def __init__(self, client: HostBridgeClient, lease: TunnelLease, prefix: str) -> None:
+        self._client = client
+        self._lease = lease
+        self._codec = TunnelCodec(prefix, inbound_direction=Direction.DOWNSTREAM)
+        self._pending: deque[TunnelFrame] = deque()
+
+    async def next(self) -> TunnelFrame:
+        while not self._pending:
+            result = await asyncio.to_thread(
+                self._client.shell_read,
+                self._lease.session_id,
+                timeout=0.2,
+                max_bytes=TUNNEL_READ_BYTES,
+                owner="mock-ssh",
+            )
+            if result.data:
+                try:
+                    self._pending.extend(self._codec.feed(result.data))
+                except TunnelProtocolError as exc:
+                    raise HostBridgeError("protocol_mismatch", str(exc)) from exc
+            if not result.alive and not self._pending:
+                raise HostBridgeError(
+                    "connection_lost",
+                    "remote TCP tunnel closed without a terminal frame",
+                    retryable=True,
+                )
+        return self._pending.popleft()
+
+
+class TunnelConnector:
+    def __init__(
+        self,
+        client: HostBridgeClient,
+        session_pool: SessionPool,
+        *,
+        connect_timeout: int,
+    ) -> None:
+        if connect_timeout <= 0:
+            raise ValueError("connect_timeout must be positive")
+        self._client = client
+        self._session_pool = session_pool
+        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)
+        try:
+            lease = await TunnelLease.acquire(self._session_pool)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
+
+        prefix = f"__HB_TUNNEL_{uuid.uuid4().hex}__"
+        frames = _TunnelFrameStream(self._client, lease, prefix)
+        lifecycle = TunnelLifecycle()
+        try:
+            await asyncio.to_thread(
+                self._client.shell_write,
+                lease.session_id,
+                build_remote_tunnel_command(dest_host, dest_port, prefix),
+                owner="mock-ssh",
+            )
+            first = await asyncio.wait_for(frames.next(), timeout=self._connect_timeout)
+            if first.kind is FrameKind.ERROR:
+                raise HostBridgeError(
+                    "connection_lost",
+                    first.payload.decode("utf-8", errors="replace") or "remote TCP connection failed",
+                    retryable=True,
+                )
+            if first.kind is not FrameKind.READY:
+                raise HostBridgeError(
+                    "protocol_mismatch",
+                    f"remote TCP tunnel returned {first.kind.value} before READY",
+                )
+            lifecycle.mark_ready()
+        except asyncio.CancelledError:
+            lifecycle.abort()
+            await lease.close(reusable=False)
+            lifecycle.finish_abort()
+            raise
+        except Exception as exc:
+            lifecycle.abort()
+            await lease.close(reusable=False)
+            lifecycle.finish_abort()
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
+        return DirectTCPIPSession(
+            self._client,
+            lease,
+            prefix,
+            frames,
+            lifecycle,
+            dest_host=dest_host,
+            dest_port=dest_port,
+        )
+
+
+_INPUT_EOF = object()
+
+
+class DirectTCPIPSession(asyncssh.SSHTCPSession[bytes]):
+    def __init__(
+        self,
+        client: HostBridgeClient,
+        lease: TunnelLease,
+        prefix: str,
+        frames: _TunnelFrameStream,
+        lifecycle: TunnelLifecycle,
+        *,
+        dest_host: str,
+        dest_port: int,
+    ) -> None:
+        self._client = client
+        self._lease = lease
+        self._codec = TunnelCodec(prefix, inbound_direction=Direction.DOWNSTREAM)
+        self._frames = frames
+        self._lifecycle = lifecycle
+        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._input_buffered_bytes = 0
+        self._input_paused = False
+        self._upstream_sequence = 0
+        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._finished = asyncio.Event()
+
+    @property
+    def state(self) -> TunnelState:
+        return self._lifecycle.state
+
+    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 not data:
+            return
+        payload = bytes(data)
+        self._input.put_nowait(payload)
+        self._input_buffered_bytes += len(payload)
+        if (
+            self._input_buffered_bytes >= TUNNEL_HIGH_WATER_BYTES
+            and not self._input_paused
+            and self._channel is not None
+        ):
+            self._channel.pause_reading()
+            self._input_paused = True
+
+    def eof_received(self) -> bool:
+        self._input.put_nowait(_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:
+        self._lifecycle.abort()
+        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:
+        self._lifecycle.abort()
+        await self._lease.close(reusable=False)
+        if self._lifecycle.state is TunnelState.ABORTING:
+            self._lifecycle.finish_abort()
+        self._close_channel()
+        self._finished.set()
+
+    async def _run(self) -> None:
+        input_task = asyncio.create_task(self._pump_input())
+        output_task = asyncio.create_task(self._pump_output())
+        try:
+            await asyncio.gather(input_task, output_task)
+        except asyncio.CancelledError:
+            self._lifecycle.abort()
+            raise
+        except Exception as exc:
+            self._lifecycle.abort()
+            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._lease.close(reusable=self._lifecycle.reusable)
+            if self._lifecycle.state is TunnelState.ABORTING:
+                self._lifecycle.finish_abort()
+            self._close_channel()
+            self._finished.set()
+
+    async def _pump_input(self) -> None:
+        while True:
+            item = await self._input.get()
+            if item is _INPUT_EOF:
+                await self._send_eof()
+                return
+
+            if not isinstance(item, bytes):
+                raise RuntimeError("TCP tunnel input queue contained an invalid item")
+            batch = bytearray(item)
+            saw_eof = False
+            while not self._input.empty():
+                queued = self._input.get_nowait()
+                if queued is _INPUT_EOF:
+                    saw_eof = True
+                    break
+                if not isinstance(queued, bytes):
+                    raise RuntimeError("TCP tunnel input queue contained an invalid item")
+                batch.extend(queued)
+            try:
+                for offset in range(0, len(batch), TUNNEL_WRITE_BATCH_BYTES):
+                    await self._send_data(bytes(batch[offset : offset + TUNNEL_WRITE_BATCH_BYTES]))
+            finally:
+                self._input_buffered_bytes -= len(batch)
+                self._resume_input_if_ready()
+            if saw_eof:
+                await self._send_eof()
+                return
+
+    async def _send_data(self, payload: bytes) -> None:
+        frames = []
+        for offset in range(0, len(payload), UPSTREAM_FRAME_BYTES):
+            frames.append(
+                TunnelFrame(
+                    Direction.UPSTREAM,
+                    self._upstream_sequence,
+                    FrameKind.DATA,
+                    payload=payload[offset : offset + UPSTREAM_FRAME_BYTES],
+                )
+            )
+            self._upstream_sequence += 1
+        await self._write_frames(frames)
+
+    async def _send_eof(self) -> None:
+        frame = TunnelFrame(Direction.UPSTREAM, self._upstream_sequence, FrameKind.EOF)
+        self._upstream_sequence += 1
+        await self._write_frames([frame])
+        self._lifecycle.mark_local_eof()
+
+    async def _write_frames(self, frames: list[TunnelFrame]) -> None:
+        payload = b"".join(self._codec.encode(frame) for frame in frames)
+        await asyncio.to_thread(
+            self._client.shell_write,
+            self._lease.session_id,
+            payload,
+            owner="mock-ssh",
+        )
+
+    async def _pump_output(self) -> None:
+        if self._channel is None:
+            raise RuntimeError("TCP tunnel channel was not initialized")
+        while True:
+            frame = await self._frames.next()
+            if frame.kind is FrameKind.DATA:
+                if self._lifecycle.state not in {TunnelState.OPEN, TunnelState.LOCAL_EOF}:
+                    raise HostBridgeError("protocol_mismatch", "remote TCP data arrived after EOF")
+                await self._write_ready.wait()
+                self._channel.write(frame.payload)
+            elif frame.kind is FrameKind.EOF:
+                self._lifecycle.mark_remote_eof()
+                self._channel.write_eof()
+            elif frame.kind is FrameKind.ERROR:
+                raise HostBridgeError(
+                    "connection_lost",
+                    frame.payload.decode("utf-8", errors="replace") or "remote TCP tunnel failed",
+                    retryable=True,
+                )
+            elif frame.kind is FrameKind.CLOSED:
+                assert frame.status is not None
+                self._lifecycle.mark_closed(frame.status)
+                if frame.status != 0:
+                    raise HostBridgeError(
+                        "connection_lost",
+                        f"remote TCP tunnel exited with status {frame.status}",
+                        retryable=True,
+                    )
+                return
+            else:
+                raise HostBridgeError(
+                    "protocol_mismatch",
+                    f"unexpected {frame.kind.value} frame after READY",
+                )
+
+    def _resume_input_if_ready(self) -> None:
+        if self._input_paused and self._input_buffered_bytes <= TUNNEL_LOW_WATER_BYTES and self._channel is not None:
+            self._channel.resume_reading()
+            self._input_paused = False
+
+    def _close_channel(self) -> None:
+        if self._channel is not None:
+            with suppress(Exception):
+                self._channel.close()