@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/mock_ssh_tunnel.py

@@ -0,0 +1,289 @@
+"""AsyncSSH direct-tcpip adaptation for persistent binary mux streams."""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from collections import deque
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
+from typing import Protocol
+
+import asyncssh
+
+TUNNEL_HIGH_WATER_BYTES = 256 * 1024
+TUNNEL_LOW_WATER_BYTES = 64 * 1024
+
+
+class MuxTunnelStream(Protocol):
+    stream_id: int
+
+    async def send(self, data: bytes) -> None: ...
+
+    async def send_eof(self) -> None: ...
+
+    async def receive(self, max_bytes: int = 64 * 1024) -> bytes | None: ...
+
+    async def finish(self) -> dict[str, object]: ...
+
+    async def reset(self, message: str = "stream reset") -> None: ...
+
+
+class MuxTunnelClient(Protocol):
+    @property
+    def closed(self) -> bool: ...
+
+    async def open_stream(self, method: str, params: dict[str, object]) -> MuxTunnelStream: ...
+
+    async def close(self) -> None: ...
+
+
+MuxClientFactory = Callable[[], Awaitable[MuxTunnelClient]]
+
+
+class TunnelConnector:
+    def __init__(
+        self,
+        host_id: str,
+        client_factory: MuxClientFactory,
+        *,
+        connect_timeout: float,
+    ) -> None:
+        if not isinstance(host_id, str) or not host_id:
+            raise ValueError("host_id must be a non-empty string")
+        if connect_timeout <= 0:
+            raise ValueError("connect_timeout must be positive")
+        self._host_id = host_id
+        self._client_factory = client_factory
+        self._connect_timeout = float(connect_timeout)
+        self._client: MuxTunnelClient | None = None
+        self._lock = asyncio.Lock()
+        self._closing = False
+
+    @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 start(self) -> None:
+        await self._get_client()
+
+    async def connect(self, dest_host: str, dest_port: int) -> asyncssh.SSHTCPSession[bytes]:
+        self.validate_destination(dest_host, dest_port)
+        client: MuxTunnelClient | None = None
+        try:
+            client = await self._get_client()
+            async with asyncio.timeout(self._connect_timeout):
+                stream = await client.open_stream(
+                    "tcp.connect",
+                    {
+                        "host_id": self._host_id,
+                        "destination_host": dest_host,
+                        "destination_port": dest_port,
+                    },
+                )
+        except asyncio.CancelledError:
+            raise
+        except (Exception, TimeoutError) as exc:
+            if client is not None and client.closed:
+                await self._discard_client(client)
+            message = str(exc) or "TCP tunnel stream open timed out"
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, message) from exc
+        return DirectTCPIPSession(stream, dest_host=dest_host, dest_port=dest_port)
+
+    async def close(self) -> None:
+        async with self._lock:
+            self._closing = True
+            client = self._client
+            if client is not None:
+                await client.close()
+                if self._client is client:
+                    self._client = None
+
+    async def _get_client(self) -> MuxTunnelClient:
+        async with self._lock:
+            if self._closing:
+                raise RuntimeError("TCP tunnel connector is closed")
+            if self._client is not None and not self._client.closed:
+                return self._client
+            stale = self._client
+            async with asyncio.timeout(self._connect_timeout):
+                client = await self._client_factory()
+            self._client = client
+        if stale is not None and stale is not client:
+            with suppress(Exception):
+                await stale.close()
+        return client
+
+    async def _discard_client(self, client: MuxTunnelClient) -> None:
+        async with self._lock:
+            if self._client is client:
+                self._client = None
+
+
+_INPUT_EOF = object()
+
+
+class DirectTCPIPSession(asyncssh.SSHTCPSession[bytes]):
+    def __init__(self, stream: MuxTunnelStream, *, dest_host: str, dest_port: int) -> None:
+        self._stream = stream
+        self._dest_host = dest_host
+        self._dest_port = dest_port
+        self._channel: asyncssh.SSHTCPChannel[bytes] | None = None
+        self._input: deque[bytes | object] = deque()
+        self._input_ready = asyncio.Event()
+        self._input_buffered_bytes = 0
+        self._input_paused = False
+        self._input_eof = False
+        self._write_ready = asyncio.Event()
+        self._write_ready.set()
+        self._runner: asyncio.Task[None] | None = None
+        self._cleanup_task: asyncio.Task[None] | None = None
+        self._finished = asyncio.Event()
+        self._aborted = False
+        self._reset_sent = False
+
+    def connection_made(self, chan: asyncssh.SSHTCPChannel[bytes]) -> None:
+        self._channel = chan
+
+    def session_started(self) -> None:
+        if self._runner is not None or self._finished.is_set():
+            return
+        self._runner = asyncio.create_task(
+            self._run(),
+            name=f"hostbridge-mock-ssh-tcp-{self._stream.stream_id}",
+        )
+        self._runner.add_done_callback(self._runner_done)
+
+    def data_received(self, data: bytes, datatype: asyncssh.DataType) -> None:  # noqa: ARG002
+        if not data or self._input_eof or self._finished.is_set():
+            return
+        payload = bytes(data)
+        self._input.append(payload)
+        self._input_buffered_bytes += len(payload)
+        self._input_ready.set()
+        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:
+        if not self._input_eof and not self._finished.is_set():
+            self._input_eof = True
+            self._input.append(_INPUT_EOF)
+            self._input_ready.set()
+        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._finished.is_set():
+            return
+        self._aborted = True
+        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 and self._cleanup_task is None:
+            self._cleanup_task = asyncio.create_task(self._finish_orphaned_session())
+
+    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() and self._cleanup_task is None:
+            self._cleanup_task = asyncio.create_task(self._finish_orphaned_session())
+
+    async def _finish_orphaned_session(self) -> None:
+        await self._reset_stream()
+        self._finish_channel()
+
+    async def _run(self) -> None:
+        input_task = asyncio.create_task(self._pump_input())
+        output_task = asyncio.create_task(self._pump_output())
+        clean = False
+        try:
+            await asyncio.gather(input_task, output_task)
+            await self._stream.finish()
+            clean = 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)
+            if not clean:
+                await self._reset_stream()
+            self._finish_channel()
+
+    async def _pump_input(self) -> None:
+        while True:
+            item = await self._next_input()
+            if item is _INPUT_EOF:
+                await self._stream.send_eof()
+                return
+            if not isinstance(item, bytes):
+                raise TypeError("tunnel input queue contains a non-bytes item")
+            try:
+                await self._stream.send(item)
+            finally:
+                self._input_buffered_bytes -= len(item)
+                self._resume_input_if_ready()
+
+    async def _next_input(self) -> bytes | object:
+        while not self._input:
+            self._input_ready.clear()
+            if not self._input:
+                await self._input_ready.wait()
+        item = self._input.popleft()
+        if not self._input:
+            self._input_ready.clear()
+        return item
+
+    async def _pump_output(self) -> None:
+        if self._channel is None:
+            raise RuntimeError("TCP tunnel channel was not initialized")
+        while True:
+            await self._write_ready.wait()
+            chunk = await self._stream.receive()
+            if chunk is None:
+                self._channel.write_eof()
+                return
+            self._channel.write(chunk)
+
+    async def _reset_stream(self) -> None:
+        if self._reset_sent:
+            return
+        self._reset_sent = True
+        with suppress(Exception):
+            await self._stream.reset("mock-ssh TCP channel closed")
+
+    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 _finish_channel(self) -> None:
+        if self._finished.is_set():
+            return
+        self._finished.set()
+        if self._channel is not None:
+            with suppress(Exception):
+                self._channel.close()

