neuron-tai/packages/node/meshnet_node/relay_bridge.py

"""Outbound relay bridge for NAT-safe node HTTP requests."""

from __future__ import annotations

import base64
import json
import logging
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass

log = logging.getLogger(__name__)


@dataclass(frozen=True)
class RelayBridgeInfo:
    peer_id: str
    relay_addr: str


def _make_envelope(topic: str, payload: dict, peer_id: str) -> dict:
    return {
        "topic": topic,
        "version": 1,
        "from_peer": peer_id,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "msg_id": f"{peer_id}-{time.time_ns():x}",
        "ttl": 1,
        "payload": payload,
    }


class RelayHttpBridge:
    """Connect outbound to a relay and proxy relay HTTP requests to localhost."""

    def __init__(
        self,
        relay_url: str,
        peer_id: str,
        local_base_url: str,
        advertised_addr: str,
        reconnect_interval: float = 3.0,
    ) -> None:
        self.relay_url = relay_url.rstrip("/")
        self.peer_id = peer_id
        self.local_base_url = local_base_url.rstrip("/")
        self.advertised_addr = advertised_addr
        self.reconnect_interval = reconnect_interval
        self._running = False
        self._thread: threading.Thread | None = None
        self._connected = threading.Event()

    @property
    def relay_addr(self) -> str:
        base = self.relay_url
        if base.endswith("/ws"):
            base = base[:-3]
        return f"{base}/rpc/{self.peer_id}"

    def start(self) -> RelayBridgeInfo:
        self._running = True
        self._thread = threading.Thread(target=self._run, daemon=True, name="relay-http-bridge")
        self._thread.start()
        return RelayBridgeInfo(peer_id=self.peer_id, relay_addr=self.relay_addr)

    def wait_connected(self, timeout: float = 5.0) -> bool:
        return self._connected.wait(timeout)

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=3.0)

    def _run(self) -> None:
        import websockets.sync.client as wsc  # type: ignore[import]

        while self._running:
            try:
                with wsc.connect(self.relay_url, open_timeout=5) as ws:
                    self._connected.set()
                    ws.send(json.dumps(_make_envelope(
                        "peer-register",
                        {"peer_id": self.peer_id, "addr": self.advertised_addr},
                        self.peer_id,
                    )))
                    while self._running:
                        try:
                            raw = ws.recv(timeout=1)
                        except TimeoutError:
                            continue
                        try:
                            envelope = json.loads(raw)
                        except (TypeError, json.JSONDecodeError):
                            continue
                        if envelope.get("topic") != "relay-http-request":
                            continue
                        payload = envelope.get("payload", {})
                        if payload.get("target_peer") not in {None, self.peer_id}:
                            continue
                        response = self._handle_request(payload)
                        ws.send(json.dumps(_make_envelope(
                            "relay-http-response",
                            response,
                            self.peer_id,
                        )))
            except Exception as exc:
                self._connected.clear()
                if self._running:
                    log.debug("relay bridge disconnected: %s", exc)
                    time.sleep(self.reconnect_interval)

    def _handle_request(self, payload: dict) -> dict:
        request_id = str(payload.get("request_id") or "")
        method = str(payload.get("method") or "POST").upper()
        path = str(payload.get("path") or "/")
        headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}

        # body_base64 carries binary data (e.g. bfloat16 activation tensors) safely.
        # Fallback to text "body" for backward-compat with non-binary requests.
        body_b64 = payload.get("body_base64")
        if body_b64:
            data = base64.b64decode(body_b64)
        else:
            body_text = payload.get("body") or ""
            data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)

        url = f"{self.local_base_url}{path}"
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(req, timeout=300.0) as resp:
                resp_bytes = resp.read()
                resp_headers = dict(resp.headers)
                # Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
                is_binary = "octet-stream" in resp.headers.get("Content-Type", "")
                result: dict = {
                    "request_id": request_id,
                    "status": resp.status,
                    "headers": resp_headers,
                }
                if is_binary:
                    result["body_base64"] = base64.b64encode(resp_bytes).decode()
                else:
                    result["body"] = resp_bytes.decode(errors="replace")
                return result
        except urllib.error.HTTPError as exc:
            return {
                "request_id": request_id,
                "status": exc.code,
                "headers": {"Content-Type": exc.headers.get("Content-Type", "application/json")},
                "body": exc.read().decode(errors="replace"),
            }
        except Exception as exc:
            return {
                "request_id": request_id,
                "status": 503,
                "headers": {"Content-Type": "application/json"},
                "body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
            }


def peer_id_from_wallet(wallet_address: str) -> str:
    return wallet_address[:16] if len(wallet_address) >= 16 else wallet_address