neuron-tai/packages/relay/meshnet_relay/server.py

"""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
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

    @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 == "/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

                # 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 _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