Skip to content

Commit b2f6c89

Browse files
bdracobdraco
committed
feat: add async_update_interfaces to rescan network interfaces at runtime
1 parent 5333fc5 commit b2f6c89

7 files changed

Lines changed: 607 additions & 14 deletions

File tree

‎src/zeroconf/_core.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,11 @@ def __init__(
199199

200200
self.unicast = unicast
201201
self._use_asyncio = use_asyncio
202+
# Retained so async_update_interfaces can re-run create_sockets /
203+
# normalize_interface_choice against the live interface set later.
204+
self._interfaces = interfaces
205+
self._ip_version = ip_version
206+
self._apple_p2p = apple_p2p
202207
listen_socket, respond_sockets = create_sockets(interfaces, unicast, ip_version, apple_p2p=apple_p2p)
203208
log.debug("Listen socket %s, respond sockets %s", listen_socket, respond_sockets)
204209

@@ -406,6 +411,59 @@ async def async_update_service(self, info: ServiceInfo) -> Awaitable:
406411
self.registry.async_update(info)
407412
return asyncio.ensure_future(self._async_broadcast_service(info, _REGISTER_TIME, None))
408413

414+
def update_interfaces(
415+
self,
416+
interfaces: InterfacesType | None = None,
417+
ip_version: IPVersion | None = None,
418+
apple_p2p: bool | None = None,
419+
) -> None:
420+
"""Rescan network interfaces and reconcile the sockets in use.
421+
422+
While it is not expected during normal operation,
423+
this function may raise EventLoopBlocked if the underlying
424+
call to `async_update_interfaces` cannot be completed.
425+
"""
426+
assert self.loop is not None
427+
run_coro_with_timeout(
428+
self.async_update_interfaces(interfaces, ip_version, apple_p2p),
429+
self.loop,
430+
_REGISTER_TIME * _REGISTER_BROADCASTS,
431+
)
432+
433+
async def async_update_interfaces(
434+
self,
435+
interfaces: InterfacesType | None = None,
436+
ip_version: IPVersion | None = None,
437+
apple_p2p: bool | None = None,
438+
) -> None:
439+
"""Rescan network interfaces and reconcile the sockets in use.
440+
441+
Adds sockets for interfaces that appeared, drops sockets for
442+
interfaces that disappeared, and re-announces existing
443+
registrations when a new sender appeared. ``interfaces``,
444+
``ip_version`` and ``apple_p2p`` each default to the value passed at
445+
construction; pass a new value to switch it. When the resulting
446+
interface set is unchanged this is a no-op (no sockets touched,
447+
nothing re-announced). The shared listen socket's family and unicast
448+
mode are fixed at construction.
449+
"""
450+
if interfaces is not None:
451+
self._interfaces = interfaces
452+
if ip_version is not None:
453+
self._ip_version = ip_version
454+
if apple_p2p is not None:
455+
self._apple_p2p = apple_p2p
456+
await self.async_wait_for_start()
457+
added = await self.engine.async_update_interfaces(self._interfaces, self._ip_version, self._apple_p2p)
458+
if not added:
459+
return
460+
await asyncio.gather(
461+
*[
462+
self._async_broadcast_service(info, _REGISTER_TIME, None)
463+
for info in self.registry.async_get_service_infos()
464+
]
465+
)
466+
409467
async def async_get_service_info(
410468
self,
411469
type_: str,

‎src/zeroconf/_engine.py‎

Lines changed: 120 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,14 @@
3030

3131
from ._record_update import RecordUpdate
3232
from ._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+
)
3341
from ._utils.time import current_time_millis
3442
from .const import _CACHE_CLEANUP_INTERVAL
3543

@@ -38,17 +46,30 @@
3846

3947

4048
from ._listener import AsyncListener
41-
from ._transport import _WrappedTransport, make_wrapped_transport
49+
from ._transport import _strip_zone, _WrappedTransport, make_wrapped_transport
4250

4351
_CLOSE_TIMEOUT = 3000 # ms
4452

4553

54+
def _interface_key(interface: str | tuple[tuple[str, int, int], int]) -> tuple[str, int]:
55+
"""Return the (address, scope_id) an interface choice maps to, for diffing.
56+
57+
Must produce the same key shape as ``_WrappedTransport.interface_key`` so
58+
the desired set (from ``normalize_interface_choice``) and the current set
59+
(from the bound senders) diff against each other.
60+
"""
61+
if isinstance(interface, tuple):
62+
return (_strip_zone(interface[0][0]), interface[0][2])
63+
return (interface, 0)
64+
65+
4666
class AsyncEngine:
4767
"""An engine wraps sockets in the event loop."""
4868

4969
__slots__ = (
5070
"_cleanup_timer",
5171
"_listen_socket",
72+
"_listen_transport",
5273
"_respond_sockets",
5374
"_setup_task",
5475
"loop",
@@ -72,6 +93,7 @@ def __init__(
7293
self.senders: list[_WrappedTransport] = []
7394
self.running_future: asyncio.Future[bool | None] | None = None
7495
self._listen_socket = listen_socket
96+
self._listen_transport: _WrappedTransport | None = None
7597
self._respond_sockets = respond_sockets
7698
self._cleanup_timer: asyncio.TimerHandle | None = None
7799
self._setup_task: asyncio.Task[None] | None = None
@@ -98,8 +120,6 @@ async def _async_setup(self, loop_thread_ready: threading.Event | None) -> None:
98120

99121
async def _async_create_endpoints(self) -> None:
100122
"""Create endpoints to send and receive."""
101-
assert self.loop is not None
102-
loop = self.loop
103123
reader_sockets = []
104124
sender_sockets = []
105125
if self._listen_socket:
@@ -110,22 +130,108 @@ async def _async_create_endpoints(self) -> None:
110130
sender_sockets.append(s)
111131

112132
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)))
133+
reader = await self._async_wrap_socket(s, s in sender_sockets)
134+
# The wrap above does not await before returning, so releasing
135+
# the engine's pending handle here keeps ``s`` in exactly one
136+
# place from a concurrent shutdown's point of view.
124137
if s is self._listen_socket:
138+
# Keep a handle to the shared listen socket so interface
139+
# rescans can add/drop multicast memberships on it.
140+
self._listen_transport = reader
125141
self._listen_socket = None
126142
if s in self._respond_sockets:
127143
self._respond_sockets.remove(s)
128144

