Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/selkies/ice/mux.py: 23%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at https://mozilla.org/MPL/2.0/.
4"""Shared ICE sockets: one UDP port and one TCP port for every session.
6A deployment that can expose only a port or two forwards one UDP port and one
7TCP port and has every session's host candidates share them, instead of an
8ephemeral port (or a `webrtc_port_range` window) per session. A `UdpMux`
9binds the port on each host address, a `TcpMux` listens on it, and each ICE
10`Connection` attaches one protocol per address under its local username
11fragment; the protocol sees a datagram transport and never learns that the
12socket is shared.
14Demultiplexing follows the ICE handshake itself. A datagram from a remote
15address already known belongs to the connection it was learned for. One from
16an unknown address is only ever STUN: a request names its connection in the
17USERNAME (`local:remote` fragments); a response answers a transaction that an
18attached protocol registered when it sent the request, a connectivity check or
19a server-reflexive query. A peer's address is learned from its request once
20the connection has authenticated it, or from its answer to a check, never from
21a STUN server's answer, so a STUN server shared by every session maps to none
22of them. Anything else is dropped.
24ICE-TCP (RFC 6544) rides the TCP port. Each accepted stream carries RFC 4571
25frames, a 16-bit length ahead of every packet, and its first frame must be the
26peer's STUN binding request, whose USERNAME hands the stream to the connection
27it belongs to; from then on the stream is that connection's datagram path to
28the peer's address, and losing it while it carries the nominated pair ends the
29connection at once rather than after consent expires. A stream that sends
30anything else first, or nothing within `FIRST_FRAME_TIMEOUT`, is closed.
31Writes are bounded: once the socket's send buffer backs up past asyncio's
32high-water mark, packets are dropped the way a full UDP socket drops them, so
33a slow link costs frames rather than memory, and the media layers recover from
34loss as they do on UDP.
35"""
36import asyncio
37import logging
38import os
39import struct
40from typing import Any, Callable, Optional
42from . import stun
43from .turn import size_udp_socket
44from .utils import bind_errno
46logger = logging.getLogger(__name__)
48# Seconds an accepted TCP stream has to send its STUN binding request.
49FIRST_FRAME_TIMEOUT = 30.0
50# Send backlog, per stream, past which packets are dropped instead of queued.
51TCP_WRITE_HIGH_WATER = 1 << 20
52_FRAME_HEADER = struct.Struct("!H")
53_MAX_FRAME = 0xFFFF
54_STUN_COOKIE = struct.pack("!I", stun.COOKIE)
57def is_stun(data: bytes) -> bool:
58 """Whether `data` starts like a STUN message: zero top bits and the magic cookie."""
59 return (
60 len(data) >= stun.HEADER_LENGTH
61 and data[0] & 0xC0 == 0
62 and data[4:8] == _STUN_COOKIE
63 )
66def username_fragment(message: stun.Message) -> Optional[str]:
67 """The local username fragment an ICE check's USERNAME (`local:remote`) names."""
68 username = message.attributes.get("USERNAME")
69 if not isinstance(username, str) or not username:
70 return None
71 return username.split(":", 1)[0]
74class _MuxPoint:
75 """Bookkeeping of one shared socket or listener on one local address.
77 Attached protocols are keyed by their local username fragment, learned
78 peers by remote address, and STUN transactions in flight by transaction id
79 together with whether their answer identifies a peer.
80 """
82 def __init__(self, address: str, port: int) -> None:
83 self.address = address
84 self.port = port
85 self.by_ufrag: dict[str, Any] = {}
86 self.by_addr: dict[tuple, Any] = {}
87 self.by_tid: dict[bytes, tuple] = {}
89 @property
90 def sockname(self) -> tuple:
91 return (self.address, self.port)
93 def attach(self, ufrag: str, protocol: Any) -> None:
94 if ufrag in self.by_ufrag:
95 raise ValueError(
96 f"username fragment {ufrag!r} is already attached on {self.sockname}"
97 )
98 self.by_ufrag[ufrag] = protocol
100 def detach(self, ufrag: str, protocol: Any) -> None:
101 if self.by_ufrag.get(ufrag) is protocol:
102 del self.by_ufrag[ufrag]
103 for addr in [a for a, p in self.by_addr.items() if p is protocol]:
104 del self.by_addr[addr]
105 for tid in [t for t, (p, _) in self.by_tid.items() if p is protocol]:
106 del self.by_tid[tid]
108 def learn(self, addr: tuple, protocol: Any) -> None:
109 if self.by_addr.get(addr) is not protocol:
110 self.by_addr[addr] = protocol
111 logger.debug("%s learned peer %s for ufrag of %r", self.sockname, addr, protocol)
113 def route(self, data: bytes, addr: tuple) -> Any:
114 """The attached protocol a datagram from `addr` belongs to, or None.
116 A learned peer's datagrams are routed without a look inside; only STUN
117 from an unknown address is parsed, to find the connection it names.
118 """
119 protocol = self.by_addr.get(addr)
120 if protocol is not None or not is_stun(data):
121 return protocol
122 try:
123 message = stun.parse_message(data)
124 except ValueError:
125 return None
126 if message.message_class in (stun.Class.RESPONSE, stun.Class.ERROR):
127 pending = self.by_tid.get(message.transaction_id)
128 if pending is not None:
129 protocol, names_peer = pending
130 if names_peer:
131 self.learn(addr, protocol)
132 elif message.message_class == stun.Class.REQUEST:
133 # Learned only once the connection authenticates it (`learn_peer`).
134 protocol = self.by_ufrag.get(username_fragment(message) or "")
135 return protocol
138class _MuxTransport:
139 """The datagram-transport face one attached protocol sends and closes through."""
141 def __init__(self, point: _MuxPoint, ufrag: str) -> None:
142 self._point = point
143 self._ufrag = ufrag
144 self.protocol: Any = None
145 self._closed = False
147 def get_extra_info(self, name: str, default: Any = None) -> Any:
148 if name == "sockname":
149 return self._point.sockname
150 return default
152 def expect_response(self, transaction_id: bytes, names_peer: bool) -> None:
153 """Route the answer to a STUN request this protocol sends back to it.
155 `names_peer` says the request is an ICE check, so the answer's source
156 is the peer and is learned; a STUN server's answer is not.
157 """
158 self._point.by_tid[transaction_id] = (self.protocol, names_peer)
160 def forget_response(self, transaction_id: bytes) -> None:
161 self._point.by_tid.pop(transaction_id, None)
163 def learn_peer(self, addr: tuple) -> None:
164 """Route `addr`'s datagrams to this protocol: a check from it authenticated."""
165 self._point.learn(addr, self.protocol)
167 def close(self) -> None:
168 if self._closed:
169 return
170 self._closed = True
171 self._point.detach(self._ufrag, self.protocol)
172 self._release()
173 self.protocol.connection_lost(None)
175 def _release(self) -> None:
176 pass
179class _UdpMuxTransport(_MuxTransport):
180 def __init__(self, point: _MuxPoint, ufrag: str, socket_transport: Any) -> None:
181 super().__init__(point, ufrag)
182 self._socket_transport = socket_transport
184 def get_extra_info(self, name: str, default: Any = None) -> Any:
185 if name == "sockname":
186 return self._point.sockname
187 return self._socket_transport.get_extra_info(name, default)
189 def sendto(self, data: bytes, addr: tuple) -> None:
190 self._socket_transport.sendto(data, addr)
193class _UdpMuxSocket(asyncio.DatagramProtocol):
194 """One shared UDP socket, bound to one local address on the mux port."""
196 def __init__(self, point: _MuxPoint) -> None:
197 self.point = point
198 self.transport: Any = None
199 self.closed: Optional[asyncio.Future] = None
201 def connection_made(self, transport: asyncio.BaseTransport) -> None:
202 self.transport = transport
203 self.closed = asyncio.get_running_loop().create_future()
205 def datagram_received(self, data: bytes, addr: tuple) -> None:
206 if not data:
207 return
208 addr = (addr[0], addr[1])
209 protocol = self.point.route(data, addr)
210 if protocol is None:
211 logger.debug("%s dropped %d bytes from unknown peer %s", self.point.sockname, len(data), addr)
212 return
213 protocol.datagram_received(data, addr)
215 def error_received(self, exc: Exception) -> None:
216 logger.debug("%s error_received(%s)", self.point.sockname, exc)
218 def connection_lost(self, exc: Optional[Exception]) -> None:
219 for protocol in list(self.point.by_ufrag.values()):
220 protocol.transport.close()
221 if self.closed is not None and not self.closed.done():
222 self.closed.set_result(None)
225class UdpMux:
226 """Every session's UDP host candidates on one port, one socket per local address.
228 `open` binds the port on the addresses given, so a port in use fails the
229 service at startup rather than the first session; `attach` binds lazily for
230 an address that appeared later.
231 """
233 def __init__(self, port: int) -> None:
234 self.port = port
235 self._sockets: dict[str, _UdpMuxSocket] = {}
236 self._closed = False
238 @property
239 def listen_addresses(self) -> list:
240 return [sock.point.sockname for sock in self._sockets.values()]
242 def attached(self, address: Optional[str] = None) -> int:
243 """Protocols attached on `address`, or on every address."""
244 sockets = self._sockets.values() if address is None else [self._sockets[address]]
245 return sum(len(sock.point.by_ufrag) for sock in sockets)
247 async def open(self, addresses: list) -> None:
248 """Bind the port on each address; on failure nothing stays bound.
250 Raises:
251 OSError: A bind failed, typically because the port is in use.
252 """
253 try:
254 for address in addresses:
255 await self._bind(address)
256 except OSError:
257 await self.close()
258 raise
260 async def _bind(self, address: str) -> _UdpMuxSocket:
261 """The socket on `address`, bound on first use.
263 Raises:
264 OSError: The bind failed, carrying the kernel's code and reason
265 even where the loop's own wrapper reports neither.
266 """
267 if self._closed:
268 raise OSError("UDP mux is closed")
269 sock = self._sockets.get(address)
270 if sock is not None:
271 return sock
272 loop = asyncio.get_running_loop()
273 point = _MuxPoint(address, self.port)
274 try:
275 transport, sock = await loop.create_datagram_endpoint(
276 lambda: _UdpMuxSocket(point), local_addr=(address, self.port)
277 )
278 except OSError as exc:
279 code = bind_errno(exc)
280 if exc.errno is None and code is not None:
281 raise OSError(code, os.strerror(code)) from exc
282 raise
283 raw = transport.get_extra_info("socket")
284 if raw is not None:
285 size_udp_socket(raw)
286 self._sockets[address] = sock
287 return sock
289 async def attach(self, address: str, ufrag: str, protocol_factory: Callable[[], Any]) -> Any:
290 """A protocol of a connection with local fragment `ufrag` on `address`.
292 Raises:
293 OSError: The address could not be bound on the mux port.
294 ValueError: The fragment is already attached on that address.
295 """
296 sock = await self._bind(address)
297 transport = _UdpMuxTransport(sock.point, ufrag, sock.transport)
298 protocol = protocol_factory()
299 transport.protocol = protocol
300 sock.point.attach(ufrag, protocol)
301 protocol.connection_made(transport)
302 return protocol
304 async def close(self) -> None:
305 """Detach every protocol and release the sockets, which are free to rebind on return."""
306 self._closed = True
307 sockets = list(self._sockets.values())
308 self._sockets.clear()
309 for sock in sockets:
310 if sock.transport is not None:
311 sock.transport.close()
312 for sock in sockets:
313 if sock.closed is not None:
314 await sock.closed
317class _TcpMuxTransport(_MuxTransport):
318 """A connection's face of the TCP listener: its accepted streams by peer address."""
320 def __init__(self, point: _MuxPoint, ufrag: str) -> None:
321 super().__init__(point, ufrag)
322 self._streams: dict[tuple, "_TcpStream"] = {}
324 def sendto(self, data: bytes, addr: tuple) -> None:
325 stream = self._streams.get(addr)
326 if stream is None:
327 logger.debug("%s no stream to %s, dropping %d bytes", self._point.sockname, addr, len(data))
328 return
329 stream.send(data)
331 def add_stream(self, stream: "_TcpStream") -> None:
332 previous = self._streams.get(stream.remote_addr)
333 self._streams[stream.remote_addr] = stream
334 if previous is not None and previous is not stream:
335 previous.close()
337 def remove_stream(self, stream: "_TcpStream") -> None:
338 if self._streams.get(stream.remote_addr) is stream:
339 del self._streams[stream.remote_addr]
340 if not self._closed:
341 self.protocol.receiver.peer_stream_lost(stream.remote_addr)
343 def _release(self) -> None:
344 for stream in list(self._streams.values()):
345 stream.close()
346 self._streams.clear()
349class _TcpStream(asyncio.BufferedProtocol):
350 """One accepted ICE-TCP stream: RFC 4571 frames in, bounded frames out.
352 The receive buffer is handed to asyncio directly, so a packet is copied
353 once, out of the buffer into the frame the connection receives.
354 """
356 def __init__(self, listener: "_TcpMuxListener") -> None:
357 self._listener = listener
358 self._owner: Optional[_TcpMuxTransport] = None
359 self._buffer = bytearray(2 * (_FRAME_HEADER.size + _MAX_FRAME))
360 self._view = memoryview(self._buffer)
361 self._start = 0
362 self._end = 0
363 self._paused = False
364 self._first_timer: Optional[asyncio.TimerHandle] = None
365 self.transport: Any = None
366 self.remote_addr: tuple = ("", 0)
367 self.dropped = 0
369 def connection_made(self, transport: asyncio.BaseTransport) -> None:
370 self.transport = transport
371 peer = transport.get_extra_info("peername") or ("", 0)
372 self.remote_addr = (peer[0], peer[1])
373 transport.set_write_buffer_limits(high=TCP_WRITE_HIGH_WATER)
374 self._first_timer = asyncio.get_running_loop().call_later(
375 FIRST_FRAME_TIMEOUT, self._first_frame_timeout
376 )
378 def _first_frame_timeout(self) -> None:
379 if self._owner is None:
380 logger.info("%s closing %s: no STUN binding request within %.0fs",
381 self._listener.point.sockname, self.remote_addr, FIRST_FRAME_TIMEOUT)
382 self.close()
384 def get_buffer(self, sizehint: int) -> memoryview:
385 if self._end == len(self._buffer):
386 self._view[: self._end - self._start] = self._view[self._start : self._end]
387 self._end -= self._start
388 self._start = 0
389 return self._view[self._end :]
391 def buffer_updated(self, nbytes: int) -> None:
392 self._end += nbytes
393 while True:
394 pending = self._end - self._start
395 if pending < _FRAME_HEADER.size:
396 break
397 (length,) = _FRAME_HEADER.unpack_from(self._buffer, self._start)
398 if pending < _FRAME_HEADER.size + length:
399 break
400 begin = self._start + _FRAME_HEADER.size
401 frame = bytes(self._view[begin : begin + length])
402 self._start = begin + length
403 if length and not self._frame(frame):
404 return
405 if self._start == self._end:
406 self._start = self._end = 0
408 def _frame(self, data: bytes) -> bool:
409 """Deliver one frame; the first must be a STUN request naming a connection."""
410 if self._owner is None:
411 owner = self._claim(data)
412 if owner is None:
413 self.close()
414 return False
415 self._owner = owner
416 if self._first_timer is not None:
417 self._first_timer.cancel()
418 owner.add_stream(self)
419 self._owner.protocol.datagram_received(data, self.remote_addr)
420 return True
422 def _claim(self, data: bytes) -> Optional[_TcpMuxTransport]:
423 point = self._listener.point
424 if not is_stun(data):
425 logger.info("%s closing %s: first frame is not STUN", point.sockname, self.remote_addr)
426 return None
427 try:
428 message = stun.parse_message(data)
429 except ValueError as exc:
430 logger.info("%s closing %s: malformed STUN (%s)", point.sockname, self.remote_addr, exc)
431 return None
432 ufrag = username_fragment(message)
433 owner = point.by_ufrag.get(ufrag or "") if message.message_class == stun.Class.REQUEST else None
434 if owner is None:
435 logger.info("%s closing %s: no connection for username fragment %r",
436 point.sockname, self.remote_addr, ufrag)
437 return None
438 return owner.transport
440 def pause_writing(self) -> None:
441 self._paused = True
443 def resume_writing(self) -> None:
444 self._paused = False
446 def send(self, data: bytes) -> None:
447 if self.transport is None or self.transport.is_closing():
448 return
449 if self._paused:
450 self.dropped += 1
451 return
452 self.transport.writelines((_FRAME_HEADER.pack(len(data)), data))
454 def close(self) -> None:
455 if self.transport is not None and not self.transport.is_closing():
456 self.transport.close()
458 def connection_lost(self, exc: Optional[Exception]) -> None:
459 if self._first_timer is not None:
460 self._first_timer.cancel()
461 if self._owner is not None:
462 owner, self._owner = self._owner, None
463 owner.remove_stream(self)
464 if self.dropped:
465 logger.info("%s stream from %s dropped %d packets to a backed-up link",
466 self._listener.point.sockname, self.remote_addr, self.dropped)
469class _TcpMuxListener:
470 """One listening socket on the mux port, bound to one local address."""
472 def __init__(self, point: _MuxPoint) -> None:
473 self.point = point
474 self.server: Any = None
477class TcpMux:
478 """Every session's passive TCP host candidates on one port, one listener per address."""
480 def __init__(self, port: int) -> None:
481 self.port = port
482 self._listeners: dict[str, _TcpMuxListener] = {}
483 self._closed = False
485 @property
486 def listen_addresses(self) -> list:
487 return [listener.point.sockname for listener in self._listeners.values()]
489 def attached(self, address: Optional[str] = None) -> int:
490 """Protocols attached on `address`, or on every address."""
491 listeners = self._listeners.values() if address is None else [self._listeners[address]]
492 return sum(len(listener.point.by_ufrag) for listener in listeners)
494 async def open(self, addresses: list) -> None:
495 """Listen on each address; on failure nothing stays bound.
497 Raises:
498 OSError: A bind failed, typically because the port is in use.
499 """
500 try:
501 for address in addresses:
502 await self._listen(address)
503 except OSError:
504 await self.close()
505 raise
507 async def _listen(self, address: str) -> _TcpMuxListener:
508 if self._closed:
509 raise OSError("TCP mux is closed")
510 listener = self._listeners.get(address)
511 if listener is not None:
512 return listener
513 loop = asyncio.get_running_loop()
514 listener = _TcpMuxListener(_MuxPoint(address, self.port))
515 listener.server = await loop.create_server(
516 lambda: _TcpStream(listener), host=address, port=self.port, reuse_address=True
517 )
518 self._listeners[address] = listener
519 return listener
521 async def attach(self, address: str, ufrag: str, protocol_factory: Callable[[], Any]) -> Any:
522 """A protocol of a connection with local fragment `ufrag` on `address`.
524 Raises:
525 OSError: The address could not be bound on the mux port.
526 ValueError: The fragment is already attached on that address.
527 """
528 listener = await self._listen(address)
529 transport = _TcpMuxTransport(listener.point, ufrag)
530 protocol = protocol_factory()
531 transport.protocol = protocol
532 listener.point.attach(ufrag, protocol)
533 protocol.connection_made(transport)
534 return protocol
536 async def close(self) -> None:
537 """Detach every protocol, end their streams, and stop listening; the port is free on return."""
538 self._closed = True
539 listeners = list(self._listeners.values())
540 self._listeners.clear()
541 for listener in listeners:
542 for protocol in list(listener.point.by_ufrag.values()):
543 protocol.transport.close()
544 if listener.server is not None:
545 listener.server.close()