dropp baes64 use binary

packages/node/meshnet_node/relay_bridge.py

@@ -19,7 +19,7 @@ log = logging.getLogger(__name__)
 
 DEFAULT_MAX_CONCURRENCY = 8
 
-# Activation tensors ride the relay as base64 inside one JSON frame, so the
+# Activation tensors ride the relay as one WebSocket frame per hop, so the
 # websockets default of 1 MiB rejects any real prefill (close code 1009).
 DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024
 
@@ -36,6 +36,26 @@ def ws_max_size() -> int | None:
     return None if value <= 0 else value
 
 
+# Binary relay frame: JSON header + raw body in one WebSocket binary message,
+# so activation bodies travel as bytes instead of base64 inside JSON. Same wire
+# format as meshnet_relay.server — duplicated because node and relay ship as
+# independent distributions.
+BINARY_FRAME_MAGIC = b"MRF1"
+
+
+def encode_binary_frame(header: dict, body: bytes) -> bytes:
+    header_bytes = json.dumps(header, separators=(",", ":")).encode()
+    return 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, bytes(frame[8 + header_len:])
+
+
 @dataclass(frozen=True)
 class RelayBridgeInfo:
     peer_id: str
@@ -125,7 +145,9 @@ class RelayHttpBridge:
 
         while self._running:
             try:
-                with wsc.connect(self.relay_url, open_timeout=5, max_size=ws_max_size()) as ws:
+                with wsc.connect(
+                    self.relay_url, open_timeout=5, max_size=ws_max_size(), compression=None,
+                ) as ws:
                     self._ws = ws
                     self._connected.set()
                     ws.send(json.dumps(_make_envelope(
@@ -138,6 +160,17 @@ class RelayHttpBridge:
                             raw = ws.recv(timeout=1)
                         except TimeoutError:
                             continue
+                        if isinstance(raw, (bytes, bytearray)):
+                            try:
+                                payload, body = decode_binary_frame(bytes(raw))
+                            except (ValueError, json.JSONDecodeError):
+                                continue
+                            if payload.get("target_peer") not in {None, self.peer_id}:
+                                continue
+                            if self._executor is None:
+                                break
+                            self._executor.submit(self._process_request, payload, body)
+                            continue
                         try:
                             envelope = json.loads(raw)
                         except (TypeError, json.JSONDecodeError):
@@ -176,23 +209,41 @@ class RelayHttpBridge:
             log.debug("relay bridge send failed (request orphaned): %s", exc)
             return False
 
-    def _process_request(self, payload: dict) -> None:
+    def _send_binary_response_frame(self, header: dict, body: bytes) -> bool:
+        """Send one binary response frame; False if the socket is gone."""
+        ws = self._ws
+        if ws is None:
+            return False
+        frame = encode_binary_frame(header, body)
+        try:
+            with self._send_lock:
+                ws.send(frame)
+            return True
+        except Exception as exc:
+            log.debug("relay bridge binary send failed (request orphaned): %s", exc)
+            return False
+
+    def _process_request(self, payload: dict, binary_body: bytes | None = None) -> None:
         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 {}
+        binary_mode = binary_body is not None
 
         req_suffix = f" request_id={request_id}" if request_id else ""
         print(f"  [node] relay {method} {path}{req_suffix}", flush=True)
 
-        # 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)
+        if binary_mode:
+            data = binary_body
         else:
-            body_text = payload.get("body") or ""
-            data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
+            # Legacy JSON request: body_base64 carries binary data, text "body"
+            # covers 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)
@@ -205,6 +256,13 @@ class RelayHttpBridge:
                     return
                 resp_bytes = resp.read()
                 # Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
+                if binary_mode:
+                    self._send_binary_response_frame({
+                        "request_id": request_id,
+                        "status": resp.status,
+                        "headers": resp_headers,
+                    }, resp_bytes)
+                    return
                 is_binary = "octet-stream" in content_type
                 result: dict = {
                     "request_id": request_id,