485 lines
14 KiB
Python
485 lines
14 KiB
Python
"""TCP tunnel for accessing services running inside sandboxes.
|
|
|
|
Establishes a WebSocket connection to the daemon's ``/tunnel`` endpoint,
|
|
runs a yamux multiplexing session on top, and forwards local TCP connections
|
|
through yamux streams to the target port inside the sandbox.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import socket
|
|
import struct
|
|
import threading
|
|
import time
|
|
from collections.abc import Mapping
|
|
from typing import TYPE_CHECKING, Any, Optional
|
|
|
|
from langsmith.sandbox._helpers import merge_headers
|
|
|
|
if TYPE_CHECKING:
|
|
from langsmith.sandbox._yamux import YamuxSession, YamuxStream
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tunnel connect-header protocol (layered on top of yamux streams)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
PROTOCOL_VERSION = 0x01
|
|
|
|
STATUS_OK = 0x00
|
|
STATUS_PORT_NOT_ALLOWED = 0x01
|
|
STATUS_DIAL_FAILED = 0x02
|
|
STATUS_UNSUPPORTED_VERSION = 0x03
|
|
|
|
_CONNECT_HEADER_FMT = ">BH" # version(1) + port(2, big-endian)
|
|
|
|
|
|
def _write_connect_header(stream: YamuxStream, port: int) -> None:
|
|
"""Write the 3-byte connect header on a freshly opened yamux stream."""
|
|
stream.write(struct.pack(_CONNECT_HEADER_FMT, PROTOCOL_VERSION, port))
|
|
|
|
|
|
def _read_status(stream: YamuxStream) -> int:
|
|
"""Read the 1-byte status response from the daemon."""
|
|
data = stream.read(1)
|
|
if not data:
|
|
raise ConnectionError("tunnel: connection closed before status")
|
|
return data[0]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WebSocket adapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _WSAdapter:
|
|
"""Adapts the ``websockets`` message API to a byte-stream interface.
|
|
|
|
yamux requires a plain read/write/close byte stream. WebSocket is
|
|
message-based, so this adapter buffers partially consumed messages on
|
|
reads and sends one binary message per write.
|
|
"""
|
|
|
|
def __init__(self, ws: Any) -> None:
|
|
self._ws = ws
|
|
self._buf = bytearray()
|
|
self._write_lock = threading.Lock()
|
|
|
|
def read(self, n: int) -> bytes:
|
|
while len(self._buf) < n:
|
|
msg = self._ws.recv()
|
|
if isinstance(msg, str):
|
|
msg = msg.encode()
|
|
self._buf.extend(msg)
|
|
|
|
result = bytes(self._buf[:n])
|
|
del self._buf[:n]
|
|
return result
|
|
|
|
def write(self, data: bytes) -> int:
|
|
with self._write_lock:
|
|
self._ws.send(data)
|
|
return len(data)
|
|
|
|
def close(self) -> None:
|
|
try:
|
|
self._ws.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bridge: bidirectional copy between yamux stream and TCP socket
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_BRIDGE_BUF_SIZE = 16384
|
|
|
|
|
|
def _bridge(stream: YamuxStream, tcp_conn: socket.socket) -> None:
|
|
"""Copy data bidirectionally until one side closes or errors."""
|
|
done = threading.Event()
|
|
|
|
def _stream_to_tcp() -> None:
|
|
try:
|
|
while True:
|
|
data = stream.read(_BRIDGE_BUF_SIZE)
|
|
if not data:
|
|
break
|
|
tcp_conn.sendall(data)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
done.set()
|
|
|
|
def _tcp_to_stream() -> None:
|
|
try:
|
|
while True:
|
|
data = tcp_conn.recv(_BRIDGE_BUF_SIZE)
|
|
if not data:
|
|
break
|
|
stream.write(data)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
done.set()
|
|
|
|
t1 = threading.Thread(target=_stream_to_tcp, daemon=True)
|
|
t2 = threading.Thread(target=_tcp_to_stream, daemon=True)
|
|
t1.start()
|
|
t2.start()
|
|
|
|
done.wait()
|
|
|
|
try:
|
|
stream.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
tcp_conn.shutdown(socket.SHUT_RDWR)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
tcp_conn.close()
|
|
except OSError:
|
|
pass
|
|
|
|
t1.join(timeout=5)
|
|
t2.join(timeout=5)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tunnel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _ensure_websockets():
|
|
"""Import websockets sync client or raise a clear error."""
|
|
try:
|
|
from websockets.sync.client import connect as ws_connect
|
|
|
|
return ws_connect
|
|
except ImportError:
|
|
raise ImportError(
|
|
"TCP tunnel requires the 'websockets' package. "
|
|
"Install it with: pip install 'langsmith[sandbox]'"
|
|
) from None
|
|
|
|
|
|
class Tunnel:
|
|
"""TCP tunnel to a port inside a sandbox.
|
|
|
|
Opens a local TCP listener and forwards each accepted connection through
|
|
a yamux-multiplexed WebSocket to the daemon, which dials the target port
|
|
inside the sandbox.
|
|
|
|
Typically used as a context manager::
|
|
|
|
with sandbox.tunnel(remote_port=5432) as t:
|
|
conn = psycopg2.connect(host="127.0.0.1", port=t.local_port)
|
|
|
|
Or with explicit lifecycle::
|
|
|
|
t = sandbox.tunnel(remote_port=5432)
|
|
# ... use tunnel ...
|
|
t.close()
|
|
"""
|
|
|
|
_BACKOFF_BASE = 0.5
|
|
_BACKOFF_MAX = 8.0
|
|
|
|
def __init__(
|
|
self,
|
|
dataplane_url: str,
|
|
api_key: Optional[str],
|
|
remote_port: int,
|
|
*,
|
|
local_port: int = 0,
|
|
max_reconnects: int = 3,
|
|
headers: Optional[Mapping[str, str]] = None,
|
|
) -> None:
|
|
self._dataplane_url = dataplane_url
|
|
self._api_key = api_key
|
|
self._headers = headers
|
|
self._remote_port = remote_port
|
|
self._requested_local_port = local_port or remote_port
|
|
self._local_port = self._requested_local_port
|
|
self._max_reconnects = max_reconnects
|
|
|
|
self._ws: object = None
|
|
self._yamux: Optional[YamuxSession] = None
|
|
self._server_socket: Optional[socket.socket] = None
|
|
self._accept_thread: Optional[threading.Thread] = None
|
|
self._reconnect_lock = threading.Lock()
|
|
self._closed = False
|
|
self._started = False
|
|
|
|
@property
|
|
def local_port(self) -> int:
|
|
"""Local port the tunnel is listening on."""
|
|
return self._local_port
|
|
|
|
@property
|
|
def remote_port(self) -> int:
|
|
"""Port inside the sandbox that the tunnel connects to."""
|
|
return self._remote_port
|
|
|
|
# -- Context manager ----------------------------------------------------
|
|
|
|
def __enter__(self) -> Tunnel:
|
|
return self
|
|
|
|
def __exit__(self, *args: object) -> None:
|
|
self.close()
|
|
|
|
# -- Lifecycle ----------------------------------------------------------
|
|
|
|
def _start(self) -> None:
|
|
if self._started:
|
|
return
|
|
self._started = True
|
|
|
|
try:
|
|
self._do_start()
|
|
except Exception:
|
|
self.close()
|
|
raise
|
|
|
|
def _do_start(self) -> None:
|
|
self._connect()
|
|
|
|
self._server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self._server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
|
|
# Check if another process is actively listening on this port.
|
|
# SO_REUSEADDR lets us rebind over TIME_WAIT, but we don't want to
|
|
# silently steal a port from a running service.
|
|
port = self._requested_local_port
|
|
if port != 0:
|
|
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
try:
|
|
probe.settimeout(0.5)
|
|
probe.connect(("127.0.0.1", port))
|
|
probe.close()
|
|
raise OSError(
|
|
f"Port {port} is already in use by another service. "
|
|
f"Choose a different local_port."
|
|
)
|
|
except ConnectionRefusedError:
|
|
pass # nothing listening — safe to bind
|
|
except OSError as e:
|
|
if "Connection refused" in str(e):
|
|
pass # same as above, different OS error message
|
|
elif "already in use" in str(e).lower():
|
|
raise
|
|
else:
|
|
pass # TIME_WAIT or other transient state — safe to bind
|
|
finally:
|
|
try:
|
|
probe.close()
|
|
except OSError:
|
|
pass
|
|
|
|
self._server_socket.bind(("127.0.0.1", self._requested_local_port))
|
|
self._server_socket.listen(128)
|
|
self._local_port = self._server_socket.getsockname()[1]
|
|
|
|
self._accept_thread = threading.Thread(
|
|
target=self._accept_loop, daemon=True, name="tunnel-accept"
|
|
)
|
|
self._accept_thread.start()
|
|
|
|
def _connect(self) -> None:
|
|
"""Establish (or re-establish) the WebSocket + yamux session."""
|
|
from langsmith.sandbox._yamux import YamuxSession
|
|
|
|
old_yamux = self._yamux
|
|
if old_yamux:
|
|
try:
|
|
old_yamux.close()
|
|
except Exception:
|
|
pass
|
|
|
|
ws_connect = _ensure_websockets()
|
|
ws_url = self._build_ws_url()
|
|
headers = merge_headers(
|
|
{"X-Api-Key": self._api_key} if self._api_key else None,
|
|
self._headers,
|
|
)
|
|
|
|
self._ws = ws_connect(
|
|
ws_url,
|
|
additional_headers=headers,
|
|
open_timeout=15,
|
|
close_timeout=5,
|
|
ping_interval=None, # yamux handles keepalive
|
|
)
|
|
|
|
adapter = _WSAdapter(self._ws)
|
|
self._yamux = YamuxSession(adapter)
|
|
|
|
def _ensure_session(self) -> YamuxSession:
|
|
"""Return a live yamux session, reconnecting if needed."""
|
|
from langsmith.sandbox._exceptions import TunnelError
|
|
|
|
if self._yamux and not self._yamux.is_closed:
|
|
return self._yamux
|
|
|
|
with self._reconnect_lock:
|
|
if self._yamux and not self._yamux.is_closed:
|
|
return self._yamux
|
|
|
|
last_err: Optional[Exception] = None
|
|
for attempt in range(self._max_reconnects):
|
|
try:
|
|
self._connect()
|
|
logger.debug("tunnel: reconnected (attempt %d)", attempt + 1)
|
|
return self._yamux # type: ignore[return-value]
|
|
except Exception as exc:
|
|
last_err = exc
|
|
if attempt < self._max_reconnects - 1:
|
|
delay = min(
|
|
self._BACKOFF_BASE * (2**attempt),
|
|
self._BACKOFF_MAX,
|
|
)
|
|
time.sleep(delay)
|
|
|
|
raise TunnelError(
|
|
f"tunnel: reconnect failed after {self._max_reconnects} attempts"
|
|
) from last_err
|
|
|
|
def close(self) -> None:
|
|
"""Shut down the tunnel, closing all connections."""
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
|
|
if self._server_socket:
|
|
try:
|
|
self._server_socket.close()
|
|
except OSError:
|
|
pass
|
|
|
|
if self._yamux:
|
|
self._yamux.close()
|
|
|
|
# -- Internal -----------------------------------------------------------
|
|
|
|
def _accept_loop(self) -> None:
|
|
while not self._closed:
|
|
try:
|
|
conn, _ = self._server_socket.accept() # type: ignore[union-attr]
|
|
except OSError:
|
|
break
|
|
threading.Thread(
|
|
target=self._handle_conn,
|
|
args=(conn,),
|
|
daemon=True,
|
|
name="tunnel-bridge",
|
|
).start()
|
|
|
|
def _handle_conn(self, tcp_conn: socket.socket) -> None:
|
|
try:
|
|
session = self._ensure_session()
|
|
stream = session.open_stream()
|
|
_write_connect_header(stream, self._remote_port)
|
|
status = _read_status(stream)
|
|
|
|
if status == STATUS_OK:
|
|
_bridge(stream, tcp_conn)
|
|
return
|
|
|
|
stream.close()
|
|
tcp_conn.close()
|
|
|
|
if status == STATUS_PORT_NOT_ALLOWED:
|
|
logger.warning(
|
|
"tunnel: port %d not allowed by daemon",
|
|
self._remote_port,
|
|
)
|
|
elif status == STATUS_DIAL_FAILED:
|
|
logger.warning(
|
|
"tunnel: nothing listening on port %d inside sandbox",
|
|
self._remote_port,
|
|
)
|
|
elif status == STATUS_UNSUPPORTED_VERSION:
|
|
logger.warning(
|
|
"tunnel: protocol version mismatch (client v%d)",
|
|
PROTOCOL_VERSION,
|
|
)
|
|
else:
|
|
logger.warning("tunnel: unknown status %d", status)
|
|
|
|
except Exception as exc:
|
|
logger.debug("tunnel: connection handler error: %s", exc)
|
|
try:
|
|
tcp_conn.close()
|
|
except OSError:
|
|
pass
|
|
|
|
def _build_ws_url(self) -> str:
|
|
url = self._dataplane_url.rstrip("/")
|
|
url = url.replace("https://", "wss://").replace("http://", "ws://")
|
|
return f"{url}/tunnel"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AsyncTunnel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AsyncTunnel:
|
|
"""Async wrapper around :class:`Tunnel`.
|
|
|
|
The underlying tunnel runs in background threads (TCP listener + bridges);
|
|
async context-manager methods delegate to the sync tunnel via the event
|
|
loop's executor.
|
|
|
|
Usage::
|
|
|
|
async with await sandbox.tunnel(remote_port=5432) as t:
|
|
conn = await asyncpg.connect(host="127.0.0.1", port=t.local_port)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataplane_url: str,
|
|
api_key: Optional[str],
|
|
remote_port: int,
|
|
*,
|
|
local_port: int = 0,
|
|
max_reconnects: int = 3,
|
|
headers: Optional[Mapping[str, str]] = None,
|
|
) -> None:
|
|
self._tunnel = Tunnel(
|
|
dataplane_url,
|
|
api_key,
|
|
remote_port,
|
|
local_port=local_port,
|
|
max_reconnects=max_reconnects,
|
|
headers=headers,
|
|
)
|
|
|
|
@property
|
|
def local_port(self) -> int:
|
|
return self._tunnel.local_port
|
|
|
|
@property
|
|
def remote_port(self) -> int:
|
|
return self._tunnel.remote_port
|
|
|
|
async def __aenter__(self) -> AsyncTunnel:
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(None, self._tunnel._start)
|
|
return self
|
|
|
|
async def __aexit__(self, *args: object) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(None, self._tunnel.close)
|
|
|
|
def close(self) -> None:
|
|
"""Shut down the tunnel (sync, safe to call from any context)."""
|
|
self._tunnel.close()
|