fix: bound DNSCache record count to prevent unbounded LAN-driven growth

bdraco · bdraco · commit c24fa39d860e · 2026-05-17T19:36:12.000-07:00
Closes #1715. A LAN-local peer multicasting unique-name mDNS responses could grow ``DNSCache.cache`` / ``_expirations`` / ``service_cache`` / ``_expire_heap`` without bound. ``_DNS_PTR_MIN_TTL = 1125`` (RFC 6762 recommendation for ServiceBrowser refresh-floor) actually extends attacker-injected PTR records to ~19 min before they expire naturally, and ``async_expire`` only runs every ``_CACHE_CLEANUP_INTERVAL = 10`` s — too slow to keep up with a sustained flood. On memory-constrained deployments (Home Assistant on a Raspberry Pi is the canonical victim) this trivially OOMs the process. Add a hard cap (``_MAX_CACHE_RECORDS = 10_000`` in ``const.py``) on the total number of records the cache will hold. ``DNSCache`` now tracks ``_total_records`` (incremented on genuinely-new inserts in ``_async_add``, decremented in ``_async_remove``). When the cap is hit, ``_async_evict_oldest`` heappops the closest-to-expiration record and removes it before inserting the new one. Reuses the existing ``_expire_heap`` for victim selection — O(log n) — and skips stale heap entries via the same expiration-equality check ``async_expire`` already uses. All four touched lines compile to direct C int ops in the Cython build (score-0 in ``cython -a`` annotation); the eviction call is a C-level vtable dispatch fired only on overflow. Wall-clock check: below-cap add stays at ~160 ns/record; over-cap add with constant eviction is ~195 ns/record (+36 ns for the heappop + cache delete). Per-source-IP record quotas — the reporter's "Better" suggestion — are deferred to a follow-up PR. This commit closes the unbounded- growth bug class on its own; the quota work changes the fairness properties but not the memory bound. CWE-400 (Uncontrolled Resource Consumption). LAN-local attack surface only.
diff --git a/src/zeroconf/_cache.pxd b/src/zeroconf/_cache.pxd
@@ -19,6 +19,7 @@ cdef object _UNIQUE_RECORD_TYPES
 cdef unsigned int _TYPE_PTR
 cdef cython.uint _ONE_SECOND
 cdef unsigned int _MIN_SCHEDULED_RECORD_EXPIRATION
+cdef unsigned int _MAX_CACHE_RECORDS
 
 
 @cython.locals(record_cache=dict)
@@ -31,6 +32,7 @@ cdef class DNSCache:
     cdef public cython.dict service_cache
     cdef public list _expire_heap
     cdef public dict _expirations
+    cdef public unsigned int _total_records
 
     cpdef bint async_add_records(self, object entries)
 
@@ -60,10 +62,14 @@ cdef class DNSCache:
         service_store=cython.dict,
         service_record=DNSService,
         when=object,
-        new=bint
+        new=bint,
+        is_new=bint
     )
     cdef bint _async_add(self, DNSRecord record)
 
+    @cython.locals(record=DNSRecord, when_record=tuple)
+    cdef void _async_evict_oldest(self)
+
     @cython.locals(service_record=DNSService)
     cdef void _async_remove(self, DNSRecord record)
 
diff --git a/src/zeroconf/_cache.py b/src/zeroconf/_cache.py
@@ -37,7 +37,7 @@
     DNSText,
 )
 from ._utils.time import current_time_millis
-from .const import _ONE_SECOND, _TYPE_PTR
+from .const import _MAX_CACHE_RECORDS, _ONE_SECOND, _TYPE_PTR
 
 _UNIQUE_RECORD_TYPES = (DNSAddress, DNSHinfo, DNSPointer, DNSText, DNSService)
 _UniqueRecordsType = DNSAddress | DNSHinfo | DNSPointer | DNSText | DNSService
@@ -72,6 +72,7 @@ def __init__(self) -> None:
         self._expire_heap: list[tuple[float, DNSRecord]] = []
         self._expirations: dict[DNSRecord, float] = {}
         self.service_cache: _DNSRecordCacheType = {}
+        self._total_records: int = 0
 
     # Functions prefixed with async_ are NOT threadsafe and must
     # be run in the event loop.
@@ -91,7 +92,16 @@ def _async_add(self, record: _DNSRecord) -> bool:
         # direction would return the old incorrect entry.
         if (store := self.cache.get(record.key)) is None:
             store = self.cache[record.key] = {}
-        new = record not in store and not isinstance(record, DNSNsec)
+        is_new = record not in store
+        # Bound total cache size; evict closest-to-expiration entry to
+        # make room before inserting a new record. Prevents a LAN-local
+        # flood of unique-name records from growing the cache without
+        # bound (RFC 6762 §10 advisory caching, defense-in-depth).
+        if is_new and self._total_records >= _MAX_CACHE_RECORDS:
+            self._async_evict_oldest()
+        new = is_new and not isinstance(record, DNSNsec)
+        if is_new:
+            self._total_records += 1
         store[record] = record
         when = record.created + (record.ttl * 1000)
         if self._expirations.get(record) != when:
@@ -106,6 +116,22 @@ def _async_add(self, record: _DNSRecord) -> bool:
             service_store[service_record] = service_record
         return new
 
