Skip to content

Commit f4c17eb

Browse files
bdracobdraco
authored
chore: split _engine.py into _transport.py and _listener.py (#1222)
1 parent e9cc5c8 commit f4c17eb

7 files changed

Lines changed: 518 additions & 443 deletions

File tree

‎src/zeroconf/_core.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929

3030
from ._cache import DNSCache
3131
from ._dns import DNSQuestion, DNSQuestionType
32-
from ._engine import AsyncEngine, _WrappedTransport
32+
from ._engine import AsyncEngine
3333
from ._exceptions import NonUniqueNameException, NotRunningException
3434
from ._handlers import (
3535
MulticastOutgoingQueue,
@@ -46,6 +46,7 @@
4646
from ._services.browser import ServiceBrowser
4747
from ._services.info import ServiceInfo, instance_name_from_service_info
4848
from ._services.registry import ServiceRegistry
49+
from ._transport import _WrappedTransport
4950
from ._updates import RecordUpdateListener
5051
from ._utils.asyncio import (
5152
await_awaitable,

‎src/zeroconf/_engine.py‎

Lines changed: 8 additions & 235 deletions
Original file line numberDiff line numberDiff line change
@@ -22,71 +22,23 @@
2222

2323
import asyncio
2424
import itertools
25-
import logging
26-
import random
2725
import socket
2826
import threading
29-
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
27+
from typing import TYPE_CHECKING, List, Optional, cast
3028

31-
from ._logger import QuietLogger, log
32-
from ._protocol.incoming import DNSIncoming
3329
from ._updates import RecordUpdate
3430
from ._utils.asyncio import get_running_loop, run_coro_with_timeout
35-
from ._utils.time import current_time_millis, millis_to_seconds
36-
from .const import (
37-
_CACHE_CLEANUP_INTERVAL,
38-
_DUPLICATE_PACKET_SUPPRESSION_INTERVAL,
39-
_MAX_MSG_ABSOLUTE,
40-
)
31+
from ._utils.time import current_time_millis
32+
from .const import _CACHE_CLEANUP_INTERVAL
4133

4234
if TYPE_CHECKING:
4335
from ._core import Zeroconf
4436

45-
_TC_DELAY_RANDOM_INTERVAL = (400, 500)
46-
47-
_CLOSE_TIMEOUT = 3000 # ms
48-
49-
50-
class _WrappedTransport:
51-
"""A wrapper for transports."""
52-
53-
__slots__ = (
54-
'transport',
55-
'is_ipv6',
56-
'sock',
57-
'fileno',
58-
'sock_name',
59-
)
60-
61-
def __init__(
62-
self,
63-
transport: asyncio.DatagramTransport,
64-
is_ipv6: bool,
65-
sock: socket.socket,
66-
fileno: int,
67-
sock_name: Any,
68-
) -> None:
69-
"""Initialize the wrapped transport.
70-
71-
These attributes are used when sending packets.
72-
"""
73-
self.transport = transport
74-
self.is_ipv6 = is_ipv6
75-
self.sock = sock
76-
self.fileno = fileno
77-
self.sock_name = sock_name
7837

38+
from ._listener import AsyncListener
39+
from ._transport import _WrappedTransport, make_wrapped_transport
7940

80-
def _make_wrapped_transport(transport: asyncio.DatagramTransport) -> _WrappedTransport:
81-
"""Make a wrapped transport."""
82-
sock: socket.socket = transport.get_extra_info('socket')
83-
return _WrappedTransport(
84-
transport=transport,
85-
is_ipv6=sock.family == socket.AF_INET6,
86-
sock=sock,
87-
fileno=sock.fileno(),
88-
sock_name=sock.getsockname(),
89-
)
41+
_CLOSE_TIMEOUT = 3000 # ms
9042

9143

9244
class AsyncEngine:
@@ -154,9 +106,9 @@ async def _async_create_endpoints(self) -> None:
154106
lambda: AsyncListener(self.zc), sock=s # type: ignore[arg-type, return-value]
155107
)
156108
self.protocols.append(cast(AsyncListener, protocol))
157-
self.readers.append(_make_wrapped_transport(cast(asyncio.DatagramTransport, transport)))
109+
self.readers.append(make_wrapped_transport(cast(asyncio.DatagramTransport, transport)))
158110
if s in sender_sockets:
159-
self.senders.append(_make_wrapped_transport(cast(asyncio.DatagramTransport, transport)))
111+
self.senders.append(make_wrapped_transport(cast(asyncio.DatagramTransport, transport)))
160112

161113
def _async_cache_cleanup(self) -> None:
162114
"""Periodic cache cleanup."""
@@ -198,182 +150,3 @@ def close(self) -> None:
198150
if not self.loop.is_running():
199151
return
200152
run_coro_with_timeout(self._async_close(), self.loop, _CLOSE_TIMEOUT)
201-
202-
203-
class AsyncListener:
204-
205-
"""A Listener is used by this module to listen on the multicast
206-
group to which DNS messages are sent, allowing the implementation
207-
to cache information as it arrives.
208-
209-
It requires registration with an Engine object in order to have
210-
the read() method called when a socket is available for reading."""
211-
212-
__slots__ = (
213-
'zc',
214-
'data',
215-
'last_time',
216-
'last_message',
217-
'transport',
218-
'sock_description',
219-
'_deferred',
220-
'_timers',
221-
)
222-
223-
def __init__(self, zc: 'Zeroconf') -> None:
224-
self.zc = zc
225-
self.data: Optional[bytes] = None
226-
self.last_time: float = 0
227-
self.last_message: Optional[DNSIncoming] = None
228-
self.transport: Optional[_WrappedTransport] = None
229-
self.sock_description: Optional[str] = None
230-
self._deferred: Dict[str, List[DNSIncoming]] = {}
231-
self._timers: Dict[str, asyncio.TimerHandle] = {}
232-
super().__init__()
233-
234-
def datagram_received(
235-
self, data: bytes, addrs: Union[Tuple[str, int], Tuple[str, int, int, int]]
236-
) -> None:
237-
assert self.transport is not None
238-
data_len = len(data)
239-
debug = log.isEnabledFor(logging.DEBUG)
240-
241-
if data_len > _MAX_MSG_ABSOLUTE:
242-
# Guard against oversized packets to ensure bad implementations cannot overwhelm
243-
# the system.
244-
if debug:
245-
log.debug(
246-
"Discarding incoming packet with length %s, which is larger "
247-
"than the absolute maximum size of %s",
248-
data_len,
249-
_MAX_MSG_ABSOLUTE,
250-
)
251-
return
252-
253-
now = current_time_millis()
254-
if (
255-
self.data == data
256-
and (now - _DUPLICATE_PACKET_SUPPRESSION_INTERVAL) < self.last_time
257-
and self.last_message is not None
258-
and not self.last_message.has_qu_question()
259-
):
260-
# Guard against duplicate packets
261-
if debug:
262-
log.debug(
263-
'Ignoring duplicate message with no unicast questions received from %s [socket %s] (%d bytes) as [%r]',
264-
addrs,
265-
self.sock_description,
266-
data_len,
267-
data,
268-
)
269-
return
270-
271-
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = ()
272-
if len(addrs) == 2:
273-
# https://github.com/python/mypy/issues/1178
274-
addr, port = addrs # type: ignore
275-
scope = None
276-
else:
277-
# https://github.com/python/mypy/issues/1178
278-
addr, port, flow, scope = addrs # type: ignore
279-
if debug: # pragma: no branch
280-
log.debug('IPv6 scope_id %d associated to the receiving interface', scope)
281-
v6_flow_scope = (flow, scope)
282-
283-
msg = DNSIncoming(data, (addr, port), scope, now)
284-
self.data = data
285-
self.last_time = now
286-
self.last_message = msg
287-
if msg.valid:
288-
if debug:
289-
log.debug(
290-
'Received from %r:%r [socket %s]: %r (%d bytes) as [%r]',
291-
addr,
292-
port,
293-
self.sock_description,
294-
msg,
295-
data_len,
296-
data,
297-
)
298-
else:
299-
if debug:
300-
log.debug(
301-
'Received from %r:%r [socket %s]: (%d bytes) [%r]',
302-
addr,
303-
port,
304-
self.sock_description,
305-
data_len,
306-
data,
307-
)
308-
return
309-
310-
if not msg.is_query():
311-
self.zc.handle_response(msg)
312-
return
313-
314-
self.handle_query_or_defer(msg, addr, port, self.transport, v6_flow_scope)
315-
316-
def handle_query_or_defer(
317-
self,
318-
msg: DNSIncoming,
319-
addr: str,
320-
port: int,
321-
transport: _WrappedTransport,
322-
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
323-
) -> None:
324-
"""Deal with incoming query packets. Provides a response if
325-
possible."""
326-
if not msg.truncated:
327-
self._respond_query(msg, addr, port, transport, v6_flow_scope)
328-
return
329-
330-
deferred = self._deferred.setdefault(addr, [])
331-
# If we get the same packet we ignore it
332-
for incoming in reversed(deferred):
333-
if incoming.data == msg.data:
334-
return
335-
deferred.append(msg)
336-
delay = millis_to_seconds(random.randint(*_TC_DELAY_RANDOM_INTERVAL))
337-
assert self.zc.loop is not None
338-
self._cancel_any_timers_for_addr(addr)
339-
self._timers[addr] = self.zc.loop.call_later(
340-
delay, self._respond_query, None, addr, port, transport, v6_flow_scope
341-
)
342-
343-
def _cancel_any_timers_for_addr(self, addr: str) -> None:
344-
"""Cancel any future truncated packet timers for the address."""
345-
if addr in self._timers:
346-
self._timers.pop(addr).cancel()
347-
348-
def _respond_query(
349-
self,
350-
msg: Optional[DNSIncoming],
351-
addr: str,
352-
port: int,
353-
transport: _WrappedTransport,
354-
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
355-
) -> None:
356-
"""Respond to a query and reassemble any truncated deferred packets."""
357-
self._cancel_any_timers_for_addr(addr)
358-
packets = self._deferred.pop(addr, [])
359-
if msg:
360-
packets.append(msg)
361-
362-
self.zc.handle_assembled_query(packets, addr, port, transport, v6_flow_scope)
363-
364-
def error_received(self, exc: Exception) -> None:
365-
"""Likely socket closed or IPv6."""
366-
# We preformat the message string with the socket as we want
367-
# log_exception_once to log a warrning message once PER EACH
368-
# different socket in case there are problems with multiple
369-
# sockets
370-
msg_str = f"Error with socket {self.sock_description}): %s"
371-
QuietLogger.log_exception_once(exc, msg_str, exc)
372-
373-
def connection_made(self, transport: asyncio.BaseTransport) -> None:
374-
wrapped_transport = _make_wrapped_transport(cast(asyncio.DatagramTransport, transport))
375-
self.transport = wrapped_transport
376-
self.sock_description = f"{wrapped_transport.fileno} ({wrapped_transport.sock_name})"
377-
378-
def connection_lost(self, exc: Optional[Exception]) -> None:
379-
"""Handle connection lost."""

0 commit comments

Comments
 (0)