Skip to content

Commit f7b6624

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

6 files changed

Lines changed: 435 additions & 13 deletions

File tree

src/zeroconf/_core.py

Lines changed: 38 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,39 @@ 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(self, interfaces: InterfacesType | None = None) -> None:
415+
"""Rescan network interfaces and reconcile the sockets in use.
416+
417+
While it is not expected during normal operation,
418+
this function may raise EventLoopBlocked if the underlying
419+
call to `async_update_interfaces` cannot be completed.
420+
"""
421+
assert self.loop is not None
422+
run_coro_with_timeout(
423+
self.async_update_interfaces(interfaces),
424+
self.loop,
425+
_REGISTER_TIME * _REGISTER_BROADCASTS,
426+
)
427+
428+
async def async_update_interfaces(self, interfaces: InterfacesType | None = None) -> None:
429+
"""Rescan network interfaces and reconcile the sockets in use.
430+
431+
Adds sockets for interfaces that appeared, drops sockets for
432+
interfaces that disappeared, and re-announces existing
433+
registrations on the resulting senders. ``interfaces`` defaults to
434+
the choice passed at construction; pass a new value to switch.
435+
"""
436+
if interfaces is not None:
437+
self._interfaces = interfaces
438+
await self.async_wait_for_start()
439+
await self.engine.async_update_interfaces(self._interfaces, self._ip_version, self._apple_p2p)
440+
await asyncio.gather(
441+
*[
442+
self._async_broadcast_service(info, _REGISTER_TIME, None)
443+
for info in self.registry.async_get_service_infos()
444+
]
445+
)
446+
409447
async def async_get_service_info(
410448
self,
411449
type_: str,

src/zeroconf/_engine.py

Lines changed: 126 additions & 13 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

@@ -43,12 +51,40 @@
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+
4681
class 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()

src/zeroconf/_utils/net.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -401,6 +401,34 @@ 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+
if get_errno(e) in {errno.EADDRNOTAVAIL, errno.EINVAL, errno.ENODEV, errno.ENOPROTOOPT}:
427+
return False
428+
raise
429+
return True
430+
431+
404432
def new_respond_socket(
405433
interface: str | tuple[tuple[str, int, int], int],
406434
apple_p2p: bool = False,

src/zeroconf/asyncio.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,15 @@ 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(self, interfaces: InterfacesType | None = None) -> None:
228+
"""Rescan network interfaces and reconcile the sockets in use.
229+
230+
Adds sockets for interfaces that appeared, drops sockets for
231+
interfaces that disappeared, and re-announces existing
232+
registrations on the resulting senders.
233+
"""
234+
await self.zeroconf.async_update_interfaces(interfaces)
235+
227236
async def async_close(self) -> None:
228237
"""Ends the background threads, and prevent this instance from
229238
servicing further queries."""

0 commit comments

Comments
 (0)