Coverage for src/local_deep_research/security/safe_requests.py: 97%
441 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Safe HTTP Requests Wrapper
4Wraps requests library to add SSRF protection and security best practices.
5"""
7import datetime
8import email.utils
9import functools
10import inspect
11import threading
12import time
13import types
14import weakref
15from typing import TYPE_CHECKING, Any, Callable, Optional
16from urllib.parse import urljoin
18import requests
19from loguru import logger
21from . import dns_pinning
22from . import ssrf_validator
23from ..constants import USER_AGENT
24from ..utilities.resource_utils import safe_close
26if TYPE_CHECKING:
27 from .egress.policy import EgressContext
29# Log lines and raised exception texts below carry only scheme://host:port:
30# the texts flow into caller logs (loguru renders ``ExcType: str(e)``) and
31# into DB error columns, and the URLs can hold userinfo passwords or
32# presigned query tokens.
33_redact = ssrf_validator.redact_url_for_log
36# Default timeout for all HTTP requests (prevents hanging)
37DEFAULT_TIMEOUT = 30 # seconds
39# Maximum response size to prevent memory exhaustion (1GB)
40# Set high to accommodate large documents (annual reports, PDFs, datasets).
41# This is a local research tool — users intentionally download these files.
42MAX_RESPONSE_SIZE = 1024 * 1024 * 1024
44# HTTP status codes that indicate a redirect
45_REDIRECT_STATUS_CODES = frozenset({301, 302, 303, 307, 308})
47# Maximum number of redirects to follow
48_MAX_REDIRECTS = 10
50# Prefix of the requests library's default User-Agent
51# (e.g. "python-requests/2.32.5"). Used by `SafeSession.request` to detect
52# whether the session is using the upstream default — if so, we override
53# it with the project `USER_AGENT` so academic API endpoints (arXiv,
54# OpenAlex, PubMed, …) can identify us. Promoting the literal to a
55# constant so a future requests rename only requires a one-line edit.
56_DEFAULT_REQUESTS_UA_PREFIX = "python-requests"
58# Instance attribute of a size-guarded raw: the list of _BodyGuard objects
59# its class-level read()/read_chunked() count against. See
60# _install_body_guard.
61_BODY_GUARDS_ATTR = "_ldr_body_guards"
64class _BodyGuard:
65 """One installed body guard: the byte budget of a single response.
67 It holds the response only through a weak reference and never refers
68 to ``raw``, so storing it on ``raw`` forms no reference cycle.
69 """
71 __slots__ = ("bytes_read", "response_ref")
73 def __init__(self, response: requests.Response) -> None:
74 self.bytes_read = 0
75 self.response_ref = weakref.ref(response)
77 def limit(self, amt):
78 """Return the read size to request, at most one byte past the cap.
80 raw.read(None) can consume an entire hostile response in one call.
81 Limit even direct raw readers to one byte past the remaining cap,
82 enough to detect overflow without buffering the whole body.
83 """
84 allowance = MAX_RESPONSE_SIZE - self.bytes_read + 1
85 if amt is None or amt < 0 or amt > allowance:
86 return allowance
87 return amt
89 def count(self, raw: Any, data: bytes) -> None:
90 """Add ``data`` to the budget; past the cap, close and raise.
92 Raises:
93 ValueError: Once more than MAX_RESPONSE_SIZE bytes were read.
94 """
95 self.bytes_read += len(data)
96 if self.bytes_read > MAX_RESPONSE_SIZE:
97 owner = self.response_ref()
98 if owner is not None:
99 owner.close()
100 else:
101 # Only raw is still referenced. Close it the way
102 # Response.close() closes an unread body.
103 raw.close()
104 release_conn = getattr(raw, "release_conn", None)
105 if release_conn is not None: 105 ↛ 107line 105 didn't jump to line 107 because the condition on line 105 was always true
106 release_conn()
107 raise ValueError(
108 f"Response body too large: >{self.bytes_read} bytes "
109 f"(max {MAX_RESPONSE_SIZE})"
110 )
113def _is_plain_method(cls: type, name: str) -> bool:
114 """Whether ``cls`` defines ``name`` as an ordinary instance method."""
115 attr = inspect.getattr_static(cls, name, None)
116 return isinstance(attr, (types.FunctionType, types.MethodDescriptorType))
119@functools.lru_cache(maxsize=32)
120def _size_guarded_class(base: type[Any], with_read_chunked: bool) -> type:
121 """Return a subclass of ``base`` whose reads count against its guards.
123 Its read(), and read_chunked() when ``with_read_chunked``, call
124 ``base``'s method (looked up when called) and count the result
125 against every _BodyGuard in the instance's ``_BODY_GUARDS_ATTR``
126 list, in install order, as nested wrappers would. It adds no slots,
127 so an instance of ``base`` can be switched to it by assigning
128 ``__class__``. Cached, so the raws of one class share one subclass.
129 """
131 class SizeGuarded(base):
132 __slots__ = ()
133 _ldr_size_guarded = True
135 def read(self, amt=None, *args, **kwargs):
136 guards = tuple(getattr(self, _BODY_GUARDS_ATTR))
137 # The last guard installed clamps first, as the outermost
138 # wrapper would.
139 for guard in reversed(guards):
140 amt = guard.limit(amt)
141 data = super().read(amt, *args, **kwargs)
142 for guard in guards:
143 guard.count(self, data)
144 return data
146 if with_read_chunked:
148 def read_chunked(self, amt=None, *args, **kwargs):
149 guards = tuple(getattr(self, _BODY_GUARDS_ATTR))
150 for guard in reversed(guards):
151 amt = guard.limit(amt)
152 for chunk in super().read_chunked(amt, *args, **kwargs):
153 for guard in guards:
154 guard.count(self, chunk)
155 yield chunk
157 SizeGuarded.__name__ = "".join(("SizeGuarded", base.__name__))
158 SizeGuarded.__qualname__ = SizeGuarded.__name__
159 return SizeGuarded
162def _class_guard_list(raw: Any) -> Optional[list]:
163 """Return the guard list that raw's class-level reads count against.
165 On the first guard this switches raw's class to the
166 ``_size_guarded_class`` of its class and gives raw an empty list. It
167 returns None, leaving raw unchanged, when raw cannot take that
168 subclass: read() or read_chunked() is set on the instance, the class
169 does not define read() as an ordinary method or defines
170 read_chunked() as something else, raw has no instance ``__dict__``
171 or its class overrides ``__setattr__``, or the ``__class__``
172 assignment is refused. That is the case for test doubles such as
173 mocks, ``SimpleNamespace`` and ``io.BytesIO``. requests' own
174 adapters always return urllib3 responses, which take the subclass.
175 """
176 cls: type = type(raw)
177 if getattr(cls, "_ldr_size_guarded", False):
178 existing: list = vars(raw).setdefault(_BODY_GUARDS_ATTR, [])
179 return existing
180 # A class with its own __setattr__ (a mock, for one) may not treat the
181 # __class__ assignment below as a plain attribute write.
182 if any("__setattr__" in vars(klass) for klass in cls.__mro__[:-1]):
183 return None
184 try:
185 own = vars(raw)
186 except TypeError:
187 return None
188 if "read" in own or "read_chunked" in own:
189 return None
190 if not _is_plain_method(cls, "read"): 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true
191 return None
192 with_read_chunked = (
193 inspect.getattr_static(cls, "read_chunked", None) is not None
194 )
195 if with_read_chunked and not _is_plain_method(cls, "read_chunked"): 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 return None
197 try:
198 raw.__class__ = _size_guarded_class(cls, with_read_chunked)
199 except TypeError:
200 return None
201 guards: list = []
202 setattr(raw, _BODY_GUARDS_ATTR, guards)
203 return guards
206def _wrap_instance_reads(raw: Any, guard: _BodyGuard) -> None:
207 """Wrap raw's read() and read_chunked() on the instance.
209 The fallback for a raw that cannot take ``_size_guarded_class`` (see
210 ``_class_guard_list``). A second guard wraps the first guard's
211 wrappers, so each counts every byte. The wrappers hold ``raw``, as
212 the bound read() they replace does.
213 """
214 original_read = raw.read
216 def bounded_read(amt=None, *args, **kwargs):
217 data = original_read(guard.limit(amt), *args, **kwargs)
218 guard.count(raw, data)
219 return data
221 raw.read = bounded_read
223 original_read_chunked = getattr(raw, "read_chunked", None)
224 if original_read_chunked is None:
225 return
227 def bounded_read_chunked(amt=None, *args, **kwargs):
228 for chunk in original_read_chunked(guard.limit(amt), *args, **kwargs):
229 guard.count(raw, chunk)
230 yield chunk
232 raw.read_chunked = bounded_read_chunked
235def _install_body_guard(response: requests.Response) -> None:
236 """Install a bounded reader that enforces MAX_RESPONSE_SIZE.
238 Counts the bytes response.raw.read() returns and raises ValueError
239 once more than MAX_RESPONSE_SIZE have been read during body
240 consumption. This transparently protects both streamed
241 (.iter_content) and non-streamed (.text, .json(), .content) access
242 patterns, including decompressed bodies whose wire Content-Length
243 was small.
245 A Transfer-Encoding: chunked body never reaches raw.read():
246 urllib3's stream() dispatches it to read_chunked(), which pulls the
247 socket through raw._fp. That generator is counted against the same
248 budget, so the two paths share one.
250 The guard (``_BodyGuard``) refers to the response only weakly and
251 never to ``raw``. A urllib3 response's ``raw`` is switched to a
252 cached subclass of its own class whose read() and read_chunked()
253 apply the guards listed on ``raw`` (``_size_guarded_class``), so
254 nothing stored on ``raw`` refers back to ``raw`` or to the response.
255 A dropped response, its ``raw`` and its cached body are therefore
256 freed by reference counting, however many guards are stacked. A
257 cycle through the guard would keep them, and urllib3's decoder
258 state, in memory until the cyclic garbage collector ran. Because
259 read() is found on the class, a call such as
260 ``safe_get(url, stream=True).raw.read()`` keeps ``raw`` alive while
261 it runs. If the response is gone when the cap is crossed, ``raw`` is
262 closed directly instead.
264 A raw that cannot take the subclass, such as a test double (a mock,
265 ``SimpleNamespace`` or ``io.BytesIO``; see ``_class_guard_list``),
266 gets instance-level wrappers instead. They hold ``raw``, as the
267 bound read() they replace does, so such a raw is freed only by the
268 cyclic garbage collector.
270 Always installs. The callers (_check_response_size, and
271 BaseDownloader._ensure_decoded_body_cap for a second guard) decide
272 when to call it.
273 """
274 raw = response.raw
275 # Synthetic responses used by callers/tests may have a cached body
276 # without a transport. Real streamed responses always have raw set.
277 if raw is None:
278 return
280 guard = _BodyGuard(response)
281 guards = _class_guard_list(raw)
282 if guards is not None:
283 guards.append(guard)
284 else:
285 _wrap_instance_reads(raw, guard)
288def _check_response_size(response: requests.Response) -> None:
289 """Reject responses whose Content-Length exceeds MAX_RESPONSE_SIZE.
291 Handles comma-separated values per RFC 7230 §3.3.2: identical
292 duplicates (from proxies) are normalized; differing values are
293 rejected as invalid framing. Empty parts from malformed headers
294 (trailing/doubled commas) are filtered before parsing. Non-integer
295 or negative values cause the header to be treated as absent.
297 Every accepted response gets a body guard. Content-Length counts wire
298 bytes; decompression can make the delivered body much larger, and a
299 server may understate its length.
301 Must be called before returning a response to the caller. On
302 rejection the response is closed to avoid leaking the connection.
304 Raises:
305 ValueError: If Content-Length values conflict or exceed
306 MAX_RESPONSE_SIZE.
307 """
308 content_length = response.headers.get("Content-Length")
309 if content_length:
310 try:
311 # Handle comma-separated Content-Length values (RFC 7230 §3.3.2).
312 # Multiple identical values may be sent by proxies; differing
313 # values indicate invalid framing and must be rejected.
314 raw_parts = [v.strip() for v in content_length.split(",")]
315 parts = [p for p in raw_parts if p]
316 if not parts:
317 _install_body_guard(response)
318 return # Only commas/whitespace — treat as absent
319 sizes = [int(p) for p in parts]
320 except (ValueError, TypeError):
321 _install_body_guard(response)
322 return # Content-Length not a valid number
323 if len(set(sizes)) > 1:
324 response.close()
325 raise ValueError(
326 f"Conflicting Content-Length values: {content_length}"
327 )
328 size = sizes[0]
329 if size < 0:
330 _install_body_guard(response)
331 return # Malformed Content-Length, treat as absent
332 if size > MAX_RESPONSE_SIZE:
333 response.close()
334 raise ValueError(
335 f"Response too large: {size} bytes (max {MAX_RESPONSE_SIZE})"
336 )
337 # Valid wire length is only a precheck; decoded bytes still need
338 # the body guard during consumption.
339 _install_body_guard(response)
340 return
342 # No Content-Length header at all — install body guard
343 _install_body_guard(response)
346def _finish_response(
347 response: requests.Response, *, preload: bool
348) -> requests.Response:
349 """Install the size guard before optionally buffering the body."""
350 _check_response_size(response)
351 if preload:
352 try:
353 _ = response.content
354 except Exception:
355 safe_close(response, "response")
356 raise
357 return response
360def _discard_redirect_body(response: requests.Response) -> None:
361 """Size-check a redirect response, then drop its body unread.
363 Requests reads ``response.content`` -- the whole body, decoded -- at the
364 start of ``Session.resolve_redirects``, which ``Session.send`` runs on
365 every redirect: to follow it, and with ``allow_redirects=False`` to
366 prepare ``Response.next``. A redirect body is not needed to follow
367 ``Location``, so close the connection and cache an empty body before
368 that read can happen.
370 Raises:
371 ValueError: If the redirect's Content-Length exceeds
372 MAX_RESPONSE_SIZE or conflicts (see ``_check_response_size``).
373 """
374 _check_response_size(response)
375 # Close before marking the body consumed: Response.close() releases an
376 # unread connection only while this flag is false. A synthetic response
377 # with no transport (raw=None) has no connection to release and
378 # Response.close() would crash on it; skip straight to discarding the
379 # (absent) body.
380 if response.raw is not None:
381 response.close()
382 response._content = b""
383 response._content_consumed = True
386def _discard_redirect_body_hook(response: requests.Response, **_kwargs):
387 """Requests response hook: discard a redirect's body (see above).
389 ``Session.send`` dispatches response hooks before it calls
390 ``resolve_redirects``, so this runs ahead of the body read there.
391 """
392 if response.is_redirect:
393 _discard_redirect_body(response)
396def _with_redirect_body_hook(hooks: Optional[dict]) -> dict:
397 """Return ``hooks`` with ``_discard_redirect_body_hook`` run first.
399 ``safe_get``/``safe_post`` send each hop through a plain
400 ``requests.get``/``requests.post``, whose throwaway ``requests.Session``
401 would otherwise buffer every redirect body (see
402 ``_discard_redirect_body``). Caller-supplied hooks are kept and run
403 after it, so they see the redirect's status and headers but an empty
404 body. The caller's dict is not mutated.
405 """
406 merged = dict(hooks or {})
407 caller_response_hooks = merged.get("response")
408 if caller_response_hooks is None:
409 caller_response_hooks = []
410 elif callable(caller_response_hooks): 410 ↛ 412line 410 didn't jump to line 412 because the condition on line 410 was always true
411 caller_response_hooks = [caller_response_hooks]
412 merged["response"] = [
413 _discard_redirect_body_hook,
414 *caller_response_hooks,
415 ]
416 return merged
419def _resolve_redirect_method(method: str, status_code: int) -> str:
420 """Determine HTTP method after redirect, per RFC 7231."""
421 if status_code == 303 and method != "HEAD":
422 method = "GET"
423 elif status_code == 302 and method == "POST":
424 method = "GET"
425 elif status_code == 301 and method == "POST":
426 method = "GET"
427 # 307, 308: preserve original method (no change needed)
428 return method
431# Credential headers: the HTTP-standard set, plus the vendor key/token family.
432# The suffixes cover the shapes vendors ship (X-Api-Key, Zotero-API-Key,
433# X-Goog-Api-Key, X-Subscription-Token, Ocp-Apim-Subscription-Key, Private-Token),
434# so a new vendor needs no edit here. No non-credential header this project
435# sends ends in one of them.
436_CREDENTIAL_HEADERS = frozenset(
437 {"authorization", "proxy-authorization", "cookie"}
438)
439_CREDENTIAL_HEADER_SUFFIXES = ("key", "token", "secret")
441# requests owns the "does this hop leave the credential's scope" rule, so
442# reuse it rather than restate it in the hand-rolled loops below.
443_redirect_scope = requests.sessions.SessionRedirectMixin()
446def _is_credential_header(name: str) -> bool:
447 """Whether a request header carries a caller credential."""
448 lowered = name.lower()
449 return lowered in _CREDENTIAL_HEADERS or lowered.endswith(
450 _CREDENTIAL_HEADER_SUFFIXES
451 )
454def _headers_for_redirect(
455 headers: Optional[dict], from_url: str, to_url: str
456) -> Optional[dict]:
457 """Return the headers to send on a redirect hop.
459 Credential headers are dropped when the hop leaves the scope they were
460 issued for, so a key stays with the host the caller addressed. Returns a
461 new dict when it strips, so the caller's own headers object is untouched.
462 A chain that leaves scope and later returns to the original host does not
463 get the credential back, which is what requests does on its own hops.
464 """
465 if not headers or not _redirect_scope.should_strip_auth(from_url, to_url):
466 return headers
467 return {k: v for k, v in headers.items() if not _is_credential_header(k)}
470def safe_get(
471 url: str,
472 params: Optional[dict] = None,
473 timeout: int = DEFAULT_TIMEOUT,
474 allow_localhost: bool = False,
475 allow_private_ips: bool = False,
476 require_https: bool = False,
477 **kwargs,
478) -> requests.Response:
479 """
480 Make a safe HTTP GET request with SSRF protection.
482 Args:
483 url: URL to request
484 params: URL parameters
485 timeout: Request timeout in seconds
486 allow_localhost: Whether to allow localhost/loopback addresses.
487 Set to True for trusted internal services like self-hosted
488 search engines (e.g., searxng). Default False.
489 allow_private_ips: Whether to allow all private/internal IPs plus localhost.
490 This includes RFC1918 (10.x, 172.16-31.x, 192.168.x), CGNAT (100.64.x.x
491 used by Podman/rootless containers), link-local (169.254.x.x), and IPv6
492 private ranges (fc00::/7, fe80::/10). Use for trusted self-hosted services
493 like SearXNG or Ollama in containerized environments.
494 Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
495 AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
496 ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
497 require_https: Refuse the fetch (and any redirect hop) whose URL
498 is not https://. Opt-in per caller for integrity-sensitive
499 downloads (e.g. the journal-quality datasets): a hop to
500 cleartext passes SSRF validation, letting a compromised
501 upstream or open-redirect move the fetch onto a tamperable
502 path. Redirect hops are judged after resolving the Location
503 against the current URL, so relative and protocol-relative
504 Locations inherit https. Combining it with ``verify=False``
505 raises ValueError — https without certificate verification
506 is not an authenticated transport. NOT global — localhost
507 SearXNG instances are legitimately http://.
508 **kwargs: Additional arguments to pass to requests.get()
510 Returns:
511 Response object. A redirect response's body (any 301/302/303/307/308
512 with a ``Location``) is discarded unread, whether the hop is followed
513 or returned because ``allow_redirects=False``: its status and headers
514 stay available, ``content`` is ``b""``. Caller ``hooks`` still run,
515 after that discard.
517 Raises:
518 ValueError: If URL fails SSRF validation, or (with
519 ``require_https``) the URL or a redirect hop is not https, or
520 ``verify=False`` is passed, or a response's Content-Length
521 (a redirect's included) exceeds MAX_RESPONSE_SIZE or conflicts,
522 or the decoded body read before returning exceeds
523 MAX_RESPONSE_SIZE (with ``stream=True`` the caller's own body
524 read raises that instead)
525 requests.RequestException: If request fails
526 """
527 if require_https:
528 # Pinning the scheme is meaningless if the certificate is not
529 # checked: an on-path attacker can terminate TLS with any cert.
530 # requests disables verification for any falsy ``verify``
531 # (``False``, ``0``, ``""``); ``None`` merges to the session
532 # default (verified) and is allowed.
533 verify = kwargs.get("verify")
534 if verify is not None and not verify:
535 raise ValueError(
536 "require_https cannot be combined with verify=False"
537 )
538 if not url.lower().startswith("https://"):
539 raise ValueError(
540 f"URL must use https when require_https is set: {_redact(url)}"
541 )
542 # Validate URL to prevent SSRF
543 if not ssrf_validator.validate_url(
544 url,
545 allow_localhost=allow_localhost,
546 allow_private_ips=allow_private_ips,
547 ):
548 raise ValueError(
549 f"URL failed security validation (possible SSRF): {_redact(url)}"
550 )
552 # Ensure timeout is set
553 if "timeout" not in kwargs: 553 ↛ 559line 553 didn't jump to line 559 because the condition on line 553 was always true
554 kwargs["timeout"] = timeout
556 # Inject the project User-Agent if the caller didn't supply one.
557 # Mutates a copy of any caller-supplied headers dict so we never
558 # touch their object.
559 headers = dict(kwargs.get("headers") or {})
560 if not any(k.lower() == "user-agent" for k in headers):
561 headers["User-Agent"] = USER_AGENT
562 kwargs["headers"] = headers
564 # Intercept allow_redirects — we handle redirects manually to validate
565 # each redirect target against SSRF rules
566 caller_wants_redirects = kwargs.pop("allow_redirects", True)
567 kwargs["allow_redirects"] = False
568 caller_stream = kwargs.get("stream", False)
569 # Requests otherwise buffers the body inside requests.get(), before
570 # _check_response_size can install its guard. Preserve the caller's
571 # normal preloaded-response contract after that guard is in place.
572 kwargs["stream"] = True
573 # Each hop's plain requests call would still buffer a redirect body
574 # while preparing Response.next; drop it unread instead.
575 kwargs["hooks"] = _with_redirect_body_hook(kwargs.get("hooks"))
577 current_url = url
578 try:
579 # Pin the validated IP for the actual connection: the request below
580 # connects to the address resolved+validated here, closing the
581 # resolve-vs-connect gap where requests/urllib3 would otherwise
582 # re-resolve the hostname independently.
583 with dns_pinning.pinned_request(
584 url, allow_localhost, allow_private_ips
585 ):
586 # The wrapper sets timeout in kwargs before the first request and retains it for redirects.
587 response = requests.get(url, params=params, **kwargs) # nosec B113
589 # Follow redirects manually with SSRF validation on each hop.
590 # Each hop uses a fresh requests.get() call without a session,
591 # so cookies set by intermediate responses are not carried
592 # forward. This is acceptable for current callers (all stateless).
593 # Callers needing cookie persistence across redirects should use
594 # SafeSession instead, which preserves cookies via its cookie jar.
595 if caller_wants_redirects:
596 redirects_followed = 0
597 while (
598 response.status_code in _REDIRECT_STATUS_CODES
599 and redirects_followed < _MAX_REDIRECTS
600 ):
601 redirect_url = (response.headers.get("Location") or "").strip()
602 if not redirect_url:
603 break
605 # Resolve relative redirects
606 redirect_url = urljoin(
607 response.url or current_url, redirect_url
608 )
610 # Scheme pin: refuse a hop that would move the fetch to
611 # cleartext. Checked after urljoin so relative and
612 # protocol-relative Locations inherit the https scheme
613 # before being judged.
614 if require_https and not redirect_url.lower().startswith(
615 "https://"
616 ):
617 logger.warning(
618 f"Redirect to {_redact(redirect_url)} would downgrade "
619 f"to cleartext (from {_redact(url)}, hop "
620 f"{redirects_followed + 1})"
621 )
622 response.close()
623 raise ValueError(
624 "Redirect would downgrade to cleartext: "
625 f"{_redact(redirect_url)}"
626 )
628 # Validate redirect target against SSRF rules
629 if not ssrf_validator.validate_url(
630 redirect_url,
631 allow_localhost=allow_localhost,
632 allow_private_ips=allow_private_ips,
633 ):
634 logger.warning(
635 f"Redirect to {_redact(redirect_url)} blocked by SSRF "
636 f"validation (from {_redact(url)}, hop "
637 f"{redirects_followed + 1})"
638 )
639 response.close()
640 raise ValueError(
641 f"Redirect target failed SSRF validation: {_redact(redirect_url)}"
642 )
644 # Credentials do not follow the hop out of their own scope.
645 kwargs["headers"] = _headers_for_redirect(
646 kwargs.get("headers"), current_url, redirect_url
647 )
648 current_url = redirect_url
649 response.close()
650 # Note: params are intentionally NOT forwarded to redirect
651 # hops. Per HTTP spec, the server's Location header contains
652 # the complete target URL. Re-appending original query params
653 # would corrupt it.
654 # Re-pin per hop: the address validated for this redirect
655 # target is the address connected to.
656 with dns_pinning.pinned_request(
657 redirect_url, allow_localhost, allow_private_ips
658 ):
659 # The wrapper sets timeout in kwargs before the first request and retains it for redirects.
660 response = requests.get(redirect_url, **kwargs) # nosec B113
661 redirects_followed += 1
663 if (
664 response.status_code in _REDIRECT_STATUS_CODES
665 and redirects_followed >= _MAX_REDIRECTS
666 ):
667 response.close()
668 # Note: raises ValueError here, while SafeSession raises
669 # requests.TooManyRedirects (delegated to the base class).
670 # Callers should catch ValueError for standalone functions.
671 raise ValueError(
672 f"Too many redirects ({_MAX_REDIRECTS}) from {_redact(url)}"
673 )
675 return _finish_response(response, preload=not caller_stream)
677 except requests.Timeout:
678 logger.warning(
679 f"Request timeout after {timeout}s: {_redact(current_url)}"
680 )
681 raise
682 except requests.RequestException:
683 logger.warning(f"Request failed for {_redact(current_url)}")
684 raise
687def safe_post(
688 url: str,
689 data: Optional[Any] = None,
690 json: Optional[dict] = None,
691 timeout: int = DEFAULT_TIMEOUT,
692 allow_localhost: bool = False,
693 allow_private_ips: bool = False,
694 **kwargs,
695) -> requests.Response:
696 """
697 Make a safe HTTP POST request with SSRF protection.
699 Args:
700 url: URL to request
701 data: Data to send in request body
702 json: JSON data to send in request body
703 timeout: Request timeout in seconds
704 allow_localhost: Whether to allow localhost/loopback addresses.
705 Set to True for trusted internal services like self-hosted
706 search engines (e.g., searxng). Default False.
707 allow_private_ips: Whether to allow all private/internal IPs plus localhost.
708 This includes RFC1918 (10.x, 172.16-31.x, 192.168.x), CGNAT (100.64.x.x
709 used by Podman/rootless containers), link-local (169.254.x.x), and IPv6
710 private ranges (fc00::/7, fe80::/10). Use for trusted self-hosted services
711 like SearXNG or Ollama in containerized environments.
712 Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
713 AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
714 ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
715 **kwargs: Additional arguments to pass to requests.post()
717 Returns:
718 Response object. Redirect bodies are discarded unread, as in
719 ``safe_get``.
721 Raises:
722 ValueError: If the URL fails SSRF validation, if a 307/308
723 redirect would re-send the request body outside the scope of
724 the URL the caller addressed, if a response's
725 Content-Length (a redirect's included) exceeds
726 MAX_RESPONSE_SIZE or conflicts, or if the decoded body read
727 before returning exceeds MAX_RESPONSE_SIZE (with
728 ``stream=True`` the caller's own body read raises that
729 instead)
730 requests.RequestException: If request fails
731 """
732 # Validate URL to prevent SSRF
733 if not ssrf_validator.validate_url(
734 url,
735 allow_localhost=allow_localhost,
736 allow_private_ips=allow_private_ips,
737 ):
738 raise ValueError(
739 f"URL failed security validation (possible SSRF): {_redact(url)}"
740 )
742 # Ensure timeout is set
743 if "timeout" not in kwargs: 743 ↛ 749line 743 didn't jump to line 749 because the condition on line 743 was always true
744 kwargs["timeout"] = timeout
746 # Inject the project User-Agent if the caller didn't supply one.
747 # Mutates a copy of any caller-supplied headers dict so we never
748 # touch their object.
749 headers = dict(kwargs.get("headers") or {})
750 if not any(k.lower() == "user-agent" for k in headers):
751 headers["User-Agent"] = USER_AGENT
752 kwargs["headers"] = headers
754 # Intercept allow_redirects — we handle redirects manually to validate
755 # each redirect target against SSRF rules
756 caller_wants_redirects = kwargs.pop("allow_redirects", True)
757 kwargs["allow_redirects"] = False
758 caller_stream = kwargs.get("stream", False)
759 kwargs["stream"] = True
760 # Each hop's plain requests call would still buffer a redirect body
761 # while preparing Response.next; drop it unread instead.
762 kwargs["hooks"] = _with_redirect_body_hook(kwargs.get("hooks"))
764 current_url = url
765 try:
766 # Pin the validated IP for the actual connection (see safe_get).
767 with dns_pinning.pinned_request(
768 url, allow_localhost, allow_private_ips
769 ):
770 # The wrapper sets timeout in kwargs before the first request and retains it for redirects.
771 response = requests.post(url, data=data, json=json, **kwargs) # nosec B113
773 # Follow redirects manually with SSRF validation on each hop.
774 # Each hop uses a fresh request without a session, so cookies
775 # set by intermediate responses are not carried forward. Callers
776 # needing cookie persistence should use SafeSession instead.
777 if caller_wants_redirects:
778 redirect_method = "POST"
779 redirects_followed = 0
780 while (
781 response.status_code in _REDIRECT_STATUS_CODES
782 and redirects_followed < _MAX_REDIRECTS
783 ):
784 redirect_url = (response.headers.get("Location") or "").strip()
785 if not redirect_url:
786 break
788 # Resolve relative redirects
789 redirect_url = urljoin(
790 response.url or current_url, redirect_url
791 )
793 # Validate redirect target against SSRF rules
794 if not ssrf_validator.validate_url(
795 redirect_url,
796 allow_localhost=allow_localhost,
797 allow_private_ips=allow_private_ips,
798 ):
799 logger.warning(
800 f"Redirect to {_redact(redirect_url)} blocked by SSRF "
801 f"validation (from {_redact(url)}, hop "
802 f"{redirects_followed + 1})"
803 )
804 response.close()
805 raise ValueError(
806 f"Redirect target failed SSRF validation: {_redact(redirect_url)}"
807 )
809 redirect_method = _resolve_redirect_method(
810 redirect_method, response.status_code
811 )
812 # 307/308 keep the body, and it cannot be dropped the way the
813 # credential headers below are without sending a request the
814 # server did not ask for, so refuse the hop instead.
815 body_would_follow = (
816 data is not None
817 or json is not None
818 or kwargs.get("files") is not None
819 )
820 if (
821 redirect_method != "GET"
822 and body_would_follow
823 and _redirect_scope.should_strip_auth(
824 current_url, redirect_url
825 )
826 ):
827 logger.warning(
828 f"Redirect to {_redact(redirect_url)} not followed: it "
829 f"would re-send the request body outside the scope "
830 f"of {_redact(current_url)} (hop "
831 f"{redirects_followed + 1})"
832 )
833 response.close()
834 raise ValueError(
835 "Redirect would re-send the request body outside "
836 f"the requested scope: {_redact(redirect_url)}"
837 )
838 # Credentials do not follow the hop out of their own scope.
839 kwargs["headers"] = _headers_for_redirect(
840 kwargs.get("headers"), current_url, redirect_url
841 )
842 current_url = redirect_url
843 response.close()
845 # Re-pin per hop: the address validated for this redirect
846 # target is the address connected to.
847 with dns_pinning.pinned_request(
848 redirect_url, allow_localhost, allow_private_ips
849 ):
850 if redirect_method == "GET":
851 # 301/302/303: convert to GET, drop body
852 data = None
853 json = None
854 # The wrapper sets timeout in kwargs before the first request and retains it for redirects.
855 response = requests.get(redirect_url, **kwargs) # nosec B113
856 else:
857 # 307/308: preserve current method and body
858 # The wrapper sets timeout in kwargs before the first request and retains it for redirects.
859 response = requests.post( # nosec B113
860 redirect_url, data=data, json=json, **kwargs
861 )
862 redirects_followed += 1
864 if (
865 response.status_code in _REDIRECT_STATUS_CODES
866 and redirects_followed >= _MAX_REDIRECTS
867 ):
868 response.close()
869 # Note: raises ValueError here, while SafeSession raises
870 # requests.TooManyRedirects (delegated to the base class).
871 # Callers should catch ValueError for standalone functions.
872 raise ValueError(
873 f"Too many redirects ({_MAX_REDIRECTS}) from {_redact(url)}"
874 )
876 return _finish_response(response, preload=not caller_stream)
878 except requests.Timeout:
879 logger.warning(
880 f"Request timeout after {timeout}s: {_redact(current_url)}"
881 )
882 raise
883 except requests.RequestException:
884 logger.warning(f"Request failed for {_redact(current_url)}")
885 raise
888# Create a safe session class
889class SafeSession(requests.Session):
890 """
891 Session with built-in SSRF protection.
893 Redirect validation relies on ``requests.Session.resolve_redirects()``
894 calling ``self.send()`` for each hop — an internal implementation detail
895 of the ``requests`` library. This is simpler than re-implementing the
896 redirect loop (as ``safe_get``/``safe_post`` do) and keeps session-level
897 features (cookies, auth) working. The trade-off is coupling to the
898 ``requests`` internals; if a future version stops routing hops through
899 ``send()``, redirect targets would no longer be validated.
901 Redirect response bodies are discarded before Requests' redirect loop
902 would buffer them, including its ``Response.next`` preparation when
903 ``allow_redirects=False``. Callers can inspect redirect status and
904 headers, but not its body.
906 Usage:
907 with SafeSession() as session:
908 response = session.get(url)
910 # For trusted internal services (e.g., searxng on localhost):
911 with SafeSession(allow_localhost=True) as session:
912 response = session.get(url)
914 # For trusted internal services on any private network IP:
915 with SafeSession(allow_private_ips=True) as session:
916 response = session.get(url)
918 Raises:
919 ValueError: If a URL (initial or redirect target) fails SSRF
920 validation, an enabled HTTPS downgrade guard refuses a redirect
921 or disabled certificate verification, a 307/308 redirect would
922 re-send the request body outside the scope of the URL the caller
923 addressed, the response Content-Length exceeds MAX_RESPONSE_SIZE,
924 or the decoded body read before a request returns exceeds it
925 (with ``stream=True`` the caller's own body read raises that
926 instead). Note: ``safe_get``/``safe_post`` also raise
927 ``ValueError`` for too-many-redirects, but ``SafeSession`` raises
928 ``requests.TooManyRedirects`` for that case since it delegates
929 redirect counting to the ``requests`` library.
930 requests.RequestException: On transport-level failures.
931 """
933 def __init__(
934 self,
935 allow_localhost: bool = False,
936 allow_private_ips: bool = False,
937 block_link_local: bool = False,
938 block_https_downgrade: bool = False,
939 egress_context: Optional["EgressContext"] = None,
940 ):
941 """
942 Initialize SafeSession.
944 Args:
945 allow_localhost: Whether to allow localhost/loopback addresses.
946 allow_private_ips: Whether to allow all private/internal IPs plus localhost.
947 This includes RFC1918, CGNAT (100.64.x.x used by Podman), link-local, and
948 IPv6 private ranges. Use for trusted self-hosted services like SearXNG or
949 Ollama in containerized environments.
950 Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
951 AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
952 ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
953 block_link_local: Keep the whole link-local range (IPv4
954 ``169.254.0.0/16``, IPv6 ``fe80::/10`` and their NAT64-wrapped
955 forms) blocked EVEN under ``allow_private_ips=True``, on the
956 initial request, every redirect hop, and the connect-time DNS
957 re-validation. Set it when relaxing ``allow_private_ips`` for a
958 trusted target that is never link-local (cloud metadata lives
959 there beyond the always-blocked literals). No effect without
960 ``allow_private_ips``; forwarded to ``validate_url``.
961 block_https_downgrade: Refuse a redirect from HTTPS to HTTP,
962 including a downgrade after earlier redirect hops. HTTP-only
963 starting URLs remain supported. HTTPS requests must verify
964 their certificates when this is enabled.
965 egress_context: Optional per-run policy enforced on each initial
966 request, redirect hop, and pinned connect-time address.
967 """
968 super().__init__()
969 self.max_redirects = _MAX_REDIRECTS
970 self.allow_localhost = allow_localhost
971 self.allow_private_ips = allow_private_ips
972 self.block_link_local = block_link_local
973 self.url_guard: Optional[Callable[[str], None]] = None
974 self.block_https_downgrade = block_https_downgrade
975 self.egress_context = egress_context
976 # Per-thread, because one session is shared across search engines.
977 self._redirect_local = threading.local()
979 @property
980 def _preparing_next_only(self) -> bool:
981 """True while ``resolve_redirects`` is only building ``Response.next``."""
982 return getattr(self._redirect_local, "preparing_next_only", False)
984 def resolve_redirects(
985 self,
986 resp,
987 req,
988 stream=False,
989 timeout=None,
990 verify=True,
991 cert=None,
992 proxies=None,
993 yield_requests=False,
994 **adapter_kwargs,
995 ):
996 """Record the redirect mode and discard bodies before Requests drains them.
998 ``Session.send`` calls this with ``yield_requests=True`` even when the
999 caller passed ``allow_redirects=False``, to populate ``Response.next``.
1000 ``rebuild_auth`` runs before that yield and cannot see which mode it is
1001 in, so the flag is published here for it to read.
1003 Requests reads ``resp.content`` at the start of its redirect loop,
1004 before control returns to ``SafeSession.send`` for the size check.
1005 A redirect can therefore buffer an unbounded body (including decoded
1006 compressed content) unless it is closed here. Redirect bodies are not
1007 needed to follow ``Location``; keep an empty cached body so Requests
1008 does not read the closed socket while preparing the next request.
1010 The parameter list mirrors ``Session.resolve_redirects`` positionally,
1011 so a caller passing ``stream`` third still gets streamed responses.
1012 """
1013 previous = self._preparing_next_only
1014 self._redirect_local.preparing_next_only = yield_requests
1015 try:
1016 if self.get_redirect_target(resp):
1017 _discard_redirect_body(resp)
1018 yield from super().resolve_redirects(
1019 resp,
1020 req,
1021 stream=stream,
1022 timeout=timeout,
1023 verify=verify,
1024 cert=cert,
1025 proxies=proxies,
1026 yield_requests=yield_requests,
1027 **adapter_kwargs,
1028 )
1029 finally:
1030 self._redirect_local.preparing_next_only = previous
1032 def rebuild_auth(
1033 self, prepared_request: requests.PreparedRequest, response
1034 ) -> None:
1035 """Drop vendor API-key headers too when a hop leaves auth scope, and
1036 refuse the hop when it would carry the request body along.
1038 The base implementation removes ``Authorization`` and nothing else, so
1039 a session default such as the Semantic Scholar engine's ``x-api-key``
1040 would otherwise be replayed to the redirect target. The strip runs
1041 before ``super()`` because the base method reapplies netrc credentials
1042 for the new host, and those belong to it.
1044 ``resolve_redirects`` clears ``body`` before calling this method for
1045 every status that downgrades the method, and keeps it for 307 and 308,
1046 so a body still present here is one the hop would re-send.
1048 The refusal is raised only when the hop is about to be followed. With
1049 ``allow_redirects=False`` nothing is sent, so raising would deny the
1050 caller the 3xx response it asked for; the credential strip still runs,
1051 which is what ``Response.next`` needs.
1052 """
1053 following_the_hop = not self._preparing_next_only
1054 if (
1055 self.block_https_downgrade
1056 and following_the_hop
1057 and (response.request.url or "").lower().startswith("https://")
1058 and (prepared_request.url or "").lower().startswith("http://")
1059 ):
1060 raise ValueError(
1061 "Redirect would downgrade to cleartext: "
1062 f"{ssrf_validator.redact_url_for_log(prepared_request.url)}"
1063 )
1065 if self.should_strip_auth(response.request.url, prepared_request.url):
1066 if prepared_request.body is not None and following_the_hop:
1067 logger.warning(
1068 f"Redirect to {_redact(prepared_request.url)} not "
1069 f"followed: it would re-send the request body outside "
1070 f"the scope of {_redact(response.request.url)}"
1071 )
1072 raise ValueError(
1073 "Redirect would re-send the request body outside "
1074 f"the requested scope: {_redact(prepared_request.url)}"
1075 )
1076 for name in [
1077 header
1078 for header in prepared_request.headers
1079 if _is_credential_header(header)
1080 ]:
1081 del prepared_request.headers[name]
1082 super().rebuild_auth(prepared_request, response)
1084 def request(self, method: str, url: str, **kwargs) -> requests.Response: # type: ignore[override]
1085 """Override request method to add SSRF validation."""
1086 # Validate URL
1087 if not ssrf_validator.validate_url(
1088 url,
1089 allow_localhost=self.allow_localhost,
1090 allow_private_ips=self.allow_private_ips,
1091 block_link_local=self.block_link_local,
1092 ):
1093 raise ValueError(
1094 f"URL failed security validation (possible SSRF): {_redact(url)}"
1095 )
1097 # Ensure timeout is set
1098 if "timeout" not in kwargs:
1099 kwargs["timeout"] = DEFAULT_TIMEOUT
1101 # Inject project User-Agent if the caller didn't already set one.
1102 # Session-level User-Agent (self.headers) is left alone — only
1103 # per-request headers are copied so we never mutate caller state.
1104 headers = dict(kwargs.get("headers") or {})
1105 session_ua = self.headers.get("User-Agent", "")
1106 has_per_request_ua = any(k.lower() == "user-agent" for k in headers)
1107 if not has_per_request_ua and (
1108 not session_ua or session_ua.startswith(_DEFAULT_REQUESTS_UA_PREFIX)
1109 ):
1110 headers["User-Agent"] = USER_AGENT
1111 kwargs["headers"] = headers
1113 return super().request(method, url, **kwargs)
1115 def send(
1116 self, request: requests.PreparedRequest, **kwargs
1117 ) -> requests.Response:
1118 """Validate every outgoing request against SSRF and run egress scope.
1120 This runs on **all** calls — both the initial request (routed
1121 here by ``requests.Session.request()``) and each redirect hop
1122 (routed here by ``resolve_redirects()``). The initial URL is
1123 therefore validated twice (once in ``request()``, once here);
1124 this is intentional defense-in-depth.
1125 """
1126 url_guard = getattr(self, "url_guard", None)
1127 if url_guard is not None:
1128 url_guard(request.url or "")
1129 if (
1130 self.block_https_downgrade
1131 and request.url
1132 and request.url.lower().startswith("https://")
1133 ):
1134 # Session.send inherits only when the argument is absent.
1135 # An explicit None reaches the adapter with verification disabled.
1136 verify = kwargs.get("verify", self.verify)
1137 if not verify:
1138 raise ValueError(
1139 "block_https_downgrade cannot be combined with verify=False"
1140 )
1141 if request.url and not ssrf_validator.validate_url(
1142 request.url,
1143 allow_localhost=self.allow_localhost,
1144 allow_private_ips=self.allow_private_ips,
1145 block_link_local=self.block_link_local,
1146 ):
1147 logger.warning(
1148 f"Request to {_redact(request.url)} blocked by SSRF validation"
1149 )
1150 # Note: This error says "security validation" while safe_get/
1151 # safe_post say "SSRF validation". The difference indicates the
1152 # source (session vs standalone function) in logs.
1153 raise ValueError(
1154 "Redirect target failed security validation (possible "
1155 f"SSRF): {_redact(request.url)}"
1156 )
1158 if request.url and self.egress_context is not None:
1159 from .egress.policy import PolicyDeniedError, evaluate_url
1161 decision = evaluate_url(request.url, self.egress_context)
1162 if not decision.allowed:
1163 raise PolicyDeniedError(decision, request.url)
1165 # Pin the validated IP for the connection this send performs. send()
1166 # runs for the initial request AND for every redirect hop
1167 # (resolve_redirects re-enters here), so each hop is independently
1168 # resolved, validated, and pinned — the address validated is the
1169 # address connected to.
1170 # Restore the resolve_redirects flag here rather than relying on the
1171 # generator's own finally: Session.send takes one item from it and
1172 # drops it, so that finally waits on collection. A leaked True would
1173 # disarm the body-scope refusal for the rest of the thread.
1174 preparing = self._preparing_next_only
1175 caller_stream = kwargs.get("stream", self.stream)
1176 # Session.send() preloads non-stream responses before returning.
1177 # Force header-first transport, then read through our size guard.
1178 kwargs["stream"] = True
1179 try:
1180 with dns_pinning.pinned_request(
1181 request.url or "",
1182 allow_localhost=self.allow_localhost,
1183 allow_private_ips=self.allow_private_ips,
1184 block_link_local=self.block_link_local,
1185 ) as pinned_addresses:
1186 if pinned_addresses and self.egress_context is not None:
1187 from .egress.policy import (
1188 Decision,
1189 PolicyDeniedError,
1190 evaluate_url,
1191 url_host_is_declared_local,
1192 )
1194 # The URL-level decision above honors an operator-declared
1195 # local hostname. Replacing that name with its pinned IP
1196 # would lose the declaration, so inspect only the
1197 # unconditional metadata block in that case. Otherwise
1198 # enforce the run scope on every address that can connect.
1199 declared_local = url_host_is_declared_local(
1200 request.url, self.egress_context
1201 )
1202 for address_info in pinned_addresses:
1203 address = str(address_info[4][0])
1204 if declared_local:
1205 if ssrf_validator.is_ip_blocked(
1206 address,
1207 allow_localhost=True,
1208 allow_private_ips=True,
1209 ):
1210 # The pinning layer normally rejects this
1211 # first. Audit without incrementing the run's
1212 # quota: this is one already-refused request,
1213 # not an additional user URL evaluation.
1214 logger.bind(policy_audit=True).warning(
1215 "policy denied URL fetch (pinned address) "
1216 f"pinned_address={address} "
1217 "reason=blocked_metadata_ip "
1218 f"scope={self.egress_context.scope.value} "
1219 "counted=False"
1220 )
1221 raise PolicyDeniedError(
1222 Decision(False, "blocked_metadata_ip"),
1223 request.url or "",
1224 )
1225 continue
1226 host = f"[{address}]" if ":" in address else address
1227 decision = evaluate_url(
1228 f"http://{host}/", self.egress_context
1229 )
1230 if not decision.allowed: 1230 ↛ 1202line 1230 didn't jump to line 1202 because the condition on line 1230 was always true
1231 raise PolicyDeniedError(decision, request.url or "")
1232 response = super().send(request, **kwargs)
1233 finally:
1234 self._redirect_local.preparing_next_only = preparing
1235 return _finish_response(response, preload=not caller_stream)
1238# Exponential backoff schedule (seconds). Kept short: journal-quality
1239# downloads are run from a user request or a scheduled job, not from a
1240# time-sensitive hot path, so three retries over ~7 seconds is plenty
1241# without adding real latency.
1242_RETRY_BACKOFF_SECONDS = (1, 2, 4)
1244# HTTP status codes worth retrying (transient server / rate-limit errors).
1245_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})
1247# Upper bound on a honored Retry-After (seconds). RFC 7231 puts no
1248# ceiling on the header, so a hostile or misconfigured upstream could
1249# pin a worker via an arbitrarily large value. Cap here to bound the
1250# damage; legitimate waits (seconds to low minutes) pass through.
1251_MAX_RETRY_AFTER_SECONDS = 300
1254def _parse_retry_after(retry_after_raw: Optional[str]) -> Optional[int]:
1255 """Parse a ``Retry-After`` header value, clamped to ``[0, MAX]``.
1257 Returns ``None`` if the header is missing or unparseable, so the
1258 caller can fall back to the exponential-backoff schedule. Accepts
1259 both RFC 7231 forms: delay-seconds (integer) and HTTP-date. An
1260 HTTP-date that names no zone is read as UTC.
1261 """
1262 if retry_after_raw is None:
1263 return None
1264 try:
1265 seconds = int(retry_after_raw)
1266 except ValueError:
1267 try:
1268 retry_dt = email.utils.parsedate_to_datetime(retry_after_raw)
1269 except (ValueError, TypeError, OverflowError):
1270 # OverflowError: a numeric field too large for datetime().
1271 logger.debug(
1272 f"Unparseable Retry-After {retry_after_raw!r}; "
1273 f"using backoff schedule"
1274 )
1275 return None
1276 if retry_dt.tzinfo is None:
1277 # "-0000", a missing zone and the asctime form parse naive,
1278 # and naive minus aware raises TypeError. An HTTP-date is
1279 # always UTC (RFC 9110 section 5.6.7).
1280 retry_dt = retry_dt.replace(tzinfo=datetime.timezone.utc)
1281 now_utc = datetime.datetime.now(datetime.timezone.utc)
1282 seconds = int((retry_dt - now_utc).total_seconds())
1283 return max(0, min(seconds, _MAX_RETRY_AFTER_SECONDS))
1286def safe_get_with_retries(
1287 url: str,
1288 params: Optional[dict] = None,
1289 timeout: int = DEFAULT_TIMEOUT,
1290 allow_localhost: bool = False,
1291 allow_private_ips: bool = False,
1292 max_retries: int = 3,
1293 backoff_times: tuple = _RETRY_BACKOFF_SECONDS,
1294 consume_body: bool = False,
1295 **kwargs,
1296) -> requests.Response:
1297 """`safe_get` plus exponential-backoff retry on transient errors.
1299 Retries on:
1300 * ``requests.ConnectionError``
1301 * ``requests.Timeout``
1302 * HTTP ``429`` (rate limit) and ``5xx`` (server error)
1303 * (when ``consume_body=True``) body-read failures —
1304 ``ChunkedEncodingError``, ``ReadTimeout``, mid-stream
1305 ``ConnectionError``
1307 Honors the ``Retry-After`` header when present (falls back to the
1308 backoff schedule otherwise). SSRF-validation errors (``ValueError``)
1309 and non-retryable HTTP 4xx responses are not retried.
1311 Without ``consume_body``, ``safe_get`` reads the guarded body before
1312 it returns (unless the caller passes ``stream=True``), so a body-read
1313 failure is raised by ``safe_get`` itself and retried only if it is a
1314 ``ConnectionError`` or ``Timeout``; a ``ChunkedEncodingError`` is
1315 neither and propagates. With ``consume_body=True``, this wrapper
1316 passes ``stream=True`` and reads the body inside the retry loop, so
1317 transient mid-body failures such as ``ChunkedEncodingError`` are also
1318 retried. The cached body is still available via ``response.content``
1319 after return.
1321 Args:
1322 url: Target URL.
1323 params: Query parameters.
1324 timeout: Per-attempt socket timeout.
1325 allow_localhost: Forwarded to ``safe_get``.
1326 allow_private_ips: Forwarded to ``safe_get``.
1327 max_retries: Maximum retry attempts after the initial try.
1328 backoff_times: Per-attempt sleep seconds.
1329 consume_body: If True, read ``response.content`` inside the
1330 retry loop so body-read transients are retried. Use for
1331 large or chunk-transferred bodies (~MB+) where mid-stream
1332 disconnects are realistic. The body-guard's ``ValueError``
1333 (oversized body) is NOT retried — it propagates immediately.
1334 **kwargs: Forwarded to ``safe_get``.
1336 Returns:
1337 The first successful (or final-attempt) ``requests.Response``.
1338 When ``consume_body=True``, the body has already been read and
1339 is cached on the response.
1341 Raises:
1342 ValueError: If SSRF validation fails or the body guard rejects an
1343 oversized response. Retries do not help in either case.
1344 requests.RequestException: If every attempt fails.
1345 """
1346 attempt = 0
1347 if consume_body:
1348 # Keep body-read failures inside this retry loop. Without
1349 # stream=True, safe_get reads the body before returning, and a
1350 # ChunkedEncodingError raised there is neither a ConnectionError
1351 # nor a Timeout, so it would escape the loop unretried.
1352 kwargs["stream"] = True
1353 while True:
1354 try:
1355 response = safe_get(
1356 url,
1357 params=params,
1358 timeout=timeout,
1359 allow_localhost=allow_localhost,
1360 allow_private_ips=allow_private_ips,
1361 **kwargs,
1362 )
1363 except (requests.ConnectionError, requests.Timeout) as exc:
1364 if attempt >= max_retries:
1365 raise
1366 wait = backoff_times[min(attempt, len(backoff_times) - 1)]
1367 logger.warning(
1368 f"{exc.__class__.__name__} on {_redact(url)}; "
1369 f"retrying in {wait}s "
1370 f"(attempt {attempt + 1}/{max_retries})"
1371 )
1372 time.sleep(wait)
1373 attempt += 1
1374 continue
1376 # On the last attempt a retryable status falls through, so with
1377 # consume_body its body is read below like any other response's.
1378 if (
1379 response.status_code in _RETRYABLE_STATUS_CODES
1380 and attempt < max_retries
1381 ):
1382 parsed = _parse_retry_after(response.headers.get("Retry-After"))
1383 wait = (
1384 parsed
1385 if parsed is not None
1386 else backoff_times[min(attempt, len(backoff_times) - 1)]
1387 )
1388 logger.warning(
1389 f"HTTP {response.status_code} on {_redact(url)}; "
1390 f"retrying in {wait}s "
1391 f"(attempt {attempt + 1}/{max_retries})"
1392 )
1393 response.close()
1394 time.sleep(wait)
1395 attempt += 1
1396 continue
1398 if consume_body:
1399 try:
1400 # Force body read while still inside the retry loop.
1401 # ChunkedEncodingError / ReadTimeout / mid-stream
1402 # ConnectionError can fire here on large responses
1403 # from flaky upstreams. ReadTimeout is a Timeout
1404 # subclass but NOT a ConnectionError subclass, so the
1405 # except must list both ConnectionError and Timeout
1406 # — listing Timeout alone would miss ConnectError, and
1407 # listing ConnectionError alone would miss ReadTimeout.
1408 _ = response.content
1409 except (
1410 requests.exceptions.ChunkedEncodingError,
1411 requests.exceptions.ConnectionError,
1412 requests.exceptions.Timeout,
1413 ) as exc:
1414 # safe_close instead of bare close: a close() that
1415 # raises here would mask the original body-read error
1416 # we actually want to surface / retry on.
1417 safe_close(response, "response")
1418 if attempt >= max_retries:
1419 raise
1420 wait = backoff_times[min(attempt, len(backoff_times) - 1)]
1421 logger.warning(
1422 f"{exc.__class__.__name__} reading body of {_redact(url)}; "
1423 f"retrying in {wait}s "
1424 f"(attempt {attempt + 1}/{max_retries})"
1425 )
1426 time.sleep(wait)
1427 attempt += 1
1428 continue
1430 return response