diff --git a/pyatv/core/mdns.py b/pyatv/core/mdns.py index f48f6ff85..d146db7dc 100644 --- a/pyatv/core/mdns.py +++ b/pyatv/core/mdns.py @@ -5,6 +5,7 @@ import logging import math import socket +import struct from types import SimpleNamespace import typing import weakref @@ -247,7 +248,13 @@ def datagram_received(self, data: bytes, _) -> None: Total=len(self.queries), ) - self.parser.add_message(DnsMessage().unpack(data)) + try: + decoded_msg = DnsMessage().unpack(data) + except (UnicodeDecodeError, struct.error): + log_binary(_LOGGER, "Failed to decode message", Msg=data) + return + + self.parser.add_message(decoded_msg) self.received_responses += 1 if self.received_responses == len(self.queries): @@ -423,17 +430,13 @@ def datagram_received(self, data, addr) -> None: Data=data, ) - query_resp = self.query_responses.setdefault( - addr[0], QueryResponse(count=0, deep_sleep=False, parser=ServiceParser()) - ) - # Suppress decode errors for now (but still log) try: decoded_msg = DnsMessage().unpack(data) parser = ServiceParser() services = parser.add_message(decoded_msg).parse() - except UnicodeDecodeError: + except (UnicodeDecodeError, struct.error): log_binary(_LOGGER, "Failed to decode message", Msg=data) return @@ -448,6 +451,9 @@ def datagram_received(self, data, addr) -> None: ): return + query_resp = self.query_responses.setdefault( + addr[0], QueryResponse(count=0, deep_sleep=False, parser=ServiceParser()) + ) is_sleep_proxy = all(service.port == 0 for service in services) query_resp.count += 1 query_resp.deep_sleep |= is_sleep_proxy diff --git a/tests/core/test_mdns.py b/tests/core/test_mdns.py index 2c4891a55..fe81b9863 100644 --- a/tests/core/test_mdns.py +++ b/tests/core/test_mdns.py @@ -264,3 +264,58 @@ def test_parse_ignore_duplicate_records(): records = parser.table["service._abc._tcp.local"] assert mdns.QueryType.SRV in records assert len(records[mdns.QueryType.SRV]) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("length", range(12)) +@pytest.mark.parametrize("unicast", [False, True]) +async def test_truncated_dns_header_then_valid_response(length, unicast): + if unicast: + protocol = mdns.UnicastDnsSdClientProtocol( + [MEDIAREMOTE_SERVICE], "127.0.0.1", 1 + ) + else: + protocol = mdns.MulticastDnsSdClientProtocol( + asyncio.get_running_loop(), [MEDIAREMOTE_SERVICE], "224.0.0.251", 5353, None + ) + address = ("127.0.0.1", 5353) + protocol.datagram_received(bytes(length), address) + if unicast: + assert protocol.received_responses == 0 + else: + assert not protocol.query_responses + query = mdns.create_service_queries([MEDIAREMOTE_SERVICE], dns.QueryType.PTR)[0] + response = fake_udns.create_response(query, TEST_SERVICES) + protocol.datagram_received(response.pack(), address) + if unicast: + assert protocol.received_responses == 1 + services = protocol.parser.parse() + else: + assert protocol.query_responses[address[0]].count == 1 + services = protocol.query_responses[address[0]].parser.parse() + assert any(service.type == MEDIAREMOTE_SERVICE for service in services) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unrelated_service", [False, True], ids=["empty", "unrelated"]) +async def test_filtered_multicast_reply_then_valid_response(unrelated_service): + protocol = mdns.MulticastDnsSdClientProtocol( + asyncio.get_running_loop(), [MEDIAREMOTE_SERVICE], "224.0.0.251", 5353, None + ) + address = ("127.0.0.1", 5353) + response = dns.DnsMessage() + if unrelated_service: + service_type = "_unrelated._tcp.local" + query = mdns.create_service_queries([service_type], dns.QueryType.PTR)[0] + response = fake_udns.create_response( + query, {service_type: TEST_SERVICES[MEDIAREMOTE_SERVICE]} + ) + protocol.datagram_received(response.pack(), address) + assert not protocol.query_responses + + query = mdns.create_service_queries([MEDIAREMOTE_SERVICE], dns.QueryType.PTR)[0] + response = fake_udns.create_response(query, TEST_SERVICES) + protocol.datagram_received(response.pack(), address) + assert protocol.query_responses[address[0]].count == 1 + services = protocol.query_responses[address[0]].parser.parse() + assert any(service.type == MEDIAREMOTE_SERVICE for service in services)