Deduplicate the outbound connect guard's sync and async paths

The two connect_tcp implementations each repeated the per-attempt budget
arithmetic verbatim, and only the sync one delegated resolution and error
mapping to a helper, so a drift between the copies could silently give one
stack a different timeout policy from the other. In the IMAP client,
_connect_pinned re-asserted a fact its only caller had already established,
which reads as a runtime invariant check on a security-relevant path when it
is only a narrowing aid.

Move the budget arithmetic into one helper called from both loops, add an
async twin of the resolve helper so the two loop bodies differ only by await,
and pass the narrowed address tuple into _connect_pinned instead of asserting
it. The PinnedIMAP4 docstring now spells out that no pinning and pinning that
yielded nothing are different things, and the monotonic clock seam says why it
exists.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
stumpylog
2026-09-22 12:49:06 -07:00
co-authored by Claude Opus 5
parent 58777076c5
commit 46c150eb67
2 changed files with 49 additions and 24 deletions
+37 -20
View File
@@ -124,6 +124,9 @@ def is_public_ip(ip: IPAddress) -> bool:
# guard dials.
_getaddrinfo = socket.getaddrinfo
_agetaddrinfo = anyio.getaddrinfo
# The clock is a seam because time-machine does not mock monotonic clocks, and
# patching time.monotonic globally would also replace the asyncio event loop's
# own clock, hanging or misfiring its timers for the rest of the test.
_monotonic = time.monotonic
@@ -251,6 +254,19 @@ def _budget_exhausted(host: str, tried: int, total: int) -> httpcore.ConnectTime
)
def _next_attempt_budget(
host: str,
deadline: float,
candidates: list[IPAddress],
index: int,
) -> float:
"""Budget for the attempt at index, or a timeout if none is left."""
remaining = deadline - _monotonic()
if remaining <= 0:
raise _budget_exhausted(host, index, len(candidates))
return _attempt_timeout(remaining, len(candidates) - index)
def _resolve_for_connect(host: str, port: int) -> tuple[IPAddress, ...]:
try:
return resolve_public_addresses(host, port)
@@ -261,6 +277,24 @@ def _resolve_for_connect(host: str, port: int) -> tuple[IPAddress, ...]:
raise httpcore.ConnectError(str(e)) from e
async def _aresolve_for_connect(
host: str,
port: int,
timeout: float | None,
) -> tuple[IPAddress, ...]:
# The scope closes before dialling; attempts are not nested inside it.
try:
with anyio.fail_after(timeout):
return await aresolve_public_addresses(host, port)
except TimeoutError as e:
raise httpcore.ConnectTimeout(f"Timed out resolving {host}") from e
except OutboundRequestBlockedError as e:
_log_block(e)
raise
except HostResolutionError as e:
raise httpcore.ConnectError(str(e)) from e
class _GuardedSyncBackend(httpcore.NetworkBackend):
"""
Wraps httpcore's sync backend. With internal addresses disallowed, it
@@ -295,10 +329,7 @@ class _GuardedSyncBackend(httpcore.NetworkBackend):
deadline = _deadline(timeout)
last_error: httpcore.ConnectError | httpcore.ConnectTimeout | None = None
for index, address in enumerate(candidates):
remaining = deadline - _monotonic()
if remaining <= 0:
raise _budget_exhausted(host, index, len(candidates))
budget = _attempt_timeout(remaining, len(candidates) - index)
budget = _next_attempt_budget(host, deadline, candidates, index)
try:
return self._inner.connect_tcp(
str(address),
@@ -361,25 +392,11 @@ class _GuardedAsyncBackend(httpcore.AsyncNetworkBackend):
)
_require_positive_timeout(host, timeout)
# Resolution counts against the budget, matching the stock backend.
# This scope closes before dialling; attempts are not nested inside it.
deadline = _deadline(timeout)
try:
with anyio.fail_after(timeout):
addresses = await aresolve_public_addresses(host, port)
except TimeoutError as e:
raise httpcore.ConnectTimeout(f"Timed out resolving {host}") from e
except OutboundRequestBlockedError as e:
_log_block(e)
raise
except HostResolutionError as e:
raise httpcore.ConnectError(str(e)) from e
candidates = _attempt_order(addresses)
candidates = _attempt_order(await _aresolve_for_connect(host, port, timeout))
last_error: httpcore.ConnectError | httpcore.ConnectTimeout | None = None
for index, address in enumerate(candidates):
remaining = deadline - _monotonic()
if remaining <= 0:
raise _budget_exhausted(host, index, len(candidates))
budget = _attempt_timeout(remaining, len(candidates) - index)
budget = _next_attempt_budget(host, deadline, candidates, index)
try:
return await self._inner.connect_tcp(
str(address),
+12 -4
View File
@@ -448,6 +448,11 @@ class PinnedIMAP4(imaplib.IMAP4):
Without pinned addresses, and with the ssl_context of the matching imaplib
class, this behaves exactly like imaplib.IMAP4 / imaplib.IMAP4_SSL.
``pinned_ips`` of ``None`` means no pinning was requested and the stock
imaplib connection path is used. An empty tuple means pinning was requested
and yielded nothing, and the connection fails without opening a socket
rather than falling back to a hostname lookup.
"""
def __init__(
@@ -462,10 +467,13 @@ class PinnedIMAP4(imaplib.IMAP4):
self.ssl_context = ssl_context
super().__init__(host, port, timeout=timeout)
def _connect_pinned(self, timeout: float | None) -> socket.socket:
assert self._pinned_ips is not None
def _connect_pinned(
self,
pinned_ips: tuple[IPAddress, ...],
timeout: float | None,
) -> socket.socket:
last_error: OSError | None = None
for ip in self._pinned_ips:
for ip in pinned_ips:
try:
address = (str(ip), self.port)
if timeout is not None:
@@ -477,7 +485,7 @@ class PinnedIMAP4(imaplib.IMAP4):
def _create_socket(self, timeout: float | None) -> socket.socket:
if self._pinned_ips is not None:
sock = self._connect_pinned(timeout)
sock = self._connect_pinned(self._pinned_ips, timeout)
else:
sock = super()._create_socket(timeout)
if self.ssl_context is None: