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 ]) -> tuple [str , int ]:
61+ """Return the (address, scope_id) an interface choice maps to, for diffing.
62+
63+ The scope_id keeps link-local IPv6 addresses that repeat across
64+ interfaces (same address, different zone) from colliding to one key.
65+ """
66+ if isinstance (interface , tuple ):
67+ return (_strip_zone (interface [0 ][0 ]), interface [0 ][2 ])
68+ return (interface , 0 )
69+
70+
71+ def _wrapped_bind_address (wrapped : _WrappedTransport ) -> tuple [str , int ]:
72+ """Return the bound (address, scope_id) of a sender transport, for diffing."""
73+ sock_name = wrapped .sock_name
74+ if wrapped .is_ipv6 :
75+ scope_id = cast (int , sock_name [3 ]) if len (sock_name ) > 3 else 0
76+ return (_strip_zone (sock_name [0 ]), scope_id )
77+ return (cast (str , sock_name [0 ]), 0 )
78+
79+
80+ def _wrapped_interface (wrapped : _WrappedTransport ) -> str | tuple [tuple [str , int , int ], int ]:
81+ """Reconstruct the interface representation for a sender, for membership drops."""
82+ sock_name = wrapped .sock_name
83+ if wrapped .is_ipv6 :
84+ scope_id = cast (int , sock_name [3 ]) if len (sock_name ) > 3 else 0
85+ return ((_strip_zone (sock_name [0 ]), 0 , 0 ), scope_id )
86+ return cast (str , sock_name [0 ])
87+
88+
4689class AsyncEngine :
4790 """An engine wraps sockets in the event loop."""
4891
4992 __slots__ = (
5093 "_cleanup_timer" ,
5194 "_listen_socket" ,
95+ "_listen_transport" ,
5296 "_respond_sockets" ,
5397 "_setup_task" ,
5498 "loop" ,
@@ -72,6 +116,7 @@ def __init__(
72116 self .senders : list [_WrappedTransport ] = []
73117 self .running_future : asyncio .Future [bool | None ] | None = None
74118 self ._listen_socket = listen_socket
119+ self ._listen_transport : _WrappedTransport | None = None
75120 self ._respond_sockets = respond_sockets
76121 self ._cleanup_timer : asyncio .TimerHandle | None = None
77122 self ._setup_task : asyncio .Task [None ] | None = None
@@ -98,8 +143,6 @@ async def _async_setup(self, loop_thread_ready: threading.Event | None) -> None:
98143
99144 async def _async_create_endpoints (self ) -> None :
100145 """Create endpoints to send and receive."""
101- assert self .loop is not None
102- loop = self .loop
103146 reader_sockets = []
104147 sender_sockets = []
105148 if self ._listen_socket :
@@ -110,22 +153,108 @@ async def _async_create_endpoints(self) -> None:
110153 sender_sockets .append (s )
111154
112155 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 )))
156+ reader = await self ._async_wrap_socket (s , s in sender_sockets )
157+ # The wrap above does not await before returning, so releasing
158+ # the engine's pending handle here keeps ``s`` in exactly one
159+ # place from a concurrent shutdown's point of view.
124160 if s is self ._listen_socket :
161+ # Keep a handle to the shared listen socket so interface
162+ # rescans can add/drop multicast memberships on it.
163+ self ._listen_transport = reader
125164 self ._listen_socket = None
126165 if s in self ._respond_sockets :
127166 self ._respond_sockets .remove (s )
128167
168+ async def _async_wrap_socket (self , sock : socket .socket , is_sender : bool ) -> _WrappedTransport :
169+ """Adopt a socket into a transport, register it, and return the reader wrapper."""
170+ assert self .loop is not None
171+ transport , protocol = await self .loop .create_datagram_endpoint ( # type: ignore[type-var]
172+ lambda : AsyncListener (self .zc ), # type: ignore[arg-type, return-value]
173+ sock = sock ,
174+ )
175+ datagram_transport = cast (asyncio .DatagramTransport , transport )
176+ reader = make_wrapped_transport (datagram_transport )
177+ # No ``await`` between wrapping and registering so a concurrent
178+ # shutdown always sees the transport in exactly one place.
179+ self .protocols .append (cast (AsyncListener , protocol ))
180+ self .readers .append (reader )
181+ if is_sender :
182+ self .senders .append (make_wrapped_transport (datagram_transport ))
183+ return reader
184+
185+ async def async_update_interfaces (
186+ self ,
187+ interfaces : InterfacesType ,
188+ ip_version : IPVersion ,
189+ apple_p2p : bool ,
190+ ) -> int :
191+ """Reconcile sender/reader sockets to the live interface set.
192+
193+ Adds a per-interface responder socket for each interface that
194+ appeared and tears down the socket for each interface that
195+ disappeared, diffing on the bound address. The shared listen
196+ socket (including the Default single-family dual-use socket) is
197+ never torn down here. Returns the number of responder sockets
198+ added so the caller can skip re-announcing when nothing appeared.
199+ """
200+ assert self .loop is not None
201+ normalized = normalize_interface_choice (interfaces , ip_version )
202+ desired = {_interface_bind_address (interface ): interface for interface in normalized }
203+ current = {_wrapped_bind_address (wrapped ): wrapped for wrapped in self .senders }
204+ listen_transport = self ._listen_transport
205+ listen_socket = listen_transport .sock if listen_transport is not None else None
206+
207+ for bind_address , wrapped in current .items ():
208+ if bind_address in desired :
209+ continue
210+ if listen_transport is not None and wrapped .transport is listen_transport .transport :
211+ # The shared listen / dual-use socket is not a per-interface
212+ # sender; leaving the group or closing it would break receive.
213+ continue
214+ self ._async_close_sender (wrapped , listen_socket )
215+
216+ added = 0
217+ for bind_address , interface in desired .items ():
218+ if bind_address in current :
219+ continue
220+ if await self ._async_add_interface (interface , listen_socket , apple_p2p ):
221+ added += 1
222+ return added
223+
224+ async def _async_add_interface (
225+ self ,
226+ interface : str | tuple [tuple [str , int , int ], int ],
227+ listen_socket : socket .socket | None ,
228+ apple_p2p : bool ,
229+ ) -> bool :
230+ """Join the multicast group and adopt a responder socket for one interface.
231+
232+ Returns whether a responder socket was actually added.
233+ """
234+ # A unicast instance has no listen socket, so membership is only
235+ # ever managed when ``listen_socket`` is present.
236+ if listen_socket is not None and not add_multicast_member (listen_socket , interface ):
237+ return False
238+ respond_socket = new_respond_socket (interface , apple_p2p = apple_p2p , unicast = self .zc .unicast )
239+ if respond_socket is None :
240+ if listen_socket is not None :
241+ drop_multicast_member (listen_socket , interface )
242+ return False
243+ await self ._async_wrap_socket (respond_socket , is_sender = True )
244+ return True
245+
246+ def _async_close_sender (self , wrapped : _WrappedTransport , listen_socket : socket .socket | None ) -> None :
247+ """Drop a per-interface sender's wrappers/protocol and close its transport."""
248+ transport = wrapped .transport
249+ self .protocols = [
250+ p for p in self .protocols if p .transport is None or p .transport .transport is not transport
251+ ]
252+ self .readers = [w for w in self .readers if w .transport is not transport ]
253+ self .senders = [w for w in self .senders if w .transport is not transport ]
254+ if listen_socket is not None :
255+ drop_multicast_member (listen_socket , _wrapped_interface (wrapped ))
256+ transport .close ()
257+
129258 def _async_cache_cleanup (self ) -> None :
130259 """Periodic cache cleanup."""
131260 now = current_time_millis ()
0 commit comments