From 73589015c3a4a07b7308277255543083d153a510 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 4 Feb 2024 16:50:08 -0600 Subject: [PATCH] Improve scalability of DHCP matchers (#109406) --- homeassistant/components/dhcp/__init__.py | 86 ++++++++--- tests/components/dhcp/test_init.py | 178 +++++++++++++++++----- 2 files changed, 202 insertions(+), 62 deletions(-) diff --git a/homeassistant/components/dhcp/__init__.py b/homeassistant/components/dhcp/__init__.py index b8a12a937e3..ad0446543db 100644 --- a/homeassistant/components/dhcp/__init__.py +++ b/homeassistant/components/dhcp/__init__.py @@ -9,6 +9,7 @@ from dataclasses import dataclass from datetime import timedelta from fnmatch import translate from functools import lru_cache +import itertools import logging import os import re @@ -89,11 +90,55 @@ class DhcpServiceInfo(BaseServiceInfo): macaddress: str +@dataclass(slots=True) +class DhcpMatchers: + """Prepared info from dhcp entries.""" + + registered_devices_domains: set[str] + no_oui_matchers: dict[str, list[DHCPMatcher]] + oui_matchers: dict[str, list[DHCPMatcher]] + + +def async_index_integration_matchers( + integration_matchers: list[DHCPMatcher], +) -> DhcpMatchers: + """Index the integration matchers. + + We have three types of matchers: + + 1. Registered devices + 2. Devices with no OUI - index by first char of lower() hostname + 3. Devices with OUI - index by OUI + """ + registered_devices_domains: set[str] = set() + no_oui_matchers: dict[str, list[DHCPMatcher]] = {} + oui_matchers: dict[str, list[DHCPMatcher]] = {} + for matcher in integration_matchers: + domain = matcher["domain"] + if REGISTERED_DEVICES in matcher: + registered_devices_domains.add(domain) + continue + + if mac_address := matcher.get(MAC_ADDRESS): + oui_matchers.setdefault(mac_address[:6], []).append(matcher) + continue + + if hostname := matcher.get(HOSTNAME): + first_char = hostname[0].lower() + no_oui_matchers.setdefault(first_char, []).append(matcher) + + return DhcpMatchers( + registered_devices_domains=registered_devices_domains, + no_oui_matchers=no_oui_matchers, + oui_matchers=oui_matchers, + ) + + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up the dhcp component.""" watchers: list[WatcherBase] = [] address_data: dict[str, dict[str, str]] = {} - integration_matchers = await async_get_dhcp(hass) + integration_matchers = async_index_integration_matchers(await async_get_dhcp(hass)) # For the passive classes we need to start listening # for state changes and connect the dispatchers before # everything else starts up or we will miss events @@ -125,7 +170,7 @@ class WatcherBase(ABC): self, hass: HomeAssistant, address_data: dict[str, dict[str, str]], - integration_matchers: list[DHCPMatcher], + integration_matchers: DhcpMatchers, ) -> None: """Initialize class.""" super().__init__() @@ -189,28 +234,29 @@ class WatcherBase(ABC): lowercase_hostname, ) - matched_domains = set() - device_domains = set() + matched_domains: set[str] = set() + matchers = self._integration_matchers + registered_devices_domains = matchers.registered_devices_domains dev_reg: DeviceRegistry = async_get(self.hass) if device := dev_reg.async_get_device( connections={(CONNECTION_NETWORK_MAC, uppercase_mac)} ): for entry_id in device.config_entries: - if entry := self.hass.config_entries.async_get_entry(entry_id): - device_domains.add(entry.domain) + if ( + entry := self.hass.config_entries.async_get_entry(entry_id) + ) and entry.domain in registered_devices_domains: + matched_domains.add(entry.domain) - for matcher in self._integration_matchers: + oui = uppercase_mac[:6] + lowercase_hostname_first_char = ( + lowercase_hostname[0] if len(lowercase_hostname) else "" + ) + for matcher in itertools.chain( + matchers.no_oui_matchers.get(lowercase_hostname_first_char, ()), + matchers.oui_matchers.get(oui, ()), + ): domain = matcher["domain"] - - if matcher.get(REGISTERED_DEVICES) and domain not in device_domains: - continue - - if ( - matcher_mac := matcher.get(MAC_ADDRESS) - ) is not None and not _memorized_fnmatch(uppercase_mac, matcher_mac): - continue - if ( matcher_hostname := matcher.get(HOSTNAME) ) is not None and not _memorized_fnmatch( @@ -241,7 +287,7 @@ class NetworkWatcher(WatcherBase): self, hass: HomeAssistant, address_data: dict[str, dict[str, str]], - integration_matchers: list[DHCPMatcher], + integration_matchers: DhcpMatchers, ) -> None: """Initialize class.""" super().__init__(hass, address_data, integration_matchers) @@ -294,7 +340,7 @@ class DeviceTrackerWatcher(WatcherBase): self, hass: HomeAssistant, address_data: dict[str, dict[str, str]], - integration_matchers: list[DHCPMatcher], + integration_matchers: DhcpMatchers, ) -> None: """Initialize class.""" super().__init__(hass, address_data, integration_matchers) @@ -349,7 +395,7 @@ class DeviceTrackerRegisteredWatcher(WatcherBase): self, hass: HomeAssistant, address_data: dict[str, dict[str, str]], - integration_matchers: list[DHCPMatcher], + integration_matchers: DhcpMatchers, ) -> None: """Initialize class.""" super().__init__(hass, address_data, integration_matchers) @@ -387,7 +433,7 @@ class DHCPWatcher(WatcherBase): self, hass: HomeAssistant, address_data: dict[str, dict[str, str]], - integration_matchers: list[DHCPMatcher], + integration_matchers: DhcpMatchers, ) -> None: """Initialize class.""" super().__init__(hass, address_data, integration_matchers) diff --git a/tests/components/dhcp/test_init.py b/tests/components/dhcp/test_init.py index a63300b1ea2..18d213a7029 100644 --- a/tests/components/dhcp/test_init.py +++ b/tests/components/dhcp/test_init.py @@ -1,6 +1,8 @@ """Test the DHCP discovery integration.""" +from collections.abc import Awaitable, Callable import datetime import threading +from typing import Any, cast from unittest.mock import MagicMock, patch import pytest @@ -130,13 +132,15 @@ RAW_DHCP_REQUEST_WITHOUT_HOSTNAME = ( ) -async def _async_get_handle_dhcp_packet(hass, integration_matchers): +async def _async_get_handle_dhcp_packet( + hass: HomeAssistant, integration_matchers: dhcp.DhcpMatchers +) -> Callable[[Any], Awaitable[None]]: dhcp_watcher = dhcp.DHCPWatcher( hass, {}, integration_matchers, ) - async_handle_dhcp_packet = None + async_handle_dhcp_packet: Callable[[Any], Awaitable[None]] | None = None def _mock_sniffer(*args, **kwargs): nonlocal async_handle_dhcp_packet @@ -158,14 +162,14 @@ async def _async_get_handle_dhcp_packet(hass, integration_matchers): ): await dhcp_watcher.async_start() - return async_handle_dhcp_packet + return cast("Callable[[Any], Awaitable[None]]", async_handle_dhcp_packet) async def test_dhcp_match_hostname_and_macaddress(hass: HomeAssistant) -> None: """Test matching based on hostname and macaddress.""" - integration_matchers = [ - {"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"} - ] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}] + ) packet = Ether(RAW_DHCP_REQUEST) async_handle_dhcp_packet = await _async_get_handle_dhcp_packet( @@ -190,9 +194,9 @@ async def test_dhcp_match_hostname_and_macaddress(hass: HomeAssistant) -> None: async def test_dhcp_renewal_match_hostname_and_macaddress(hass: HomeAssistant) -> None: """Test renewal matching based on hostname and macaddress.""" - integration_matchers = [ - {"domain": "mock-domain", "hostname": "irobot-*", "macaddress": "501479*"} - ] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "irobot-*", "macaddress": "501479*"}] + ) packet = Ether(RAW_DHCP_RENEWAL) @@ -220,10 +224,12 @@ async def test_registered_devices( hass: HomeAssistant, device_registry: dr.DeviceRegistry ) -> None: """Test discovery flows are created for registered devices.""" - integration_matchers = [ - {"domain": "not-matching", "registered_devices": True}, - {"domain": "mock-domain", "registered_devices": True}, - ] + integration_matchers = dhcp.async_index_integration_matchers( + [ + {"domain": "not-matching", "registered_devices": True}, + {"domain": "mock-domain", "registered_devices": True}, + ] + ) packet = Ether(RAW_DHCP_RENEWAL) @@ -265,7 +271,9 @@ async def test_registered_devices( async def test_dhcp_match_hostname(hass: HomeAssistant) -> None: """Test matching based on hostname only.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "connect"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "connect"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -289,7 +297,9 @@ async def test_dhcp_match_hostname(hass: HomeAssistant) -> None: async def test_dhcp_match_macaddress(hass: HomeAssistant) -> None: """Test matching based on macaddress only.""" - integration_matchers = [{"domain": "mock-domain", "macaddress": "B8B7F1*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "macaddress": "B8B7F1*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -313,10 +323,12 @@ async def test_dhcp_match_macaddress(hass: HomeAssistant) -> None: async def test_dhcp_multiple_match_only_one_flow(hass: HomeAssistant) -> None: """Test matching the domain multiple times only generates one flow.""" - integration_matchers = [ - {"domain": "mock-domain", "macaddress": "B8B7F1*"}, - {"domain": "mock-domain", "hostname": "connect"}, - ] + integration_matchers = dhcp.async_index_integration_matchers( + [ + {"domain": "mock-domain", "macaddress": "B8B7F1*"}, + {"domain": "mock-domain", "hostname": "connect"}, + ] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -340,7 +352,9 @@ async def test_dhcp_multiple_match_only_one_flow(hass: HomeAssistant) -> None: async def test_dhcp_match_macaddress_without_hostname(hass: HomeAssistant) -> None: """Test matching based on macaddress only.""" - integration_matchers = [{"domain": "mock-domain", "macaddress": "606BBD*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "macaddress": "606BBD*"}] + ) packet = Ether(RAW_DHCP_REQUEST_WITHOUT_HOSTNAME) @@ -364,7 +378,9 @@ async def test_dhcp_match_macaddress_without_hostname(hass: HomeAssistant) -> No async def test_dhcp_nomatch(hass: HomeAssistant) -> None: """Test not matching based on macaddress only.""" - integration_matchers = [{"domain": "mock-domain", "macaddress": "ABC123*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "macaddress": "ABC123*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -379,7 +395,9 @@ async def test_dhcp_nomatch(hass: HomeAssistant) -> None: async def test_dhcp_nomatch_hostname(hass: HomeAssistant) -> None: """Test not matching based on hostname only.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "nomatch*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -394,7 +412,9 @@ async def test_dhcp_nomatch_hostname(hass: HomeAssistant) -> None: async def test_dhcp_nomatch_non_dhcp_packet(hass: HomeAssistant) -> None: """Test matching does not throw on a non-dhcp packet.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "nomatch*"}] + ) packet = Ether(b"") @@ -409,7 +429,9 @@ async def test_dhcp_nomatch_non_dhcp_packet(hass: HomeAssistant) -> None: async def test_dhcp_nomatch_non_dhcp_request_packet(hass: HomeAssistant) -> None: """Test nothing happens with the wrong message-type.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "nomatch*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -433,7 +455,9 @@ async def test_dhcp_nomatch_non_dhcp_request_packet(hass: HomeAssistant) -> None async def test_dhcp_invalid_hostname(hass: HomeAssistant) -> None: """Test we ignore invalid hostnames.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "nomatch*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -457,7 +481,9 @@ async def test_dhcp_invalid_hostname(hass: HomeAssistant) -> None: async def test_dhcp_missing_hostname(hass: HomeAssistant) -> None: """Test we ignore missing hostnames.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "nomatch*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -481,7 +507,9 @@ async def test_dhcp_missing_hostname(hass: HomeAssistant) -> None: async def test_dhcp_invalid_option(hass: HomeAssistant) -> None: """Test we ignore invalid hostname option.""" - integration_matchers = [{"domain": "mock-domain", "hostname": "nomatch*"}] + integration_matchers = dhcp.async_index_integration_matchers( + [{"domain": "mock-domain", "hostname": "nomatch*"}] + ) packet = Ether(RAW_DHCP_REQUEST) @@ -628,7 +656,15 @@ async def test_device_tracker_hostname_and_macaddress_exists_before_start( device_tracker_watcher = dhcp.DeviceTrackerWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -653,7 +689,15 @@ async def test_device_tracker_registered(hass: HomeAssistant) -> None: device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -684,7 +728,15 @@ async def test_device_tracker_registered_hostname_none(hass: HomeAssistant) -> N device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -709,7 +761,15 @@ async def test_device_tracker_hostname_and_macaddress_after_start( device_tracker_watcher = dhcp.DeviceTrackerWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -748,7 +808,15 @@ async def test_device_tracker_hostname_and_macaddress_after_start_not_home( device_tracker_watcher = dhcp.DeviceTrackerWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -877,7 +945,15 @@ async def test_device_tracker_ignore_self_assigned_ips_before_start( device_tracker_watcher = dhcp.DeviceTrackerWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -902,7 +978,15 @@ async def test_aiodiscover_finds_new_hosts(hass: HomeAssistant) -> None: device_tracker_watcher = dhcp.NetworkWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -953,13 +1037,15 @@ async def test_aiodiscover_does_not_call_again_on_shorter_hostname( device_tracker_watcher = dhcp.NetworkWatcher( hass, {}, - [ - { - "domain": "mock-domain", - "hostname": "irobot-*", - "macaddress": "B8B7F1*", - } - ], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "irobot-*", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done() @@ -996,7 +1082,15 @@ async def test_aiodiscover_finds_new_hosts_after_interval(hass: HomeAssistant) - device_tracker_watcher = dhcp.NetworkWatcher( hass, {}, - [{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}], + dhcp.async_index_integration_matchers( + [ + { + "domain": "mock-domain", + "hostname": "connect", + "macaddress": "B8B7F1*", + } + ] + ), ) await device_tracker_watcher.async_start() await hass.async_block_till_done()