feat(webrtc): add UdpMux for shared-port WebRTC-Direct inbound dispatch

libp2p/transport/webrtc/_udp_mux.py

@@ -0,0 +1,315 @@
+"""
+UdpMux: shared UDP socket dispatcher for WebRTC-Direct inbound connections.
+
+Routes incoming datagrams to per-dial aioice.Connection instances:
+  - STUN packets:  by local ufrag prefix in the USERNAME attribute (pre-ICE)
+  - non-STUN data: by remote (host, port) pair (post-ICE DTLS / SCTP)
+
+This is the foundational primitive for a spec-aligned WebRTC-Direct listener
+that advertises a single fixed UDP port and demuxes concurrent inbound dials
+without spinning up a new port per peer.
+
+Validated against ``aioice`` 0.10.x. It relies on private ``aioice`` internals
+(``StunProtocol``, ``Connection._protocols`` / ``_local_candidates`` /
+``_local_candidates_end``) to inject a mux-backed protocol without aioice
+binding its own UDP socket, so it may need updating on aioice major bumps.
+Upstreaming a real ``UdpMux`` into aioice is tracked as follow-up on #1352.
+
+Refs: libp2p/specs#715 (WebRTC-Direct v2), libp2p/py-libp2p#1352 (spike).
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+import logging
+import struct
+from typing import Protocol
+
+try:
+    from aioice.candidate import Candidate
+    import aioice.ice as _ice
+    import aioice.stun as _stun
+
+    _HAS_AIOICE = True
+except ImportError:
+    _HAS_AIOICE = False
+
+
+class _HasConnectionLost(Protocol):
+    def connection_lost(self, exc: Exception | None) -> None: ...
+    def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None: ...
+
+
+logger = logging.getLogger(__name__)
+
+
+class _MuxedTransport:
+    """
+    Fake DatagramTransport backed by the mux's real shared socket.
+
+    Passed to aioice's StunProtocol so that aioice can send datagrams without
+    ever owning its own UDP socket.  close() signals the protocol rather than
+    closing the shared socket.
+    """
+
+    def __init__(
+        self,
+        real_transport: asyncio.DatagramTransport,
+        local_addr: tuple[str, int],
+    ) -> None:
+        self._real = real_transport
+        self._local_addr = local_addr
+        # Set by UdpMux after the StunProtocol is constructed (avoids circular ref).
+        self._protocol: _HasConnectionLost | None = None
+
+    def sendto(self, data: bytes, addr: tuple[str, int]) -> None:
+        self._real.sendto(data, addr)
+
+    def close(self) -> None:
+        # Let the protocol know the "transport" is gone so aioice's
+        # `await protocol.close()` can complete without blocking forever.
+        if self._protocol is not None:
+            self._protocol.connection_lost(None)
+
+    def get_extra_info(self, key: str, default: object = None) -> object:
+        if key == "sockname":
+            return self._local_addr
+        return default
+
+
+class UdpMux(asyncio.DatagramProtocol):
+    """
+    Shared UDP socket for WebRTC-Direct inbound connections.
+
+    Create one instance per listener port via :meth:`create`, then call
+    :meth:`add_ice_connection` for each inbound dial.  After ICE completes,
+    call :meth:`register_addr` with the elected remote address so that
+    subsequent DTLS/SCTP datagrams from that peer are dispatched correctly.
+
+    Example::
+
+        mux, port = await UdpMux.create("0.0.0.0", 0)
+
+        # Observe first-contact STUN for unregistered ufrags:
+        mux.set_unknown_stun_handler(on_new_dial)
+
+        # When a new STUN BINDING REQUEST arrives for ufrag "abc":
+        conn = mux.add_ice_connection("abc", "password123", host="0.0.0.0")
+        for c in remote_candidates:
+            await conn.add_remote_candidate(c)
+        await conn.add_remote_candidate(None)  # end-of-candidates
+        await conn.connect()                   # do NOT call gather_candidates()
+
+        # After ICE selects a candidate pair:
+        mux.register_addr(("203.0.113.5", 54321), conn._protocols[0])
+
+        # Teardown:
+        mux.unregister("abc")
+        mux.unregister_addr(("203.0.113.5", 54321))
+        await mux.close()
+    """
+
+    def __init__(self) -> None:
+        self._transport: asyncio.DatagramTransport | None = None
+        self._local_addr: tuple[str, int] | None = None
+        # ufrag -> StunProtocol (pre-ICE STUN dispatch)
+        self._by_ufrag: dict[str, _HasConnectionLost] = {}
+        # (host, port) -> StunProtocol (post-ICE DTLS/SCTP dispatch)
+        self._by_addr: dict[tuple[str, int], _HasConnectionLost] = {}
+        # Called for a STUN packet whose ufrag is not registered (see
+        # set_unknown_stun_handler); lets a listener observe first-contact dials.
+        self._unknown_stun_handler: (
+            Callable[[str, bytes, tuple[str, int]], None] | None
+        ) = None
+
+    # ------------------------------------------------------------------
+    # Factory
+    # ------------------------------------------------------------------
+
+    @classmethod
+    async def create(cls, host: str, port: int) -> tuple[UdpMux, int]:
+        """
+        Bind a shared UDP socket on *host*:*port* (use ``port=0`` for OS choice).
+        Returns ``(mux, bound_port)``.
+        """
+        loop = asyncio.get_running_loop()
+        mux = cls()
+        transport, _ = await loop.create_datagram_endpoint(
+            lambda: mux, local_addr=(host, port)
+        )
+        # connection_made is called synchronously inside create_datagram_endpoint,
+        # so _transport and _local_addr are set by the time we return.
+        assert mux._local_addr is not None
+        return mux, mux._local_addr[1]
+
+    # ------------------------------------------------------------------
+    # asyncio.DatagramProtocol
+    # ------------------------------------------------------------------
+
+    def connection_made(self, transport: asyncio.BaseTransport) -> None:
+        self._transport = transport  # type: ignore[assignment]
+        addr = transport.get_extra_info("sockname")
+        self._local_addr = (addr[0], addr[1])
+
+    def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
+        norm: tuple[str, int] = (addr[0], addr[1])
+        try:
+            msg = _stun.parse_message(data)
+            username: str = msg.attributes.get("USERNAME", "")
+            ufrag = username.split(":")[0]
+            protocol = self._by_ufrag.get(ufrag)
+            if protocol is not None:
+                protocol.datagram_received(data, norm)
+                return
+            # First contact: an inbound BINDING REQUEST for an unregistered
+            # ufrag. A WebRTC-Direct listener registers a handler to create the
+            # connection (add_ice_connection) and replay this datagram.
+            if self._unknown_stun_handler is not None:
+                self._unknown_stun_handler(ufrag, data, norm)
+                return
+            logger.debug("UdpMux: no handler for ufrag %r from %s", ufrag, norm)
+        except (ValueError, struct.error):
+            # Not valid STUN. aioice raises ValueError for non-STUN framing, but
+            # a STUN-shaped packet with a malformed/short fixed-width attribute
+            # makes its struct.unpack raise struct.error (NOT a ValueError
+            # subclass). Both mean "not usable STUN" and must fall through to
+            # addr-based routing rather than escape datagram_received — a remote
+            # peer must not be able to crash or log-flood the mux with crafted
+            # packets.
+            # Non-STUN (DTLS handshake, SCTP frames after ICE): route by addr.
+            protocol = self._by_addr.get(norm)
+            if protocol is not None:
+                protocol.datagram_received(data, norm)
+                return
+            logger.debug("UdpMux: no handler for non-STUN datagram from %s", norm)
+
+    def error_received(self, exc: Exception) -> None:
+        logger.warning("UdpMux socket error: %s", exc)
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        logger.debug("UdpMux connection lost: %s", exc)
+
+    # ------------------------------------------------------------------
+    # Registration
+    # ------------------------------------------------------------------
+
+    def register(self, ufrag: str, protocol: _HasConnectionLost) -> None:
+        """Route STUN packets with local ufrag *ufrag* to *protocol*."""
+        self._by_ufrag[ufrag] = protocol
+
+    def register_addr(
+        self, addr: tuple[str, int], protocol: _HasConnectionLost
+    ) -> None:
+        """Route non-STUN packets from *addr* to *protocol* (call after ICE)."""
+        self._by_addr[addr] = protocol
+
+    def unregister(self, ufrag: str) -> None:
+        self._by_ufrag.pop(ufrag, None)
+
+    def unregister_addr(self, addr: tuple[str, int]) -> None:
+        self._by_addr.pop(addr, None)
+
+    def set_unknown_stun_handler(
+        self, handler: Callable[[str, bytes, tuple[str, int]], None] | None
+    ) -> None:
+        """
+        Register a callback for STUN packets whose ufrag is not yet registered.
+
+        The callback receives ``(ufrag, data, addr)``. For WebRTC-Direct this is
+        how a listener observes the *first* inbound BINDING REQUEST from a new
+        dialer — it can create and register the connection via
+        :meth:`add_ice_connection` and replay ``data``. Called from the asyncio
+        event-loop thread; keep it non-blocking (e.g. hand off to a queue).
+        Pass ``None`` to clear.
+        """
+        self._unknown_stun_handler = handler
+
+    # ------------------------------------------------------------------
+    # Connection factory
+    # ------------------------------------------------------------------
+
+    def add_ice_connection(
+        self,
+        local_username: str,
+        local_password: str,
+        *,
+        host: str,
+    ) -> _ice.Connection:
+        """
+        Create an ``aioice.Connection`` backed by this mux (no own UDP socket).
+
+        The connection is pre-registered for *local_username* so that inbound
+        STUN connectivity checks are dispatched correctly before the caller
+        has a chance to set remote candidates. The single host candidate on the
+        shared port is injected as the connection's local candidate, so the
+        caller must **not** call ``conn.gather_candidates()`` — doing so binds
+        additional UDP sockets and defeats the shared-port design.
+
+        The caller must:
+        1. For each of the dialer's candidates, ``await conn.add_remote_candidate(c)``,
+           then ``await conn.add_remote_candidate(None)`` to signal end-of-candidates.
+        2. Await ``conn.connect()`` to complete ICE negotiation.
+        3. Call ``register_addr(remote_addr, conn._protocols[0])`` once ICE
+           selects a candidate pair so that post-ICE DTLS/SCTP is routed.
+        4. Call ``unregister(local_username)`` and
+           ``unregister_addr(remote_addr)`` on teardown.
+        """
+        if not _HAS_AIOICE:
+            raise RuntimeError("aioice is required (install py-libp2p[webrtc])")
+        assert self._transport is not None, "UdpMux not yet bound (call create() first)"
+        assert self._local_addr is not None
+
+        conn = _ice.Connection(
+            ice_controlling=False,
+            local_username=local_username,
+            local_password=local_password,
+        )
+
+        muxed_transport = _MuxedTransport(self._transport, self._local_addr)
+        protocol = _ice.StunProtocol(conn)
+        # Wire the fake transport so aioice can send responses.
+        protocol.transport = muxed_transport  # type: ignore[assignment]
+        # Back-reference so _MuxedTransport.close() can resolve protocol.close().
+        # aioice types connection_lost(exc: Exception) but asyncio protocol is Optional.
+        muxed_transport._protocol = protocol  # type: ignore[assignment]
+
+        local_candidate = Candidate(
+            foundation=_ice.candidate_foundation("host", "udp", host),
+            component=1,
+            transport="udp",
+            priority=_ice.candidate_priority(1, "host"),
+            host=host,
+            port=self._local_addr[1],
+            type="host",
+        )
+        protocol.local_candidate = local_candidate
+
+        # Inject into aioice's internal protocol list, bypassing gather_candidates
+        # (which would bind extra UDP sockets). Also mark gathering complete and
+        # expose the candidate so Connection.connect() proceeds instead of raising
+        # "Local candidates gathering was not performed".
+        conn._protocols.append(protocol)
+        conn._local_candidates = [local_candidate]
+        conn._local_candidates_end = True
+        self.register(local_username, protocol)  # type: ignore[arg-type]
+        return conn
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    @property
+    def local_addr(self) -> tuple[str, int] | None:
+        return self._local_addr
+
+    async def close(self) -> None:
+        """Close the shared UDP socket and drop all dispatch registrations."""
+        if self._transport is not None:
+            self._transport.close()
+            self._transport = None
+        # Clear routing tables so a datagram racing with close() can't be
+        # dispatched to a half-torn-down protocol.
+        self._by_ufrag.clear()
+        self._by_addr.clear()
+        self._unknown_stun_handler = None

newsfragments/1352.feature.rst

@@ -0,0 +1 @@
+Added ``libp2p.transport.webrtc._udp_mux.UdpMux`` — a shared UDP socket dispatcher for WebRTC-Direct inbound connections. Routes pre-ICE STUN datagrams by ``USERNAME`` ufrag prefix and post-ICE DTLS/SCTP frames by remote address, enabling a single fixed port to demultiplex concurrent inbound dials without spinning up a separate socket per peer (prerequisite for a spec-aligned WebRTC-Direct v2 listener, libp2p/specs#715).

pyproject.toml

@@ -337,6 +337,7 @@ project_excludes = [
     "./tests/core/transport/webrtc",
     "./libp2p/transport/webrtc/_asyncio_bridge.py",
     "./libp2p/transport/webrtc/_aiortc_helpers.py",
+    "./libp2p/transport/webrtc/_udp_mux.py",
     "./libp2p/transport/webrtc/certificate.py",
     "./libp2p/transport/webrtc/transport.py",
     "./libp2p/transport/webrtc/listener.py",