Skip to content

Commit 74251cc

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

7 files changed

Lines changed: 716 additions & 16 deletions

File tree

‎src/zeroconf/_core.py‎

Lines changed: 65 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

@@ -216,6 +221,9 @@ def __init__(
216221
self.record_manager = RecordManager(self)
217222

218223
self._notify_futures: set[asyncio.Future] = set()
224+
# Serializes async_update_interfaces so overlapping calls (a bursty
225+
# adapter-change source) don't diff against a stale sender snapshot.
226+
self._interface_update_lock = asyncio.Lock()
219227
self.loop: asyncio.AbstractEventLoop | None = None
220228
self._loop_thread: threading.Thread | None = None
221229

@@ -406,6 +414,63 @@ async def async_update_service(self, info: ServiceInfo) -> Awaitable:
406414
self.registry.async_update(info)
407415
return asyncio.ensure_future(self._async_broadcast_service(info, _REGISTER_TIME, None))
408416

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

‎src/zeroconf/_engine.py‎

Lines changed: 123 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,17 @@
2828
import threading
2929
from typing import TYPE_CHECKING, cast
3030

31+
from ._logger import log
3132
from ._record_update import RecordUpdate
3233
from ._utils.asyncio import get_running_loop, run_coro_with_timeout
34+
from ._utils.net import (
35+
InterfacesType,
36+
IPVersion,
37+
add_multicast_member,
38+
drop_multicast_member,
39+
new_respond_socket,
40+
normalize_interface_choice,
41+
)
3342
from ._utils.time import current_time_millis
3443
from .const import _CACHE_CLEANUP_INTERVAL
3544

@@ -38,17 +47,30 @@
3847

3948

4049
from ._listener import AsyncListener
41-
from ._transport import _WrappedTransport, make_wrapped_transport
50+
from ._transport import _strip_zone, _WrappedTransport, make_wrapped_transport
4251

4352
_CLOSE_TIMEOUT = 3000 # ms
4453

4554

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

4970
__slots__ = (
5071
"_cleanup_timer",
5172
"_listen_socket",
73+
"_listen_transport",
5274
"_respond_sockets",
5375
"_setup_task",
5476
"loop",
@@ -72,6 +94,7 @@ def __init__(
7294
self.senders: list[_WrappedTransport] = []
7395
self.running_future: asyncio.Future[bool | None] | None = None
7496
self._listen_socket = listen_socket
97+
self._listen_transport: _WrappedTransport | None = None
7598
self._respond_sockets = respond_sockets
7699
self._cleanup_timer: asyncio.TimerHandle | None = None
77100
self._setup_task: asyncio.Task[None] | None = None
@@ -98,8 +121,6 @@ async def _async_setup(self, loop_thread_ready: threading.Event | None) -> None:
98121

99122
async def _async_create_endpoints(self) -> None:
100123
"""Create endpoints to send and receive."""
101-
assert self.loop is not None
102-
loop = self.loop
103124
reader_sockets = []
104125
sender_sockets = []
105126
if self._listen_socket:
@@ -110,22 +131,110 @@ async def _async_create_endpoints(self) -> None:
110131
sender_sockets.append(s)
111132

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

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

‎src/zeroconf/_transport.py‎

Lines changed: 52 additions & 2 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:
@@ -32,6 +39,7 @@ class _WrappedTransport:
3239
__slots__ = (
3340
"fileno",
3441
"is_ipv6",
42+
"multicast_index",
3543
"sock",
3644
"sock_name",
3745
"transport",
@@ -44,25 +52,67 @@ def __init__(
4452
sock: socket.socket,
4553
fileno: int,
4654
sock_name: tuple,
55+
multicast_index: int = 0,
4756
) -> None:
4857
"""Initialize the wrapped transport.
4958
50-
These attributes are used when sending packets.
59+
``multicast_index`` is the IPV6_MULTICAST_IF interface index the
60+
sender joined the group with, carried so a group leave uses the same
61+
index the join did (the bound socket's scope_id is 0 for global IPv6).
5162
"""
5263
self.transport = transport
5364
self.is_ipv6 = is_ipv6
5465
self.sock = sock
5566
self.fileno = fileno
5667
self.sock_name = sock_name
68+
self.multicast_index = multicast_index
69+
70+
@property
71+
def interface_key(self) -> tuple[str, int]:
72+
"""The bound (address, scope_id) identifying this sender's interface.
73+
74+
Used to diff senders against the desired interface set. The scope_id
75+
keeps link-local IPv6 addresses that repeat across interfaces (same
76+
address, different zone) from colliding to one key.
77+
"""
78+
sock_name = self.sock_name
79+
if self.is_ipv6:
80+
scope_id = cast(int, sock_name[3]) if len(sock_name) > 3 else 0
81+
return (_strip_zone(sock_name[0]), scope_id)
82+
return (cast(str, sock_name[0]), 0)
83+
84+
@property
85+
def multicast_interface(self) -> str | tuple[tuple[str, int, int], int]:
86+
"""The interface value a group leave takes for this transport.
87+
88+
For IPv6 this carries ``multicast_index`` (the index the join used),
89+
not the bound scope_id, so leave and join stay symmetric.
90+
"""
91+
address, _scope_id = self.interface_key
92+
if self.is_ipv6:
93+
return ((address, 0, 0), self.multicast_index)
94+
return address
5795

5896

5997
def make_wrapped_transport(transport: asyncio.DatagramTransport) -> _WrappedTransport:
6098
"""Make a wrapped transport."""
6199
sock: socket.socket = transport.get_extra_info("socket")
100+
is_ipv6 = sock.family == socket.AF_INET6
101+
multicast_index = 0
102+
if is_ipv6:
103+
# IPV6_MULTICAST_IF holds the interface index new_respond_socket
104+
# joined the group with; capture it so a later group leave uses the
105+
# same index. Windows rejects reading the option (WSAEINVAL); there
106+
# the leave falls back to the default interface as it did before.
107+
try:
108+
multicast_index = sock.getsockopt(socket.IPPROTO_IPV6, socket.IPV6_MULTICAST_IF)
109+
except OSError:
110+
multicast_index = 0
62111
return _WrappedTransport(
63112
transport=transport,
64-
is_ipv6=sock.family == socket.AF_INET6,
113+
is_ipv6=is_ipv6,
65114
sock=sock,
66115
fileno=sock.fileno(),
67116
sock_name=sock.getsockname(),
117+
multicast_index=multicast_index,
68118
)

0 commit comments

Comments
 (0)