+    def _async_evict_oldest(self) -> None:
+        """Drop the closest-to-expiration record to make room for a new one.
+
+        Skips stale heap entries (records re-added with a different TTL),
+        which mirrors the staleness check in ``async_expire``.
+
+        This function must be run in from event loop.
+        """
+        while self._expire_heap:
+            when_record = heappop(self._expire_heap)
+            record = when_record[1]
+            if self._expirations.get(record) != when_record[0]:
+                continue
+            self._async_remove(record)
+            return
+
     def async_add_records(self, entries: Iterable[DNSRecord]) -> bool:
         """Add multiple records.
 
@@ -129,6 +155,7 @@ def _async_remove(self, record: _DNSRecord) -> None:
             _remove_key(self.service_cache, service_record.server_key, service_record)
         _remove_key(self.cache, record.key, record)
         self._expirations.pop(record, None)
+        self._total_records -= 1
 
     def async_remove_records(self, entries: Iterable[DNSRecord]) -> None:
         """Remove multiple records.
diff --git a/src/zeroconf/const.py b/src/zeroconf/const.py
@@ -59,6 +59,12 @@
 # level of rate limit and safe guards so we use 1/4 of the recommended value
 _DNS_PTR_MIN_TTL = 1125
 
+# Upper bound on the number of records the DNSCache will hold before it
+# starts evicting the closest-to-expiration entry to make room for new
+# arrivals. Bounds the memory a malicious LAN peer can force the cache
+# to retain by multicasting many unique-name records.
+_MAX_CACHE_RECORDS = 10000
+
 _DNS_PACKET_HEADER_LEN = 12
 
 _MAX_MSG_TYPICAL = 1460  # unused
diff --git a/tests/benchmarks/test_cache_bound.py b/tests/benchmarks/test_cache_bound.py
@@ -0,0 +1,44 @@
+"""Benchmark for the DNSCache record-count bound + overflow eviction."""
+
+from __future__ import annotations
+
+from pytest_codspeed import BenchmarkFixture
+
+from zeroconf import DNSAddress, DNSCache, current_time_millis
+from zeroconf.const import _CLASS_IN, _MAX_CACHE_RECORDS, _TYPE_A
+
+
+def _make_records(count: int, now: float) -> list[DNSAddress]:
+    return [
+        DNSAddress(
+            f"bench-{i}.local.",
+            _TYPE_A,
+            _CLASS_IN,
+            120,
+            bytes((i & 0xFF, (i >> 8) & 0xFF, 0, 1)),
+            created=now + i,
+        )
+        for i in range(count)
+    ]
+
+
+def test_cache_add_below_cap(benchmark: BenchmarkFixture) -> None:
+    """Adding records while the cache is well below the cap (no eviction)."""
+    now = current_time_millis()
+    records = _make_records(1000, now)
+
+    @benchmark
+    def _add() -> None:
+        cache = DNSCache()
+        cache.async_add_records(records)
+
+
+def test_cache_add_at_cap_evicts(benchmark: BenchmarkFixture) -> None:
+    """Steady-state add at the cap: every new record forces one eviction."""
+    now = current_time_millis()
+    overflow_records = _make_records(_MAX_CACHE_RECORDS + 1000, now)
+
+    @benchmark
+    def _flood() -> None:
+        cache = DNSCache()
+        cache.async_add_records(overflow_records)
diff --git a/tests/test_cache.py b/tests/test_cache.py
@@ -482,3 +482,60 @@ async def test_cache_heap_pops_order() -> None:
         ts, _ = heappop(cache._expire_heap)
         assert ts >= start_ts
         start_ts = ts
+
+
+def test_cache_size_is_bounded() -> None:
+    """A flood of unique-name records is capped at ``_MAX_CACHE_RECORDS``."""
+    cache = r.DNSCache()
+    now = r.current_time_millis()
+    overflow = 1000
+    flood_size = const._MAX_CACHE_RECORDS + overflow
+
+    cache.async_add_records(
+        r.DNSAddress(
+            f"flood-{i}.local.",
+            const._TYPE_A,
+            const._CLASS_IN,
+            120,
+            bytes((i & 0xFF, (i >> 8) & 0xFF, 0, 1)),
+            created=now + i,
+        )
+        for i in range(flood_size)
+    )
+
+    total = sum(len(store) for store in cache.cache.values())
+    assert total == const._MAX_CACHE_RECORDS
+    assert cache._total_records == const._MAX_CACHE_RECORDS
+    # FIFO-ish: the earliest-created records (closest to expiration) get
+    # evicted first, so the names that remain are from the tail.
+    for i in range(overflow):
+        assert f"flood-{i}.local." not in cache.cache
+    for i in range(flood_size - overflow, flood_size):
+        assert f"flood-{i}.local." in cache.cache
+
+
+def test_cache_eviction_decrements_total_records() -> None:
+    """Natural removal (goodbyes, expirations) must keep ``_total_records`` in sync."""
+    cache = r.DNSCache()
+    now = r.current_time_millis()
+    records = [
+        r.DNSAddress(
+            f"sync-{i}.local.",
+            const._TYPE_A,
+            const._CLASS_IN,
+            120,
+            bytes((i & 0xFF, (i >> 8) & 0xFF, 0, 1)),
+            created=now,
+        )
+        for i in range(50)
+    ]
+    cache.async_add_records(records)
+    assert cache._total_records == 50
+
+    cache.async_remove_records(records[:20])
+    assert cache._total_records == 30
+
+    # async_expire on a future time should drop the rest
+    cache.async_expire(now + (200 * 1000))
+    assert cache._total_records == 0
+    assert not cache.cache
