"""Validated upstream requests and bounded redirect reuse, independent of API state.""" import asyncio import re import time from urllib.parse import unquote, urljoin, urlparse import aiohttp ALLOWED_DOWNLOAD_HOSTS = {"huggingface.co", "hf-mirror.com", "hf.co"} ALLOWED_DOWNLOAD_HOST_SUFFIXES = (".hf.co", ".huggingface.co", ".xethub.hf.co") # Redirects are a narrower trust boundary than user-supplied source URLs. Keep # this exact-host list in sync with the signed storage targets we actually use; # accepting every future *.hf.co host would let a redirect escape the dataset # scope checked on the initial URL. TRUSTED_REDIRECT_HOSTS = frozenset({ "cdn-lfs.hf.co", "cdn-lfs.huggingface.co", "cdn-lfs-us-1.hf.co", "cdn-lfs-us-1.huggingface.co", "us.aws.cdn.hf.co", "cas-bridge.xethub.hf.co", "cas-server.xethub.hf.co", "transfer.xethub.hf.co", }) MAX_REDIRECTS = 5 READER_ASSET_PATH_RE = re.compile( r"objects/[0-9a-f]{2}/[0-9a-f]{64}/(?:[0-9a-f]{16}/)?(?:page-manifest\.json|ocr-manifest\.json|pages/page-[0-9]{6}\.(?:webp|jxl)|ocr/(?:page-[0-9]{6}\.json\.gz|book-text\.json\.gz)|(?:[a-z0-9-]+/)?(?:linearized\.pdf|chapter-manifest\.json|document\.(?:pdf|epub|mobi|azw|azw3|fb2|docx|html|txt|md|webp|jpg|jpeg|png|gif|bmp|swf)|book\.epub|audio\.(?:mp3|wav|m4a|flac|mpga)|video\.(?:mp4|mov)|epub-chapters/(?:chapter-manifest\.json|chapters/chapter-[0-9]{4}\.xhtml|resources/[A-Za-z0-9._~%+\-/]+|epub-search-index\.json\.gz)))" ) # Converted Reader assets are served from buckets; the former # vomebook/Reader-Assets dataset is deleted and is not a Reader source. VOICEOFML_READER_SOURCE_RE = re.compile(r"^/datasets/VoiceOfML/[A-Za-z0-9._-]+/(?:resolve|raw)/main/.+$") UPSTREAM_REDIRECT_TTL_SECONDS = 300 UPSTREAM_REDIRECT_MAX_ENTRIES = 2000 # Statuses from a cached signed target that mean it expired or moved; the # cached entry is evicted and the redirect chain is walked once from scratch. STALE_REDIRECT_TARGET_STATUSES = frozenset({401, 403, 404, 410}) upstream_redirect_cache: dict[tuple[str, str], tuple[float, str]] = {} def is_allowed_download_host(hostname: str | None) -> bool: return bool(hostname) and ( hostname in ALLOWED_DOWNLOAD_HOSTS or hostname.endswith(ALLOWED_DOWNLOAD_HOST_SUFFIXES) ) def validate_download_url(url: str) -> str: parsed = urlparse(url) if ( parsed.scheme != "https" or not is_allowed_download_host(parsed.hostname) or parsed.port not in (None, 443) or parsed.username is not None or parsed.password is not None ): raise ValueError(f"不允许的下载跳转地址: {parsed.hostname or 'unknown'}") return url def normalize_download_url(url: str) -> str: parsed = urlparse(url) if parsed.hostname == "hf-mirror.com": return parsed._replace(netloc="huggingface.co").geturl() return url def decode_url_path(path: str) -> str: decoded_path = path for _ in range(8): next_path = unquote(decoded_path) if next_path == decoded_path: return decoded_path decoded_path = next_path raise ValueError("不允许的下载来源路径") def validate_path_safety(path: str) -> str: decoded_path = decode_url_path(path) if ( "\\" in decoded_path or "\x00" in decoded_path or any(part in (".", "..") for part in decoded_path.split("/")) ): raise ValueError("不允许的下载来源路径") return decoded_path def validate_voiceofml_source_url(url: str) -> str: url = normalize_download_url(url) validate_download_url(url) parsed = urlparse(url) decoded_path = validate_path_safety(parsed.path) if ( parsed.hostname != "huggingface.co" or parsed.fragment or parsed.query or not decoded_path.startswith("/datasets/VoiceOfML/") ): raise ValueError("只允许 VoiceOfML 数据集文件") return url def validate_reader_source_url(url: str) -> str: url = normalize_download_url(url) validate_download_url(url) parsed = urlparse(url) if parsed.hostname != "huggingface.co" or parsed.fragment or parsed.query: raise ValueError("不允许的阅读来源") decoded_path = validate_path_safety(parsed.path) if VOICEOFML_READER_SOURCE_RE.fullmatch(decoded_path): return url raise ValueError("不允许的阅读来源") def validate_voiceofml_redirect_url(url: str) -> str: normalized = normalize_download_url(url) validate_download_url(normalized) if urlparse(normalized).hostname == "huggingface.co": parsed = urlparse(normalized) cache_source = reader_cache_source_path(parsed.path) if cache_source is not None and not parsed.fragment: # Hugging Face redirects existing files to commit-pinned cache # paths. Validate the mapped dataset source while preserving the # redirect URL and its cache metadata for the upstream request. validate_voiceofml_source_url("https://huggingface.co" + cache_source) return normalized return validate_voiceofml_source_url(normalized) if urlparse(normalized).hostname not in TRUSTED_REDIRECT_HOSTS: raise ValueError("上游重定向超出允许的存储范围") return normalized def reader_cache_source_path(path: str) -> str | None: """Map HF's commit-pinned cache route back to its dataset source for validation.""" match = re.fullmatch( r"/api/resolve-cache/datasets/([^/]+/[^/]+)/[0-9a-f]{40}/(.+)", validate_path_safety(path), ) if match: return f"/datasets/{match[1]}/resolve/main/{match[2]}" return None def validate_reader_redirect_url(url: str) -> str: normalized = normalize_download_url(url) validate_download_url(normalized) if urlparse(normalized).hostname == "huggingface.co": parsed = urlparse(normalized) cache_source = reader_cache_source_path(parsed.path) if cache_source is not None and not parsed.fragment: # Only redirects may use this route and its HF query metadata. # Apply the same dataset and path restrictions as the source. validate_reader_source_url("https://huggingface.co" + cache_source) return normalized return validate_reader_source_url(normalized) if urlparse(normalized).hostname not in TRUSTED_REDIRECT_HOSTS: raise ValueError("上游重定向超出允许的存储范围") return normalized def validate_fixed_source_redirect_url(url: str, source_url: str) -> str: """Accept only the exact fixed source URL or a known signed-storage target.""" validate_download_url(url) if url == source_url: return url parsed = urlparse(url) if parsed.hostname not in TRUSTED_REDIRECT_HOSTS or parsed.fragment: raise ValueError("上游重定向超出允许的存储范围") return url def source_url_scope(url: str) -> str | None: parsed = urlparse(url) if parsed.hostname != "huggingface.co": return None decoded_path = reader_cache_source_path(parsed.path) or validate_path_safety(parsed.path) if decoded_path.startswith("/datasets/VoiceOfML/"): return "voiceofml" return None def _redirect_cache_get(method: str, validated_url: str, url_validator, initial_scope: str | None) -> str | None: entry = upstream_redirect_cache.get((method, validated_url)) if not entry: return None cached_at, target = entry if time.monotonic() - cached_at >= UPSTREAM_REDIRECT_TTL_SECONDS: upstream_redirect_cache.pop((method, validated_url), None) return None try: checked = url_validator(target) except ValueError: upstream_redirect_cache.pop((method, validated_url), None) return None next_scope = source_url_scope(checked) if initial_scope is not None and next_scope is not None and next_scope != initial_scope: upstream_redirect_cache.pop((method, validated_url), None) return None return checked def _redirect_cache_put(method: str, validated_url: str, target: str) -> None: upstream_redirect_cache[(method, validated_url)] = (time.monotonic(), target) while len(upstream_redirect_cache) > UPSTREAM_REDIRECT_MAX_ENTRIES: oldest = min(upstream_redirect_cache, key=lambda key: upstream_redirect_cache[key][0]) upstream_redirect_cache.pop(oldest, None) async def _request(session, method, url, **options): request = getattr(session, "request", None) if request is not None: return await request(method, url, **options) request = getattr(session, method.lower(), None) if request is None: raise ValueError(f"上游会话不支持 {method} 请求") return await request(url, **options) async def open_download_response( session, url: str, timeout: aiohttp.ClientTimeout, *, method: str = "GET", request_headers: dict[str, str] | None = None, url_validator=validate_download_url, redirect_validator=None, auto_decompress: bool | None = None, cache_redirects: bool = True, ): method = method.upper() if method not in ("GET", "HEAD"): raise ValueError("不支持的上游请求方法") validated_url = url_validator(url) redirect_validator = redirect_validator or url_validator initial_scope = source_url_scope(validated_url) headers = {"Accept-Encoding": "identity"} headers.update(request_headers or {}) request_options = {} if auto_decompress is None else {"auto_decompress": auto_decompress} for attempt in range(2): current_url = validated_url cached = None if attempt == 0 and cache_redirects: cached = _redirect_cache_get(method, validated_url, redirect_validator, initial_scope) if cached is not None: current_url = cached walked = current_url != validated_url # Only a walk that started from a cached signed target may be retried; # a fresh walk's failure is returned or raised as-is. from_cache = cached is not None for _ in range(MAX_REDIRECTS + 1): try: response = await _request( session, method, current_url, allow_redirects=False, timeout=timeout, headers=headers, **request_options, ) except (aiohttp.ClientError, asyncio.TimeoutError): if not from_cache: raise upstream_redirect_cache.pop((method, validated_url), None) break if response.status not in (301, 302, 303, 307, 308): if response.status in STALE_REDIRECT_TARGET_STATUSES and from_cache: upstream_redirect_cache.pop((method, validated_url), None) response.release() break # Reusing a signed target must not renew its fixed lifetime. if (cache_redirects and walked and current_url != cached and response.status in (200, 206, 304)): _redirect_cache_put(method, validated_url, current_url) return response location = response.headers.get("Location") response.release() if not location: raise ValueError("上游重定向缺少地址") next_url = redirect_validator(urljoin(current_url, location)) next_scope = source_url_scope(next_url) if initial_scope is not None and next_scope is not None and next_scope != initial_scope: raise ValueError("上游重定向超出允许的数据集范围") current_url = next_url walked = True else: raise ValueError("上游重定向次数过多")