145+
async def _async_wrap_socket(self, sock: socket.socket, is_sender: bool) -> _WrappedTransport:
146+
"""Adopt a socket into a transport, register it, and return the reader wrapper."""
147+
assert self.loop is not None
148+
transport, protocol = await self.loop.create_datagram_endpoint( # type: ignore[type-var]
149+
lambda: AsyncListener(self.zc), # type: ignore[arg-type, return-value]
150+
sock=sock,
151+
)
152+
datagram_transport = cast(asyncio.DatagramTransport, transport)
153+
reader = make_wrapped_transport(datagram_transport)
154+
# No ``await`` between wrapping and registering so a concurrent
155+
# shutdown always sees the transport in exactly one place.
156+
self.protocols.append(cast(AsyncListener, protocol))
157+
self.readers.append(reader)
158+
if is_sender:
159+
self.senders.append(make_wrapped_transport(datagram_transport))
160+
return reader
161+
162+
async def async_update_interfaces(
163+
self,
164+
interfaces: InterfacesType,
165+
ip_version: IPVersion,
166+
apple_p2p: bool,
167+
) -> bool:
168+
"""Reconcile sender/reader sockets to the live interface set.
169+
170+
Adds a per-interface responder socket for each interface that
171+
appeared and tears down the socket for each interface that
172+
disappeared, diffing on the bound address. The shared listen
173+
socket (including the Default single-family dual-use socket) is
174+
never torn down here. Returns whether any responder socket was
175+
added, so the caller can skip re-announcing when nothing appeared.
176+
"""
177+
assert self.loop is not None
178+
normalized = normalize_interface_choice(interfaces, ip_version)
179+
desired = {_interface_key(interface): interface for interface in normalized}
180+
current = {wrapped.interface_key: wrapped for wrapped in self.senders}
181+
listen_transport = self._listen_transport
182+
listen_socket = listen_transport.sock if listen_transport is not None else None
183+
184+
for bind_address, wrapped in current.items():
185+
if bind_address in desired:
186+
continue
187+
if listen_transport is not None and wrapped.transport is listen_transport.transport:
188+
# The shared listen / dual-use socket is not a per-interface
189+
# sender; leaving the group or closing it would break receive.
190+
continue
191+
self._async_close_sender(wrapped, listen_socket)
192+
193+
added = False
194+
for bind_address, interface in desired.items():
195+
if bind_address in current:
196+
continue
197+
if await self._async_add_interface(interface, listen_socket, apple_p2p):
198+
added = True
199+
return added
200+
201+
async def _async_add_interface(
202+
self,
203+
interface: str | tuple[tuple[str, int, int], int],
204+
listen_socket: socket.socket | None,
205+
apple_p2p: bool,
206+
) -> bool:
207+
"""Join the multicast group and adopt a responder socket for one interface.
208+
209+
Returns whether a responder socket was actually added.
210+
"""
211+
# A unicast instance has no listen socket, so membership is only
212+
# ever managed when ``listen_socket`` is present.
213+
if listen_socket is not None and not add_multicast_member(listen_socket, interface):
214+
return False
215+
respond_socket = new_respond_socket(interface, apple_p2p=apple_p2p, unicast=self.zc.unicast)
216+
if respond_socket is None:
217+
if listen_socket is not None:
218+
drop_multicast_member(listen_socket, interface)
219+
return False
220+
await self._async_wrap_socket(respond_socket, is_sender=True)
221+
return True
222+
223+
def _async_close_sender(self, wrapped: _WrappedTransport, listen_socket: socket.socket | None) -> None:
224+
"""Drop a per-interface sender's wrappers/protocol and close its transport."""
225+
transport = wrapped.transport
226+
self.protocols = [
227+
p for p in self.protocols if p.transport is None or p.transport.transport is not transport
228+
]
229+
self.readers = [w for w in self.readers if w.transport is not transport]
230+
self.senders = [w for w in self.senders if w.transport is not transport]
231+
if listen_socket is not None:
232+
drop_multicast_member(listen_socket, wrapped.multicast_interface)
233+
transport.close()
234+
129235
def _async_cache_cleanup(self) -> None:
130236
"""Periodic cache cleanup."""
131237
now = current_time_millis()

