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

[yashksaini-coder]

Fixes #1352. Recreates #1381, which was auto-closed when the source fork was recreated — same diff, no code changes.

Summary

• Adds libp2p/transport/webrtc/_udp_mux.py — a UdpMux class (a shared asyncio.DatagramProtocol) that owns one OS UDP socket and demultiplexes concurrent inbound WebRTC-Direct dials
• Pre-ICE: routes STUN BINDING REQUESTs by USERNAME ufrag prefix (parsed from the STUN message)
• Post-ICE: routes DTLS/SCTP frames by remote (host, port) (registered after ICE nomination)
• _MuxedTransport adapter gives each aioice.Connection a fake transport backed by the shared socket, bypassing _gather_candidates entirely so aioice never binds its own port
• add_ice_connection(ufrag, pwd, host=) creates an aioice.Connection with preset local_username/local_password and injects a mux-backed StunProtocol into connection._protocols
• 11 tests covering all dispatch paths, transport adapter, and connection factory

Why

Closes the ICE mux gap identified in the spike (#1352). A spec-aligned WebRTC-Direct listener (libp2p/specs#715) must advertise a single fixed UDP port and demux concurrent dials on that port — impossible with the current aioice architecture where each Connection binds its own ephemeral socket. UdpMux solves this without modifying aioice.

Architecture note

StunProtocol is an internal aioice class. We access it directly because we're injecting into aioice's _protocols list — intentional monkey-patching. The two # type: ignore annotations mark aioice's connection_lost(exc: Exception) annotation gap (should be Optional[Exception] per asyncio docs).

What's next

Follow-up PR: wire UdpMux into listener.py to replace the current HTTP-signaling harness with a spec-aligned STUN-dispatch listener (WebRTC-Direct v2, libp2p/specs#715).

Test plan

• uv run pytest tests/core/transport/webrtc/test_udp_mux.py -v → 11/11 pass (skipped when aioice not installed)
• uv run make lint → all pre-commit hooks pass (mypy + pyrefly clean)

[acul71]

AI PR Review: #1397

PR: feat(webrtc): add UdpMux for shared-port WebRTC-Direct inbound dispatch
Author: yashksaini-coder
Reviewer: AI-assisted maintainer review
Date: 2026-08-12
Base: main ← feat/webrtc-udpmux-1352
Related: Recreates closed #1381; Fixes #1352

---

1. Summary of Changes

This PR adds a foundational WebRTC transport primitive: UdpMux, a shared asyncio.DatagramProtocol that owns one OS UDP socket and demultiplexes concurrent inbound WebRTC-Direct dials without binding a new port per aioice.Connection.

It addresses #1352 (spike: aiortc ICE mux / STUN USERNAME exposure). That spike concluded ICE mux was the hard blocker for a spec-aligned WebRTC-Direct listener (libp2p/specs#715) and recommended Path A (UdpMux, ideally upstream in aioice). This PR implements an in-tree mux via intentional monkey-patching of aioice internals (StunProtocol, Connection._protocols) rather than an upstream aioice change.

Behavior:

• Pre-ICE: route STUN by USERNAME local-ufrag prefix
• Post-ICE: route non-STUN (DTLS/SCTP) by remote (host, port)
• _MuxedTransport fakes a per-connection datagram transport on the shared socket
• add_ice_connection() presets ICE credentials and injects a mux-backed StunProtocol

Files affected:

File	Role
libp2p/transport/webrtc/_udp_mux.py	New shared-port ICE mux (transport / WebRTC)
tests/core/transport/webrtc/test_udp_mux.py	11 unit tests for dispatch / transport / factory
newsfragments/1352.feature.rst	Towncrier newsfragment
pyproject.toml	Add _udp_mux.py to pyrefly project_excludes (consistent with other webrtc modules)

Breaking changes / deprecations: None. Not wired into listener.py yet (explicit follow-up).

Discussions referenced: None (specs PR libp2p/specs#715 and issues #546 / #1309 referenced from #1352).

Maintainer feedback on PR: No review comments. Maintainer acul71 merged main into the branch multiple times (latest c2630b0e). CI mostly green; Windows core jobs were still pending at review time; Linux tox/docs checks passed.

---

2. Branch Sync Status and Merge Conflicts

Branch Sync Status

• Status: ℹ️ Ahead of origin/main (not behind)
• Details: branch_sync_status.txt reports 0 8 → 0 commits behind, 8 commits ahead of origin/main
• Recent merges from main by acul71 keep the branch current

Merge Conflict Analysis

• Conflicts Detected: ✅ No conflicts
• Test merge of origin/main reported Already up to date. / === NO MERGE CONFLICTS DETECTED ===

✅ **No merge conflicts detected.** The PR branch can be merged cleanly into origin/main.

---

3. Strengths

• Correct problem framing: Matches the #1352 spike gap (aioice binds one ephemeral UDP socket per Connection; WebRTC-Direct needs one advertised port).
• Clear module docs and example explaining pre-ICE vs post-ICE routing and teardown responsibilities.
• Optional aioice import with a clear RuntimeError when add_ice_connection is used without the webrtc extra.
• Solid unit coverage of the dispatch table (known/unknown ufrag, known/unknown addr, unregister, muxed sendto/sockname/close) — 11/11 pass locally with aioice installed.
• Process hygiene: Linked issue + 1352.feature.rst newsfragment present; PR body and commit message are high quality.
• Local validation clean: make lint, make typecheck, make linux-docs all exit 0; full make test green (3211 passed, 15 skipped).

---

4. Issues Found

Critical

None for merge of an unwired primitive (no production listener path calls this yet). The Major items below will block the follow-up listener PR and should be fixed here so the API is actually usable with aioice 0.10.x.

Major

• File: libp2p/transport/webrtc/_udp_mux.py
• Line(s): 185–239 (esp. 212–238) and docstring 86–88 / 199–205
• Issue: add_ice_connection() injects a StunProtocol into conn._protocols but does not mark local gathering complete. On aioice 0.10.2, Connection.connect() immediately raises ConnectionError: Local candidates gathering was not performed unless _local_candidates_end is true. Callers who follow the documented await conn.connect() path therefore cannot complete ICE. Calling real gather_candidates() afterward binds additional ephemeral sockets (probe observed mux port plus five extra ports), defeating the shared-port goal. Docstring also tells callers to use conn.set_remote_candidates([...]), which does not exist on aioice 0.10.2 (add_remote_candidate is the API).
• Suggestion: After injecting the protocol, also set conn._local_candidates = [protocol.local_candidate] and conn._local_candidates_end = True (and document the private-API dependency). Fix the docstring to use add_remote_candidate / end-of-candidates (None). Add a regression test that connect() is reachable without opening extra UDP sockets (mock remote or assert len(conn._protocols) == 1 and no new ports after a guarded gather-skip path).

libp2p/transport/webrtc/_udp_mux.py — lines 212–238

        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)
        protocol.transport = muxed_transport  # type: ignore[assignment]
        muxed_transport._protocol = protocol  # type: ignore[assignment]

        protocol.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",
        )

        conn._protocols.append(protocol)
        self.register(local_username, protocol)  # type: ignore[arg-type]
        return conn

• File: libp2p/transport/webrtc/_udp_mux.py
• Line(s): 136–153
• Issue: Unknown-ufrag STUN is logged and dropped. For WebRTC-Direct (the stated use case), the first inbound BINDING REQUEST is how the listener learns the dialer’s ufrag and must create/register a connection. With only register / add_ice_connection, there is a chicken-and-egg: you cannot register until you know the ufrag, but you only see the ufrag inside datagram_received after drop. The spike (#1352) explicitly required inbound STUN USERNAME exposure for listener dispatch.
• Suggestion: Add an optional unhandled-STUN callback / queue (e.g. on_unknown_stun(ufrag, data, addr) or await mux.next_stun()) invoked when parse succeeds but ufrag is unregistered, so the follow-up listener can create the connection and optionally replay the datagram. Cover with a test.

libp2p/transport/webrtc/_udp_mux.py — lines 136–147

    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
            logger.debug("UdpMux: no handler for ufrag %r from %s", ufrag, norm)

• File: tests/core/transport/webrtc/test_udp_mux.py

• Line(s): 240–286 (and missing coverage overall)

• Issue: Tests never drive a real (or lightly stubbed) ICE handshake through the mux. They assert registration and dispatch tables only. That left the _local_candidates_end / gather_candidates footgun undetected. There is also no test that gather_candidates is not required / must not be called.

• Suggestion: Add an integration-style test (two muxes or mux + raw UDP peer) that completes STUN connectivity checks on the shared port, or at minimum asserts post-add_ice_connection flags allow connect() without extra sockets.

• File: libp2p/transport/webrtc/_udp_mux.py / issue #1352

• Line(s): n/a (process / architecture)

• Issue: Spike #1352’s recommendation was Path A upstream aioice UdpMux (or Path B alternate backend). This lands an in-tree monkey-patch with no version pin note, no upstream tracking issue, and reliance on private StunProtocol / _protocols. Fragile across aioice releases.

• Suggestion: Document supported aioice version(s) in the module docstring; open a tracking issue/PR for upstreaming; consider pinning or testing against the aioice version pulled by aiortc. Not a newsfragment blocker, but should be explicit before calling #1352 “closed” as fully resolved.

Minor

• File: libp2p/transport/webrtc/_udp_mux.py

• Line(s): 117, 249–253

• Issue: create() uses asyncio.get_event_loop() (prefer get_running_loop() in async context). close() closes the shared transport but does not clear _by_ufrag / _by_addr or signal per-connection protocols, so late datagrams could hit half-dead handlers if close races with dispatch.

• Suggestion: Use get_running_loop(); on close(), clear maps and/or call connection_lost on registered protocols.

• File: tests/core/transport/webrtc/test_udp_mux.py

• Line(s): 199

• Issue: await asyncio.sleep(0.05) for UDP delivery can flake under load.

• Suggestion: Poll with a short timeout until received is non-empty (same de-flake pattern used elsewhere in the repo).

• File: newsfragments/1352.feature.rst

• Line(s): 1

• Issue: Valid and ends with newline, but very long and implementation-heavy for towncrier user-facing notes.

• Suggestion: Optional shorten to user impact only (shared-port WebRTC-Direct inbound demux prerequisite).

---

5. Security Review

• Risk: Unregistered STUN/non-STUN datagrams are dropped (good fail-closed default). No authentication at the mux layer — expected; ICE MESSAGE-INTEGRITY and later DTLS/Noise remain responsible.

• Impact: Low for this PR (primitive unused by listener yet).

• Mitigation: When wiring the listener, rate-limit unknown-ufrag STUN handling to avoid CPU DoS from parse_message on floods; do not log full USERNAME/password material at info level (current debug logs of ufrag only are acceptable).

• Risk: Post-ICE dispatch keys only on (host, port). NAT rebinding / address change could mis-route or drop DTLS until re-registered; STUN consent still needs the ufrag mapping to remain registered for the connection lifetime (docstring teardown order is easy to get wrong).

• Impact: Medium once listener lands; Low now.

• Mitigation: Keep ufrag registration for the full connection lifetime; document that _by_addr is non-STUN only; consider falling through unknown-ufrag STUN to addr map as a secondary key.

• Risk: Monkey-patching private aioice types could break unexpectedly on dependency bumps (availability / correctness, not direct exploit).

• Impact: Low–Medium operationally.

• Mitigation: Pin/test aioice version; upstream the mux API.

Security Impact: Low (for this unwired change).

---

6. Documentation and Examples

• Module docstring and class example are strong for developers.
• Example / add_ice_connection docstring are incorrect against aioice 0.10.2 (set_remote_candidates, incomplete connect prerequisites) — must be fixed (see Major).
• No public README/tutorial update needed yet (not user-wired); newsfragment covers changelog.
• Follow-up listener PR should document the STUN-first discovery flow once the unknown-ufrag hook exists.

---

7. Newsfragment Requirement

• Issue reference: ✅ PR body Fixes #1352 (commit also Closes #1352)
• Newsfragment: ✅ newsfragments/1352.feature.rst present, .feature type appropriate, ends with newline
• Note: Closing a spike issue with an implementation is reasonable if maintainers accept in-tree Path A; still call out upstream tracking as follow-up (Major process note above). Not a newsfragment blocker.

---

8. Tests and Validation

New tests

• 11 focused unit tests; good table coverage; missing end-to-end ICE-on-mux and “no extra bind” assertions (Major).
• Sync asyncio.run wrappers are appropriate under project-wide trio mode.

Local commands

Command	Result
make lint	✅ Passed (all pre-commit hooks)
make typecheck	✅ Passed (mypy + pyrefly)
make test	✅ 3211 passed, 15 skipped, 6 warnings in ~173s; plus serial identify/host/quic subset 6 passed
make linux-docs	✅ build succeeded / EXIT 0
pytest tests/core/transport/webrtc/test_udp_mux.py -v	✅ 11 passed (with .[webrtc] installed)

Lint: No errors/warnings.

Typecheck: Clean. _udp_mux.py excluded from pyrefly (same pattern as other webrtc modules) — acceptable but means pyrefly does not guard this file.

Tests: Full suite green. Warnings are pre-existing RuntimeWarning: coroutine 'AsyncMockMixin._execute_mock_call' was never awaited in unrelated webrtc/pubsub tests — not introduced by this PR. Skips are optional DB backends / platform-specific path tests.

Docs: Sphinx build succeeded; no new doc errors attributable to this PR.

Manual probe (aioice 0.10.2): Confirmed connect() fails without _local_candidates_end; gather_candidates() after add_ice_connection opens extra UDP ports.

GitHub CI: Linux tox (core/lint/docs/interop/…) and Read the Docs passed; Windows core jobs were pending at review time.

---

9. Recommendations for Improvement

1. Make add_ice_connection() set _local_candidates + _local_candidates_end so connect() works without gather_candidates().
2. Fix docstring API (add_remote_candidate, not set_remote_candidates) and document that callers must not call gather_candidates().
3. Add unknown-ufrag STUN hook for WebRTC-Direct discovery (or document an alternate peek path and implement it before listener work).
4. Add a regression test proving no extra UDP binds and that ICE connect is reachable.
5. Prefer get_running_loop(); clear registrations on mux close().
6. Track upstreaming to aioice / pin tested aioice version in docs.
7. Optional: shorten newsfragment; replace sleep(0.05) with a poll in the sendto test.

---

10. Questions for the Author

1. Was add_ice_connection() validated against a live conn.connect() on aioice 0.10.x, or only against the unit registration tests?
2. How should the follow-up listener observe the first STUN packet’s ufrag if unknown ufrags are dropped today?
3. Is the intent to keep this in-tree long-term, or still pursue the #1352 Path A upstream aioice PR?
4. After ICE succeeds, should ufrag registration remain until full teardown (needed for STUN consent), and should that lifetime be documented more explicitly?
5. Should _MuxedTransport implement additional DatagramTransport surface (is_closing, get_protocol, etc.) for forward compatibility with aioice?

---

11. Overall Assessment

• Quality Rating: Needs Work
• Security Impact: Low
• Merge Readiness: Needs fixes (correct ICE lifecycle wiring + docstring; strongly prefer unknown-STUN hook + regression test before merge, or accept as “dispatch-table only” with explicit WIP labeling and do not close #1352 as fully done)
• Confidence: High

Bottom line: Directionally right and well-scoped for the spike follow-through, with good unit tests and clean CI/lint. As written, add_ice_connection() does not produce an aioice Connection that can connect() without re-binding ports, and the mux cannot support STUN-first WebRTC-Direct discovery. Fix those before treating this as the foundational primitive the follow-up listener will rely on.

[yashksaini-coder]

Thanks, this was a really useful review — you were right that the primitive didn't actually connect. Fixed all of it and validated against aioice 0.10.2 locally:

ICE lifecycle (the main bug): add_ice_connection() now sets conn._local_candidates and _local_candidates_end = True, so connect() clears the "Local candidates gathering was not performed" guard without calling gather_candidates(). I checked the socket count before/after — no extra UDP binds, and connect() now gets past the gather check (fails only on missing remote creds / no real peer, as expected). Added a regression test that asserts exactly this.

Docstring: fixed to the real API — await conn.add_remote_candidate(c) then add_remote_candidate(None) to end, and an explicit "do not call gather_candidates()" note.

Unknown-STUN discovery: added set_unknown_stun_handler(cb). Unregistered-ufrag STUN now goes to the callback with (ufrag, data, addr) so the listener can create the connection and replay the datagram, instead of dropping it.

Minors: get_running_loop(); close() clears the dispatch maps; the sendto test polls instead of sleep(0.05).

Answering your questions directly:

1. It was only validated against unit registration before — that's exactly how the gather footgun slipped through. It's now driven through connect() in a test.
2. Via the new set_unknown_stun_handler hook — that's the first-BINDING-REQUEST path.
3. In-tree for now; I documented the aioice-0.10.x version dependency in the module and noted upstreaming as a spike: aiortc ICE mux / STUN USERNAME exposure for webrtc-direct v2 listener #1352 follow-up.
4. Yes — ufrag registration is kept for the connection lifetime (teardown calls unregister); documented.
5. Deferred until the listener needs it, to avoid guessing at surface.

13/13 tests green, ruff + mypy clean.

[yashksaini-coder]

Pushed 993219de: hardened UdpMux.datagram_received against malformed STUN. aioice parses attributes with struct.unpack, which raises struct.error (not a ValueError) on a short/malformed fixed-width attribute — the demux only caught ValueError, so a crafted packet escaped the handler (mis-classification + traceback log-flood under a flood). Now catches struct.error too and routes by addr. Added a regression test (green with fix, red without).

[acul71]

Re-reviewed after the follow-up commits (8154dafd, 993219de).

All items from the previous CHANGES_REQUESTED review are addressed:

• add_ice_connection() marks _local_candidates / _local_candidates_end so connect() works without gather_candidates() / extra UDP binds
• Docstrings use the real aioice API (add_remote_candidate / end-of-candidates)
• set_unknown_stun_handler covers first-contact STUN discovery
• Regression coverage for connect-reachable + malformed STUN (struct.error)
• Minors: get_running_loop(), map clear on close(), aioice 0.10.x docs

Approved. Not merging yet — waiting for CI to go green; will merge manually afterward.