Coverage for src/local_deep_research/security/safe_requests.py: 98%
246 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2Safe HTTP Requests Wrapper
4Wraps requests library to add SSRF protection and security best practices.
5"""
7import datetime
8import email.utils
9import time
10from typing import Any, Optional
11from urllib.parse import urljoin
13import requests
14from loguru import logger
16from . import dns_pinning
17from . import ssrf_validator
18from ..constants import USER_AGENT
19from ..utilities.resource_utils import safe_close
22# Default timeout for all HTTP requests (prevents hanging)
23DEFAULT_TIMEOUT = 30 # seconds
25# Maximum response size to prevent memory exhaustion (1GB)
26# Set high to accommodate large documents (annual reports, PDFs, datasets).
27# This is a local research tool — users intentionally download these files.
28MAX_RESPONSE_SIZE = 1024 * 1024 * 1024
30# HTTP status codes that indicate a redirect
31_REDIRECT_STATUS_CODES = frozenset({301, 302, 303, 307, 308})
33# Maximum number of redirects to follow
34_MAX_REDIRECTS = 10
36# Prefix of the requests library's default User-Agent
37# (e.g. "python-requests/2.32.5"). Used by `SafeSession.request` to detect
38# whether the session is using the upstream default — if so, we override
39# it with the project `USER_AGENT` so academic API endpoints (arXiv,
40# OpenAlex, PubMed, …) can identify us. Promoting the literal to a
41# constant so a future requests rename only requires a one-line edit.
42_DEFAULT_REQUESTS_UA_PREFIX = "python-requests"
45def _install_body_guard(response: requests.Response) -> None:
46 """Install a bounded reader that enforces MAX_RESPONSE_SIZE.
48 Wraps response.raw.read() to track cumulative bytes and raise
49 ValueError if MAX_RESPONSE_SIZE is exceeded during body consumption.
50 This transparently protects both streamed (.iter_content) and
51 non-streamed (.text, .json(), .content) access patterns.
53 Always installs — callers (currently only _check_response_size)
54 are responsible for deciding when to call this function.
55 """
56 original_read = response.raw.read
57 bytes_read = 0
59 def bounded_read(amt=None, *args, **kwargs):
60 nonlocal bytes_read
61 data = original_read(amt, *args, **kwargs)
62 bytes_read += len(data)
63 if bytes_read > MAX_RESPONSE_SIZE:
64 response.close()
65 raise ValueError(
66 f"Response body too large: >{bytes_read} bytes "
67 f"(max {MAX_RESPONSE_SIZE}, Content-Length absent or invalid)"
68 )
69 return data
71 response.raw.read = bounded_read # type: ignore[method-assign]
74def _check_response_size(response: requests.Response) -> None:
75 """Reject responses whose Content-Length exceeds MAX_RESPONSE_SIZE.
77 Handles comma-separated values per RFC 7230 §3.3.2: identical
78 duplicates (from proxies) are normalized; differing values are
79 rejected as invalid framing. Empty parts from malformed headers
80 (trailing/doubled commas) are filtered before parsing. Non-integer
81 or negative values cause the header to be treated as absent.
83 When Content-Length is absent, unparseable, negative, or consists
84 only of commas/whitespace, installs a body guard that enforces
85 the size limit during body consumption.
87 Must be called before returning a response to the caller. On
88 rejection the response is closed to avoid leaking the connection.
90 Raises:
91 ValueError: If Content-Length values conflict or exceed
92 MAX_RESPONSE_SIZE.
93 """
94 content_length = response.headers.get("Content-Length")
95 if content_length:
96 try:
97 # Handle comma-separated Content-Length values (RFC 7230 §3.3.2).
98 # Multiple identical values may be sent by proxies; differing
99 # values indicate invalid framing and must be rejected.
100 raw_parts = [v.strip() for v in content_length.split(",")]
101 parts = [p for p in raw_parts if p]
102 if not parts:
103 _install_body_guard(response)
104 return # Only commas/whitespace — treat as absent
105 sizes = [int(p) for p in parts]
106 except (ValueError, TypeError):
107 _install_body_guard(response)
108 return # Content-Length not a valid number
109 if len(set(sizes)) > 1:
110 response.close()
111 raise ValueError(
112 f"Conflicting Content-Length values: {content_length}"
113 )
114 size = sizes[0]
115 if size < 0:
116 _install_body_guard(response)
117 return # Malformed Content-Length, treat as absent
118 if size > MAX_RESPONSE_SIZE:
119 response.close()
120 raise ValueError(
121 f"Response too large: {size} bytes (max {MAX_RESPONSE_SIZE})"
122 )
123 # Valid Content-Length within limit — no body guard needed
124 return
126 # No Content-Length header at all — install body guard
127 _install_body_guard(response)
130def _resolve_redirect_method(method: str, status_code: int) -> str:
131 """Determine HTTP method after redirect, per RFC 7231."""
132 if status_code == 303 and method != "HEAD":
133 method = "GET"
134 elif status_code == 302 and method == "POST":
135 method = "GET"
136 elif status_code == 301 and method == "POST":
137 method = "GET"
138 # 307, 308: preserve original method (no change needed)
139 return method
142# Credential headers: the HTTP-standard set, plus the vendor key/token family.
143# The suffixes cover the shapes vendors ship (X-Api-Key, Zotero-API-Key,
144# X-Goog-Api-Key, X-Subscription-Token, Ocp-Apim-Subscription-Key, Private-Token),
145# so a new vendor needs no edit here. No non-credential header this project
146# sends ends in one of them.
147_CREDENTIAL_HEADERS = frozenset(
148 {"authorization", "proxy-authorization", "cookie"}
149)
150_CREDENTIAL_HEADER_SUFFIXES = ("key", "token", "secret")
152# requests owns the "does this hop leave the credential's scope" rule, so
153# reuse it rather than restate it in the hand-rolled loops below.
154_redirect_scope = requests.sessions.SessionRedirectMixin()
157def _is_credential_header(name: str) -> bool:
158 """Whether a request header carries a caller credential."""
159 lowered = name.lower()
160 return lowered in _CREDENTIAL_HEADERS or lowered.endswith(
161 _CREDENTIAL_HEADER_SUFFIXES
162 )
165def _headers_for_redirect(
166 headers: Optional[dict], from_url: str, to_url: str
167) -> Optional[dict]:
168 """Return the headers to send on a redirect hop.
170 Credential headers are dropped when the hop leaves the scope they were
171 issued for, so a key stays with the host the caller addressed. Returns a
172 new dict when it strips, so the caller's own headers object is untouched.
173 A chain that leaves scope and later returns to the original host does not
174 get the credential back, which is what requests does on its own hops.
175 """
176 if not headers or not _redirect_scope.should_strip_auth(from_url, to_url):
177 return headers
178 return {k: v for k, v in headers.items() if not _is_credential_header(k)}
181def safe_get(
182 url: str,
183 params: Optional[dict] = None,
184 timeout: int = DEFAULT_TIMEOUT,
185 allow_localhost: bool = False,
186 allow_private_ips: bool = False,
187 **kwargs,
188) -> requests.Response:
189 """
190 Make a safe HTTP GET request with SSRF protection.
192 Args:
193 url: URL to request
194 params: URL parameters
195 timeout: Request timeout in seconds
196 allow_localhost: Whether to allow localhost/loopback addresses.
197 Set to True for trusted internal services like self-hosted
198 search engines (e.g., searxng). Default False.
199 allow_private_ips: Whether to allow all private/internal IPs plus localhost.
200 This includes RFC1918 (10.x, 172.16-31.x, 192.168.x), CGNAT (100.64.x.x
201 used by Podman/rootless containers), link-local (169.254.x.x), and IPv6
202 private ranges (fc00::/7, fe80::/10). Use for trusted self-hosted services
203 like SearXNG or Ollama in containerized environments.
204 Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
205 AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
206 ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
207 **kwargs: Additional arguments to pass to requests.get()
209 Returns:
210 Response object
212 Raises:
213 ValueError: If URL fails SSRF validation
214 requests.RequestException: If request fails
215 """
216 # Validate URL to prevent SSRF
217 if not ssrf_validator.validate_url(
218 url,
219 allow_localhost=allow_localhost,
220 allow_private_ips=allow_private_ips,
221 ):
222 raise ValueError(
223 f"URL failed security validation (possible SSRF): {url}"
224 )
226 # Ensure timeout is set
227 if "timeout" not in kwargs: 227 ↛ 233line 227 didn't jump to line 233 because the condition on line 227 was always true
228 kwargs["timeout"] = timeout
230 # Inject the project User-Agent if the caller didn't supply one.
231 # Mutates a copy of any caller-supplied headers dict so we never
232 # touch their object.
233 headers = dict(kwargs.get("headers") or {})
234 if not any(k.lower() == "user-agent" for k in headers):
235 headers["User-Agent"] = USER_AGENT
236 kwargs["headers"] = headers
238 # Intercept allow_redirects — we handle redirects manually to validate
239 # each redirect target against SSRF rules
240 caller_wants_redirects = kwargs.pop("allow_redirects", True)
241 kwargs["allow_redirects"] = False
243 current_url = url
244 try:
245 # Pin the validated IP for the actual connection: the request below
246 # connects to the address resolved+validated here, closing the
247 # resolve-vs-connect gap where requests/urllib3 would otherwise
248 # re-resolve the hostname independently.
249 with dns_pinning.pinned_request(
250 url, allow_localhost, allow_private_ips
251 ):
252 response = requests.get(url, params=params, **kwargs)
254 # Follow redirects manually with SSRF validation on each hop.
255 # Each hop uses a fresh requests.get() call without a session,
256 # so cookies set by intermediate responses are not carried
257 # forward. This is acceptable for current callers (all stateless).
258 # Callers needing cookie persistence across redirects should use
259 # SafeSession instead, which preserves cookies via its cookie jar.
260 if caller_wants_redirects:
261 redirects_followed = 0
262 while (
263 response.status_code in _REDIRECT_STATUS_CODES
264 and redirects_followed < _MAX_REDIRECTS
265 ):
266 redirect_url = (response.headers.get("Location") or "").strip()
267 if not redirect_url:
268 break
270 # Resolve relative redirects
271 redirect_url = urljoin(
272 response.url or current_url, redirect_url
273 )
275 # Validate redirect target against SSRF rules
276 if not ssrf_validator.validate_url(
277 redirect_url,
278 allow_localhost=allow_localhost,
279 allow_private_ips=allow_private_ips,
280 ):
281 logger.warning(
282 f"Redirect to {redirect_url} blocked by SSRF validation "
283 f"(from {url}, hop {redirects_followed + 1})"
284 )
285 response.close()
286 raise ValueError(
287 f"Redirect target failed SSRF validation: {redirect_url}"
288 )
290 # Credentials do not follow the hop out of their own scope.
291 kwargs["headers"] = _headers_for_redirect(
292 kwargs.get("headers"), current_url, redirect_url
293 )
294 current_url = redirect_url
295 response.close()
296 # Note: params are intentionally NOT forwarded to redirect
297 # hops. Per HTTP spec, the server's Location header contains
298 # the complete target URL. Re-appending original query params
299 # would corrupt it.
300 # Re-pin per hop: the address validated for this redirect
301 # target is the address connected to.
302 with dns_pinning.pinned_request(
303 redirect_url, allow_localhost, allow_private_ips
304 ):
305 response = requests.get(redirect_url, **kwargs)
306 redirects_followed += 1
308 if (
309 response.status_code in _REDIRECT_STATUS_CODES
310 and redirects_followed >= _MAX_REDIRECTS
311 ):
312 response.close()
313 # Note: raises ValueError here, while SafeSession raises
314 # requests.TooManyRedirects (delegated to the base class).
315 # Callers should catch ValueError for standalone functions.
316 raise ValueError(
317 f"Too many redirects ({_MAX_REDIRECTS}) from {url}"
318 )
320 _check_response_size(response)
322 return response
324 except requests.Timeout:
325 logger.warning(f"Request timeout after {timeout}s: {current_url}")
326 raise
327 except requests.RequestException:
328 logger.warning(f"Request failed for {current_url}")
329 raise
332def safe_post(
333 url: str,
334 data: Optional[Any] = None,
335 json: Optional[dict] = None,
336 timeout: int = DEFAULT_TIMEOUT,
337 allow_localhost: bool = False,
338 allow_private_ips: bool = False,
339 **kwargs,
340) -> requests.Response:
341 """
342 Make a safe HTTP POST request with SSRF protection.
344 Args:
345 url: URL to request
346 data: Data to send in request body
347 json: JSON data to send in request body
348 timeout: Request timeout in seconds
349 allow_localhost: Whether to allow localhost/loopback addresses.
350 Set to True for trusted internal services like self-hosted
351 search engines (e.g., searxng). Default False.
352 allow_private_ips: Whether to allow all private/internal IPs plus localhost.
353 This includes RFC1918 (10.x, 172.16-31.x, 192.168.x), CGNAT (100.64.x.x
354 used by Podman/rootless containers), link-local (169.254.x.x), and IPv6
355 private ranges (fc00::/7, fe80::/10). Use for trusted self-hosted services
356 like SearXNG or Ollama in containerized environments.
357 Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
358 AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
359 ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
360 **kwargs: Additional arguments to pass to requests.post()
362 Returns:
363 Response object
365 Raises:
366 ValueError: If URL fails SSRF validation
367 requests.RequestException: If request fails
368 """
369 # Validate URL to prevent SSRF
370 if not ssrf_validator.validate_url(
371 url,
372 allow_localhost=allow_localhost,
373 allow_private_ips=allow_private_ips,
374 ):
375 raise ValueError(
376 f"URL failed security validation (possible SSRF): {url}"
377 )
379 # Ensure timeout is set
380 if "timeout" not in kwargs: 380 ↛ 386line 380 didn't jump to line 386 because the condition on line 380 was always true
381 kwargs["timeout"] = timeout
383 # Inject the project User-Agent if the caller didn't supply one.
384 # Mutates a copy of any caller-supplied headers dict so we never
385 # touch their object.
386 headers = dict(kwargs.get("headers") or {})
387 if not any(k.lower() == "user-agent" for k in headers):
388 headers["User-Agent"] = USER_AGENT
389 kwargs["headers"] = headers
391 # Intercept allow_redirects — we handle redirects manually to validate
392 # each redirect target against SSRF rules
393 caller_wants_redirects = kwargs.pop("allow_redirects", True)
394 kwargs["allow_redirects"] = False
396 current_url = url
397 try:
398 # Pin the validated IP for the actual connection (see safe_get).
399 with dns_pinning.pinned_request(
400 url, allow_localhost, allow_private_ips
401 ):
402 response = requests.post(url, data=data, json=json, **kwargs)
404 # Follow redirects manually with SSRF validation on each hop.
405 # Each hop uses a fresh request without a session, so cookies
406 # set by intermediate responses are not carried forward. Callers
407 # needing cookie persistence should use SafeSession instead.
408 if caller_wants_redirects:
409 redirect_method = "POST"
410 redirects_followed = 0
411 while (
412 response.status_code in _REDIRECT_STATUS_CODES
413 and redirects_followed < _MAX_REDIRECTS
414 ):
415 redirect_url = (response.headers.get("Location") or "").strip()
416 if not redirect_url:
417 break
419 # Resolve relative redirects
420 redirect_url = urljoin(
421 response.url or current_url, redirect_url
422 )
424 # Validate redirect target against SSRF rules
425 if not ssrf_validator.validate_url(
426 redirect_url,
427 allow_localhost=allow_localhost,
428 allow_private_ips=allow_private_ips,
429 ):
430 logger.warning(
431 f"Redirect to {redirect_url} blocked by SSRF validation "
432 f"(from {url}, hop {redirects_followed + 1})"
433 )
434 response.close()
435 raise ValueError(
436 f"Redirect target failed SSRF validation: {redirect_url}"
437 )
439 redirect_method = _resolve_redirect_method(
440 redirect_method, response.status_code
441 )
442 # Credentials do not follow the hop out of their own scope.
443 kwargs["headers"] = _headers_for_redirect(
444 kwargs.get("headers"), current_url, redirect_url
445 )
446 current_url = redirect_url
447 response.close()
449 # Re-pin per hop: the address validated for this redirect
450 # target is the address connected to.
451 with dns_pinning.pinned_request(
452 redirect_url, allow_localhost, allow_private_ips
453 ):
454 if redirect_method == "GET":
455 # 301/302/303: convert to GET, drop body
456 data = None
457 json = None
458 response = requests.get(redirect_url, **kwargs)
459 else:
460 # 307/308: preserve current method and body
461 response = requests.post(
462 redirect_url, data=data, json=json, **kwargs
463 )
464 redirects_followed += 1
466 if (
467 response.status_code in _REDIRECT_STATUS_CODES
468 and redirects_followed >= _MAX_REDIRECTS
469 ):
470 response.close()
471 # Note: raises ValueError here, while SafeSession raises
472 # requests.TooManyRedirects (delegated to the base class).
473 # Callers should catch ValueError for standalone functions.
474 raise ValueError(
475 f"Too many redirects ({_MAX_REDIRECTS}) from {url}"
476 )
478 _check_response_size(response)
480 return response
482 except requests.Timeout:
483 logger.warning(f"Request timeout after {timeout}s: {current_url}")
484 raise
485 except requests.RequestException:
486 logger.warning(f"Request failed for {current_url}")
487 raise
490# Create a safe session class
491class SafeSession(requests.Session):
492 """
493 Session with built-in SSRF protection.
495 Redirect validation relies on ``requests.Session.resolve_redirects()``
496 calling ``self.send()`` for each hop — an internal implementation detail
497 of the ``requests`` library. This is simpler than re-implementing the
498 redirect loop (as ``safe_get``/``safe_post`` do) and keeps session-level
499 features (cookies, auth) working. The trade-off is coupling to the
500 ``requests`` internals; if a future version stops routing hops through
501 ``send()``, redirect targets would no longer be validated.
503 Usage:
504 with SafeSession() as session:
505 response = session.get(url)
507 # For trusted internal services (e.g., searxng on localhost):
508 with SafeSession(allow_localhost=True) as session:
509 response = session.get(url)
511 # For trusted internal services on any private network IP:
512 with SafeSession(allow_private_ips=True) as session:
513 response = session.get(url)
515 Raises:
516 ValueError: If a URL (initial or redirect target) fails SSRF
517 validation, or if the response Content-Length exceeds
518 MAX_RESPONSE_SIZE. Note: ``safe_get``/``safe_post`` also raise
519 ``ValueError`` for too-many-redirects, but ``SafeSession`` raises
520 ``requests.TooManyRedirects`` for that case since it delegates
521 redirect counting to the ``requests`` library.
522 requests.RequestException: On transport-level failures.
523 """
525 def __init__(
526 self, allow_localhost: bool = False, allow_private_ips: bool = False
527 ):
528 """
529 Initialize SafeSession.
531 Args:
532 allow_localhost: Whether to allow localhost/loopback addresses.
533 allow_private_ips: Whether to allow all private/internal IPs plus localhost.
534 This includes RFC1918, CGNAT (100.64.x.x used by Podman), link-local, and
535 IPv6 private ranges. Use for trusted self-hosted services like SearXNG or
536 Ollama in containerized environments.
537 Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
538 AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
539 ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
540 """
541 super().__init__()
542 self.max_redirects = _MAX_REDIRECTS
543 self.allow_localhost = allow_localhost
544 self.allow_private_ips = allow_private_ips
546 def rebuild_auth(
547 self, prepared_request: requests.PreparedRequest, response
548 ) -> None:
549 """Drop vendor API-key headers too when a hop leaves auth scope.
551 The base implementation removes ``Authorization`` and nothing else, so
552 a session default such as the Semantic Scholar engine's ``x-api-key``
553 would otherwise be replayed to the redirect target. The strip runs
554 before ``super()`` because the base method reapplies netrc credentials
555 for the new host, and those belong to it.
556 """
557 if self.should_strip_auth(response.request.url, prepared_request.url):
558 for name in [
559 header
560 for header in prepared_request.headers
561 if _is_credential_header(header)
562 ]:
563 del prepared_request.headers[name]
564 super().rebuild_auth(prepared_request, response)
566 def request(self, method: str, url: str, **kwargs) -> requests.Response: # type: ignore[override]
567 """Override request method to add SSRF validation."""
568 # Validate URL
569 if not ssrf_validator.validate_url(
570 url,
571 allow_localhost=self.allow_localhost,
572 allow_private_ips=self.allow_private_ips,
573 ):
574 raise ValueError(
575 f"URL failed security validation (possible SSRF): {url}"
576 )
578 # Ensure timeout is set
579 if "timeout" not in kwargs:
580 kwargs["timeout"] = DEFAULT_TIMEOUT
582 # Inject project User-Agent if the caller didn't already set one.
583 # Session-level User-Agent (self.headers) is left alone — only
584 # per-request headers are copied so we never mutate caller state.
585 headers = dict(kwargs.get("headers") or {})
586 session_ua = self.headers.get("User-Agent", "")
587 has_per_request_ua = any(k.lower() == "user-agent" for k in headers)
588 if not has_per_request_ua and (
589 not session_ua or session_ua.startswith(_DEFAULT_REQUESTS_UA_PREFIX)
590 ):
591 headers["User-Agent"] = USER_AGENT
592 kwargs["headers"] = headers
594 return super().request(method, url, **kwargs)
596 def send(
597 self, request: requests.PreparedRequest, **kwargs
598 ) -> requests.Response:
599 """Override send to validate every outgoing request against SSRF.
601 This runs on **all** calls — both the initial request (routed
602 here by ``requests.Session.request()``) and each redirect hop
603 (routed here by ``resolve_redirects()``). The initial URL is
604 therefore validated twice (once in ``request()``, once here);
605 this is intentional defense-in-depth.
606 """
607 if request.url and not ssrf_validator.validate_url(
608 request.url,
609 allow_localhost=self.allow_localhost,
610 allow_private_ips=self.allow_private_ips,
611 ):
612 logger.warning(
613 f"Request to {request.url} blocked by SSRF validation"
614 )
615 # Note: This error says "security validation" while safe_get/
616 # safe_post say "SSRF validation". The difference indicates the
617 # source (session vs standalone function) in logs.
618 raise ValueError(
619 f"Redirect target failed security validation (possible SSRF): {request.url}"
620 )
622 # Pin the validated IP for the connection this send performs. send()
623 # runs for the initial request AND for every redirect hop
624 # (resolve_redirects re-enters here), so each hop is independently
625 # resolved, validated, and pinned — the address validated is the
626 # address connected to.
627 with dns_pinning.pinned_request(
628 request.url or "",
629 allow_localhost=self.allow_localhost,
630 allow_private_ips=self.allow_private_ips,
631 ):
632 response = super().send(request, **kwargs)
633 _check_response_size(response)
634 return response
637# Exponential backoff schedule (seconds). Kept short: journal-quality
638# downloads are run from a user request or a scheduled job, not from a
639# time-sensitive hot path, so three retries over ~7 seconds is plenty
640# without adding real latency.
641_RETRY_BACKOFF_SECONDS = (1, 2, 4)
643# HTTP status codes worth retrying (transient server / rate-limit errors).
644_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})
646# Upper bound on a honored Retry-After (seconds). RFC 7231 puts no
647# ceiling on the header, so a hostile or misconfigured upstream could
648# pin a worker via an arbitrarily large value. Cap here to bound the
649# damage; legitimate waits (seconds to low minutes) pass through.
650_MAX_RETRY_AFTER_SECONDS = 300
653def _parse_retry_after(retry_after_raw: Optional[str]) -> Optional[int]:
654 """Parse a ``Retry-After`` header value, clamped to ``[0, MAX]``.
656 Returns ``None`` if the header is missing or unparseable, so the
657 caller can fall back to the exponential-backoff schedule. Accepts
658 both RFC 7231 forms: delay-seconds (integer) and HTTP-date.
659 """
660 if retry_after_raw is None:
661 return None
662 try:
663 seconds = int(retry_after_raw)
664 except ValueError:
665 try:
666 retry_dt = email.utils.parsedate_to_datetime(retry_after_raw)
667 except (ValueError, TypeError):
668 logger.debug(
669 f"Unparseable Retry-After {retry_after_raw!r}; "
670 f"using backoff schedule"
671 )
672 return None
673 now_utc = datetime.datetime.now(datetime.timezone.utc)
674 seconds = int((retry_dt - now_utc).total_seconds())
675 return max(0, min(seconds, _MAX_RETRY_AFTER_SECONDS))
678def safe_get_with_retries(
679 url: str,
680 params: Optional[dict] = None,
681 timeout: int = DEFAULT_TIMEOUT,
682 allow_localhost: bool = False,
683 allow_private_ips: bool = False,
684 max_retries: int = 3,
685 backoff_times: tuple = _RETRY_BACKOFF_SECONDS,
686 consume_body: bool = False,
687 **kwargs,
688) -> requests.Response:
689 """`safe_get` plus exponential-backoff retry on transient errors.
691 Retries on:
692 * ``requests.ConnectionError``
693 * ``requests.Timeout``
694 * HTTP ``429`` (rate limit) and ``5xx`` (server error)
695 * (when ``consume_body=True``) body-read failures —
696 ``ChunkedEncodingError``, ``ReadTimeout``, mid-stream
697 ``ConnectionError``
699 Honors the ``Retry-After`` header when present (falls back to the
700 backoff schedule otherwise). SSRF-validation errors (``ValueError``)
701 and non-retryable HTTP 4xx responses are not retried.
703 Without ``consume_body``, only failures raised inside ``safe_get``
704 itself (DNS, connect, header timeout, retryable status) trigger a
705 retry. The body isn't read until the caller touches ``.content`` /
706 ``.text`` / ``.json()``, by which point this wrapper has already
707 returned — so a mid-stream S3 hiccup (``ChunkedEncodingError``)
708 propagates uncaught. ``consume_body=True`` reads the body inside
709 the retry loop so those transient body-read failures are also
710 retried. The cached body is still available to the caller via
711 ``response.content`` after the wrapper returns.
713 Args:
714 url: Target URL.
715 params: Query parameters.
716 timeout: Per-attempt socket timeout.
717 allow_localhost: Forwarded to ``safe_get``.
718 allow_private_ips: Forwarded to ``safe_get``.
719 max_retries: Maximum retry attempts after the initial try.
720 backoff_times: Per-attempt sleep seconds.
721 consume_body: If True, read ``response.content`` inside the
722 retry loop so body-read transients are retried. Use for
723 large or chunk-transferred bodies (~MB+) where mid-stream
724 disconnects are realistic. The body-guard's ``ValueError``
725 (oversized body) is NOT retried — it propagates immediately.
726 **kwargs: Forwarded to ``safe_get``.
728 Returns:
729 The first successful (or final-attempt) ``requests.Response``.
730 When ``consume_body=True``, the body has already been read and
731 is cached on the response.
733 Raises:
734 ValueError: If SSRF validation fails or, with
735 ``consume_body=True``, the body-guard rejects an oversized
736 response. Retries do not help in either case.
737 requests.RequestException: If every attempt fails.
738 """
739 attempt = 0
740 while True:
741 try:
742 response = safe_get(
743 url,
744 params=params,
745 timeout=timeout,
746 allow_localhost=allow_localhost,
747 allow_private_ips=allow_private_ips,
748 **kwargs,
749 )
750 except (requests.ConnectionError, requests.Timeout) as exc:
751 if attempt >= max_retries:
752 raise
753 wait = backoff_times[min(attempt, len(backoff_times) - 1)]
754 logger.warning(
755 f"{exc.__class__.__name__} on {url}; "
756 f"retrying in {wait}s "
757 f"(attempt {attempt + 1}/{max_retries})"
758 )
759 time.sleep(wait)
760 attempt += 1
761 continue
763 if response.status_code in _RETRYABLE_STATUS_CODES:
764 if attempt >= max_retries: 764 ↛ 765line 764 didn't jump to line 765 because the condition on line 764 was never true
765 return response
766 parsed = _parse_retry_after(response.headers.get("Retry-After"))
767 wait = (
768 parsed
769 if parsed is not None
770 else backoff_times[min(attempt, len(backoff_times) - 1)]
771 )
772 logger.warning(
773 f"HTTP {response.status_code} on {url}; "
774 f"retrying in {wait}s "
775 f"(attempt {attempt + 1}/{max_retries})"
776 )
777 response.close()
778 time.sleep(wait)
779 attempt += 1
780 continue
782 if consume_body:
783 try:
784 # Force body read while still inside the retry loop.
785 # ChunkedEncodingError / ReadTimeout / mid-stream
786 # ConnectionError can fire here on large responses
787 # from flaky upstreams. ReadTimeout is a Timeout
788 # subclass but NOT a ConnectionError subclass, so the
789 # except must list both ConnectionError and Timeout
790 # — listing Timeout alone would miss ConnectError, and
791 # listing ConnectionError alone would miss ReadTimeout.
792 _ = response.content
793 except (
794 requests.exceptions.ChunkedEncodingError,
795 requests.exceptions.ConnectionError,
796 requests.exceptions.Timeout,
797 ) as exc:
798 # safe_close instead of bare close: a close() that
799 # raises here would mask the original body-read error
800 # we actually want to surface / retry on.
801 safe_close(response, "response")
802 if attempt >= max_retries:
803 raise
804 wait = backoff_times[min(attempt, len(backoff_times) - 1)]
805 logger.warning(
806 f"{exc.__class__.__name__} reading body of {url}; "
807 f"retrying in {wait}s "
808 f"(attempt {attempt + 1}/{max_retries})"
809 )
810 time.sleep(wait)
811 attempt += 1
812 continue
814 return response