package/src/server_control_mcp/mux_connection.py

@@ -0,0 +1,326 @@
+"""Lifecycle and bounded resource primitives for multiplexed connections."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import Enum, StrEnum
+
+from .mux_protocol import FrameType, MuxProtocolError
+
+
+class MuxStateError(RuntimeError):
+    """Raised when an operation is invalid for the current lifecycle state."""
+
+
+class StreamState(StrEnum):
+    OPENING = "OPENING"
+    OPEN = "OPEN"
+    LOCAL_EOF = "LOCAL_EOF"
+    REMOTE_EOF = "REMOTE_EOF"
+    DRAINING = "DRAINING"
+    CLOSING = "CLOSING"
+    CLOSED = "CLOSED"
+    RESET = "RESET"
+
+
+class TerminalKind(StrEnum):
+    CLOSED = "CLOSED"
+    RESET = "RESET"
+
+
+class ConnectionState(StrEnum):
+    DISCONNECTED = "DISCONNECTED"
+    CONNECTING = "CONNECTING"
+    HANDSHAKING = "HANDSHAKING"
+    READY = "READY"
+    DRAINING = "DRAINING"
+    CLOSED = "CLOSED"
+    FAILED = "FAILED"
+
+
+class FrameDisposition(Enum):
+    UNKNOWN_STREAM = "unknown_stream"
+    DISCARD_LATE_DATA = "discard_late_data"
+    IGNORE_TERMINAL = "ignore_terminal"
+    PROTOCOL_ERROR = "protocol_error"
+
+
+class ByteCreditWindow:
+    def __init__(self, *, initial_credit: int, maximum_credit: int) -> None:
+        _positive_or_zero(initial_credit, "initial_credit")
+        _positive(maximum_credit, "maximum_credit")
+        if initial_credit > maximum_credit:
+            raise ValueError("initial credit exceeds maximum credit")
+        self._available = initial_credit
+        self._maximum = maximum_credit
+        self._closed = False
+
+    @property
+    def available(self) -> int:
+        return self._available
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def consume(self, amount: int) -> None:
+        _positive(amount, "amount")
+        if self._closed:
+            raise MuxStateError("credit window is closed")
+        if amount > self._available:
+            raise MuxStateError("insufficient byte credit")
+        self._available -= amount
+
+    def grant(self, amount: int) -> None:
+        _positive(amount, "amount")
+        if self._closed:
+            return
+        if self._available + amount > self._maximum:
+            raise MuxProtocolError("credit grant exceeds negotiated maximum")
+        self._available += amount
+
+    def close(self) -> None:
+        self._closed = True
+
+
+class StreamLifecycle:
+    def __init__(self, *, open_initiator: bool) -> None:
+        if not isinstance(open_initiator, bool):
+            raise TypeError("open_initiator must be a boolean")
+        self.open_initiator = open_initiator
+        self.state = StreamState.OPENING
+        self.terminal_kind: TerminalKind | None = None
+        self._opened = False
+        self._local_eof = False
+        self._remote_eof = False
+        self._close_sent = False
+        self._close_received = False
+
+    @property
+    def opened(self) -> bool:
+        return self._opened
+
+    @property
+    def local_eof(self) -> bool:
+        return self._local_eof
+
+    @property
+    def remote_eof(self) -> bool:
+        return self._remote_eof
+
+    @property
+    def close_sent(self) -> bool:
+        return self._close_sent
+
+    @property
+    def close_received(self) -> bool:
+        return self._close_received
+
+    @property
+    def terminal(self) -> bool:
+        return self.state in {StreamState.CLOSED, StreamState.RESET}
+
+    def mark_open(self) -> None:
+        self._require(StreamState.OPENING)
+        self._opened = True
+        self.state = StreamState.OPEN
+
+    def send_eof(self) -> None:
+        if self._local_eof:
+            return
+        self._require_active("send EOF")
+        self._local_eof = True
+        self._update_eof_state()
+
+    def receive_eof(self) -> None:
+        if self._remote_eof:
+            return
+        self._require_active("receive EOF")
+        self._remote_eof = True
+        self._update_eof_state()
+
+    def validate_send_data(self) -> None:
+        if self.state not in {StreamState.OPEN, StreamState.REMOTE_EOF}:
+            raise MuxStateError(f"cannot send DATA while stream is {self.state.value}")
+
+    def validate_receive_data(self) -> None:
+        if self.state not in {StreamState.OPEN, StreamState.LOCAL_EOF}:
+            raise MuxStateError(f"cannot receive DATA while stream is {self.state.value}")
+
+    def send_close(self) -> None:
+        if self._close_sent:
+            return
+        if not self.open_initiator:
+            raise MuxStateError("only the OPEN initiator may send CLOSE")
+        self._require(StreamState.DRAINING)
+        self._close_sent = True
+        self.state = StreamState.CLOSING
+
+    def receive_close(self) -> None:
+        if self._close_received:
+            return
+        if self.open_initiator:
+            raise MuxStateError("OPEN initiator cannot receive CLOSE")
+        self._require(StreamState.DRAINING)
+        self._close_received = True
+        self.state = StreamState.CLOSING
+
+    def send_close_ack(self) -> None:
+        if self.state is StreamState.CLOSED:
+            return
+        if self.open_initiator or not self._close_received:
+            raise MuxStateError("CLOSE_ACK requires a received CLOSE on the acceptor")
+        self._finish_closed()
+
+    def receive_close_ack(self) -> None:
+        if self.state is StreamState.CLOSED:
+            return
+        if not self.open_initiator or not self._close_sent:
+            raise MuxStateError("CLOSE_ACK requires a sent CLOSE on the initiator")
+        self._finish_closed()
+
+    def send_reset(self) -> None:
+        self._finish_reset()
+
+    def receive_reset(self) -> None:
+        self._finish_reset()
+
+    def _update_eof_state(self) -> None:
+        if self._local_eof and self._remote_eof:
+            self.state = StreamState.DRAINING
+        elif self._local_eof:
+            self.state = StreamState.LOCAL_EOF
+        else:
+            self.state = StreamState.REMOTE_EOF
+
+    def _require_active(self, action: str) -> None:
+        if self.state not in {StreamState.OPEN, StreamState.LOCAL_EOF, StreamState.REMOTE_EOF}:
+            raise MuxStateError(f"cannot {action} while stream is {self.state.value}")
+
+    def _require(self, state: StreamState) -> None:
+        if self.state is not state:
+            raise MuxStateError(f"stream state must be {state.value}, found {self.state.value}")
+
+    def _finish_closed(self) -> None:
+        self.state = StreamState.CLOSED
+        self.terminal_kind = TerminalKind.CLOSED
+
+    def _finish_reset(self) -> None:
+        if self.state in {StreamState.CLOSED, StreamState.RESET}:
+            return
+        self.state = StreamState.RESET
+        self.terminal_kind = TerminalKind.RESET
+
+
+class ConnectionLifecycle:
+    def __init__(self) -> None:
+        self.state = ConnectionState.DISCONNECTED
+
+    def start_connecting(self) -> None:
+        self._require_any(ConnectionState.DISCONNECTED, ConnectionState.CLOSED, ConnectionState.FAILED)
+        self.state = ConnectionState.CONNECTING
+
+    @property
+    def terminal(self) -> bool:
+        return self.state in {ConnectionState.CLOSED, ConnectionState.FAILED}
+
+    @property
+    def draining(self) -> bool:
+        return self.state is ConnectionState.DRAINING
+
+    def start_handshake(self) -> None:
+        self._require(ConnectionState.CONNECTING)
+        self.state = ConnectionState.HANDSHAKING
+
+    def mark_ready(self) -> None:
+        self._require(ConnectionState.HANDSHAKING)
+        self.state = ConnectionState.READY
+
+    def start_draining(self) -> None:
+        if self.state is ConnectionState.DRAINING:
+            return
+        self._require(ConnectionState.READY)
+        self.state = ConnectionState.DRAINING
+
+    def mark_closed(self) -> None:
+        if self.state is ConnectionState.CLOSED:
+            return
+        self._require_any(ConnectionState.CONNECTING, ConnectionState.HANDSHAKING, ConnectionState.DRAINING)
+        self.state = ConnectionState.CLOSED
+
+    def mark_failed(self) -> None:
+        if self.state is ConnectionState.FAILED:
+            return
+        if self.state is ConnectionState.CLOSED:
+            raise MuxStateError("cannot fail connection while CLOSED")
+        self.state = ConnectionState.FAILED
+
+    def _require(self, state: ConnectionState) -> None:
+        if self.state is not state:
+            raise MuxStateError(f"connection state must be {state.value}, found {self.state.value}")
+
+    def _require_any(self, *states: ConnectionState) -> None:
+        if self.state not in states:
+            allowed = ", ".join(state.value for state in states)
+            raise MuxStateError(f"connection state must be one of {allowed}, found {self.state.value}")
+
+
+@dataclass(frozen=True, slots=True)
+class _Tombstone:
+    terminal_kind: TerminalKind
+    expires_at: float
+
+
+class StreamTombstones:
+    def __init__(
+        self,
+        *,
+        retention_seconds: float,
+        max_entries: int = 4096,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if retention_seconds <= 0:
+            raise ValueError("retention_seconds must be positive")
+        _positive(max_entries, "max_entries")
+        self._retention_seconds = float(retention_seconds)
+        self._max_entries = max_entries
+        self._clock = clock
+        self._entries: dict[int, _Tombstone] = {}
+
+    def add(self, stream_id: int, terminal_kind: TerminalKind) -> None:
+        if not isinstance(stream_id, int) or isinstance(stream_id, bool) or stream_id <= 0:
+            raise ValueError("stream_id must be a positive integer")
+        if not isinstance(terminal_kind, TerminalKind):
+            raise TypeError("terminal_kind must be a TerminalKind")
+        self.expire()
+        if stream_id not in self._entries and len(self._entries) >= self._max_entries:
+            raise MuxStateError("tombstone capacity is exhausted")
+        self._entries[stream_id] = _Tombstone(terminal_kind, self._clock() + self._retention_seconds)
+
+    def classify(self, stream_id: int, frame_type: FrameType) -> FrameDisposition:
+        self.expire()
+        if stream_id not in self._entries:
+            return FrameDisposition.UNKNOWN_STREAM
+        if frame_type is FrameType.OPEN:
+            return FrameDisposition.PROTOCOL_ERROR
+        if frame_type is FrameType.DATA:
+            return FrameDisposition.DISCARD_LATE_DATA
+        return FrameDisposition.IGNORE_TERMINAL
+
+    def expire(self) -> None:
+        now = self._clock()
+        expired = [stream_id for stream_id, tombstone in self._entries.items() if tombstone.expires_at <= now]
+        for stream_id in expired:
+            del self._entries[stream_id]
+
+
+def _positive(value: int, name: str) -> None:
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+
+
+def _positive_or_zero(value: int, name: str) -> None:
+    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+        raise ValueError(f"{name} must be a non-negative integer")