Skip to content

Commit 1c46e23

Browse files
bdracobdraco
committed
use transport
1 parent c92c928 commit 1c46e23

2 files changed

Lines changed: 47 additions & 36 deletions

File tree

‎tests/test_core.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -695,7 +695,6 @@ def test_guard_against_oversized_packets():
695695
# will guard against the oversized packet and we won't see it.
696696
listener = _core.AsyncListener(zc)
697697
listener.transport = unittest.mock.MagicMock()
698-
listener.sock_fileno = 1
699698

700699
listener.datagram_received(ok_packet, ('127.0.0.1', const._MDNS_PORT))
701700
assert zc.cache.async_get_unique(okpacket_record) is not None

‎zeroconf/_core.py‎

Lines changed: 47 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,6 @@ def datagram_received(
234234
self, data: bytes, addrs: Union[Tuple[str, int], Tuple[str, int, int, int]]
235235
) -> None:
236236
assert self.transport is not None
237-
assert self.sock_fileno is not None
238237
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = ()
239238
if len(addrs) == 2:
240239
# https://github.com/python/mypy/issues/1178
@@ -296,20 +295,20 @@ def datagram_received(
296295
self.zc.handle_response(msg)
297296
return
298297

299-
self.handle_query_or_defer(msg, addr, port, self.sock_fileno, v6_flow_scope)
298+
self.handle_query_or_defer(msg, addr, port, self.transport, v6_flow_scope)
300299

301300
def handle_query_or_defer(
302301
self,
303302
msg: DNSIncoming,
304303
addr: str,
305304
port: int,
306-
sock_fileno: int,
305+
transport: asyncio.DatagramTransport,
307306
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
308307
) -> None:
309308
"""Deal with incoming query packets. Provides a response if
310309
possible."""
311310
if not msg.truncated:
312-
self._respond_query(msg, addr, port, sock_fileno, v6_flow_scope)
311+
self._respond_query(msg, addr, port, transport, v6_flow_scope)
313312
return
314313

315314
deferred = self._deferred.setdefault(addr, [])
@@ -322,7 +321,7 @@ def handle_query_or_defer(
322321
assert self.zc.loop is not None
323322
self._cancel_any_timers_for_addr(addr)
324323
self._timers[addr] = self.zc.loop.call_later(
325-
delay, self._respond_query, None, addr, port, sock_fileno, v6_flow_scope
324+
delay, self._respond_query, None, addr, port, transport, v6_flow_scope
326325
)
327326

328327
def _cancel_any_timers_for_addr(self, addr: str) -> None:
@@ -335,7 +334,7 @@ def _respond_query(
335334
msg: Optional[DNSIncoming],
336335
addr: str,
337336
port: int,
338-
sock_fileno: int,
337+
transport: asyncio.DatagramTransport,
339338
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
340339
) -> None:
341340
"""Respond to a query and reassemble any truncated deferred packets."""
@@ -344,7 +343,7 @@ def _respond_query(
344343
if msg:
345344
packets.append(msg)
346345

347-
self.zc.handle_assembled_query(packets, addr, port, sock_fileno, v6_flow_scope)
346+
self.zc.handle_assembled_query(packets, addr, port, transport, v6_flow_scope)
348347

349348
@property
350349
def _socket_description(self) -> str:
@@ -740,7 +739,7 @@ def handle_assembled_query(
740739
packets: List[DNSIncoming],
741740
addr: str,
742741
port: int,
743-
sock_fileno: int,
742+
transport: asyncio.DatagramTransport,
744743
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
745744
) -> None:
746745
"""Respond to a (re)assembled query.
@@ -761,7 +760,7 @@ def handle_assembled_query(
761760
# When sending unicast, only send back the reply
762761
# via the same socket that it was recieved from
763762
# as we know its reachable from that socket
764-
self.async_send(out, addr, port, v6_flow_scope, sock_fileno)
763+
self.async_send(out, addr, port, v6_flow_scope, transport)
765764
if question_answers.mcast_now:
766765
self.async_send(construct_outgoing_multicast_answers(question_answers.mcast_now))
767766
if question_answers.mcast_aggregate:
@@ -778,51 +777,64 @@ def send(
778777
addr: Optional[str] = None,
779778
port: int = _MDNS_PORT,
780779
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
781-
sock_fileno: Optional[int] = None,
780+
transport: Optional[asyncio.DatagramTransport] = None,
782781
) -> None:
783782
"""Sends an outgoing packet threadsafe."""
784783
assert self.loop is not None
785-
self.loop.call_soon_threadsafe(self.async_send, out, addr, port, v6_flow_scope, sock_fileno)
784+
self.loop.call_soon_threadsafe(self.async_send, out, addr, port, v6_flow_scope, transport)
786785

787786
def async_send(
788787
self,
789788
out: DNSOutgoing,
790789
addr: Optional[str] = None,
791790
port: int = _MDNS_PORT,
792791
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
793-
sock_fileno: Optional[int] = None,
792+
transport: Optional[asyncio.DatagramTransport] = None,
794793
) -> None:
795794
"""Sends an outgoing packet."""
796795
if self._GLOBAL_DONE:
797796
return
798797

798+
# If no transport is specified, we send to all the ones
799+
# with the same address family
800+
transports = [transport] if transport else self.engine.senders
801+
799802
for packet_num, packet in enumerate(out.packets()):
800803
if len(packet) > _MAX_MSG_ABSOLUTE:
801804
self.log_warning_once("Dropping %r over-sized packet (%d bytes) %r", out, len(packet), packet)
802805
return
803-
for transport in self.engine.senders:
804-
s = transport.get_extra_info('socket')
805-
fileno = s.fileno()
806-
if addr is None:
807-
real_addr = _MDNS_ADDR6 if s.family == socket.AF_INET6 else _MDNS_ADDR
808-
else:
809-
real_addr = addr
810-
if not can_send_to(s, real_addr):
811-
continue
812-
if not self.unicast and sock_fileno is not None and sock_fileno != fileno:
813-
continue
814-
log.debug(
815-
'Sending to (%s, %d) via [socket %s (%s)] (%d bytes #%d) %r as %r...',
816-
real_addr,
817-
port or _MDNS_PORT,
818-
fileno,
819-
transport.get_extra_info('sockname'),
820-
len(packet),
821-
packet_num + 1,
822-
out,
823-
packet,
824-
)
825-
transport.sendto(packet, (real_addr, port or _MDNS_PORT, *v6_flow_scope))
806+
for send_transport in transports:
807+
self._async_send_transport(send_transport, packet, packet_num, out, addr, port, v6_flow_scope)
808+
809+
def _async_send_transport(
810+
self,
811+
transport: asyncio.DatagramTransport,
812+
packet: bytes,
813+
packet_num: int,
814+
out: DNSOutgoing,
815+
addr: Optional[str],
816+
port: int,
817+
v6_flow_scope: Union[Tuple[()], Tuple[int, int]] = (),
818+
) -> None:
819+
s = transport.get_extra_info('socket')
820+
if addr is None:
821+
real_addr = _MDNS_ADDR6 if s.family == socket.AF_INET6 else _MDNS_ADDR
822+
else:
823+
real_addr = addr
824+
if not can_send_to(s, real_addr):
825+
return
826+
log.debug(
827+
'Sending to (%s, %d) via [socket %s (%s)] (%d bytes #%d) %r as %r...',
828+
real_addr,
829+
port or _MDNS_PORT,
830+
s.fileno(),
831+
transport.get_extra_info('sockname'),
832+
len(packet),
833+
packet_num + 1,
834+
out,
835+
packet,
836+
)
837+
transport.sendto(packet, (real_addr, port or _MDNS_PORT, *v6_flow_scope))
826838

827839
def _close(self) -> None:
828840
"""Set global done and remove all service listeners."""

0 commit comments

Comments
 (0)