‎src/zeroconf/_transport.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,13 @@
2424

2525
import asyncio
2626
import socket
27+
from typing import cast
28+
29+
30+
def _strip_zone(address: str) -> str:
31+
"""Drop a ``%zone`` suffix from an IPv6 address string."""
32+
percent = address.find("%")
33+
return address[:percent] if percent != -1 else address
2734

2835

2936
class _WrappedTransport:
@@ -55,6 +62,28 @@ def __init__(
5562
self.fileno = fileno
5663
self.sock_name = sock_name
5764

65+
@property
66+
def interface_key(self) -> tuple[str, int]:
67+
"""The bound (address, scope_id) identifying this sender's interface.
68+
69+
Used to diff senders against the desired interface set. The scope_id
70+
keeps link-local IPv6 addresses that repeat across interfaces (same
71+
address, different zone) from colliding to one key.
72+
"""
73+
sock_name = self.sock_name
74+
if self.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+
@property
80+
def multicast_interface(self) -> str | tuple[tuple[str, int, int], int]:
81+
"""The interface value membership calls take for this transport's group."""
82+
address, scope_id = self.interface_key
83+
if self.is_ipv6:
84+
return ((address, 0, 0), scope_id)
85+
return address
86+
5887

5988
def make_wrapped_transport(transport: asyncio.DatagramTransport) -> _WrappedTransport:
6089
"""Make a wrapped transport."""

‎src/zeroconf/_utils/net.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -401,6 +401,38 @@ def add_multicast_member(
401401
return True
402402

403403

404+
def drop_multicast_member(
405+
listen_socket: socket.socket,
406+
interface: str | tuple[tuple[str, int, int], int],
407+
) -> bool:
408+
"""Leave the mDNS multicast group on an interface; inverse of add_multicast_member."""
409+
# This is based on assumptions in normalize_interface_choice
410+
is_v6 = isinstance(interface, tuple)
411+
log.debug("Dropping %r (socket %d) from multicast group", interface, listen_socket.fileno())
412+
try:
413+
if is_v6:
414+
try:
415+
mdns_addr6_bytes = socket.inet_pton(socket.AF_INET6, _MDNS_ADDR6)
416+
except OSError:
417+
return False
418+
iface_bin = struct.pack("@I", cast(int, interface[1]))
419+
listen_socket.setsockopt(_IPPROTO_IPV6, socket.IPV6_LEAVE_GROUP, mdns_addr6_bytes + iface_bin)
420+
else:
421+
_value = socket.inet_aton(_MDNS_ADDR) + socket.inet_aton(cast(str, interface))
422+
listen_socket.setsockopt(socket.IPPROTO_IP, socket.IP_DROP_MEMBERSHIP, _value)
423+
except OSError as e:
424+
# The kernel drops memberships automatically when an interface
425+
# disappears, so a stale leave is expected to fail benignly.
426+
benign = {errno.EADDRNOTAVAIL, errno.EINVAL, errno.ENODEV, errno.ENOPROTOOPT}
427+
if sys.platform == "win32":
428+
# No WSAEINVAL definition in typeshed
429+
benign.add(cast(Any, errno).WSAEINVAL) # pylint: disable=no-member
430+
if get_errno(e) in benign:
431+
return False
432+
raise
433+
return True
434+
435+
404436
def new_respond_socket(
405437
interface: str | tuple[tuple[str, int, int], int],
406438
apple_p2p: bool = False,

‎src/zeroconf/asyncio.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,22 @@ async def async_update_service(self, info: ServiceInfo) -> Awaitable:
224224
"""
225225
return await self.zeroconf.async_update_service(info)
226226

227+
async def async_update_interfaces(
228+
self,
229+
interfaces: InterfacesType | None = None,
230+
ip_version: IPVersion | None = None,
231+
apple_p2p: bool | None = None,
232+
) -> None:
233+
"""Rescan network interfaces and reconcile the sockets in use.
234+
235+
Adds sockets for interfaces that appeared, drops sockets for
236+
interfaces that disappeared, and re-announces existing
237+
registrations on the resulting senders. ``interfaces``,
238+
``ip_version`` and ``apple_p2p`` each default to the construction-time
239+
value.
240+
"""
241+
await self.zeroconf.async_update_interfaces(interfaces, ip_version, apple_p2p)
242+
227243
async def async_close(self) -> None:
228244
"""Ends the background threads, and prevent this instance from
229245
servicing further queries."""

0 commit comments

Comments
 (0)