| 1 | from __future__ import annotations |
| 2 | |
| 3 | from dataclasses import dataclass |
| 4 | import ipaddress |
| 5 | import os |
| 6 | import socket |
| 7 | import struct |
| 8 | from urllib.parse import urljoin, urlparse |
| 9 | |
| 10 | import requests |
| 11 | |
| 12 | |
| 13 | SAFE_HTTP_SCHEMES = frozenset({"http", "https"}) |
| 14 | DEFAULT_FETCH_TIMEOUT = (3.05, 10.0) |
| 15 | DEFAULT_HTTP_USER_AGENT = "@mixedbread-ai/unstructured" |
| 16 | |
| 17 | |
| 18 | @dataclass(frozen=True) |
| 19 | class HttpFetchResult: |
| 20 | url: str |
| 21 | content: bytes |
| 22 | content_type: str | None |
| 23 | encoding: str | None |
| 24 | |
| 25 | |
| 26 | class UnsafeUrlError(ValueError): |
| 27 | """Raised when a remote URL resolves to a non-public destination.""" |
| 28 | |
| 29 | |
| 30 | def _build_request_headers() -> dict[str, str]: |
| 31 | user_agent = ( |
| 32 | os.getenv("USER_AGENT") |
| 33 | or os.getenv("user_agent") |
| 34 | or DEFAULT_HTTP_USER_AGENT |
| 35 | ).strip() |
| 36 | return {"User-Agent": user_agent or DEFAULT_HTTP_USER_AGENT} |
| 37 | |
| 38 | |
| 39 | def _normalize_content_type(content_type: str | None) -> str | None: |
| 40 | if not content_type: |
| 41 | return None |
| 42 | return content_type.split(";", 1)[0].strip().lower() or None |
| 43 | |
| 44 | |
| 45 | def resolve_host_ips(hostname: str) -> tuple[ipaddress._BaseAddress, ...]: |
| 46 | try: |
| 47 | results = socket.getaddrinfo( |
| 48 | hostname, |
| 49 | None, |
| 50 | family=socket.AF_UNSPEC, |
| 51 | type=socket.SOCK_STREAM, |
| 52 | ) |
| 53 | except socket.gaierror as exc: |
| 54 | raise UnsafeUrlError(f"Unable to resolve hostname '{hostname}'") from exc |
| 55 | |
| 56 | ips: list[ipaddress._BaseAddress] = [] |
| 57 | seen: set[str] = set() |
| 58 | for _family, _type, _proto, _canonname, sockaddr in results: |
| 59 | address = sockaddr[0] |
| 60 | if "%" in address: |
| 61 | address = address.split("%", 1)[0] |
| 62 | ip = ipaddress.ip_address(address) |
| 63 | key = ip.compressed |
| 64 | if key in seen: |
| 65 | continue |
| 66 | seen.add(key) |
| 67 | ips.append(ip) |
| 68 | |
| 69 | if not ips: |
| 70 | raise UnsafeUrlError(f"Hostname '{hostname}' did not resolve to an IP address") |
| 71 | |
| 72 | return tuple(ips) |
| 73 | |
| 74 | |
| 75 | def validate_public_http_url(url: str) -> tuple[ipaddress._BaseAddress, ...]: |
| 76 | parsed = urlparse(url) |
| 77 | |
| 78 | if parsed.scheme not in SAFE_HTTP_SCHEMES: |
| 79 | raise UnsafeUrlError("Only http:// and https:// URLs are supported") |
| 80 | if not parsed.hostname: |
| 81 | raise UnsafeUrlError("URL hostname is required") |
| 82 | if parsed.username or parsed.password: |
| 83 | raise UnsafeUrlError("URLs with embedded credentials are not allowed") |
| 84 | |
| 85 | hostname = parsed.hostname.rstrip(".").lower() |
| 86 | if hostname == "localhost" or hostname.endswith(".localhost"): |
| 87 | raise UnsafeUrlError(f"Blocked local hostname '{hostname}'") |
| 88 | |
| 89 | ips = resolve_host_ips(hostname) |
| 90 | blocked = [str(ip) for ip in ips if not ip.is_global] |
| 91 | if blocked: |
| 92 | raise UnsafeUrlError( |
| 93 | f"Blocked non-public address resolution for '{hostname}': {', '.join(blocked)}" |
| 94 | ) |
| 95 | |
| 96 | return ips |
| 97 | |
| 98 | |
| 99 | def fetch_public_http_resource( |
| 100 | url: str, |
| 101 | *, |
| 102 | max_bytes: int, |
| 103 | max_redirects: int = 5, |
| 104 | timeout: tuple[float, float] = DEFAULT_FETCH_TIMEOUT, |
| 105 | ) -> HttpFetchResult: |
| 106 | current_url = url |
| 107 | session = requests.Session() |
| 108 | session.trust_env = False |
| 109 | |
| 110 | for redirect_count in range(max_redirects + 1): |
| 111 | validate_public_http_url(current_url) |
| 112 | |
| 113 | try: |
| 114 | with session.get( |
| 115 | current_url, |
| 116 | stream=True, |
| 117 | allow_redirects=False, |
| 118 | headers=_build_request_headers(), |
| 119 | timeout=timeout, |
| 120 | ) as response: |
| 121 | if 300 <= response.status_code < 400: |
| 122 | location = response.headers.get("Location") |
| 123 | if not location: |
| 124 | raise ValueError( |
| 125 | f"Remote URL redirect is missing a Location header: {current_url}" |
| 126 | ) |
| 127 | if redirect_count >= max_redirects: |
| 128 | raise ValueError( |
| 129 | f"Remote URL exceeded redirect limit ({max_redirects}): {url}" |
| 130 | ) |
| 131 | current_url = urljoin(current_url, location) |
| 132 | continue |
| 133 | |
| 134 | if response.status_code >= 400: |
| 135 | raise ValueError( |
| 136 | f"Remote URL returned HTTP {response.status_code}: {current_url}" |
| 137 | ) |
| 138 | |
| 139 | content_length = response.headers.get("Content-Length") |
| 140 | if content_length: |
| 141 | try: |
| 142 | declared_length = int(content_length) |
| 143 | except ValueError: |
| 144 | declared_length = None |
| 145 | if declared_length is not None and declared_length > max_bytes: |
| 146 | raise ValueError( |
| 147 | f"Remote document exceeds max size {max_bytes} bytes: {current_url}" |
| 148 | ) |
| 149 | |
| 150 | body = bytearray() |
| 151 | for chunk in response.iter_content(chunk_size=64 * 1024): |
| 152 | if not chunk: |
| 153 | continue |
| 154 | body.extend(chunk) |
| 155 | if len(body) > max_bytes: |
| 156 | raise ValueError( |
| 157 | f"Remote document exceeds max size {max_bytes} bytes: {current_url}" |
| 158 | ) |
| 159 | |
| 160 | return HttpFetchResult( |
| 161 | url=current_url, |
| 162 | content=bytes(body), |
| 163 | content_type=_normalize_content_type( |
| 164 | response.headers.get("Content-Type") |
| 165 | ), |
| 166 | encoding=response.encoding, |
| 167 | ) |
| 168 | except requests.RequestException as exc: |
| 169 | raise ValueError( |
| 170 | f"Remote document fetch failed for {current_url}: {exc}" |
| 171 | ) from exc |
| 172 | |
| 173 | raise ValueError(f"Remote URL exceeded redirect limit ({max_redirects}): {url}") |
| 174 | |
| 175 | |
| 176 | def is_loopback_address(address: str) -> bool: |
| 177 | """Check whether *address* resolves to a loopback interface.""" |
| 178 | _checkers = { |
| 179 | socket.AF_INET: lambda x: ( |
| 180 | struct.unpack("!I", socket.inet_aton(x))[0] >> (32 - 8) |
| 181 | ) == 127, |
| 182 | socket.AF_INET6: lambda x: x == "::1", |
| 183 | } |
| 184 | try: |
| 185 | socket.inet_pton(socket.AF_INET6, address) |
| 186 | return _checkers[socket.AF_INET6](address) |
| 187 | except socket.error: |
| 188 | pass |
| 189 | try: |
| 190 | socket.inet_pton(socket.AF_INET, address) |
| 191 | return _checkers[socket.AF_INET](address) |
| 192 | except socket.error: |
| 193 | pass |
| 194 | for family in (socket.AF_INET, socket.AF_INET6): |
| 195 | try: |
| 196 | r = socket.getaddrinfo(address, None, family, socket.SOCK_STREAM) |
| 197 | except socket.gaierror: |
| 198 | return False |
| 199 | for fam, _, _, _, sockaddr in r: |
| 200 | if not _checkers[fam](sockaddr[0]): |
| 201 | return False |
| 202 | return True |