406 lines
16 KiB
Python
406 lines
16 KiB
Python
"""Relay server — WebSocket gossip hub + circuit relay proxy.
|
|
|
|
HTTP API (served via asyncio-based handler on same port):
|
|
GET /health → {"status": "ok", "peers": N}
|
|
GET /v1/peers → [{peer_id, addr, last_seen}]
|
|
POST /v1/gossip → accept a gossip envelope, fan out to connected peers
|
|
|
|
WebSocket endpoints:
|
|
ws[s]://host:port/ws → gossip PubSub connection
|
|
ws[s]://host:port/relay/{peer_id} → circuit relay to that peer
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import threading
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from .peer_registry import PeerRegistry
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# Activation tensors ride the relay as base64 inside one JSON frame, so the
|
|
# websockets default of 1 MiB rejects any real prefill (close code 1009).
|
|
DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024
|
|
|
|
|
|
def ws_max_size() -> int | None:
|
|
"""Max inbound WebSocket frame size; MESHNET_WS_MAX_BYTES<=0 means unlimited."""
|
|
raw = os.environ.get("MESHNET_WS_MAX_BYTES", "").strip()
|
|
if not raw:
|
|
return DEFAULT_WS_MAX_BYTES
|
|
try:
|
|
value = int(raw)
|
|
except ValueError:
|
|
return DEFAULT_WS_MAX_BYTES
|
|
return None if value <= 0 else value
|
|
|
|
|
|
# Binary relay frame: JSON header (request/response metadata) + raw body in one
|
|
# WebSocket binary message. Activation bodies stay bytes end to end — no base64
|
|
# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the
|
|
# control plane (gossip, peer-register, streamed SSE responses).
|
|
BINARY_FRAME_MAGIC = b"MRF1"
|
|
_RPC_PEER_DISCONNECTED = object()
|
|
|
|
|
|
def encode_binary_frame(header: dict, body: bytes) -> bytes:
|
|
header_bytes = json.dumps(header, separators=(",", ":")).encode()
|
|
return b"".join((BINARY_FRAME_MAGIC, len(header_bytes).to_bytes(4, "big"), header_bytes, body))
|
|
|
|
|
|
def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
|
|
if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC:
|
|
raise ValueError("not a meshnet binary relay frame")
|
|
header_len = int.from_bytes(frame[4:8], "big")
|
|
header = json.loads(frame[8:8 + header_len].decode())
|
|
return header, frame[8 + header_len:]
|
|
|
|
|
|
class RelayServer:
|
|
"""Async WebSocket relay server that runs in a background thread.
|
|
|
|
Usage::
|
|
|
|
server = RelayServer(host="0.0.0.0", port=8765)
|
|
port = server.start() # returns actual port
|
|
...
|
|
server.stop()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "0.0.0.0",
|
|
port: int = 8765,
|
|
ssl_cert: Path | None = None,
|
|
ssl_key: Path | None = None,
|
|
max_peers: int = 500,
|
|
rpc_timeout: float = 310.0,
|
|
rpc_idle_timeout: float = 120.0,
|
|
):
|
|
self.host = host
|
|
self.port = port
|
|
self.ssl_cert = ssl_cert
|
|
self.ssl_key = ssl_key
|
|
self.max_peers = max_peers
|
|
self.rpc_timeout = rpc_timeout
|
|
self.rpc_idle_timeout = rpc_idle_timeout
|
|
|
|
self._registry = PeerRegistry()
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self._server = None
|
|
self._actual_port: int = port
|
|
self._ready = threading.Event()
|
|
self._running = False
|
|
self._stop_event: asyncio.Event | None = None
|
|
# relay request id → (target peer, requester request id, response queue).
|
|
# The relay-generated id prevents two Route Sessions using the same
|
|
# legacy request_id from overwriting each other's pending response.
|
|
self._pending_rpc: dict[str, tuple[str, str, asyncio.Queue]] = {}
|
|
|
|
@property
|
|
def registry(self) -> PeerRegistry:
|
|
return self._registry
|
|
|
|
def start(self) -> int:
|
|
"""Start server in background thread. Returns actual bound port."""
|
|
self._running = True
|
|
self._loop = asyncio.new_event_loop()
|
|
self._thread = threading.Thread(target=self._run, daemon=True, name="relay")
|
|
self._thread.start()
|
|
self._ready.wait(timeout=5)
|
|
return self._actual_port
|
|
|
|
def stop(self) -> None:
|
|
self._running = False
|
|
if self._loop and self._stop_event is not None:
|
|
self._loop.call_soon_threadsafe(self._stop_event.set)
|
|
if self._thread:
|
|
self._thread.join(timeout=3.0)
|
|
|
|
def _fail_pending_for_peer(self, peer_id: str) -> None:
|
|
"""Wake requesters immediately when their selected bridge disconnects."""
|
|
for target, _, queue in tuple(self._pending_rpc.values()):
|
|
if target == peer_id:
|
|
queue.put_nowait(_RPC_PEER_DISCONNECTED)
|
|
|
|
def _run(self) -> None:
|
|
asyncio.set_event_loop(self._loop)
|
|
self._loop.run_until_complete(self._serve())
|
|
|
|
async def _serve(self) -> None:
|
|
import websockets # type: ignore[import]
|
|
import websockets.server # type: ignore[import]
|
|
|
|
ssl_ctx = None
|
|
if self.ssl_cert and self.ssl_key:
|
|
import ssl
|
|
ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
|
ssl_ctx.load_cert_chain(str(self.ssl_cert), str(self.ssl_key))
|
|
|
|
server = await websockets.serve(
|
|
self._handle_connection,
|
|
self.host,
|
|
self.port,
|
|
ssl=ssl_ctx,
|
|
max_size=ws_max_size(),
|
|
# Bulk payloads are zstd-compressed at the pipeline layer already;
|
|
# per-message deflate would recompress them on every hop for nothing.
|
|
compression=None,
|
|
)
|
|
# Record actual port after bind
|
|
for sock in server.sockets or []:
|
|
self._actual_port = sock.getsockname()[1]
|
|
break
|
|
|
|
self._stop_event = asyncio.Event()
|
|
self._server = server
|
|
self._ready.set()
|
|
log.info("Relay listening on %s:%d", self.host, self._actual_port)
|
|
|
|
await self._stop_event.wait()
|
|
server.close()
|
|
await server.wait_closed()
|
|
|
|
async def _handle_connection(self, ws) -> None:
|
|
"""Dispatch incoming WebSocket to gossip hub or circuit relay."""
|
|
try:
|
|
path = ws.request.path
|
|
except AttributeError:
|
|
path = getattr(ws, "path", "/ws")
|
|
|
|
if path.startswith("/relay/"):
|
|
peer_id = path[len("/relay/"):]
|
|
await self._handle_circuit_relay(ws, peer_id)
|
|
elif path.startswith("/rpc/"):
|
|
peer_id = path[len("/rpc/"):]
|
|
await self._handle_rpc(ws, peer_id)
|
|
elif path == "/health":
|
|
await ws.send(json.dumps({"status": "ok", "peers": len(self._registry)}))
|
|
await ws.close()
|
|
elif path == "/v1/peers":
|
|
await ws.send(json.dumps(self._registry.list_peers()))
|
|
await ws.close()
|
|
else:
|
|
await self._handle_gossip(ws)
|
|
|
|
async def _handle_gossip(self, ws) -> None:
|
|
"""Accept a gossip peer connection, register it, and fan out messages."""
|
|
peer_id: str | None = None
|
|
peer_addr: str = ""
|
|
|
|
try:
|
|
async for raw in ws:
|
|
# Binary frames are relay-http-response bodies from a bridge —
|
|
# route them to the waiting rpc requester as-is, never fan out.
|
|
if isinstance(raw, (bytes, bytearray)):
|
|
try:
|
|
header, _ = decode_binary_frame(bytes(raw))
|
|
except (ValueError, json.JSONDecodeError):
|
|
continue
|
|
pending = self._pending_rpc.get(header.get("request_id"))
|
|
if pending is not None:
|
|
pending[2].put_nowait(bytes(raw))
|
|
continue
|
|
try:
|
|
envelope = json.loads(raw)
|
|
except (json.JSONDecodeError, TypeError):
|
|
continue
|
|
|
|
topic = envelope.get("topic", "")
|
|
from_peer = envelope.get("from_peer", "")
|
|
|
|
# Handle peer registration message
|
|
if topic == "peer-register":
|
|
payload = envelope.get("payload", {})
|
|
peer_id = payload.get("peer_id") or from_peer
|
|
peer_addr = payload.get("addr", "")
|
|
if len(self._registry) >= self.max_peers:
|
|
await ws.close(1008, "relay at capacity")
|
|
return
|
|
self._registry.register(peer_id, peer_addr, ws)
|
|
log.debug("Peer registered: %s", peer_id)
|
|
# Send current peer list back
|
|
await ws.send(json.dumps({
|
|
"topic": "peer-list",
|
|
"version": 1,
|
|
"from_peer": "relay",
|
|
"payload": {"peers": self._registry.list_peers()},
|
|
}))
|
|
continue
|
|
|
|
if topic == "relay-http-response":
|
|
payload = envelope.get("payload", {})
|
|
request_id = payload.get("request_id")
|
|
pending = self._pending_rpc.get(request_id)
|
|
if pending is not None:
|
|
pending[2].put_nowait(payload)
|
|
continue
|
|
|
|
# Fan out to all other registered peers
|
|
if peer_id:
|
|
self._registry.touch(peer_id)
|
|
fan_out_peers = self._registry.all_except(peer_id or "")
|
|
await _broadcast(raw, fan_out_peers)
|
|
|
|
except Exception as exc:
|
|
log.debug("Gossip connection error: %s", exc)
|
|
finally:
|
|
if peer_id:
|
|
self._registry.unregister(peer_id)
|
|
# Do not leave a requester waiting for its full timeout after
|
|
# the selected bridge goes away.
|
|
self._fail_pending_for_peer(peer_id)
|
|
log.debug("Peer unregistered: %s", peer_id)
|
|
|
|
async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None:
|
|
"""Proxy WebSocket traffic between ws_requester and target_peer_id's ws."""
|
|
target = self._registry.get(target_peer_id)
|
|
if target is None:
|
|
try:
|
|
await ws_requester.send(json.dumps({
|
|
"error": f"peer {target_peer_id!r} not connected to relay"
|
|
}))
|
|
await ws_requester.close(1011, "target peer not found")
|
|
except Exception:
|
|
pass
|
|
return
|
|
|
|
log.debug("Circuit relay: ??? → %s", target_peer_id)
|
|
|
|
async def pipe(src, dst) -> None:
|
|
try:
|
|
async for msg in src:
|
|
await dst.send(msg)
|
|
except Exception:
|
|
pass
|
|
|
|
await asyncio.gather(
|
|
pipe(ws_requester, target.ws),
|
|
pipe(target.ws, ws_requester),
|
|
return_exceptions=True,
|
|
)
|
|
|
|
async def _handle_rpc(self, ws_requester, target_peer_id: str) -> None:
|
|
"""Relay sequential HTTP-shaped requests over one requester connection."""
|
|
target = self._registry.get(target_peer_id)
|
|
if target is None:
|
|
await ws_requester.send(json.dumps({
|
|
"status": 503,
|
|
"headers": {"Content-Type": "application/json"},
|
|
"body": json.dumps({"error": f"peer {target_peer_id!r} not connected to relay"}),
|
|
}))
|
|
await ws_requester.close()
|
|
return
|
|
|
|
while True:
|
|
try:
|
|
raw = await ws_requester.recv()
|
|
except Exception:
|
|
return
|
|
|
|
requester_request_id: str | None = None
|
|
relay_request_id = uuid.uuid4().hex
|
|
try:
|
|
if isinstance(raw, (bytes, bytearray)):
|
|
header, body = decode_binary_frame(bytes(raw))
|
|
requester_request_id = str(header.get("request_id") or uuid.uuid4())
|
|
header["request_id"] = relay_request_id
|
|
header["target_peer"] = target_peer_id
|
|
outbound: str | bytes = encode_binary_frame(header, body)
|
|
else:
|
|
payload = json.loads(raw)
|
|
requester_request_id = str(payload.get("request_id") or uuid.uuid4())
|
|
payload["request_id"] = relay_request_id
|
|
payload["target_peer"] = target_peer_id
|
|
outbound = json.dumps({
|
|
"topic": "relay-http-request",
|
|
"version": 1,
|
|
"from_peer": "relay",
|
|
"payload": payload,
|
|
})
|
|
except Exception:
|
|
await ws_requester.close(1003, "invalid relay rpc request")
|
|
return
|
|
|
|
queue: asyncio.Queue = asyncio.Queue()
|
|
self._pending_rpc[relay_request_id] = (
|
|
target_peer_id, requester_request_id, queue,
|
|
)
|
|
overall_timeout = self.rpc_timeout
|
|
idle_timeout = self.rpc_idle_timeout
|
|
loop = asyncio.get_running_loop()
|
|
deadline = loop.time() + overall_timeout
|
|
target = self._registry.get(target_peer_id)
|
|
try:
|
|
if target is None:
|
|
await ws_requester.send(json.dumps({
|
|
"request_id": requester_request_id,
|
|
"status": 503,
|
|
"headers": {"Content-Type": "application/json"},
|
|
"body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}),
|
|
}))
|
|
continue
|
|
await target.ws.send(outbound)
|
|
# Streamed responses end with done=true. Binary and legacy JSON
|
|
# responses are complete in one frame.
|
|
while True:
|
|
remaining = deadline - loop.time()
|
|
if remaining <= 0:
|
|
raise asyncio.TimeoutError
|
|
frame = await asyncio.wait_for(
|
|
queue.get(), timeout=min(idle_timeout, remaining)
|
|
)
|
|
if frame is _RPC_PEER_DISCONNECTED:
|
|
raise ConnectionError(f"peer {target_peer_id!r} disconnected")
|
|
if isinstance(frame, (bytes, bytearray)):
|
|
header, body = decode_binary_frame(bytes(frame))
|
|
header["request_id"] = requester_request_id
|
|
await ws_requester.send(encode_binary_frame(header, body))
|
|
break
|
|
response = dict(frame)
|
|
response["request_id"] = requester_request_id
|
|
await ws_requester.send(json.dumps(response))
|
|
if not response.get("stream") or response.get("done"):
|
|
break
|
|
except asyncio.TimeoutError:
|
|
await ws_requester.send(json.dumps({
|
|
"request_id": requester_request_id,
|
|
"status": 504,
|
|
"headers": {"Content-Type": "application/json"},
|
|
"body": json.dumps({"error": "relay rpc timed out"}),
|
|
}))
|
|
except ConnectionError as exc:
|
|
await ws_requester.send(json.dumps({
|
|
"request_id": requester_request_id,
|
|
"status": 503,
|
|
"headers": {"Content-Type": "application/json"},
|
|
"body": json.dumps({"error": str(exc)}),
|
|
}))
|
|
finally:
|
|
self._pending_rpc.pop(relay_request_id, None)
|
|
|
|
|
|
async def _broadcast(raw: str | bytes, peers: list) -> None:
|
|
"""Send raw message to all peers; ignore individual send failures."""
|
|
if not peers:
|
|
return
|
|
import asyncio
|
|
await asyncio.gather(
|
|
*[_safe_send(p.ws, raw) for p in peers],
|
|
return_exceptions=True,
|
|
)
|
|
|
|
|
|
async def _safe_send(ws, msg) -> None:
|
|
try:
|
|
await ws.send(msg)
|
|
except Exception:
|
|
pass
|