Download tests/test_async_proxy.py from VoiceOfML/Search: direct link, hf CLI and curl.
- Browser
- Download file 87.4 kB
-
https://huggingface.co/spaces/VoiceOfML/Search/resolve/main/tests/test_async_proxy.py
- Command line
-
hf download hf://spaces/VoiceOfML/Search/tests/test_async_proxy.py
-
curl -L -o test_async_proxy.py https://huggingface.co/spaces/VoiceOfML/Search/resolve/main/tests/test_async_proxy.py
87.4 kB
| from __future__ import annotations | |
| import asyncio | |
| import copy | |
| import gzip | |
| import json | |
| import tempfile | |
| import threading | |
| import time | |
| import unittest | |
| from collections import OrderedDict | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| from unittest.mock import AsyncMock, Mock, patch | |
| IMPORT_ERROR = None | |
| try: | |
| import aiohttp | |
| from backend import app | |
| from backend import upstream_http | |
| except ModuleNotFoundError as exc: # pragma: no cover | |
| app = None | |
| IMPORT_ERROR = exc | |
| class FakeResponse: | |
| def __init__(self, status=200, headers=None, chunks=(b"ok",)): | |
| self.status = status | |
| self.headers = headers or {} | |
| self.content = self | |
| self.chunks = chunks | |
| self.released = False | |
| def release(self): | |
| self.released = True | |
| async def read(self, amount=-1): | |
| return b"".join(self.chunks)[:amount if amount >= 0 else None] | |
| def iter_chunked(self, size): | |
| async def generator(): | |
| for chunk in self.chunks: | |
| yield chunk | |
| return generator() | |
| async def __aenter__(self): | |
| return self | |
| async def __aexit__(self, *args): | |
| self.release() | |
| class DownloadConnectionWarmTests(unittest.IsolatedAsyncioTestCase): | |
| async def asyncSetUp(self): | |
| self.response = FakeResponse() | |
| self.state = SimpleNamespace( | |
| download_warm_enabled=True, download_warm_task=None, | |
| download_warm_at=float("-inf"), upstream_semaphore=asyncio.Semaphore(2), | |
| http_session=SimpleNamespace(head=Mock(return_value=self.response)), | |
| ) | |
| self.application = SimpleNamespace(state=self.state) | |
| async def asyncTearDown(self): | |
| if self.state.download_warm_task is not None: | |
| self.state.download_warm_task.cancel() | |
| await asyncio.gather(self.state.download_warm_task, return_exceptions=True) | |
| async def test_ping_returns_before_warmup_and_coalesces_repeated_visits(self): | |
| with patch.object(app, "app", self.application): | |
| first = await app.api_ping() | |
| task = self.state.download_warm_task | |
| second = await app.api_ping() | |
| self.assertEqual([first.status_code, second.status_code], [204, 204]) | |
| self.assertEqual(first.headers["cache-control"], "no-store") | |
| self.state.http_session.head.assert_not_called() | |
| self.assertIs(self.state.download_warm_task, task) | |
| await task | |
| self.state.http_session.head.assert_called_once() | |
| args, kwargs = self.state.http_session.head.call_args | |
| self.assertEqual(args, ("https://huggingface.co/robots.txt",)) | |
| self.assertFalse(kwargs["allow_redirects"]) | |
| self.assertEqual(kwargs["timeout"].total, 4) | |
| self.assertTrue(self.response.released) | |
| app.schedule_download_connection_warm(self.application) | |
| self.assertIs(self.state.download_warm_task, task) | |
| self.assertEqual(self.state.upstream_semaphore._value, 2) | |
| async def test_failed_warmup_backs_off_then_allows_retry(self): | |
| self.state.http_session.head.side_effect = aiohttp.ClientConnectionError() | |
| app.schedule_download_connection_warm(self.application) | |
| await self.state.download_warm_task | |
| app.schedule_download_connection_warm(self.application) | |
| self.assertEqual(self.state.http_session.head.call_count, 1) | |
| self.state.download_warm_at -= app.DOWNLOAD_CONNECTION_WARM_INTERVAL_SECONDS | |
| self.state.http_session.head.side_effect = None | |
| app.schedule_download_connection_warm(self.application) | |
| await self.state.download_warm_task | |
| self.assertEqual(self.state.http_session.head.call_count, 2) | |
| self.assertEqual(self.state.upstream_semaphore._value, 2) | |
| async def test_saturated_pool_skips_warmup_without_queuing(self): | |
| self.state.upstream_semaphore = asyncio.Semaphore(0) | |
| await app.warm_download_connection(self.application) | |
| self.state.http_session.head.assert_not_called() | |
| async def test_cancellation_releases_pool_and_disabled_state_stops_scheduling(self): | |
| entered = asyncio.Event() | |
| class StalledResponse(FakeResponse): | |
| async def __aenter__(self): | |
| entered.set() | |
| await asyncio.Future() | |
| self.state.http_session.head.return_value = StalledResponse() | |
| app.schedule_download_connection_warm(self.application) | |
| await entered.wait() | |
| task = self.state.download_warm_task | |
| task.cancel() | |
| with self.assertRaises(asyncio.CancelledError): | |
| await task | |
| self.assertEqual(self.state.upstream_semaphore._value, 2) | |
| self.state.download_warm_enabled = False | |
| self.state.download_warm_at = float("-inf") | |
| app.schedule_download_connection_warm(self.application) | |
| self.assertIs(self.state.download_warm_task, task) | |
| class FakeSession: | |
| def __init__(self, responses): | |
| self.responses = list(responses) | |
| self.urls = [] | |
| async def get(self, url, **kwargs): | |
| self.urls.append((url, kwargs)) | |
| response = self.responses.pop(0) | |
| if isinstance(response, Exception): | |
| raise response | |
| return response | |
| async def request(self, method, url, **kwargs): | |
| kwargs["method"] = method | |
| return await self.get(url, **kwargs) | |
| class AsyncProxyTests(unittest.IsolatedAsyncioTestCase): | |
| async def asyncSetUp(self): | |
| self._app_state = dict(app.app.state._state) | |
| self._assets_state = {name: getattr(app, name) for name in ( | |
| 'reader_assets_cache', 'reader_assets_cache_at', 'reader_assets_revision')} | |
| app.app.state.reader_assets_task = None | |
| app.app.state.reader_assets_lock = asyncio.Lock() | |
| app.app.state.http_session = FakeSession([]) | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.initialize_reader_bucket_state(app.app) | |
| self._redirect_cache = dict(upstream_http.upstream_redirect_cache) | |
| upstream_http.upstream_redirect_cache.clear() | |
| async def asyncTearDown(self): | |
| task = getattr(app.app.state, 'reader_assets_task', None) | |
| if task is not None and task is not self._app_state.get('reader_assets_task'): | |
| task.cancel() | |
| await asyncio.gather(task, return_exceptions=True) | |
| await app.shutdown_reader_bucket_state(app.app) | |
| for name, value in self._assets_state.items(): | |
| setattr(app, name, value) | |
| app.app.state._state.clear() | |
| app.app.state._state.update(self._app_state) | |
| upstream_http.upstream_redirect_cache.clear() | |
| upstream_http.upstream_redirect_cache.update(self._redirect_cache) | |
| def test_reader_assets_decoder_filters_invalid_entries(self): | |
| valid_path = "objects/aa/" + "b" * 64 + "/document.pdf" | |
| valid_linearized_path = "objects/aa/" + "f" * 64 + "/linearized.pdf" | |
| valid_azw_path = "objects/aa/" + "e" * 64 + "/document.azw" | |
| page_manifest = "objects/aa/" + "d" * 64 + "/page-manifest.json" | |
| media_path = "objects/aa/" + "c" * 64 + "/ffmpeg-audio-mp3-v1/audio.mp3" | |
| raw = gzip.compress(json.dumps({"v": 1, "f": { | |
| "VoiceOfML/Test\0book.doc": {"s": 2, "m": "p", "p": valid_path}, | |
| "VoiceOfML/Test\0book-linearized.pdf": {"s": 2, "m": "p", "p": valid_linearized_path}, | |
| "VoiceOfML/Test\0book.azw": {"s": 2, "m": "e", "p": valid_azw_path}, | |
| "VoiceOfML/Test\0sound.wma": {"s": 2, "m": "a", "p": media_path}, | |
| "VoiceOfML/Test\0scan.pdf": {"s": 2, "m": "p", "p": page_manifest}, | |
| "bad-path": {"s": 2, "m": "p", "p": "objects/../../secret.pdf"}, | |
| "null-path": {"s": 2, "m": "p", "p": None}, | |
| "unknown-mode": {"s": 2, "m": "unknown", "p": valid_path}, | |
| "unhashable-mode": {"s": 2, "m": [], "p": valid_path}, | |
| "failed": {"s": 4}, | |
| }}).encode()) | |
| self.assertEqual(app.decode_reader_assets(raw), {"v": 1, "f": { | |
| "VoiceOfML/Test\0book.doc": {"s": 2, "m": "p", "p": valid_path}, | |
| "VoiceOfML/Test\0book-linearized.pdf": {"s": 2, "m": "p", "p": valid_linearized_path}, | |
| "VoiceOfML/Test\0book.azw": {"s": 2, "m": "e", "p": valid_azw_path}, | |
| "VoiceOfML/Test\0sound.wma": {"s": 2, "m": "a", "p": media_path}, | |
| "VoiceOfML/Test\0scan.pdf": {"s": 2, "m": "p", "p": page_manifest}, | |
| "failed": {"s": 4}, | |
| }}) | |
| async def test_reader_assets_refresh_is_cached_and_releases_response(self): | |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"failed": {"s": 4}}}).encode()) | |
| response = FakeResponse(200, chunks=(payload,)) | |
| session = FakeSession([response]) | |
| app.app.state.http_session = session | |
| app.app.state.reader_assets_lock = asyncio.Lock() | |
| old_cache, old_at = app.reader_assets_cache, app.reader_assets_cache_at | |
| try: | |
| app.reader_assets_cache = {"v": 1, "f": {}} | |
| app.reader_assets_cache_at = 0 | |
| with patch.object(app, "load_reading_catalog", AsyncMock(return_value=None)): | |
| first = await app.get_reader_assets() | |
| second = await app.get_reader_assets() | |
| self.assertEqual(first, {"v": 1, "f": {"failed": {"s": 4}}}) | |
| self.assertIs(first, second) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertTrue(response.released) | |
| finally: | |
| app.reader_assets_cache, app.reader_assets_cache_at = old_cache, old_at | |
| async def test_search_returns_while_shared_reader_assets_refresh_is_stalled(self): | |
| started, release = asyncio.Event(), asyncio.Event() | |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"new": {"s": 4}}}).encode()) | |
| response = FakeResponse(200, chunks=(payload,)) | |
| async def get(*args, **kwargs): | |
| started.set() | |
| await release.wait() | |
| return response | |
| app.app.state.http_session = SimpleNamespace(get=AsyncMock(side_effect=get)) | |
| app.app.state.reader_assets_lock = asyncio.Lock() | |
| app.app.state.reader_assets_task = None | |
| with patch.object(app, "reader_assets_cache", {"v": 1, "f": {}}), \ | |
| patch.object(app, "load_reading_catalog", AsyncMock(return_value=None)), \ | |
| patch.object(app, "reader_assets_cache_at", 0), \ | |
| patch.object(app, "search", return_value={"results": [], "total": 0, "page": 1, "page_size": 100}): | |
| try: | |
| first = await asyncio.wait_for(app.api_search(app.SearchRequest()), 1) | |
| await asyncio.wait_for(started.wait(), 1) | |
| task = app.app.state.reader_assets_task | |
| second = await asyncio.wait_for(app.api_search(app.SearchRequest()), 1) | |
| self.assertEqual([first.status_code, second.status_code], [200, 200]) | |
| self.assertIs(task, app.app.state.reader_assets_task) | |
| self.assertFalse(task.done()) | |
| release.set() | |
| await task | |
| self.assertEqual(app.reader_assets_cache["f"], {"new": {"s": 4}}) | |
| self.assertTrue(response.released) | |
| self.assertEqual(app.app.state.http_session.get.call_count, 1) | |
| finally: | |
| release.set() | |
| if app.app.state.reader_assets_task: | |
| await app.app.state.reader_assets_task | |
| async def test_reader_assets_refresh_failure_backs_off_with_existing_cache(self): | |
| response = FakeResponse(503) | |
| session = FakeSession([response]) | |
| app.app.state.http_session = session | |
| app.app.state.reader_assets_lock = asyncio.Lock() | |
| old_cache, old_at = app.reader_assets_cache, app.reader_assets_cache_at | |
| try: | |
| app.reader_assets_cache = {"v": 1, "f": {"old": {"s": 4}}} | |
| app.reader_assets_cache_at = time.monotonic() - app.READER_ASSETS_CACHE_TTL_SECONDS - 1 | |
| first = await app.get_reader_assets() | |
| second = await app.get_reader_assets() | |
| self.assertIs(first, second) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertTrue(response.released) | |
| finally: | |
| app.reader_assets_cache, app.reader_assets_cache_at = old_cache, old_at | |
| async def test_initial_reader_assets_failure_is_not_http_cached_and_recovers(self): | |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"ready": {"s": 4}}}).encode()) | |
| session = FakeSession([FakeResponse(503), FakeResponse(200, chunks=(payload,))]) | |
| app.app.state.http_session = session | |
| app.app.state.reader_assets_lock = asyncio.Lock() | |
| with patch.object(app, "reader_assets_cache", {"v": 1, "f": {}}), \ | |
| patch.object(app, "load_reading_catalog", AsyncMock(return_value=None)), \ | |
| patch.object(app, "reader_assets_cache_at", 0), \ | |
| patch.object(app, 'time', SimpleNamespace(monotonic=Mock(return_value=100), | |
| time=time.time, perf_counter=time.perf_counter)) as clock: | |
| failed = await app.api_reader_assets() | |
| self.assertEqual(failed.headers["cache-control"], "no-store") | |
| await app.get_reader_assets() | |
| self.assertEqual(len(session.urls), 1) | |
| clock.monotonic.return_value += 6 | |
| recovered = await app.api_reader_assets() | |
| self.assertEqual(json.loads(recovered.body)["f"], {"ready": {"s": 4}}) | |
| self.assertEqual(len(session.urls), 2) | |
| async def test_reader_bucket_cache_hit_preserves_success_cache_control(self): | |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" | |
| with patch.object(app, "read_reader_bucket_resource", return_value=b"manifest") as read: | |
| first = await app.api_reader_bucket_resource(path) | |
| second = await app.api_reader_bucket_resource(path) | |
| self.assertEqual(first.body, b"manifest") | |
| self.assertEqual(second.body, b"manifest") | |
| self.assertEqual(read.call_count, 1) | |
| self.assertEqual(second.headers["cache-control"], "public, max-age=3600") | |
| async def test_random_reader_uses_complete_snapshot_without_waiting_for_refresh(self): | |
| snapshot = {"v": 1, "f": {"ready": {"s": 4}}} | |
| entered, release = asyncio.Event(), asyncio.Event() | |
| async def refresh(): | |
| entered.set() | |
| await release.wait() | |
| return snapshot | |
| with patch.object(app, "reader_assets_cache", snapshot), \ | |
| patch.object(app, "reader_assets_cache_at", time.monotonic() - app.READER_ASSETS_CACHE_TTL_SECONDS - 1), \ | |
| patch.object(app, "refresh_reader_assets", side_effect=refresh), \ | |
| patch.object(app, "ensure_reader_indexes_async", AsyncMock()), \ | |
| patch.object(app, "converted_reader_indices", return_value={}), \ | |
| patch.object(app, "reader_record_indices", [0]), \ | |
| patch.object(app, "records", [{"Repo": "VoiceOfML/Test", "File": "book", "Extension": "txt", "Folder": [], "Size": 1}]): | |
| result = await asyncio.wait_for(app.api_random_reader(None), .5) | |
| await entered.wait() | |
| self.assertFalse(release.is_set()) | |
| self.assertEqual(result.status_code, 200) | |
| self.assertEqual(json.loads(result.body)["File"], "book") | |
| status = await asyncio.wait_for(app.api_random_reader_status(None), .5) | |
| self.assertEqual(json.loads(status.body), {"available": True, "count": 1}) | |
| release.set() | |
| await app.app.state.reader_assets_task | |
| async def test_cold_random_reader_waits_for_first_complete_asset_snapshot(self): | |
| entered, release = asyncio.Event(), asyncio.Event() | |
| async def refresh(): | |
| entered.set() | |
| await release.wait() | |
| return {"v": 1, "f": {}} | |
| with patch.object(app, "reader_assets_cache", {"v": 1, "f": {}}), \ | |
| patch.object(app, "reader_assets_cache_at", 0), \ | |
| patch.object(app, "refresh_reader_assets", side_effect=refresh), \ | |
| patch.object(app, "ensure_reader_indexes_async", AsyncMock()), \ | |
| patch.object(app, "converted_reader_indices", return_value={}), \ | |
| patch.object(app, "reader_record_indices", []): | |
| task = asyncio.create_task(app.api_random_reader(None)) | |
| await entered.wait() | |
| self.assertFalse(task.done()) | |
| release.set() | |
| self.assertEqual((await task).status_code, 404) | |
| async def test_reader_asset_decode_does_not_block_ping(self): | |
| entered, release = threading.Event(), threading.Event() | |
| def decode(raw): | |
| entered.set() | |
| release.wait(3) | |
| return {"v": 1, "f": {}} | |
| app.app.state.http_session = FakeSession([FakeResponse(200, chunks=(b"sidecar",))]) | |
| with patch.object(app, "reader_assets_cache_at", 0), \ | |
| patch.object(app, "decode_reader_assets", side_effect=decode), \ | |
| patch.object(app, "load_reading_catalog", AsyncMock(return_value=None)): | |
| task = asyncio.create_task(app.refresh_reader_assets()) | |
| try: | |
| self.assertTrue(await asyncio.to_thread(entered.wait, 1)) | |
| self.assertEqual((await asyncio.wait_for(app.api_ping(), .2)).status_code, 204) | |
| self.assertFalse(task.done()) | |
| finally: | |
| release.set() | |
| await task | |
| async def test_v2_chapter_manifest_and_search_index_reuse_bounded_cache(self): | |
| prefix = "chapters/ebook/epub/" + "a" * 64 + "/" + "b" * 16 + "/" | |
| for name in ("chapter-manifest.json", "epub-search-index.json.gz"): | |
| with self.subTest(name=name), patch.object(app, "read_reader_bucket_resource", return_value=b"chapter metadata") as read: | |
| first = await app.api_reader_bucket_resource(prefix + name) | |
| second = await app.api_reader_bucket_resource(prefix + name) | |
| self.assertEqual(first.status_code, 200) | |
| self.assertEqual(second.body, first.body) | |
| self.assertEqual(read.call_count, 1) | |
| self.assertLessEqual(app.app.state.reader_bucket_cache_bytes, app.READER_BUCKET_CACHE_MAX_BYTES) | |
| async def test_bucket_shutdown_keeps_slot_owned_until_filesystem_thread_finishes(self): | |
| started, release = threading.Event(), threading.Event() | |
| path = "chapters/ebook/epub/" + "a" * 64 + "/" + "b" * 16 + "/chapter-manifest.json" | |
| semaphore = app.app.state.reader_bucket_semaphore = asyncio.Semaphore(1) | |
| def read(_path, _limit): | |
| started.set() | |
| release.wait(3) | |
| return b"manifest" | |
| with patch.object(app, "read_reader_bucket_resource", side_effect=read): | |
| loading = asyncio.create_task(app.get_reader_bucket_resource(path)) | |
| shutdown = None | |
| try: | |
| self.assertTrue(await asyncio.to_thread(started.wait, 1)) | |
| shutdown = asyncio.create_task(app.shutdown_reader_bucket_state(app.app)) | |
| await asyncio.sleep(0.02) | |
| self.assertEqual(semaphore._value, 0) | |
| self.assertEqual(len(app.app.state.reader_bucket_workers), 1) | |
| self.assertFalse(shutdown.done()) | |
| finally: | |
| release.set() | |
| await asyncio.gather(loading, return_exceptions=True) | |
| if shutdown is not None: | |
| await shutdown | |
| self.assertEqual(semaphore._value, 1) | |
| self.assertEqual(app.app.state.reader_bucket_workers, set()) | |
| self.assertEqual(app.app.state.reader_bucket_cache, OrderedDict()) | |
| async def test_reader_content_routes_both_v2_bucket_origins_to_proxy(self): | |
| request = SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}) | |
| paths = { | |
| "reader-assets-v2": "documents/pdf/ppt/" + "a" * 64 + "/document.pdf", | |
| "pdf-pages-v2": "derived/test/" + "a" * 32 + "/document.pdf", | |
| } | |
| for bucket, path in paths.items(): | |
| url = f"https://huggingface.co/buckets/vomebook/{bucket}/resolve/{path}" | |
| with self.subTest(bucket=bucket), patch.object(app, "api_reader_bucket_resource", AsyncMock(return_value=app.Response(b"ab", status_code=206))) as proxy: | |
| response = await app.api_reader_content(request, url=url, id=None) | |
| self.assertEqual(response.status_code, 206) | |
| proxy.assert_awaited_once_with(path, request) | |
| async def test_reader_bucket_page_uses_immutable_cache_control(self): | |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000002.webp" | |
| with patch.object(app, "read_reader_bucket_resource", return_value=b"page"): | |
| response = await app.api_reader_bucket_resource(path) | |
| self.assertEqual(response.headers["cache-control"], "public, max-age=31536000, immutable") | |
| async def test_bucket_pdf_streams_initial_get_head_and_range_without_buffering(self): | |
| path = "documents/pdf/test/" + "a" * 64 + "/document.pdf" | |
| for method, headers, status, upstream_headers in ( | |
| ("GET", {}, 200, {"Content-Length": "10", "Accept-Ranges": "bytes"}), | |
| ("HEAD", {}, 200, {"Content-Length": "10", "Accept-Ranges": "bytes"}), | |
| ("GET", {"Range": "bytes=0-1"}, 206, | |
| {"Content-Length": "2", "Content-Range": "bytes 0-1/10", "Accept-Ranges": "bytes"}), | |
| ): | |
| with self.subTest(method=method, headers=headers): | |
| upstream = FakeResponse(status, upstream_headers, chunks=(b"ab",)) | |
| session = FakeSession([upstream]) | |
| app.app.state.http_session = session | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| request = SimpleNamespace(method=method, headers=headers) | |
| with patch.object(app, "get_reader_bucket_resource", AsyncMock()) as whole_read: | |
| response = await app.api_reader_bucket_resource(path, request) | |
| whole_read.assert_not_awaited() | |
| self.assertEqual(response.status_code, status) | |
| self.assertEqual(session.urls[0][0], app.reader_bucket_source_url(path)) | |
| self.assertEqual(session.urls[0][1]["method"], method) | |
| self.assertEqual(session.urls[0][1]["headers"], {"Accept-Encoding": "identity", **headers}) | |
| if method != "HEAD": | |
| self.assertFalse(upstream.released) | |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"ab"]) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_bucket_media_streams_head_and_ranges_without_whole_file_reads(self): | |
| for kind, suffix, mime in (("audio", "flac", "audio/flac"), ("video", "mp4", "video/mp4")): | |
| path = f"media/{kind}/native/{'a' * 64}/{kind}.{suffix}" | |
| for method, forwarded, status, upstream_headers in ( | |
| ("GET", {}, 200, {"Content-Length": "10", "Accept-Ranges": "bytes"}), | |
| ("HEAD", {}, 200, {"Content-Length": "10", "Accept-Ranges": "bytes"}), | |
| ("GET", {"Range": "bytes=0-1"}, 206, | |
| {"Content-Length": "2", "Content-Range": "bytes 0-1/10", "Accept-Ranges": "bytes"}), | |
| ): | |
| with self.subTest(kind=kind, method=method, range=forwarded): | |
| upstream = FakeResponse(status, {"Content-Type": "application/octet-stream", **upstream_headers}, chunks=(b"ab",)) | |
| session = app.app.state.http_session = FakeSession([upstream]) | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| with patch.object(app, "get_reader_bucket_resource", AsyncMock()) as whole_read: | |
| response = await app.api_reader_bucket_resource(path, SimpleNamespace(method=method, headers=forwarded)) | |
| whole_read.assert_not_awaited() | |
| self.assertEqual(response.status_code, status) | |
| self.assertEqual(response.headers["content-type"], mime) | |
| self.assertNotIn("content-encoding", response.headers) | |
| self.assertEqual(session.urls[0][1]["headers"], {"Accept-Encoding": "identity", **forwarded}) | |
| if method != "HEAD": | |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"ab"]) | |
| if forwarded: | |
| cached = await app.api_reader_bucket_resource(path, SimpleNamespace(method=method, headers=forwarded)) | |
| self.assertEqual(cached.body, b"ab") | |
| self.assertNotIn("content-encoding", cached.headers) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_media_remains_unencoded_through_real_gzip_middleware_and_range_cache(self): | |
| path = f"media/audio/native/{'a' * 64}/audio.mp3" | |
| payload = b"media" * 400 | |
| for status in (200, 206): | |
| forwarded = {"Range": f"bytes=0-{len(payload) - 1}"} if status == 206 else {} | |
| headers = {"Content-Length": str(len(payload)), "Content-Encoding": "identity"} | |
| if status == 206: | |
| headers["Content-Range"] = f"bytes 0-{len(payload) - 1}/{len(payload) * 2}" | |
| upstream = FakeResponse(status, headers, chunks=(payload,)) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| for cached in (False, True) if status == 206 else (False,): | |
| with self.subTest(status=status, cached=cached): | |
| response = await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers=forwarded)) | |
| messages = [] | |
| async def serve(scope, receive, send): | |
| await response(scope, receive, send) | |
| async def receive(): | |
| return {"type": "http.request", "body": b"", "more_body": False} | |
| async def send(message): | |
| messages.append(message) | |
| await app.GZipMiddleware(serve, minimum_size=1)( | |
| {"type": "http", "method": "GET", "asgi": {"spec_version": "2.4"}, | |
| "headers": [(b"accept-encoding", b"gzip")]}, receive, send) | |
| start = messages[0] | |
| self.assertEqual(start["status"], status) | |
| self.assertNotIn(b"content-encoding", dict(start["headers"])) | |
| self.assertEqual(dict(start["headers"])[b"content-length"], str(len(payload)).encode()) | |
| self.assertEqual(b"".join(message.get("body", b"") for message in messages[1:]), payload) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(app.app.state.upstream_semaphore._value, 1) | |
| async def test_complete_reader_text_negotiates_gzip_without_changing_decoded_bytes(self): | |
| payload = ("正文 UTF-8 literal \ufffd / unchanged\n" * 5000).encode() | |
| for extension in ("txt", "md", "html", "xhtml", "json"): | |
| for encoding in (b"gzip", b"identity"): | |
| with self.subTest(extension=extension, encoding=encoding): | |
| upstream = FakeResponse(200, {"Content-Length": str(len(payload)), "ETag": '"same-source"'}, | |
| chunks=tuple(payload[i:i + 65536] for i in range(0, len(payload), 65536))) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| response = await app.stream_reader_document(SimpleNamespace(method="GET", headers={}), | |
| f"https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.{extension}") | |
| self.assertNotIn("content-encoding", response.headers) | |
| messages = [] | |
| async def serve(scope, receive, send): | |
| await response(scope, receive, send) | |
| async def receive(): | |
| return {"type": "http.request", "body": b"", "more_body": False} | |
| async def send(message): | |
| messages.append(message) | |
| await app.GZipMiddleware(serve, minimum_size=500)( | |
| {"type": "http", "method": "GET", "asgi": {"spec_version": "2.4"}, | |
| "headers": [(b"accept-encoding", encoding)]}, receive, send) | |
| headers = dict(messages[0]["headers"]) | |
| self.assertEqual(headers[b"etag"], b'W/"same-source"') | |
| body = b"".join(message.get("body", b"") for message in messages[1:]) | |
| self.assertEqual(gzip.decompress(body) if encoding == b"gzip" else body, payload) | |
| if encoding == b"gzip": | |
| self.assertEqual(headers[b"content-encoding"], b"gzip") | |
| self.assertIn(b"Accept-Encoding", headers[b"vary"]) | |
| self.assertLess(len(body), len(payload) // 2) | |
| else: | |
| self.assertNotIn(b"content-encoding", headers) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(app.app.state.upstream_semaphore._value, 1) | |
| async def test_text_compression_cannot_change_range_head_or_binary_policy(self): | |
| source = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/" | |
| for suffix in ("pdf", "epub", "docx", "jpg", "webp", "mp3", "mp4", "json.gz"): | |
| self.assertFalse(app.reader_body_allows_compression(source + "book." + suffix)) | |
| for method, headers, status in (("HEAD", {}, 200), ("GET", {"Range": "bytes=0-1"}, 200), | |
| ("GET", {"Range": "bytes=0-1"}, 206), ("GET", {"If-None-Match": '"A"'}, 200), | |
| ("GET", {}, 304), ("GET", {}, 416)): | |
| self.assertFalse(app.reader_body_allows_compression(source + "book.txt", | |
| SimpleNamespace(method=method, headers=headers), status)) | |
| async def test_large_media_streams_metadata_and_ranges_without_whole_file_limit(self): | |
| total = app.READER_DOCUMENT_MAX_BYTES + 100 | |
| for kind, suffix in (("audio", "mp3"), ("video", "mp4")): | |
| path = f"media/{kind}/native/{'a' * 64}/{kind}.{suffix}" | |
| for method, forwarded, status, headers in ( | |
| ("HEAD", {}, 200, {"Content-Length": str(total)}), | |
| ("GET", {"Range": "bytes=0-1"}, 206, | |
| {"Content-Length": "2", "Content-Range": f"bytes 0-1/{total}"}), | |
| ("GET", {"Range": "bytes=-2"}, 206, | |
| {"Content-Length": "2", "Content-Range": f"bytes {total - 2}-{total - 1}/{total}"}), | |
| ("GET", {"Range": f"bytes={total}-"}, 416, | |
| {"Content-Length": "0", "Content-Range": f"bytes */{total}"}), | |
| ): | |
| for source in ("bucket", "original"): | |
| with self.subTest(kind=kind, method=method, range=forwarded, source=source): | |
| upstream = FakeResponse(status, headers, chunks=() if status == 416 else (b"ab",)) | |
| session = app.app.state.http_session = FakeSession([upstream]) | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| request = SimpleNamespace(method=method, headers=forwarded) | |
| if source == "bucket": | |
| response = await app.api_reader_bucket_resource(path, request) | |
| else: | |
| url = f"https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/large.{suffix}" | |
| response = await app.api_reader_content(request, url=url, id=None) | |
| self.assertEqual(response.status_code, status) | |
| self.assertIsNone(session.urls[0][1]["timeout"].total) | |
| self.assertEqual(session.urls[0][1]["timeout"].sock_read, 60) | |
| if method != "HEAD": | |
| self.assertEqual(b"".join([chunk async for chunk in response.body_iterator]), | |
| b"" if status == 416 else b"ab") | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_media_unknown_length_stream_does_not_use_document_body_limit(self): | |
| path = f"media/audio/native/{'a' * 64}/audio.mp3" | |
| upstream = FakeResponse(200, chunks=(b"abc", b"def")) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| with patch.object(app, "READER_DOCUMENT_MAX_BYTES", 2), patch.object(app, "READER_RESOURCE_MAX_BYTES", 2): | |
| response = await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(response.status_code, 200) | |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"abc", b"def"]) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(app.app.state.upstream_semaphore._value, 1) | |
| async def test_bucket_media_preserves_range_validation_and_cancellation_cleanup(self): | |
| path = f"media/video/native/{'a' * 64}/video.mp4" | |
| for status, headers in ( | |
| (206, {"Content-Length": "2", "Content-Range": "invalid"}), | |
| (206, {"Content-Length": "2", "Content-Range": "bytes 0-2/100"}), | |
| (206, {"Content-Length": "2", "Content-Range": "bytes 99-100/100"}), | |
| ): | |
| upstream = FakeResponse(status, headers) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| response = await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"})) | |
| self.assertEqual(response.status_code, 502) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| with patch.object(app, "open_download_response", AsyncMock(side_effect=asyncio.CancelledError)): | |
| with self.assertRaises(asyncio.CancelledError): | |
| await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_bucket_pdf_validates_range_and_releases_on_failure_or_cancel(self): | |
| path = "derived/test/" + "a" * 32 + "/document.pdf" | |
| for status, headers, expected in ( | |
| (206, {"Content-Range": "invalid"}, 502), | |
| (206, {"Content-Range": f"bytes 0-1/{app.READER_DOCUMENT_MAX_BYTES + 1}"}, 413), | |
| (200, {"Content-Length": str(app.READER_DOCUMENT_MAX_BYTES + 1)}, 413), | |
| (200, {"Content-Encoding": "gzip"}, 502), | |
| (404, {}, 404), | |
| ): | |
| with self.subTest(status=status, headers=headers): | |
| upstream = FakeResponse(status, headers) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| response = await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(response.status_code, expected) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| with patch.object(app, "open_download_response", AsyncMock(side_effect=asyncio.CancelledError)): | |
| with self.assertRaises(asyncio.CancelledError): | |
| await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_bucket_pdf_redirect_cannot_switch_to_another_hf_source(self): | |
| path = "derived/test/" + "a" * 32 + "/document.pdf" | |
| redirect = FakeResponse(302, {"Location": "https://huggingface.co/datasets/Other/Private/resolve/main/book.pdf"}) | |
| session = app.app.state.http_session = FakeSession([redirect]) | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| response = await app.api_reader_bucket_resource(path, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(response.status_code, 502) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertTrue(redirect.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_accepts_resolved_v2_proxy_but_not_other_hosts(self): | |
| path = "derived/test/" + "a" * 32 + "/document.pdf" | |
| request = SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}) | |
| proxy_url = "https://voiceofml-search.hf.space/api/reader-bucket-resource?path=" + path | |
| with patch.object(app, "api_reader_bucket_resource", AsyncMock(return_value=app.Response(b"ab", status_code=206))) as bucket: | |
| with patch.object(app, "api_reader_resolve", AsyncMock(return_value=app.JSONResponse({"url": proxy_url}))): | |
| response = await app.api_reader_content(request, url=None, id="398vk0yyy8m29") | |
| self.assertEqual(response.status_code, 206) | |
| bucket.assert_awaited_once_with(path, request) | |
| bucket.reset_mock() | |
| response = await app.api_reader_content(request, url=proxy_url.replace("voiceofml-search.hf.space", "other.example"), id=None) | |
| self.assertEqual(response.status_code, 403) | |
| bucket.assert_not_awaited() | |
| async def test_reader_bucket_single_flight_coalesces_uncached_page(self): | |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000002.webp" | |
| started = threading.Event() | |
| release = threading.Event() | |
| def read(_path, max_bytes=app.READER_RESOURCE_MAX_BYTES): | |
| started.set() | |
| release.wait(1) | |
| return b"page" | |
| with patch.object(app, "read_reader_bucket_resource", side_effect=read) as bucket_read: | |
| first = asyncio.create_task(app.api_reader_bucket_resource(path)) | |
| self.assertTrue(await asyncio.to_thread(started.wait, 1)) | |
| second = asyncio.create_task(app.api_reader_bucket_resource(path)) | |
| await asyncio.sleep(0) | |
| self.assertEqual(bucket_read.call_count, 1) | |
| release.set() | |
| responses = await asyncio.gather(first, second) | |
| self.assertEqual([response.body for response in responses], [b"page", b"page"]) | |
| self.assertEqual(bucket_read.call_count, 1) | |
| self.assertNotIn(path, app.app.state.reader_bucket_cache) | |
| self.assertEqual(app.app.state.reader_bucket_inflight, {}) | |
| async def test_reader_bucket_semaphore_bounds_distinct_reads(self): | |
| app.app.state.reader_bucket_semaphore = asyncio.Semaphore(1) | |
| first_path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000002.webp" | |
| second_path = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000003.webp" | |
| first_started = threading.Event() | |
| release = threading.Event() | |
| def read(path, max_bytes=app.READER_RESOURCE_MAX_BYTES): | |
| if path == first_path: | |
| first_started.set() | |
| release.wait(1) | |
| return path.encode() | |
| with patch.object(app, "read_reader_bucket_resource", side_effect=read) as bucket_read: | |
| first = asyncio.create_task(app.api_reader_bucket_resource(first_path)) | |
| self.assertTrue(await asyncio.to_thread(first_started.wait, 1)) | |
| second = asyncio.create_task(app.api_reader_bucket_resource(second_path)) | |
| await asyncio.sleep(0.05) | |
| self.assertEqual(bucket_read.call_count, 1) | |
| release.set() | |
| await asyncio.gather(first, second) | |
| self.assertEqual(bucket_read.call_count, 2) | |
| async def test_reader_bucket_failures_are_no_store_and_not_cached(self): | |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" | |
| invalid = await app.api_reader_bucket_resource("invalid") | |
| self.assertEqual(invalid.status_code, 400) | |
| self.assertEqual(invalid.headers["cache-control"], "no-store") | |
| for error, status in ((FileNotFoundError(), 404), (RuntimeError(), 502)): | |
| with self.subTest(status=status), patch.object( | |
| app, "read_reader_bucket_resource", side_effect=error | |
| ) as bucket_read: | |
| first = await app.api_reader_bucket_resource(path) | |
| second = await app.api_reader_bucket_resource(path) | |
| self.assertEqual(first.status_code, status) | |
| self.assertEqual(second.status_code, status) | |
| self.assertEqual(first.headers["cache-control"], "no-store") | |
| self.assertEqual(bucket_read.call_count, 2) | |
| async def test_reader_bucket_cache_evicts_by_bytes_and_skips_oversize_items(self): | |
| manifest = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" | |
| first_page = f"objects/aa/{'a' * 64}/{'b' * 16}/pages/page-000001.webp" | |
| other_manifest = f"objects/bb/{'c' * 64}/{'d' * 16}/page-manifest.json" | |
| old_total = app.READER_BUCKET_CACHE_MAX_BYTES | |
| old_item = app.READER_BUCKET_CACHE_MAX_ITEM_BYTES | |
| app.READER_BUCKET_CACHE_MAX_BYTES = 5 | |
| app.READER_BUCKET_CACHE_MAX_ITEM_BYTES = 4 | |
| payloads = {manifest: b"aaa", first_page: b"bbb", other_manifest: b"12345"} | |
| try: | |
| with patch.object(app, "read_reader_bucket_resource", side_effect=lambda path, limit: payloads.get(path)) as bucket_read: | |
| await app.api_reader_bucket_resource(manifest) | |
| await app.api_reader_bucket_resource(first_page) | |
| self.assertNotIn(manifest, app.app.state.reader_bucket_cache) | |
| self.assertEqual(app.app.state.reader_bucket_cache_bytes, 3) | |
| await app.api_reader_bucket_resource(manifest) | |
| await app.api_reader_bucket_resource(other_manifest) | |
| await app.api_reader_bucket_resource(other_manifest) | |
| self.assertEqual(bucket_read.call_count, 5) | |
| self.assertNotIn(other_manifest, app.app.state.reader_bucket_cache) | |
| self.assertLessEqual(app.app.state.reader_bucket_cache_bytes, 5) | |
| finally: | |
| app.READER_BUCKET_CACHE_MAX_BYTES = old_total | |
| app.READER_BUCKET_CACHE_MAX_ITEM_BYTES = old_item | |
| async def test_reader_bucket_shutdown_cancels_flights_and_clears_cache(self): | |
| task = asyncio.create_task(asyncio.sleep(10)) | |
| app.app.state.reader_bucket_inflight["path"] = task | |
| app.app.state.reader_bucket_cache["path"] = (time.monotonic(), b"data") | |
| app.app.state.reader_bucket_cache_bytes = 4 | |
| await app.shutdown_reader_bucket_state(app.app) | |
| self.assertTrue(task.cancelled()) | |
| self.assertEqual(app.app.state.reader_bucket_inflight, {}) | |
| self.assertEqual(app.app.state.reader_bucket_cache, {}) | |
| self.assertEqual(app.app.state.reader_bucket_cache_bytes, 0) | |
| async def test_relative_redirect_is_resolved_against_current_url(self): | |
| session = FakeSession([FakeResponse(302, {"Location": "next"}), FakeResponse(200)]) | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" | |
| result = await app.open_download_response( | |
| session, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_source_url, | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertEqual([url for url, _ in session.urls], [start, "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/next"]) | |
| async def test_redirect_cannot_escape_voiceofml_dataset_scope(self): | |
| first = FakeResponse(302, {"Location": "https://huggingface.co/datasets/Other/Test/resolve/main/book.pdf"}) | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" | |
| with self.assertRaises(ValueError): | |
| await app.open_download_response( | |
| FakeSession([first]), start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_source_url, | |
| ) | |
| self.assertTrue(first.released) | |
| async def test_redirect_cannot_escape_reader_assets_scope(self): | |
| digest = "b" * 64 | |
| start = f"https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/{digest}/document.pdf" | |
| first = FakeResponse(302, {"Location": "/datasets/VoiceOfML/Test/resolve/main/book.pdf"}) | |
| with self.assertRaises(ValueError): | |
| await app.open_download_response( | |
| FakeSession([first]), start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_reader_source_url, | |
| ) | |
| self.assertTrue(first.released) | |
| async def test_multilayer_encoded_reader_assets_redirect_cannot_escape_scope(self): | |
| digest = "b" * 64 | |
| start = f"https://huggingface.co/datasets/%252576omebook/Reader-Assets/resolve/main/objects/aa/{digest}/document.pdf" | |
| first = FakeResponse(302, {"Location": "/datasets/VoiceOfML/Test/resolve/main/book.pdf"}) | |
| self.assertEqual(app.validate_reader_source_url(start), start) | |
| with self.assertRaises(ValueError): | |
| await app.open_download_response( | |
| FakeSession([first, FakeResponse(200)]), start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_reader_source_url, | |
| ) | |
| self.assertTrue(first.released) | |
| async def test_redirect_to_disallowed_host_is_rejected_and_releases_first(self): | |
| first = FakeResponse(302, {"Location": "https://evil.test/x"}) | |
| with self.assertRaises(ValueError): | |
| await app.open_download_response(FakeSession([first]), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) | |
| self.assertTrue(first.released) | |
| async def test_redirect_without_location_is_rejected_and_released(self): | |
| response = FakeResponse(302) | |
| with self.assertRaises(ValueError): | |
| await app.open_download_response(FakeSession([response]), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) | |
| self.assertTrue(response.released) | |
| async def test_too_many_redirects_are_rejected(self): | |
| responses = [FakeResponse(302, {"Location": "/next"}) for _ in range(app.MAX_REDIRECTS + 1)] | |
| with self.assertRaises(ValueError): | |
| await app.open_download_response(FakeSession(responses), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) | |
| async def test_resolved_redirect_target_is_reused_and_scoped(self): | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/cached.pdf" | |
| final = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/actual.pdf" | |
| first = FakeSession([FakeResponse(302, {"Location": "actual.pdf"}), FakeResponse(200)]) | |
| result = await app.open_download_response( | |
| first, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_source_url, | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertEqual([url for url, _ in first.urls], [start, final]) | |
| second = FakeSession([FakeResponse(200)]) | |
| result = await app.open_download_response( | |
| second, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_source_url, | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertEqual([url for url, _ in second.urls], [final]) | |
| async def test_rejected_redirect_target_falls_back_to_fresh_walk(self): | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/retry.pdf" | |
| stale = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/stale.pdf" | |
| fresh = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/fresh.pdf" | |
| upstream_http.upstream_redirect_cache[("GET", start)] = (time.monotonic(), stale) | |
| session = FakeSession([ | |
| FakeResponse(403), FakeResponse(302, {"Location": "fresh.pdf"}), FakeResponse(200), | |
| ]) | |
| result = await app.open_download_response( | |
| session, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_source_url, | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertEqual([url for url, _ in session.urls], [stale, start, fresh]) | |
| self.assertEqual(upstream_http.upstream_redirect_cache[("GET", start)][1], fresh) | |
| async def test_out_of_scope_cached_target_is_ignored(self): | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/scope.pdf" | |
| escaped = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" | |
| upstream_http.upstream_redirect_cache[("GET", start)] = (time.monotonic(), escaped) | |
| session = FakeSession([FakeResponse(200)]) | |
| result = await app.open_download_response( | |
| session, start, aiohttp.ClientTimeout(total=1), | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertEqual([url for url, _ in session.urls], [start]) | |
| async def test_non_redirect_response_is_returned_without_release(self): | |
| response = FakeResponse(404) | |
| result = await app.open_download_response(FakeSession([response]), "https://huggingface.co/x", aiohttp.ClientTimeout(total=1)) | |
| self.assertIs(result, response) | |
| self.assertFalse(response.released) | |
| async def test_download_rejects_external_source_before_session_call(self): | |
| session = FakeSession([]) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| result = await app.api_download("x.pdf", "https://evil.test/x") | |
| self.assertEqual(result.status_code, 403) | |
| self.assertEqual(session.urls, []) | |
| async def test_download_upstream_error_releases_semaphore(self): | |
| session = FakeSession([FakeResponse(502)]) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = semaphore | |
| result = await app.api_download("x.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") | |
| self.assertEqual(result.status_code, 502) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_download_cancellation_releases_semaphore(self): | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = FakeSession([]) | |
| app.app.state.upstream_semaphore = semaphore | |
| with patch.object(app, "open_download_response", AsyncMock(side_effect=asyncio.CancelledError)): | |
| with self.assertRaises(asyncio.CancelledError): | |
| await app.api_download("x.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_download_non_200_releases_upstream_and_semaphore(self): | |
| upstream = FakeResponse(404) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.upstream_semaphore = semaphore | |
| result = await app.api_download("x.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") | |
| self.assertEqual(result.status_code, 502) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_download_response_stream_releases_resources_after_consumption(self): | |
| upstream = FakeResponse(200, chunks=(b"a", b"b")) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.upstream_semaphore = semaphore | |
| response = await app.api_download("../book.pdf", "https://huggingface.co/datasets/VoiceOfML/A/x") | |
| chunks = [chunk async for chunk in response.body_iterator] | |
| self.assertEqual(chunks, [b"a", b"b"]) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| self.assertIn("filename*=UTF-8''book.pdf", response.headers["content-disposition"]) | |
| # Slow large transfers are bounded by inactivity, not a total deadline. | |
| timeout = app.app.state.http_session.urls[0][1]["timeout"] | |
| self.assertIsNone(timeout.total) | |
| self.assertEqual((timeout.connect, timeout.sock_read), (15, 60)) | |
| async def test_reader_content_forwards_range_and_preserves_partial_headers(self): | |
| upstream = FakeResponse(206, { | |
| "Content-Type": "application/pdf", | |
| "Content-Length": "2", | |
| "Content-Range": "bytes 0-1/10", | |
| "Accept-Ranges": "bytes", | |
| "ETag": '"reader"', | |
| "Set-Cookie": "secret=1", | |
| }, chunks=(b"ab",)) | |
| session = FakeSession([upstream]) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = semaphore | |
| request = SimpleNamespace(method="GET", headers={ | |
| "Range": "bytes=0-1", "If-None-Match": '"reader"', | |
| "If-Modified-Since": "Tue, 25 Aug 2026 00:00:00 GMT", "Cookie": "secret=1", | |
| }) | |
| url = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" | |
| response = await app.api_reader_content(request, url) | |
| self.assertEqual(response.status_code, 206) | |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"ab"]) | |
| self.assertEqual(session.urls[0][1]["headers"], { | |
| "Accept-Encoding": "identity", "Range": "bytes=0-1", | |
| "If-None-Match": '"reader"', "If-Modified-Since": "Tue, 25 Aug 2026 00:00:00 GMT", | |
| }) | |
| self.assertEqual(response.headers["content-range"], "bytes 0-1/10") | |
| self.assertEqual(response.headers["content-length"], "2") | |
| self.assertEqual(response.headers["content-encoding"], "identity") | |
| self.assertEqual(response.headers["cache-control"], "no-store") | |
| self.assertIs(session.urls[0][1]["auto_decompress"], False) | |
| self.assertNotIn("set-cookie", response.headers) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_reuses_completed_small_range_from_cache(self): | |
| upstream = FakeResponse(206, { | |
| "Content-Type": "application/pdf", "Content-Length": "2", | |
| "Content-Range": "bytes 0-1/10", "Accept-Ranges": "bytes", | |
| }, chunks=(b"ab",)) | |
| session = FakeSession([upstream]) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| request = SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}) | |
| url = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" | |
| first = await app.api_reader_content(request, url) | |
| self.assertEqual(first.headers["cache-control"], "no-store") | |
| self.assertEqual([chunk async for chunk in first.body_iterator], [b"ab"]) | |
| next(iter(app.app.state.reader_range_cache.values()))[2]["Cache-Control"] = "public, max-age=31536000" | |
| second = await app.api_reader_content(request, url) | |
| self.assertEqual(second.body, b"ab") | |
| self.assertEqual(second.headers["cache-control"], "no-store") | |
| self.assertEqual(len(session.urls), 1) | |
| async def test_reader_browser_cache_policy_only_changes_partial_pdf_responses(self): | |
| for suffix, mime, status, expected in ( | |
| ("pdf", "application/pdf", 206, "no-store"), | |
| ("PDF", "application/octet-stream", 206, "no-store"), | |
| ("bin", "application/pdf; charset=binary", 206, "no-store"), | |
| ("pdf", "application/pdf", 200, "public, max-age=3600"), | |
| ("pdf", "application/pdf", 304, "public, max-age=3600"), | |
| ("mp4", "video/mp4", 206, "public, max-age=3600"), | |
| ): | |
| with self.subTest(suffix=suffix, mime=mime, status=status): | |
| headers = {"Content-Type": mime, "Cache-Control": "public, max-age=3600", "ETag": '"same"'} | |
| result = app.reader_document_browser_headers( | |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book." + suffix, headers, status) | |
| self.assertEqual(result["Cache-Control"], expected) | |
| self.assertEqual(result["Content-Type"], mime) | |
| self.assertEqual(result["ETag"], '"same"') | |
| self.assertEqual(headers["Cache-Control"], "public, max-age=3600") | |
| async def test_reader_content_does_not_reuse_range_cache_for_conditional_request(self): | |
| first_upstream = FakeResponse(206, { | |
| "Content-Type": "application/pdf", "Content-Length": "2", | |
| "Content-Range": "bytes 0-1/10", "Accept-Ranges": "bytes", | |
| "ETag": '"reader"', | |
| }, chunks=(b"ab",)) | |
| second_upstream = FakeResponse(304, {"ETag": '"reader"'}) | |
| session = FakeSession([first_upstream, second_upstream]) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| url = "https://huggingface.co/datasets/vomebook/Reader-Assets/resolve/main/objects/aa/" + "b" * 64 + "/document.pdf" | |
| first = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}), url) | |
| self.assertEqual([chunk async for chunk in first.body_iterator], [b"ab"]) | |
| second = await app.api_reader_content(SimpleNamespace(method="GET", headers={ | |
| "Range": "bytes=0-1", "If-None-Match": '"reader"', | |
| }), url) | |
| self.assertEqual(second.status_code, 304) | |
| self.assertEqual(len(session.urls), 2) | |
| async def test_reader_range_cache_does_not_turn_origin_200_into_206(self): | |
| upstream = FakeResponse(200, {"Content-Type": "application/pdf", "Content-Length": "2"}, chunks=(b"ab",)) | |
| session = FakeSession([upstream]) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| url = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" | |
| response = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}), url) | |
| self.assertEqual(response.status_code, 200) | |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"ab"]) | |
| self.assertEqual(app.app.state.reader_range_cache, OrderedDict()) | |
| async def test_reader_range_cache_rejects_incomplete_body_and_retries_origin(self): | |
| headers = {"Content-Length": "2", "Content-Range": "bytes 0-1/10"} | |
| partial = FakeResponse(206, headers, chunks=(b"a",)) | |
| complete = FakeResponse(206, headers, chunks=(b"ab",)) | |
| session = app.app.state.http_session = FakeSession([partial, complete]) | |
| semaphore = app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| request = SimpleNamespace(method="GET", headers={"Range": "bytes=0-1"}) | |
| url = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/incomplete.pdf" | |
| first = await app.api_reader_content(request, url) | |
| self.assertEqual([chunk async for chunk in first.body_iterator], [b"a"]) | |
| self.assertEqual(app.app.state.reader_range_cache, OrderedDict()) | |
| second = await app.api_reader_content(request, url) | |
| self.assertEqual([chunk async for chunk in second.body_iterator], [b"ab"]) | |
| third = await app.api_reader_content(request, url) | |
| self.assertEqual(third.body, b"ab") | |
| self.assertEqual(len(session.urls), 2) | |
| self.assertTrue(partial.released and complete.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_rejects_oversized_content_length_before_streaming(self): | |
| upstream = FakeResponse(200, {"Content-Length": str(app.READER_DOCUMENT_MAX_BYTES + 1)}, chunks=(b"not-read",)) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| app.app.state.upstream_semaphore = semaphore | |
| response = await app.api_reader_content( | |
| SimpleNamespace(method="GET", headers={}), | |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf", | |
| ) | |
| self.assertEqual(response.status_code, 413) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_rejects_oversized_range_total(self): | |
| old_limit = app.READER_DOCUMENT_MAX_BYTES | |
| app.READER_DOCUMENT_MAX_BYTES = 3 | |
| upstream = FakeResponse(206, { | |
| "Content-Length": "1", | |
| "Content-Range": "bytes 0-0/4", | |
| }, chunks=(b"a",)) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| app.app.state.upstream_semaphore = semaphore | |
| try: | |
| response = await app.api_reader_content( | |
| SimpleNamespace(method="GET", headers={}), | |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf", | |
| ) | |
| finally: | |
| app.READER_DOCUMENT_MAX_BYTES = old_limit | |
| self.assertEqual(response.status_code, 413) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_enforces_actual_stream_limit_and_releases(self): | |
| old_limit = app.READER_DOCUMENT_MAX_BYTES | |
| app.READER_DOCUMENT_MAX_BYTES = 3 | |
| upstream = FakeResponse(200, {}, chunks=(b"ab", b"cd")) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| app.app.state.upstream_semaphore = semaphore | |
| try: | |
| response = await app.api_reader_content( | |
| SimpleNamespace(method="GET", headers={}), | |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf", | |
| ) | |
| with self.assertRaises(app.ReaderDocumentLimitError): | |
| [chunk async for chunk in response.body_iterator] | |
| finally: | |
| app.READER_DOCUMENT_MAX_BYTES = old_limit | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_cancellation_releases_semaphore(self): | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.upstream_semaphore = semaphore | |
| request = SimpleNamespace(method="GET", headers={}) | |
| with patch.object(app, "open_download_response", AsyncMock(side_effect=asyncio.CancelledError)): | |
| with self.assertRaises(asyncio.CancelledError): | |
| await app.api_reader_content( | |
| request, "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" | |
| ) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_head_releases_resources_without_streaming(self): | |
| upstream = FakeResponse(200, {"Content-Length": "10", "Accept-Ranges": "bytes"}) | |
| session = FakeSession([upstream]) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = semaphore | |
| request = SimpleNamespace(method="HEAD", headers={}) | |
| response = await app.api_reader_content(request, "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.epub") | |
| self.assertEqual(response.status_code, 200) | |
| self.assertEqual(session.urls[0][1]["method"], "HEAD") | |
| self.assertIs(session.urls[0][1]["auto_decompress"], False) | |
| self.assertEqual(response.headers["content-length"], "10") | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_reader_content_rejects_external_source_before_session_call(self): | |
| session = FakeSession([]) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| result = await app.api_reader_content(SimpleNamespace(method="GET", headers={}), "https://evil.test/book.pdf") | |
| self.assertEqual(result.status_code, 403) | |
| self.assertEqual(session.urls, []) | |
| async def test_download_check_upstream_error_returns_502_and_releases(self): | |
| upstream = FakeResponse(503) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| result = await app.api_download_check("https://huggingface.co/datasets/VoiceOfML/A/x") | |
| self.assertEqual(result.status_code, 502) | |
| self.assertTrue(upstream.released) | |
| async def test_download_check_uses_head_and_shared_semaphore(self): | |
| upstream = FakeResponse(200) | |
| session = FakeSession([upstream]) | |
| semaphore = asyncio.Semaphore(1) | |
| app.app.state.http_session = session | |
| app.app.state.upstream_semaphore = semaphore | |
| result = await app.api_download_check("https://huggingface.co/datasets/VoiceOfML/A/x") | |
| self.assertEqual(result.status_code, 200) | |
| self.assertEqual(session.urls[0][1]["method"], "HEAD") | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(semaphore._value, 1) | |
| async def test_local_txt_file_is_served_without_remote_request(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| old_dir = app.TXT_DIR | |
| app.TXT_DIR = Path(directory) | |
| local = Path(directory) / "local.txt" | |
| local.write_text("local", encoding="utf-8") | |
| session = FakeSession([]) | |
| app.app.state.http_session = session | |
| try: | |
| response = await app.serve_txt_proxy("local.txt") | |
| self.assertEqual(response.path, local) | |
| self.assertEqual(session.urls, []) | |
| finally: | |
| app.TXT_DIR = old_dir | |
| async def test_remote_txt_success_streams_content(self): | |
| upstream = FakeResponse(200, chunks=(b"remote",)) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| with tempfile.TemporaryDirectory() as directory, patch.object(app, "resolve_safe_txt_path", return_value=Path(directory) / "missing.txt"): | |
| response = await app.serve_txt_proxy("remote.txt") | |
| self.assertEqual([chunk async for chunk in response.body_iterator], [b"remote"]) | |
| self.assertTrue(upstream.released) | |
| async def test_txt_traversal_is_400_without_remote_request(self): | |
| session = FakeSession([]) | |
| app.app.state.http_session = session | |
| response = await app.serve_txt_proxy("../secret.txt") | |
| self.assertEqual(response.status_code, 400) | |
| self.assertEqual(session.urls, []) | |
| async def test_remote_txt_404_maps_to_404(self): | |
| upstream = FakeResponse(404) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| with tempfile.TemporaryDirectory() as directory, patch.object(app, "resolve_safe_txt_path", return_value=Path(directory) / "missing.txt"): | |
| response = await app.serve_txt_proxy("missing.txt") | |
| self.assertEqual(response.status_code, 404) | |
| self.assertTrue(upstream.released) | |
| async def test_reader_cached_containing_range_serves_exact_slice_without_upstream(self): | |
| upstream = FakeResponse(206, { | |
| "Content-Type": "application/pdf", "Content-Length": "8", | |
| "Content-Range": "bytes 2-9/10", "ETag": '"same"', | |
| }, chunks=(b"2345", b"6789")) | |
| session = app.app.state.http_session = FakeSession([upstream]) | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| url = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/book.pdf" | |
| first = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": "bytes=2-9"}), url) | |
| self.assertIn('reader_cache;desc="miss"', first.headers['server-timing']) | |
| self.assertEqual(b"".join([part async for part in first.body_iterator]), b"23456789") | |
| for requested, body, returned in (("bytes=4-6", b"456", "bytes 4-6/10"), | |
| ("bytes=7-20", b"789", "bytes 7-9/10")): | |
| response = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": requested}), url) | |
| self.assertEqual(response.body, body) | |
| self.assertEqual(response.headers["content-range"], returned) | |
| self.assertEqual(response.headers["content-length"], str(len(body))) | |
| self.assertEqual(response.headers["etag"], '"same"') | |
| self.assertEqual(response.headers["cache-control"], "no-store") | |
| self.assertIn('reader_cache;desc="slice"', response.headers['server-timing']) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertEqual(app.app.state.reader_range_cache_bytes, 8) | |
| self.assertEqual(len(app.app.state.reader_range_cache), 1) | |
| self.assertIsNone(app.reader_cached_subrange((url, 'bytes=0-3', app.READER_DOCUMENT_MAX_BYTES, None))) | |
| self.assertIsNone(app.reader_cached_subrange((url, 'bytes=4-6', 10, None))) | |
| self.assertIsNone(app.reader_cached_subrange((url, 'bytes=4-6', app.READER_DOCUMENT_MAX_BYTES, 'video/mp4'))) | |
| self.assertIsNone(app.reader_cached_subrange((url, 'bytes=-3', app.READER_DOCUMENT_MAX_BYTES, None))) | |
| key = next(iter(app.app.state.reader_range_cache)) | |
| _, payload, headers = app.app.state.reader_range_cache[key] | |
| app.app.state.reader_range_cache[key] = (float('-inf'), payload, headers) | |
| self.assertIsNone(app.reader_cached_subrange((url, 'bytes=4-6', app.READER_DOCUMENT_MAX_BYTES, None))) | |
| self.assertEqual(app.app.state.reader_range_cache_bytes, 0) | |
| async def test_reader_range_cache_entry_cap_evicts_oldest_completed_range(self): | |
| responses = [FakeResponse(206, {"Content-Length": "1", "Content-Range": f"bytes {i}-{i}/10"}, | |
| chunks=(str(i).encode(),)) for i in range(3)] | |
| app.app.state.http_session = FakeSession(responses) | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| app.app.state.reader_range_cache = OrderedDict() | |
| app.app.state.reader_range_cache_bytes = 0 | |
| url = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/range-cap.pdf" | |
| with patch.object(app, "READER_RANGE_CACHE_MAX_ENTRIES", 2): | |
| for i in range(3): | |
| response = await app.api_reader_content(SimpleNamespace(method="GET", headers={"Range": f"bytes={i}-{i}"}), url) | |
| self.assertEqual(b"".join([part async for part in response.body_iterator]), str(i).encode()) | |
| self.assertEqual([key[1] for key in app.app.state.reader_range_cache], ["bytes=1-1", "bytes=2-2"]) | |
| self.assertEqual(app.app.state.reader_range_cache_bytes, 2) | |
| # Reader archive resources (/api/reader-resource). | |
| def zip_bytes(files=None): | |
| import io | |
| import zipfile | |
| buffer = io.BytesIO() | |
| with zipfile.ZipFile(buffer, "w") as archive: | |
| for name, body in (files or {"image.svg": b"<svg/>"}).items(): | |
| archive.writestr(name, body) | |
| return buffer.getvalue() | |
| async def test_reader_archive_rejects_non_archive_formats_before_request(self): | |
| app.initialize_reader_archive_state(app.app) | |
| session = app.app.state.http_session = FakeSession([]) | |
| for name in ("book.pdf", "book.mobi", "book.azw3", "book.fb2", "audio.mp3", "cover.jpg"): | |
| with self.subTest(name=name): | |
| response = await app.api_reader_resource( | |
| "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/" + name, "image.svg") | |
| self.assertEqual(response.status_code, 400) | |
| self.assertEqual(session.urls, []) | |
| async def test_reader_archive_rejects_declared_oversize_without_reading_body(self): | |
| app.initialize_reader_archive_state(app.app) | |
| source = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/huge.epub" | |
| reads = [] | |
| class LargeResponse(FakeResponse): | |
| def iter_chunked(self, size): | |
| reads.append(size) | |
| return super().iter_chunked(size) | |
| upstream = LargeResponse(200, {"Content-Length": str(app.READER_RESOURCE_MAX_ARCHIVE_BYTES + 1)}, | |
| chunks=(self.zip_bytes(),)) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| response = await app.api_reader_resource(source, "image.svg") | |
| self.assertEqual(response.status_code, 413) | |
| self.assertEqual(reads, []) | |
| self.assertTrue(upstream.released) | |
| self.assertEqual(app.app.state.reader_archive_cache, {}) | |
| self.assertEqual(app.app.state.upstream_semaphore._value, 1) | |
| async def test_reader_archive_stops_after_non_archive_signature(self): | |
| app.initialize_reader_archive_state(app.app) | |
| source = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/fake.epub" | |
| consumed = [] | |
| class TextResponse(FakeResponse): | |
| def iter_chunked(self, size): | |
| async def generator(): | |
| for chunk in (b"plain text, not an archive", b"x" * 1024, b"y" * 1024): | |
| consumed.append(chunk) | |
| yield chunk | |
| return generator() | |
| upstream = TextResponse() | |
| app.app.state.http_session = FakeSession([upstream]) | |
| response = await app.api_reader_resource(source, "image.svg") | |
| self.assertEqual(response.status_code, 400) | |
| self.assertEqual(len(consumed), 1) | |
| self.assertTrue(upstream.released) | |
| async def test_reader_archive_split_signature_and_single_copy_roundtrip(self): | |
| app.initialize_reader_archive_state(app.app) | |
| source = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/split.epub" | |
| raw = self.zip_bytes() | |
| upstream = FakeResponse(200, {"Content-Length": str(len(raw))}, chunks=(raw[:1], raw[1:3], raw[3:])) | |
| app.app.state.http_session = FakeSession([upstream]) | |
| prepared_inputs = [] | |
| original_prepare = app.prepare_reader_archive | |
| def spy(value): | |
| prepared_inputs.append(type(value)) | |
| return original_prepare(value) | |
| with patch.object(app, "prepare_reader_archive", side_effect=spy): | |
| response = await app.api_reader_resource(source, "image.svg") | |
| self.assertEqual(response.status_code, 200) | |
| self.assertEqual(response.body, b"<svg/>") | |
| # The worker receives the single joined bytes buffer, not a bytearray | |
| # that must be copied again while both are alive. | |
| self.assertEqual(prepared_inputs, [bytes]) | |
| cached = app.app.state.reader_archive_cache[source][1] | |
| self.assertIs(type(cached), bytes) | |
| self.assertIs(app.prepare_reader_archive(cached), cached) | |
| gzipped = FakeResponse(chunks=(gzip.compress(raw),)) | |
| app.app.state.http_session = FakeSession([gzipped]) | |
| response = await app.api_reader_resource(source.replace("split", "packed"), "image.svg") | |
| self.assertEqual(response.status_code, 200) | |
| # Redirect-cache recovery (upstream_http.open_download_response). | |
| async def test_cached_target_transport_failure_evicts_and_walks_fresh(self): | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/conn.pdf" | |
| stale = "https://cdn-lfs.hf.co/stale-signed" | |
| fresh = "https://cdn-lfs.hf.co/fresh-signed" | |
| for error in (aiohttp.ClientConnectionError("reset"), asyncio.TimeoutError()): | |
| with self.subTest(error=type(error).__name__): | |
| upstream_http.upstream_redirect_cache.clear() | |
| upstream_http.upstream_redirect_cache[("GET", start)] = (time.monotonic(), stale) | |
| session = FakeSession([error, FakeResponse(302, {"Location": fresh}), FakeResponse(200)]) | |
| result = await app.open_download_response( | |
| session, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_redirect_url, | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertEqual([url for url, _ in session.urls], [stale, start, fresh]) | |
| self.assertEqual(upstream_http.upstream_redirect_cache[("GET", start)][1], fresh) | |
| async def test_cached_target_gone_status_evicts_and_walks_fresh(self): | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/gone.pdf" | |
| stale = "https://cdn-lfs.hf.co/stale-signed" | |
| fresh = "https://cdn-lfs.hf.co/fresh-signed" | |
| for status in (404, 410): | |
| with self.subTest(status=status): | |
| upstream_http.upstream_redirect_cache.clear() | |
| upstream_http.upstream_redirect_cache[("GET", start)] = (time.monotonic(), stale) | |
| gone = FakeResponse(status) | |
| session = FakeSession([gone, FakeResponse(302, {"Location": fresh}), FakeResponse(200)]) | |
| result = await app.open_download_response( | |
| session, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_redirect_url, | |
| ) | |
| self.assertEqual(result.status, 200) | |
| self.assertTrue(gone.released) | |
| self.assertEqual([url for url, _ in session.urls], [stale, start, fresh]) | |
| async def test_fresh_walk_failure_is_not_walked_twice(self): | |
| start = "https://huggingface.co/datasets/VoiceOfML/Test/resolve/main/denied.pdf" | |
| signed = "https://cdn-lfs.hf.co/signed" | |
| for status in (401, 403, 404): | |
| with self.subTest(status=status): | |
| upstream_http.upstream_redirect_cache.clear() | |
| denied = FakeResponse(status) | |
| session = FakeSession([FakeResponse(302, {"Location": signed}), denied, | |
| FakeResponse(302, {"Location": signed}), FakeResponse(200)]) | |
| result = await app.open_download_response( | |
| session, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_redirect_url, | |
| ) | |
| self.assertIs(result, denied) | |
| self.assertFalse(denied.released) | |
| self.assertEqual([url for url, _ in session.urls], [start, signed]) | |
| self.assertEqual(upstream_http.upstream_redirect_cache, {}) | |
| session = FakeSession([aiohttp.ClientConnectionError("down"), FakeResponse(200)]) | |
| with self.assertRaises(aiohttp.ClientConnectionError): | |
| await app.open_download_response(session, start, aiohttp.ClientTimeout(total=1), | |
| url_validator=app.validate_voiceofml_redirect_url) | |
| self.assertEqual(len(session.urls), 1) | |
| # Reader bucket path decoding, admission and redirect scope. | |
| async def test_reader_bucket_path_is_decoded_exactly_once(self): | |
| root = "chapters/ebook/epub/" + "a" * 64 + "/" + "b" * 16 | |
| literal = root + "/chapter-manifest.json" | |
| with patch.object(app, "read_reader_bucket_resource", return_value=b"{}") as read: | |
| response = await app.api_reader_bucket_resource(literal) | |
| self.assertEqual(response.status_code, 200) | |
| read.assert_called_once_with(literal, app.READER_RESOURCE_MAX_BYTES) | |
| resource = root + "/resources/image%41.svg" | |
| source_urls = [] | |
| async def stream(request, source_url, **kwargs): | |
| source_urls.append(source_url) | |
| return app.Response(b"svg") | |
| with patch.object(app, "stream_reader_document", side_effect=stream): | |
| response = await app.api_reader_bucket_resource(resource, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(response.status_code, 200) | |
| self.assertEqual(source_urls, [app.reader_bucket_source_url(resource)]) | |
| self.assertTrue(source_urls[0].endswith("/resources/image%2541.svg")) | |
| for unsafe in (root + "/resources/%2e%2e/x.svg", root + "/resources/%252e%252e/x.svg", | |
| root + "/resources/../x.svg"): | |
| with self.subTest(path=unsafe), patch.object(app, "stream_reader_document", side_effect=stream): | |
| rejected = await app.api_reader_bucket_resource(unsafe, SimpleNamespace(method="GET", headers={})) | |
| self.assertEqual(rejected.status_code, 400) | |
| async def test_reader_bucket_admission_limit_is_retryable_503(self): | |
| path = f"objects/aa/{'a' * 64}/{'b' * 16}/page-manifest.json" | |
| with patch.object(app, "READER_BUCKET_MAX_INFLIGHT", 0), \ | |
| patch.object(app, "read_reader_bucket_resource", return_value=b"{}") as read: | |
| response = await app.api_reader_bucket_resource(path) | |
| self.assertEqual(response.status_code, 503) | |
| self.assertEqual(response.headers["retry-after"], "1") | |
| self.assertEqual(response.headers["cache-control"], "no-store") | |
| read.assert_not_called() | |
| async def test_bucket_pdf_redirect_rejects_unlisted_hf_storage_host(self): | |
| path = "derived/test/" + "a" * 32 + "/document.pdf" | |
| for location in ("https://attacker.hf.co/secret.pdf", "https://datasets.huggingface.co/x.pdf", | |
| "https://hf-mirror.com/buckets/vomebook/pdf-pages-v2/resolve/x.pdf"): | |
| with self.subTest(location=location): | |
| upstream_http.upstream_redirect_cache.clear() | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| redirect = FakeResponse(302, {"Location": location}) | |
| session = app.app.state.http_session = FakeSession([redirect, FakeResponse(200)]) | |
| response = await asyncio.wait_for(app.api_reader_bucket_resource( | |
| path, SimpleNamespace(method="GET", headers={})), 2) | |
| self.assertEqual(response.status_code, 502) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertTrue(redirect.released) | |
| signed = "https://cas-bridge.xethub.hf.co/xet-bridge-us/signed" | |
| app.app.state.upstream_semaphore = asyncio.Semaphore(1) | |
| session = app.app.state.http_session = FakeSession([ | |
| FakeResponse(302, {"Location": signed}), | |
| FakeResponse(200, {"Content-Length": "3", "Content-Type": "application/pdf"}, chunks=(b"pdf",)), | |
| ]) | |
| response = await asyncio.wait_for(app.api_reader_bucket_resource( | |
| path, SimpleNamespace(method="GET", headers={})), 2) | |
| self.assertEqual(response.status_code, 200) | |
| self.assertEqual(b"".join([part async for part in response.body_iterator]), b"pdf") | |
| self.assertEqual([url for url, _ in session.urls][1], signed) | |
| # Reader asset sidecar and v3 catalog fetches. | |
| async def test_reader_assets_sidecar_redirect_hops_are_validated(self): | |
| payload = gzip.compress(json.dumps({"v": 1, "f": {"ready": {"s": 4}}}).encode()) | |
| signed = "https://cas-bridge.xethub.hf.co/xet-bridge-us/sidecar" | |
| cases = ( | |
| ([FakeResponse(302, {"Location": "https://attacker.hf.co/sidecar"})], {}), | |
| ([FakeResponse(302, {"Location": signed}), FakeResponse(200, chunks=(payload,))], | |
| {"ready": {"s": 4}}), | |
| ) | |
| for responses, expected in cases: | |
| with self.subTest(expected=expected): | |
| session = FakeSession(responses) | |
| app.app.state.http_session = session | |
| app.app.state.reader_assets_lock = asyncio.Lock() | |
| with patch.object(app, "reader_assets_cache", {"v": 1, "f": {}}), \ | |
| patch.object(app, "reader_assets_cache_at", 0), \ | |
| patch.object(app, "load_reading_catalog", AsyncMock(return_value=None)): | |
| assets = await app.refresh_reader_assets() | |
| self.assertEqual(assets["f"], expected) | |
| self.assertEqual(len(session.urls), len(responses)) | |
| self.assertTrue(all(not kwargs["allow_redirects"] for _, kwargs in session.urls)) | |
| self.assertTrue(all(response.released for response in responses)) | |
| self.assertEqual(upstream_http.upstream_redirect_cache, {}) | |
| async def test_v3_catalog_redirect_hops_are_validated_and_not_cached(self): | |
| from backend import reader_v3_catalog as v3 | |
| from tests.test_reader_v3_catalog import fixture | |
| pointer, raw = fixture() | |
| escaped = FakeResponse(302, {"Location": "https://attacker.hf.co/pointer"}) | |
| session = FakeSession([escaped]) | |
| with self.assertRaisesRegex(ValueError, "存储范围"): | |
| await v3.load_reading_catalog(session) | |
| self.assertEqual(len(session.urls), 1) | |
| self.assertTrue(escaped.released) | |
| signed = "https://cas-bridge.xethub.hf.co/xet-bridge-us/" | |
| session = FakeSession([ | |
| FakeResponse(302, {"Location": signed + "pointer"}), | |
| FakeResponse(200, chunks=(json.dumps(pointer).encode(),)), | |
| FakeResponse(302, {"Location": signed + "catalog"}), | |
| FakeResponse(200, chunks=(raw,)), | |
| ]) | |
| result = await v3.load_reading_catalog(session) | |
| self.assertEqual(result["generation"], pointer["generation"]) | |
| self.assertEqual([url for url, _ in session.urls], [ | |
| v3.BASE + v3.POINTER, signed + "pointer", v3.BASE + pointer["catalog"]["path"], signed + "catalog"]) | |
| self.assertTrue(all(not kwargs["allow_redirects"] for _, kwargs in session.urls)) | |
| self.assertEqual(session.urls[0][1]["headers"]["Cache-Control"], "no-cache") | |
| self.assertEqual(upstream_http.upstream_redirect_cache, {}) | |
| if __name__ == "__main__": | |
| unittest.main() | |