fleetmesh-node 0.2.0__py3-none-any.whl

node/transport/fips_rung/core.py

@@ -0,0 +1,666 @@
+"""FIPS Native Datagram API Transport Client for Fleetmesh.
+
+Spec reference: docs/fips-native-transport.md.
+Connects to the FIPS v0.5.1 native API socket (/var/run/fips/api.sock, ~/.fips/api.sock, ...),
+requests FSP service flows and listener descriptors via SCM_RIGHTS, and exchanges binary-framed
+datagrams addressed to peer npub identities without requiring TUN/IPv6 adapters or root privileges.
+"""
+from __future__ import annotations
+
+import array
+import errno
+import json
+import os
+import socket
+import struct
+import threading
+import time
+from typing import Any, Mapping, MutableMapping, Optional, Union
+
+from .types import FipsArrivalInfo
+
+# Protocol Framing Constants (docs/fips-native-transport.md section 6)
+FRAMING_MAGIC: int = 0x4D  # 'M'
+FRAMING_VERSION: int = 0x01
+FRAME_MAGIC: bytes = bytes([FRAMING_MAGIC, FRAMING_VERSION])
+HEADER_FORMAT: str = '!BBHBBH'
+HEADER_SIZE: int = struct.calcsize(HEADER_FORMAT)  # 8 bytes
+
+FLAG_FINAL: int = 0x01
+FLAG_COMPRESSED: int = 0x02
+FLAG_DIDCOMM: int = 0x04
+FLAG_ACK_REQUESTED: int = 0x08
+
+# Where a FIPS daemon may be found when the operator did not say. Every path
+# here is one only root or this user can create: /var/run and /run are
+# root-owned, ~/.fips is ours. /tmp/fips/api.sock used to be on this list and
+# is not any more, because any local user can create /tmp/fips first, and a
+# node that then connected to it would hand its flows - and, until the
+# envelope became mandatory on this rung, its peers' plaintext - to whoever
+# squatted the name. An operator with a daemon in /tmp names it explicitly.
+DEFAULT_SOCKET_PATHS: tuple[str, ...] = (
+    "/var/run/fips/api.sock",
+    "/run/fips/api.sock",
+    os.path.expanduser("~/.fips/api.sock"),
+    "/var/run/fips/control.sock",
+    os.path.expanduser("~/.fips/control.sock"),
+)
+
+# Standard FSP Service Ports
+PORT_PEERING: int = 11801
+PORT_TELEMETRY: int = 21801
+PORT_ACCOUNTING: int = 30801
+PORT_COMPUTE_MARKET: int = 30810
+PORT_COMPUTE_STREAM: int = 30811
+PORT_COMPUTE_CANCEL: int = 30812
+
+# FIPS v0.5.1 advertises max_payload per flow (1170 over a UDP transport,
+# 1362 in-process); fall back to the smaller value when a flow did not say.
+DEFAULT_MAX_PAYLOAD: int = 1170
+# Upper bound on frames per message; larger messages raise so the transport
+# ladder can cascade to a stream-capable rung (iroh / nostr) instead.
+MAX_FRAGMENTS: int = 64
+REASSEMBLY_TIMEOUT_S: float = 30.0
+
+
+def find_api_socket() -> Optional[str]:
+    """Find first existing and accessible FIPS Native API socket path, or None."""
+    for p in DEFAULT_SOCKET_PATHS:
+        if os.path.exists(p) and os.access(p, os.R_OK | os.W_OK):
+            return p
+    return None
+
+
+class FipsTransportError(Exception):
+    """Raised when FIPS control or flow operations fail."""
+
+
+def encode_frame(
+    seq: int,
+    payload: bytes,
+    is_final: bool = False,
+    is_didcomm: bool = True,
+    compressed: bool = False,
+    ack_requested: bool = False,
+    flags: Optional[int] = None,
+    reserved: int = 0,
+) -> bytes:
+    """Pack payload with 8-byte binary framing header.
+
+    ``reserved`` carries the message id used to reassemble fragmented messages.
+    """
+    if len(payload) > 65535:
+        raise FipsTransportError(f"payload exceeds maximum 65535 bytes ({len(payload)})")
+
+    if flags is not None:
+        computed_flags = flags
+        if is_final:
+            computed_flags |= FLAG_FINAL
+    else:
+        computed_flags = 0
+        if is_final:
+            computed_flags |= FLAG_FINAL
+        if compressed:
+            computed_flags |= FLAG_COMPRESSED
+        if is_didcomm:
+            computed_flags |= FLAG_DIDCOMM
+        if ack_requested:
+            computed_flags |= FLAG_ACK_REQUESTED
+
+    header = struct.pack(
+        HEADER_FORMAT,
+        FRAMING_MAGIC,
+        FRAMING_VERSION,
+        seq & 0xFFFF,
+        computed_flags,
+        reserved & 0xFF,  # message id for fragment reassembly
+        len(payload),
+    )
+    return header + payload
+
+
+def decode_frame_ex(data: bytes) -> tuple[int, int, int, bytes]:
+    """Like decode_frame but also returns the message id: (seq, flags, msg_id, payload)."""
+    seq, flags, payload = decode_frame(data)
+    return seq, flags, data[5], payload
+
+
+class Reassembler:
+    """Rebuild messages split by FipsNativeTransport.send_message on one flow.
+
+    Frames with message id 0 are legacy single-frame messages and pass through.
+    Frames arrive in order on a FIPS flow but may be lost. A message is
+    delivered only when fragments 0..N arrive contiguously for one message id
+    and fragment N has FLAG_FINAL; any gap discards the partial message.
+    """
+
+    def __init__(self, timeout_s: float = REASSEMBLY_TIMEOUT_S, max_fragments: int = MAX_FRAGMENTS,
+                 clock: Any = time.monotonic) -> None:
+        self.timeout_s = timeout_s
+        self.max_fragments = max_fragments
+        self.clock = clock
+        self._msg_id: Optional[int] = None
+        self._parts: list[bytes] = []
+        self._flags: int = 0
+        self._started: float = 0.0
+        self.dropped: int = 0
+
+    def _reset(self, dropped: bool) -> None:
+        if dropped and self._parts:
+            self.dropped += 1
+        self._msg_id, self._parts, self._flags = None, [], 0
+
+    def feed(self, seq: int, flags: int, msg_id: int, payload: bytes) -> Optional[tuple[int, bytes]]:
+        """Feed one frame; returns (flags, full_payload) when a message completes."""
+        if msg_id == 0:
+            # Pre-fragmentation sender (reserved byte was always 0): one frame = one message.
+            return flags, payload
+        now = self.clock()
+        if self._parts and now - self._started > self.timeout_s:
+            self._reset(dropped=True)
+        if seq == 0:
+            self._reset(dropped=True)
+            self._msg_id, self._flags, self._started = msg_id, flags, now
+        elif msg_id != self._msg_id or seq != len(self._parts):
+            self._reset(dropped=True)  # gap, loss, or interleaving: discard
+            return None
+        self._parts.append(payload)
+        if flags & FLAG_FINAL:
+            out = (self._flags | FLAG_FINAL, b"".join(self._parts))
+            self._reset(dropped=False)
+            return out
+        if len(self._parts) >= self.max_fragments:
+            self._reset(dropped=True)
+        return None
+
+
+def decode_frame(data: bytes) -> tuple[int, int, bytes]:
+    """Unpack 8-byte binary framing header.
+
+    Returns:
+        (seq, flags, payload)
+    """
+    if len(data) < HEADER_SIZE:
+        raise FipsTransportError(f"frame truncated ({len(data)} < {HEADER_SIZE} bytes)")
+
+    magic, ver, seq, flags, _, plen = struct.unpack(HEADER_FORMAT, data[:HEADER_SIZE])
+    if magic != FRAMING_MAGIC:
+        raise FipsTransportError(f"invalid framing magic 0x{magic:02x} (expected 0x{FRAMING_MAGIC:02x})")
+    if ver != FRAMING_VERSION:
+        raise FipsTransportError(f"unsupported framing version {ver} (expected {FRAMING_VERSION})")
+
+    payload = data[HEADER_SIZE:HEADER_SIZE + plen]
+    if len(payload) != plen:
+        raise FipsTransportError(f"payload length mismatch ({len(payload)} != {plen})")
+
+    return seq, flags, payload
+
+
+def pubkey_to_npub(pubkey: object) -> str:
+    """Normalize a pubkey (hex, did:nostr:, or npub) to standard npub format.
+
+    Takes `object` because config and CLI hand it whatever they parsed; a
+    non-string is stringified rather than refused, as before."""
+    if not isinstance(pubkey, str):
+        return str(pubkey)
+    if pubkey.startswith("npub1"):
+        return pubkey
+    if pubkey.startswith("did:nostr:"):
+        pubkey = pubkey[len("did:nostr:"):]
+    if len(pubkey) == 64:
+        try:
+            from fleetmesh_core.crypto.bech32 import bech32_encode
+            return bech32_encode("npub", bytes.fromhex(pubkey))
+        except Exception:
+            pass
+    return pubkey
+
+
+class FipsNativeTransport:
+    """Manages bilateral datagram flows over FIPS Native Control API."""
+
+    def __init__(self, sock_path: Optional[str] = None, node_pk: Optional[str] = None) -> None:
+        self.sock_path = sock_path or self._discover_socket()
+        self.node_pk: Optional[str] = node_pk
+        self.ctrl_sock: Optional[socket.socket] = None
+        self._lock: threading.Lock = threading.Lock()
+        self.open_flows: dict[str, socket.socket] = {}
+        self.listeners: dict[int, socket.socket] = {}
+        # fileno -> max datagram size the daemon reported for that flow
+        self.flow_max_payload: dict[int, int] = {}
+        self._msg_counter: int = 0
+
+    @staticmethod
+    def _discover_socket() -> str:
+        for p in DEFAULT_SOCKET_PATHS:
+            if os.path.exists(p):
+                return p
+        return DEFAULT_SOCKET_PATHS[0]
+
+    def is_available(self) -> bool:
+        """Check if local FIPS control socket exists and is accessible."""
+        return os.path.exists(self.sock_path) and os.access(self.sock_path, os.R_OK | os.W_OK)
+
+    def connect(self) -> None:
+        """Connect to local FIPS Unix control socket."""
+        if not self.is_available():
+            raise FipsTransportError(f"FIPS control socket not accessible at {self.sock_path}")
+        with self._lock:
+            if self.ctrl_sock is not None:
+                return
+            s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            s.connect(self.sock_path)
+            self.ctrl_sock = s
+
+    def open_flow(self, peer_npub: str, port: int, timeout_s: float = 5.0) -> socket.socket:
+        """Request connected FSP datagram flow descriptor from FIPS via SCM_RIGHTS."""
+        if self.ctrl_sock is None:
+            self.connect()
+        assert self.ctrl_sock is not None
+
+        norm_peer = pubkey_to_npub(peer_npub)
+        params: MutableMapping[str, Any] = {
+            "peer": norm_peer,
+            "remote_port": int(port),
+        }
+        cmd_dict: MutableMapping[str, Any] = {
+            "command": "connect",
+            "params": params,
+        }
+        if self.node_pk:
+            norm_caller = pubkey_to_npub(self.node_pk)
+            params["caller"] = norm_caller
+        req = json.dumps(cmd_dict) + "\n"
+        with self._lock:
+            try:
+                self.ctrl_sock.settimeout(timeout_s)
+                self.ctrl_sock.sendall(req.encode("utf-8"))
+
+                fds = array.array("i")
+                cmsg_space = socket.CMSG_SPACE(fds.itemsize * 4)
+                msg, ancdata, flags, addr = self.ctrl_sock.recvmsg(4096, cmsg_space)
+            except (OSError, BrokenPipeError, ConnectionResetError) as e:
+                try:
+                    self.ctrl_sock.close()
+                except Exception:
+                    pass
+                self.ctrl_sock = None
+                raise FipsTransportError(f"FIPS control socket disconnected: {e}")
+
+            if not msg and not ancdata:
+                try:
+                    self.ctrl_sock.close()
+                except Exception:
+                    pass
+                self.ctrl_sock = None
+                raise FipsTransportError("FIPS control socket closed unexpectedly (EOF)")
+
+            for cmsg_level, cmsg_type, cmsg_data in ancdata:
+                if cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS:
+                    fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
+                    if fds:
+                        raw_fd = fds[0]
+                        for extra_fd in fds[1:]:
+                            try:
+                                os.close(extra_fd)
+                            except OSError:
+                                pass
+                        try:
+                            flow_sock = socket.fromfd(raw_fd, socket.AF_UNIX, socket.SOCK_DGRAM)
+                        finally:
+                            try:
+                                os.close(raw_fd)
+                            except OSError:
+                                pass
+                        key = f"{peer_npub}:{port}"
+                        self.open_flows[key] = flow_sock
+                        self._record_max_payload(flow_sock, msg)
+                        return flow_sock
+
+        resp = {}
+        if msg:
+            try:
+                resp = json.loads(msg.decode("utf-8"))
+            except Exception:
+                pass
+        err_msg = resp.get("message") or resp.get("error") or "no descriptor received"
+        raise FipsTransportError(f"FIPS flow connect failed: {err_msg}")
+
+    def listen_port(self, port: int, timeout_s: float = 5.0) -> socket.socket:
+        """Register inbound listener on FSP port, receiving descriptor via SCM_RIGHTS."""
+        if self.ctrl_sock is None:
+            self.connect()
+        assert self.ctrl_sock is not None
+
+        params: MutableMapping[str, Any] = {
+            "local_port": int(port),
+        }
+        cmd_dict: MutableMapping[str, Any] = {
+            "command": "listen",
+            "params": params,
+        }
+        if self.node_pk:
+            norm_pk = pubkey_to_npub(self.node_pk)
+            params["node_pk"] = norm_pk
+        req = json.dumps(cmd_dict) + "\n"
+        with self._lock:
+            try:
+                self.ctrl_sock.settimeout(timeout_s)
+                self.ctrl_sock.sendall(req.encode("utf-8"))
+
+                fds = array.array("i")
+                cmsg_space = socket.CMSG_SPACE(fds.itemsize * 4)
+                msg, ancdata, flags, addr = self.ctrl_sock.recvmsg(4096, cmsg_space)
+            except (OSError, BrokenPipeError, ConnectionResetError) as e:
+                try:
+                    self.ctrl_sock.close()
+                except Exception:
+                    pass
+                self.ctrl_sock = None
+                raise FipsTransportError(f"FIPS control socket disconnected: {e}")
+
+            if not msg and not ancdata:
+                try:
+                    self.ctrl_sock.close()
+                except Exception:
+                    pass
+                self.ctrl_sock = None
+                raise FipsTransportError("FIPS control socket closed unexpectedly (EOF)")
+
+            for cmsg_level, cmsg_type, cmsg_data in ancdata:
+                if cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS:
+                    fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
+                    if fds:
+                        raw_fd = fds[0]
+                        for extra_fd in fds[1:]:
+                            try:
+                                os.close(extra_fd)
+                            except OSError:
+                                pass
+                        try:
+                            listener_sock = socket.fromfd(raw_fd, socket.AF_UNIX, socket.SOCK_DGRAM)
+                        finally:
+                            try:
+                                os.close(raw_fd)
+                            except OSError:
+                                pass
+                        self.listeners[port] = listener_sock
+                        return listener_sock
+
+        resp = {}
+        if msg:
+            try:
+                resp = json.loads(msg.decode("utf-8"))
+            except Exception:
+                pass
+        err_msg = resp.get("message") or resp.get("error") or "no descriptor received"
+        raise FipsTransportError(f"FIPS listen failed: {err_msg}")
+
+    def accept_flow(
+        self,
+        listener_sock: socket.socket,
+        timeout_s: Optional[float] = 5.0,
+    ) -> tuple[socket.socket, FipsArrivalInfo]:
+        """Accept an incoming flow on a registered listener socket.
+
+        In FIPS v0.5.1, when a remote peer initiates a flow to a listening port,
+        the listener descriptor receives an arrival notice datagram with SCM_RIGHTS
+        carrying the new connected flow datagram socket.
+
+        Returns:
+            (flow_socket, arrival_info_dict)
+        """
+        if timeout_s is not None:
+            listener_sock.settimeout(timeout_s)
+        fds = array.array("i")
+        cmsg_space = socket.CMSG_SPACE(fds.itemsize * 4)
+        msg, ancdata, flags, addr = listener_sock.recvmsg(4096, cmsg_space)
+        flow_fd = None
+        for cmsg_level, cmsg_type, cmsg_data in ancdata:
+            if cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS:
+                fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
+                if fds:
+                    flow_fd = fds[0]
+                    for extra_fd in fds[1:]:
+                        try:
+                            os.close(extra_fd)
+                        except OSError:
+                            pass
+                    break
+        if flow_fd is None:
+            raise FipsTransportError("no flow descriptor received in listener arrival")
+
+        try:
+            flow_sock = socket.fromfd(flow_fd, socket.AF_UNIX, socket.SOCK_DGRAM)
+        finally:
+            try:
+                os.close(flow_fd)
+            except OSError:
+                pass
+        arrival_info: FipsArrivalInfo = {}
+        if msg:
+            try:
+                arrival_info = json.loads(msg.decode("utf-8"))
+            except Exception:
+                pass
+        peer = arrival_info.get("peer")
+        remote_port = arrival_info.get("remote_port")
+        if peer and remote_port:
+            self.open_flows[f"{peer}:{remote_port}"] = flow_sock
+        self._record_max_payload(flow_sock, arrival_info)
+        return flow_sock, arrival_info
+
+    def accept_listener_event(
+        self,
+        listener_sock: socket.socket,
+    ) -> tuple[Optional[socket.socket], Mapping[str, Any], bytes]:
+        """Receive the next incoming event on a registered listener socket.
+
+        Returns:
+            (flow_sock, arrival_info, direct_msg)
+            If an incoming flow arrived via SCM_RIGHTS: returns (flow_sock, arrival_info, b"")
+            If a direct binary frame arrived: returns (None, {}, direct_msg)
+            If empty/timeout: returns (None, {}, b"")
+        """
+        fds = array.array("i")
+        cmsg_space = socket.CMSG_SPACE(fds.itemsize * 4)
+        msg, ancdata, flags, addr = listener_sock.recvmsg(65536, cmsg_space)
+        if not msg and not ancdata:
+            return None, {}, b""
+
+        flow_fd = None
+        for cmsg_level, cmsg_type, cmsg_data in ancdata:
+            if cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS:
+                fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
+                if fds:
+                    flow_fd = fds[0]
+                    for extra_fd in fds[1:]:
+                        try:
+                            os.close(extra_fd)
+                        except OSError:
+                            pass
+                    break
+
+        if flow_fd is not None:
+            try:
+                flow_sock = socket.fromfd(flow_fd, socket.AF_UNIX, socket.SOCK_DGRAM)
+            finally:
+                try:
+                    os.close(flow_fd)
+                except OSError:
+                    pass
+            arrival_info = {}
+            if msg:
+                try:
+                    arrival_info = json.loads(msg.decode("utf-8"))
+                    peer = arrival_info.get("peer")
+                    rport = arrival_info.get("remote_port")
+                    if peer and rport:
+                        self.open_flows[f"{peer}:{rport}"] = flow_sock
+                except (ValueError, UnicodeDecodeError, AttributeError):
+                    pass
+            self._record_max_payload(flow_sock, arrival_info)
+            return flow_sock, arrival_info, b""
+
+        return None, {}, msg
+
+    def _record_max_payload(self, flow_sock: socket.socket, reply: Any) -> None:
+        """Remember the per-flow datagram ceiling from a connect reply or arrival."""
+        info: Any = reply
+        if isinstance(reply, (bytes, bytearray)):
+            try:
+                info = json.loads(bytes(reply).decode("utf-8"))
+            except Exception:
+                info = {}
+        if isinstance(info, dict) and isinstance(info.get("data"), dict):
+            info = info["data"]
+        mp = info.get("max_payload") if isinstance(info, dict) else None
+        if isinstance(mp, int) and mp > HEADER_SIZE:
+            try:
+                self.flow_max_payload[flow_sock.fileno()] = mp
+            except Exception:
+                pass
+
+    def max_payload_for(self, flow_sock: socket.socket) -> int:
+        """Largest datagram (header included) this flow will carry."""
+        try:
+            return self.flow_max_payload.get(flow_sock.fileno(), DEFAULT_MAX_PAYLOAD)
+        except Exception:
+            return DEFAULT_MAX_PAYLOAD
+
+    def send_message(
+        self,
+        flow_sock: socket.socket,
+        payload: bytes,
+        flags: int = 0,
+        max_fragments: int = MAX_FRAGMENTS,
+    ) -> int:
+        """Send one logical message, fragmenting it to fit the flow's max_payload.
+
+        Every frame carries the same message id in the header's reserved byte
+        and a per-message fragment index in ``seq``; the last frame sets
+        FLAG_FINAL. Raises FipsTransportError before sending anything if the
+        message would need more than ``max_fragments`` frames, so callers can
+        fall back to a transport without a datagram ceiling. Returns the number
+        of frames sent.
+        """
+        raw = bytes(payload)
+        chunk = self.max_payload_for(flow_sock) - HEADER_SIZE
+        if chunk <= 0:
+            raise FipsTransportError("flow max_payload too small for framing header")
+        parts = [raw[i:i + chunk] for i in range(0, len(raw), chunk)] or [b""]
+        if len(parts) > max_fragments:
+            raise FipsTransportError(
+                f"message of {len(raw)} bytes needs {len(parts)} frames "
+                f"(limit {max_fragments} at {chunk} bytes/frame)"
+            )
+        with self._lock:
+            # ids 1..255; 0 marks a legacy/unfragmented frame
+            self._msg_counter = self._msg_counter % 255 + 1
+            msg_id = self._msg_counter
+        base = flags & ~FLAG_FINAL
+        for idx, part in enumerate(parts):
+            fl = base | (FLAG_FINAL if idx == len(parts) - 1 else 0)
+            packet = encode_frame(seq=idx, payload=part, flags=fl, is_final=False, reserved=msg_id)
+            sent = 0
+            for attempt in range(10):
+                try:
+                    sent = flow_sock.send(packet)
+                    break
+                except OSError as e:
+                    if e.errno == getattr(errno, 'ENOBUFS', 55) and attempt < 9:
+                        time.sleep(0.005 * (attempt + 1))
+                        continue
+                    raise
+            if sent != len(packet):
+                raise FipsTransportError(f"short datagram write ({sent} of {len(packet)} bytes)")
+        return len(parts)
+
+    def send_framed(
+        self,
+        flow_sock: socket.socket,
+        payload: Union[bytes, str, Mapping[str, Any]],
+        seq: int = 0,
+        is_final: bool = False,
+        flags: Optional[int] = None,
+    ) -> int:
+        """Transmit binary-framed datagram over open flow socket."""
+        if isinstance(payload, Mapping):
+            raw_bytes = json.dumps(payload, sort_keys=True).encode("utf-8")
+            is_didcomm = True
+        elif isinstance(payload, str):
+            raw_bytes = payload.encode("utf-8")
+            is_didcomm = False
+        else:
+            raw_bytes = bytes(payload)
+            is_didcomm = False
+
+        packet = encode_frame(
+            seq=seq,
+            payload=raw_bytes,
+            is_final=is_final,
+            is_didcomm=is_didcomm,
+            flags=flags,
+        )
+        limit = self.max_payload_for(flow_sock)
+        if len(packet) > limit:
+            raise FipsTransportError(
+                f"frame of {len(packet)} bytes exceeds flow max_payload {limit}; use send_message()"
+            )
+        for attempt in range(10):
+            try:
+                return flow_sock.send(packet)
+            except OSError as e:
+                if e.errno == getattr(errno, 'ENOBUFS', 55) and attempt < 9:
+                    time.sleep(0.005 * (attempt + 1))
+                    continue
+                raise
+        return flow_sock.send(packet)
+
+    def recv_framed(
+        self,
+        flow_sock: socket.socket,
+        timeout_s: Optional[float] = 5.0,
+    ) -> tuple[int, bool, bytes]:
+        """Receive and decode framed datagram. Returns (seq, is_final, payload_bytes)."""
+        if timeout_s is not None:
+            flow_sock.settimeout(timeout_s)
+        data = flow_sock.recv(65536)
+        seq, flags, payload = decode_frame(data)
+        is_final = bool(flags & FLAG_FINAL)
+        return seq, is_final, payload
+
+    def recv_raw_frame(
+        self,
+        flow_sock: socket.socket,
+        timeout_s: Optional[float] = 5.0,
+    ) -> tuple[int, int, bytes]:
+        """Receive and decode raw frame with exact integer bitflags. Returns (seq, flags, payload_bytes)."""
+        if timeout_s is not None:
+            flow_sock.settimeout(timeout_s)
+        data = flow_sock.recv(65536)
+        return decode_frame(data)
+
+    def close(self) -> None:
+        """Close all open flows, listeners, and control connection."""
+        with self._lock:
+            for s in self.open_flows.values():
+                try:
+                    s.close()
+                except Exception:
+                    pass
+            self.open_flows.clear()
+
+            for s in self.listeners.values():
+                try:
+                    s.close()
+                except Exception:
+                    pass
+            self.listeners.clear()
+
+            if self.ctrl_sock is not None:
+                try:
+                    self.ctrl_sock.close()
+                except Exception:
+                    pass
+                self.ctrl_sock = None