3030
3131from ._record_update import RecordUpdate
3232from ._utils .asyncio import get_running_loop , run_coro_with_timeout
33+ from ._utils .net import (
34+ InterfacesType ,
35+ IPVersion ,
36+ add_multicast_member ,
37+ drop_multicast_member ,
38+ new_respond_socket ,
39+ normalize_interface_choice ,
40+ )
3341from ._utils .time import current_time_millis
3442from .const import _CACHE_CLEANUP_INTERVAL
3543
4351_CLOSE_TIMEOUT = 3000 # ms
4452
4553
54+ def _strip_zone (address : str ) -> str :
55+ """Drop a ``%zone`` suffix from an IPv6 address string."""
56+ percent = address .find ("%" )
57+ return address [:percent ] if percent != - 1 else address
58+
59+
60+ def _interface_bind_address (interface : str | tuple [tuple [str , int , int ], int ]) -> str :
61+ """Return the bind address an interface choice maps to, for diffing."""
62+ if isinstance (interface , tuple ):
63+ return _strip_zone (interface [0 ][0 ])
64+ return interface
65+
66+
67+ def _wrapped_bind_address (wrapped : _WrappedTransport ) -> str :
68+ """Return the bound address of a sender transport, for diffing."""
69+ return _strip_zone (wrapped .sock_name [0 ])
70+
71+
72+ def _wrapped_interface (wrapped : _WrappedTransport ) -> str | tuple [tuple [str , int , int ], int ]:
73+ """Reconstruct the interface representation for a sender, for membership drops."""
74+ sock_name = wrapped .sock_name
75+ if wrapped .is_ipv6 :
76+ scope_id = cast (int , sock_name [3 ]) if len (sock_name ) > 3 else 0
77+ return ((_strip_zone (sock_name [0 ]), 0 , 0 ), scope_id )
78+ return cast (str , sock_name [0 ])
79+
80+
4681class AsyncEngine :
4782 """An engine wraps sockets in the event loop."""
4883
4984 __slots__ = (
5085 "_cleanup_timer" ,
5186 "_listen_socket" ,
87+ "_listen_transport" ,
5288 "_respond_sockets" ,
5389 "_setup_task" ,
5490 "loop" ,
@@ -72,6 +108,7 @@ def __init__(
72108 self .senders : list [_WrappedTransport ] = []
73109 self .running_future : asyncio .Future [bool | None ] | None = None
74110 self ._listen_socket = listen_socket
111+ self ._listen_transport : _WrappedTransport | None = None
75112 self ._respond_sockets = respond_sockets
76113 self ._cleanup_timer : asyncio .TimerHandle | None = None
77114 self ._setup_task : asyncio .Task [None ] | None = None
@@ -98,8 +135,6 @@ async def _async_setup(self, loop_thread_ready: threading.Event | None) -> None:
98135
99136 async def _async_create_endpoints (self ) -> None :
100137 """Create endpoints to send and receive."""
101- assert self .loop is not None
102- loop = self .loop
103138 reader_sockets = []
104139 sender_sockets = []
105140 if self ._listen_socket :
@@ -110,22 +145,100 @@ async def _async_create_endpoints(self) -> None:
110145 sender_sockets .append (s )
111146
112147 for s in reader_sockets :
113- transport , protocol = await loop .create_datagram_endpoint ( # type: ignore[type-var]
114- lambda : AsyncListener (self .zc ), # type: ignore[arg-type, return-value]
115- sock = s ,
116- )
117- # Register the wrapped transport before releasing the engine's
118- # handle so a concurrent shutdown always sees ``s`` in exactly
119- # one place; do not add an ``await`` between these two steps.
120- self .protocols .append (cast (AsyncListener , protocol ))
121- self .readers .append (make_wrapped_transport (cast (asyncio .DatagramTransport , transport )))
122- if s in sender_sockets :
123- self .senders .append (make_wrapped_transport (cast (asyncio .DatagramTransport , transport )))
148+ reader = await self ._async_wrap_socket (s , s in sender_sockets )
149+ # The wrap above does not await before returning, so releasing
150+ # the engine's pending handle here keeps ``s`` in exactly one
151+ # place from a concurrent shutdown's point of view.
124152 if s is self ._listen_socket :
153+ # Keep a handle to the shared listen socket so interface
154+ # rescans can add/drop multicast memberships on it.
155+ self ._listen_transport = reader
125156 self ._listen_socket = None
126157 if s in self ._respond_sockets :
127158 self ._respond_sockets .remove (s )
128159
160+ async def _async_wrap_socket (self , sock : socket .socket , is_sender : bool ) -> _WrappedTransport :
161+ """Adopt a socket into a transport, register it, and return the reader wrapper."""
162+ assert self .loop is not None
163+ transport , protocol = await self .loop .create_datagram_endpoint ( # type: ignore[type-var]
164+ lambda : AsyncListener (self .zc ), # type: ignore[arg-type, return-value]
165+ sock = sock ,
166+ )
167+ datagram_transport = cast (asyncio .DatagramTransport , transport )
168+ reader = make_wrapped_transport (datagram_transport )
169+ # No ``await`` between wrapping and registering so a concurrent
170+ # shutdown always sees the transport in exactly one place.
171+ self .protocols .append (cast (AsyncListener , protocol ))
172+ self .readers .append (reader )
173+ if is_sender :
174+ self .senders .append (make_wrapped_transport (datagram_transport ))
175+ return reader
176+
177+ async def async_update_interfaces (
178+ self ,
179+ interfaces : InterfacesType ,
180+ ip_version : IPVersion ,
181+ apple_p2p : bool ,
182+ ) -> None :
183+ """Reconcile sender/reader sockets to the live interface set.
184+
185+ Adds a per-interface responder socket for each interface that
186+ appeared and tears down the socket for each interface that
187+ disappeared, diffing on the bound address. The shared listen
188+ socket (including the Default single-family dual-use socket) is
189+ never torn down here.
190+ """
191+ assert self .loop is not None
192+ normalized = normalize_interface_choice (interfaces , ip_version )
193+ desired = {_interface_bind_address (interface ): interface for interface in normalized }
194+ current = {_wrapped_bind_address (wrapped ): wrapped for wrapped in self .senders }
195+ listen_transport = self ._listen_transport
196+ listen_socket = listen_transport .sock if listen_transport is not None else None
197+
198+ for bind_address , wrapped in current .items ():
199+ if bind_address in desired :
200+ continue
201+ if listen_transport is not None and wrapped .transport is listen_transport .transport :
202+ # The shared listen / dual-use socket is not a per-interface
203+ # sender; leaving the group or closing it would break receive.
204+ continue
205+ self ._async_close_sender (wrapped , listen_socket )
206+
207+ for bind_address , interface in desired .items ():
208+ if bind_address in current :
209+ continue
210+ await self ._async_add_interface (interface , listen_socket , apple_p2p )
211+
212+ async def _async_add_interface (
213+ self ,
214+ interface : str | tuple [tuple [str , int , int ], int ],
215+ listen_socket : socket .socket | None ,
216+ apple_p2p : bool ,
217+ ) -> None :
218+ """Join the multicast group and adopt a responder socket for one interface."""
219+ # A unicast instance has no listen socket, so membership is only
220+ # ever managed when ``listen_socket`` is present.
221+ if listen_socket is not None and not add_multicast_member (listen_socket , interface ):
222+ return
223+ respond_socket = new_respond_socket (interface , apple_p2p = apple_p2p , unicast = self .zc .unicast )
224+ if respond_socket is None :
225+ if listen_socket is not None :
226+ drop_multicast_member (listen_socket , interface )
227+ return
228+ await self ._async_wrap_socket (respond_socket , is_sender = True )
229+
230+ def _async_close_sender (self , wrapped : _WrappedTransport , listen_socket : socket .socket | None ) -> None :
231+ """Drop a per-interface sender's wrappers/protocol and close its transport."""
232+ transport = wrapped .transport
233+ self .protocols = [
234+ p for p in self .protocols if p .transport is None or p .transport .transport is not transport
235+ ]
236+ self .readers = [w for w in self .readers if w .transport is not transport ]
237+ self .senders = [w for w in self .senders if w .transport is not transport ]
238+ if listen_socket is not None :
239+ drop_multicast_member (listen_socket , _wrapped_interface (wrapped ))
240+ transport .close ()
241+
129242 def _async_cache_cleanup (self ) -> None :
130243 """Periodic cache cleanup."""
131244 now = current_time_millis ()
0 commit comments