355 lines
11 KiB
Python
355 lines
11 KiB
Python
"""Minimal yamux (Yet Another Multiplexer) client for TCP tunneling.
|
|
|
|
Implements the client side of the yamux protocol as specified at
|
|
https://github.com/hashicorp/yamux/blob/master/spec.md
|
|
|
|
Only the subset needed for tunnel client operation is implemented:
|
|
opening streams, sending/receiving data, flow control, and keepalive.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import struct
|
|
import threading
|
|
from typing import Protocol
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Protocol constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_VERSION = 0
|
|
|
|
_TYPE_DATA = 0
|
|
_TYPE_WINDOW_UPDATE = 1
|
|
_TYPE_PING = 2
|
|
_TYPE_GO_AWAY = 3
|
|
|
|
_FLAG_SYN = 0x0001
|
|
_FLAG_ACK = 0x0002
|
|
_FLAG_FIN = 0x0004
|
|
_FLAG_RST = 0x0008
|
|
|
|
_HEADER_SIZE = 12
|
|
_HEADER_FMT = ">BBHII" # version(1), type(1), flags(2), streamID(4), length(4)
|
|
|
|
_INITIAL_WINDOW_SIZE = 256 * 1024 # 256 KB
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Byte-stream interface required by the session
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _ReadWriteCloser(Protocol):
|
|
def read(self, n: int) -> bytes: ...
|
|
|
|
def write(self, data: bytes) -> int: ...
|
|
|
|
def close(self) -> None: ...
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# YamuxStream
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class YamuxStream:
|
|
"""A single multiplexed stream within a yamux session.
|
|
|
|
Streams are created via :meth:`YamuxSession.open_stream` and provide
|
|
blocking read/write/close with per-stream flow control.
|
|
"""
|
|
|
|
def __init__(self, stream_id: int, session: YamuxSession) -> None:
|
|
self._id = stream_id
|
|
self._session = session
|
|
|
|
self._recv_buf = bytearray()
|
|
self._recv_cond = threading.Condition()
|
|
self._recv_closed = False
|
|
self._recv_error = False
|
|
self._recv_window = _INITIAL_WINDOW_SIZE
|
|
|
|
self._send_window = _INITIAL_WINDOW_SIZE
|
|
self._send_cond = threading.Condition()
|
|
self._send_closed = False
|
|
|
|
@property
|
|
def stream_id(self) -> int:
|
|
return self._id
|
|
|
|
def read(self, n: int) -> bytes:
|
|
"""Read up to *n* bytes, blocking until data is available.
|
|
|
|
Returns ``b""`` on EOF (FIN received).
|
|
Raises :class:`ConnectionResetError` on RST.
|
|
"""
|
|
delta_to_send = 0
|
|
|
|
with self._recv_cond:
|
|
while not self._recv_buf and not self._recv_closed and not self._recv_error:
|
|
self._recv_cond.wait()
|
|
|
|
if self._recv_error and not self._recv_buf:
|
|
raise ConnectionResetError("yamux stream reset by peer")
|
|
|
|
if not self._recv_buf:
|
|
return b""
|
|
|
|
size = min(n, len(self._recv_buf))
|
|
data = bytes(self._recv_buf[:size])
|
|
del self._recv_buf[:size]
|
|
|
|
consumed = _INITIAL_WINDOW_SIZE - self._recv_window
|
|
if consumed >= _INITIAL_WINDOW_SIZE // 2:
|
|
delta_to_send = consumed
|
|
self._recv_window += consumed
|
|
|
|
if delta_to_send > 0:
|
|
try:
|
|
self._session._send_window_update(self._id, delta_to_send)
|
|
except Exception:
|
|
pass
|
|
|
|
return data
|
|
|
|
def write(self, data: bytes) -> int:
|
|
"""Write *data*, blocking if the send window is exhausted."""
|
|
if self._send_closed:
|
|
raise BrokenPipeError("yamux stream closed for writing")
|
|
|
|
offset = 0
|
|
mv = memoryview(data)
|
|
|
|
while offset < len(data):
|
|
with self._send_cond:
|
|
while self._send_window == 0 and not self._send_closed:
|
|
self._send_cond.wait()
|
|
if self._send_closed:
|
|
raise BrokenPipeError("yamux stream closed for writing")
|
|
chunk = min(len(data) - offset, self._send_window)
|
|
self._send_window -= chunk
|
|
|
|
self._session._send_data(self._id, bytes(mv[offset : offset + chunk]))
|
|
offset += chunk
|
|
|
|
return len(data)
|
|
|
|
def close(self) -> None:
|
|
"""Close the stream (sends FIN to the remote end)."""
|
|
if not self._send_closed:
|
|
self._send_closed = True
|
|
try:
|
|
self._session._send_frame(_TYPE_DATA, _FLAG_FIN, self._id, 0)
|
|
except Exception:
|
|
pass
|
|
|
|
with self._recv_cond:
|
|
self._recv_closed = True
|
|
self._recv_cond.notify_all()
|
|
with self._send_cond:
|
|
self._send_cond.notify_all()
|
|
|
|
# -- Internal: called by YamuxSession._read_loop ------------------------
|
|
|
|
def _receive_data(self, data: bytes) -> None:
|
|
with self._recv_cond:
|
|
self._recv_buf.extend(data)
|
|
self._recv_window -= len(data)
|
|
self._recv_cond.notify_all()
|
|
|
|
def _receive_fin(self) -> None:
|
|
with self._recv_cond:
|
|
self._recv_closed = True
|
|
self._recv_cond.notify_all()
|
|
|
|
def _receive_rst(self) -> None:
|
|
with self._recv_cond:
|
|
self._recv_error = True
|
|
self._recv_cond.notify_all()
|
|
with self._send_cond:
|
|
self._send_closed = True
|
|
self._send_cond.notify_all()
|
|
|
|
def _update_send_window(self, delta: int) -> None:
|
|
with self._send_cond:
|
|
self._send_window += delta
|
|
self._send_cond.notify_all()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# YamuxSession
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class YamuxSession:
|
|
"""Client-side yamux session over a byte-stream connection.
|
|
|
|
The connection must implement ``read(n) -> bytes``, ``write(data) -> int``,
|
|
and ``close() -> None``. Typically this is a :class:`_WSAdapter` wrapping
|
|
a WebSocket.
|
|
|
|
Usage::
|
|
|
|
session = YamuxSession(conn)
|
|
stream = session.open_stream()
|
|
stream.write(b"hello")
|
|
data = stream.read(1024)
|
|
stream.close()
|
|
session.close()
|
|
"""
|
|
|
|
def __init__(self, conn: _ReadWriteCloser) -> None:
|
|
self._conn = conn
|
|
self._streams: dict[int, YamuxStream] = {}
|
|
self._next_stream_id = 1 # client uses odd IDs
|
|
self._lock = threading.Lock()
|
|
self._write_lock = threading.Lock()
|
|
self._closed = False
|
|
self._shutdown_event = threading.Event()
|
|
|
|
self._reader_thread = threading.Thread(
|
|
target=self._read_loop, daemon=True, name="yamux-reader"
|
|
)
|
|
self._reader_thread.start()
|
|
|
|
self._keepalive_thread = threading.Thread(
|
|
target=self._keepalive_loop, daemon=True, name="yamux-keepalive"
|
|
)
|
|
self._keepalive_thread.start()
|
|
|
|
@property
|
|
def is_closed(self) -> bool:
|
|
return self._closed
|
|
|
|
def open_stream(self) -> YamuxStream:
|
|
"""Open a new multiplexed stream.
|
|
|
|
Raises :class:`RuntimeError` if the session is closed.
|
|
"""
|
|
with self._lock:
|
|
if self._closed:
|
|
raise RuntimeError("yamux session is closed")
|
|
stream_id = self._next_stream_id
|
|
self._next_stream_id += 2
|
|
stream = YamuxStream(stream_id, self)
|
|
self._streams[stream_id] = stream
|
|
|
|
self._send_frame(_TYPE_WINDOW_UPDATE, _FLAG_SYN, stream_id, 0)
|
|
return stream
|
|
|
|
def close(self) -> None:
|
|
"""Close the session and all streams."""
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
self._shutdown_event.set()
|
|
|
|
try:
|
|
self._send_frame(_TYPE_GO_AWAY, 0, 0, 0)
|
|
except Exception:
|
|
pass
|
|
|
|
with self._lock:
|
|
for stream in self._streams.values():
|
|
stream._receive_rst()
|
|
|
|
try:
|
|
self._conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
# -- Frame I/O ----------------------------------------------------------
|
|
|
|
def _send_frame(
|
|
self, msg_type: int, flags: int, stream_id: int, length: int
|
|
) -> None:
|
|
hdr = struct.pack(_HEADER_FMT, _VERSION, msg_type, flags, stream_id, length)
|
|
with self._write_lock:
|
|
self._conn.write(hdr)
|
|
|
|
def _send_data(self, stream_id: int, data: bytes) -> None:
|
|
hdr = struct.pack(_HEADER_FMT, _VERSION, _TYPE_DATA, 0, stream_id, len(data))
|
|
with self._write_lock:
|
|
self._conn.write(hdr + data)
|
|
|
|
def _send_window_update(self, stream_id: int, delta: int) -> None:
|
|
self._send_frame(_TYPE_WINDOW_UPDATE, 0, stream_id, delta)
|
|
|
|
# -- Read loop ----------------------------------------------------------
|
|
|
|
def _read_loop(self) -> None:
|
|
try:
|
|
while not self._closed:
|
|
hdr_bytes = self._conn.read(_HEADER_SIZE)
|
|
if len(hdr_bytes) < _HEADER_SIZE:
|
|
break
|
|
|
|
_ver, msg_type, flags, stream_id, length = struct.unpack(
|
|
_HEADER_FMT, hdr_bytes
|
|
)
|
|
|
|
if msg_type == _TYPE_DATA:
|
|
self._handle_data(flags, stream_id, length)
|
|
elif msg_type == _TYPE_WINDOW_UPDATE:
|
|
self._handle_window_update(flags, stream_id, length)
|
|
elif msg_type == _TYPE_PING:
|
|
self._handle_ping(flags, length)
|
|
elif msg_type == _TYPE_GO_AWAY:
|
|
break
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
if not self._closed:
|
|
self._closed = True
|
|
self._shutdown_event.set()
|
|
with self._lock:
|
|
for stream in self._streams.values():
|
|
stream._receive_rst()
|
|
|
|
def _handle_data(self, flags: int, stream_id: int, length: int) -> None:
|
|
payload = self._conn.read(length) if length > 0 else b""
|
|
|
|
with self._lock:
|
|
stream = self._streams.get(stream_id)
|
|
if stream is None:
|
|
return
|
|
|
|
if payload:
|
|
stream._receive_data(payload)
|
|
if flags & _FLAG_FIN:
|
|
stream._receive_fin()
|
|
if flags & _FLAG_RST:
|
|
stream._receive_rst()
|
|
|
|
def _handle_window_update(self, flags: int, stream_id: int, length: int) -> None:
|
|
with self._lock:
|
|
stream = self._streams.get(stream_id)
|
|
if stream is None:
|
|
return
|
|
|
|
if length > 0:
|
|
stream._update_send_window(length)
|
|
if flags & _FLAG_FIN:
|
|
stream._receive_fin()
|
|
if flags & _FLAG_RST:
|
|
stream._receive_rst()
|
|
|
|
def _handle_ping(self, flags: int, opaque: int) -> None:
|
|
if flags & _FLAG_SYN:
|
|
try:
|
|
self._send_frame(_TYPE_PING, _FLAG_ACK, 0, opaque)
|
|
except Exception:
|
|
pass
|
|
|
|
# -- Keepalive ----------------------------------------------------------
|
|
|
|
def _keepalive_loop(self) -> None:
|
|
ping_id = 0
|
|
while not self._shutdown_event.wait(30):
|
|
ping_id += 1
|
|
try:
|
|
self._send_frame(_TYPE_PING, _FLAG_SYN, 0, ping_id)
|
|
except Exception:
|
|
break
|