From d094b4d3ad66d824ebfc8900234dff7aaf4acec0 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Feb 2026 12:58:22 +0000 Subject: [PATCH] Bound DNS cache (#12106) (#12117) --------- CVE: CVE-2026-34513 Upstream-Status: Backport [https://github.com/aio-libs/aiohttp/commit/c4d77c3533122be353b8afca8e8675e3b4cbda98] Backport Changes: - Imported OrderedDict for the aiohttp 3.9.5 collections layout. - Kept Dict[str, Any] because 3.9.5 lacks ResolveResult. - Removed ResolveResult annotations from the backported tests. - Kept family: int = 0 from the 3.9.5 connector API. - Kept _throttle_dns_events from the 3.9.5 DNS flow. (cherry picked from commit 8ab84c52fe58ef34794fa9b12f00b06e626adcc0) Co-authored-by: gonas (cherry picked from commit c4d77c3533122be353b8afca8e8675e3b4cbda98) Signed-off-by: Darsh Kelaiya --- CHANGES/12106.feature.rst | 1 + aiohttp/connector.py | 26 ++++++++++---- tests/test_connector.py | 76 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 CHANGES/12106.feature.rst diff --git a/CHANGES/12106.feature.rst b/CHANGES/12106.feature.rst new file mode 100644 index 000000000..daa9088ee --- /dev/null +++ b/CHANGES/12106.feature.rst @@ -0,0 +1 @@ +Added a ``dns_cache_max_size`` parameter to ``TCPConnector`` to limit the size of the cache -- by :user:`Dreamsorcerer`. diff --git a/aiohttp/connector.py b/aiohttp/connector.py index f95ebe84c..7267d9a5c 100644 --- a/aiohttp/connector.py +++ b/aiohttp/connector.py @@ -4,7 +4,7 @@ import random import sys import traceback import warnings -from collections import defaultdict, deque +from collections import OrderedDict, defaultdict, deque from contextlib import suppress from http import HTTPStatus from http.cookies import SimpleCookie @@ -690,25 +690,33 @@ class BaseConnector: class _DNSCacheTable: - def __init__(self, ttl: Optional[float] = None) -> None: - self._addrs_rr: Dict[Tuple[str, int], Tuple[Iterator[Dict[str, Any]], int]] = {} + def __init__(self, ttl: Optional[float] = None, max_size: int = 1000) -> None: + self._addrs_rr: OrderedDict[ + Tuple[str, int], Tuple[Iterator[Dict[str, Any]], int] + ] = OrderedDict() self._timestamps: Dict[Tuple[str, int], float] = {} self._ttl = ttl + self._max_size = max_size def __contains__(self, host: object) -> bool: return host in self._addrs_rr def add(self, key: Tuple[str, int], addrs: List[Dict[str, Any]]) -> None: + if key in self._addrs_rr: + self._addrs_rr.move_to_end(key) + self._addrs_rr[key] = (cycle(addrs), len(addrs)) if self._ttl is not None: self._timestamps[key] = monotonic() + if len(self._addrs_rr) > self._max_size: + oldest_key, _ = self._addrs_rr.popitem(last=False) + self._timestamps.pop(oldest_key, None) + def remove(self, key: Tuple[str, int]) -> None: self._addrs_rr.pop(key, None) - - if self._ttl is not None: - self._timestamps.pop(key, None) + self._timestamps.pop(key, None) def clear(self) -> None: self._addrs_rr.clear() @@ -719,6 +727,7 @@ class _DNSCacheTable: addrs = list(islice(loop, length)) # Consume one more element to shift internal state of `cycle` next(loop) + self._addrs_rr.move_to_end(key) return addrs def expired(self, key: Tuple[str, int]) -> bool: @@ -760,6 +769,7 @@ class TCPConnector(BaseConnector): fingerprint: Optional[bytes] = None, use_dns_cache: bool = True, ttl_dns_cache: Optional[int] = 10, + dns_cache_max_size: int = 1000, family: int = 0, ssl_context: Optional[SSLContext] = None, ssl: Union[bool, Fingerprint, SSLContext] = True, @@ -789,7 +799,9 @@ class TCPConnector(BaseConnector): self._resolver = resolver self._use_dns_cache = use_dns_cache - self._cached_hosts = _DNSCacheTable(ttl=ttl_dns_cache) + self._cached_hosts = _DNSCacheTable( + ttl=ttl_dns_cache, max_size=dns_cache_max_size + ) self._throttle_dns_events: Dict[Tuple[str, int], EventResultOrError] = {} self._family = family self._local_addr = local_addr diff --git a/tests/test_connector.py b/tests/test_connector.py index 02e48bc10..dd5e2c9e1 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -2197,6 +2197,25 @@ async def test_named_pipe_connector( class TestDNSCacheTable: + host1 = ("localhost", 80) + host2 = ("foo", 80) + result1 = { + "hostname": "localhost", + "host": "127.0.0.1", + "port": 80, + "family": socket.AF_INET, + "proto": 0, + "flags": socket.AI_NUMERICHOST, + } + result2 = { + "hostname": "foo", + "host": "127.0.0.2", + "port": 80, + "family": socket.AF_INET, + "proto": 0, + "flags": socket.AI_NUMERICHOST, + } + @pytest.fixture def dns_cache_table(self): return _DNSCacheTable() @@ -2282,6 +2301,63 @@ class TestDNSCacheTable: addrs = dns_cache_table.next_addrs("foo") assert addrs == ["127.0.0.1"] + def test_max_size_eviction(self) -> None: + table = _DNSCacheTable(max_size=2) + + table.add(self.host1, [self.result1]) + table.add(self.host2, [self.result2]) + + host3 = ("example.com", 80) + result3 = { + **self.result1, + "hostname": "example.com", + "host": "1.2.3.4", + } + table.add(host3, [result3]) + + assert len(table._addrs_rr) == 2 + assert self.host1 not in table._addrs_rr + assert host3 in table._addrs_rr + + def test_lru_eviction(self) -> None: + table = _DNSCacheTable(max_size=2) + + table.add(self.host1, [self.result1]) + table.add(self.host2, [self.result2]) + + table.next_addrs(self.host1) + + host3 = ("example.com", 80) + result3 = { + **self.result1, + "hostname": "example.com", + "host": "1.2.3.4", + } + table.add(host3, [result3]) + + assert self.host1 in table._addrs_rr + assert self.host2 not in table._addrs_rr + + def test_lru_eviction_add(self) -> None: + table = _DNSCacheTable(max_size=2) + + table.add(self.host1, [self.result1]) + table.add(self.host2, [self.result2]) + + # Re-add, thus making host1 the most recently used. + table.add(self.host1, [self.result1]) + + host3 = ("example.com", 80) + result3 = { + **self.result1, + "hostname": "example.com", + "host": "1.2.3.4", + } + table.add(host3, [result3]) + + assert self.host1 in table._addrs_rr + assert self.host2 not in table._addrs_rr + async def test_connector_cache_trace_race(): class DummyTracer: -- 2.35.6