282 lines
9.7 KiB
Python
282 lines
9.7 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 threading
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from .peer_registry import PeerRegistry
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
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,
|
|
):
|
|
self.host = host
|
|
self.port = port
|
|
self.ssl_cert = ssl_cert
|
|
self.ssl_key = ssl_key
|
|
self.max_peers = max_peers
|
|
|
|
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
|
|
self._pending_rpc: dict[str, asyncio.Future] = {}
|
|
|
|
@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 _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,
|
|
)
|
|
# 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:
|
|
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")
|
|
fut = self._pending_rpc.pop(request_id, None)
|
|
if fut is not None and not fut.done():
|
|
fut.set_result(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)
|
|
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:
|
|
"""Send one HTTP-shaped request to a connected peer and relay its response."""
|
|
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
|
|
|
|
try:
|
|
raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0)
|
|
payload = json.loads(raw)
|
|
except Exception:
|
|
await ws_requester.close(1003, "invalid relay rpc request")
|
|
return
|
|
|
|
request_id = str(payload.get("request_id") or uuid.uuid4())
|
|
payload["request_id"] = request_id
|
|
payload["target_peer"] = target_peer_id
|
|
fut = self._loop.create_future() if self._loop is not None else asyncio.get_running_loop().create_future()
|
|
self._pending_rpc[request_id] = fut
|
|
try:
|
|
await target.ws.send(json.dumps({
|
|
"topic": "relay-http-request",
|
|
"version": 1,
|
|
"from_peer": "relay",
|
|
"payload": payload,
|
|
}))
|
|
response = await asyncio.wait_for(fut, timeout=310.0)
|
|
await ws_requester.send(json.dumps(response))
|
|
except asyncio.TimeoutError:
|
|
await ws_requester.send(json.dumps({
|
|
"request_id": request_id,
|
|
"status": 504,
|
|
"headers": {"Content-Type": "application/json"},
|
|
"body": json.dumps({"error": "relay rpc timed out"}),
|
|
}))
|
|
finally:
|
|
self._pending_rpc.pop(request_id, None)
|
|
await ws_requester.close()
|
|
|
|
|
|
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
|